diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b12d8d86..b58a3ee5 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -110,6 +110,24 @@ def normalize_tool(tool: str) -> str: return normalized +def normalize_agent(agent: str) -> str: + """Canonical name for an ``--agents`` entry. Same alias resolution as + ``normalize_tool`` (``claude-code`` -> ``claude``, ``gemini-cli`` -> ``gemini``), + plus the MCP-only ``cursor`` agent, which has no model routing and so is not a + ``normalize_tool`` value.""" + if agent.strip().lower() == "cursor": + return "cursor" + return normalize_tool(agent) + + +def normalize_agents(spec: str | None) -> set[str] | None: + """Parse a comma-separated ``--agents`` value into canonical agent names, or + ``None`` when unset or empty.""" + if spec is None: + return None + return {normalize_agent(a) for a in spec.split(",") if a.strip()} or None + + def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool: spec = TOOL_SPECS[tool] binary = spec["binary"] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index aabaa01f..f87a47b4 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -28,6 +28,8 @@ explicit_model_arg_value, install_databricks_ai_tools_for_agents, install_tool_binary, + normalize_agent, + normalize_agents, normalize_tool, provider_permission_error, resolve_gemini_provider_model, @@ -1133,7 +1135,7 @@ def _configure_agents_for_mcp( Cursor is MCP-only, so it just needs workspace state established and rides along via MCP_ONLY_CLIENTS. Interactive — prompts for the workspace URL on first run.""" - scope = {a if a == "cursor" else normalize_tool(a) for a in requested} + scope = {normalize_agent(a) for a in requested} ready = set(configured_mcp_clients(load_state(), available_mcp_clients())) to_bootstrap = scope - ready model_agents = sorted(a for a in to_bootstrap if a != "cursor") @@ -1201,11 +1203,7 @@ def mcp_add( (and, if needed, set up) specific agents. """ selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()} - requested_agents = ( - None - if agents is None - else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) - ) + requested_agents = normalize_agents(agents) try: scope = _configure_agents_for_mcp(sorted(requested_agents)) if requested_agents else None add_mcp_command(location=location, services=selected, agents=scope) @@ -1235,11 +1233,7 @@ def mcp_remove( Interactive: shows the servers you currently have configured and unregisters the ones you select. Needs no Databricks login. """ - requested_agents = ( - None - if agents is None - else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) - ) + requested_agents = normalize_agents(agents) try: remove_mcp_command(agents=requested_agents) except RuntimeError as exc: @@ -1315,11 +1309,7 @@ def skills_add( requested_skills = ( None if skills is None else {s.strip() for s in skills.split(",") if s.strip()} ) - requested_agents = ( - None - if agents is None - else ({agent.strip().lower() for agent in agents.split(",") if agent.strip()} or None) - ) + requested_agents = normalize_agents(agents) if mcp and path is not None: raise RuntimeError("--path is not supported when using --mcp") if mcp and requested_skills is not None: @@ -1421,11 +1411,7 @@ def skills_remove( "Removing downloaded skills is not supported yet. Pass --mcp to remove " "schemas from the skills MCP connection." ) - requested_agents = ( - None - if agents is None - else ({agent.strip().lower() for agent in agents.split(",") if agent.strip()} or None) - ) + requested_agents = normalize_agents(agents) remove_skills_command(agents=requested_agents) except RuntimeError as exc: print_err(str(exc)) diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index a59e40a1..16444867 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -19,6 +19,8 @@ explicit_model_arg_value, install_databricks_ai_tools_for_agents, install_tool_binary, + normalize_agent, + normalize_agents, normalize_tool, provider_permission_error, resolve_launch_model, @@ -204,6 +206,32 @@ def test_unknown_raises(self): normalize_tool("unknown-agent") +class TestNormalizeAgent: + def test_resolves_aliases_like_normalize_tool(self): + assert normalize_agent("claude-code") == "claude" + assert normalize_agent(" Gemini-CLI ") == "gemini" + + def test_accepts_cursor(self): + assert normalize_agent("cursor") == "cursor" + assert normalize_agent(" CURSOR ") == "cursor" + + def test_unknown_raises(self): + with pytest.raises(RuntimeError, match="Unsupported"): + normalize_agent("unknown-agent") + + +class TestNormalizeAgents: + def test_none_passthrough(self): + assert normalize_agents(None) is None + + def test_empty_becomes_none(self): + assert normalize_agents(" , ") is None + + def test_parses_dedupes_and_normalizes(self): + assert normalize_agents("claude-code, codex , claude") == {"claude", "codex"} + assert normalize_agents("cursor,claude-code") == {"cursor", "claude"} + + class TestCheckGatewayEndpoint: def test_claude_available_when_models_present(self): assert check_gateway_endpoint({"claude_models": {"sonnet": "s4"}}, "claude") is True diff --git a/tests/test_cli.py b/tests/test_cli.py index 36f4d842..b574b88a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1063,6 +1063,20 @@ def test_mcp_group_lists_web_search(self): assert result.exit_code == 0 assert "web-search" in result.output + def test_mcp_remove_forwards_agent_scope(self): + with patch("ucode.cli.remove_mcp_command") as remove: + result = runner.invoke(app, ["mcp", "remove", "--agents", "claude, codex"]) + + assert result.exit_code == 0, result.output + remove.assert_called_once_with(agents={"claude", "codex"}) + + def test_mcp_remove_normalizes_agent_aliases(self): + with patch("ucode.cli.remove_mcp_command") as remove: + result = runner.invoke(app, ["mcp", "remove", "--agents", "claude-code, gemini-cli"]) + + assert result.exit_code == 0, result.output + remove.assert_called_once_with(agents={"claude", "gemini"}) + class TestAuthTokenCommand: """`ucode auth-token` is the cross-platform apiKeyHelper (#116).""" @@ -1621,6 +1635,17 @@ def test_mcp_remove_forwards_agent_scope(self): assert result.exit_code == 0, result.output remove.assert_called_once_with(agents={"claude", "codex"}) + def test_mcp_remove_normalizes_agent_aliases(self): + # `skill add` resolves aliases like `claude-code`; removal must too, or a + # user who added with an alias can't remove with the same name. + with patch("ucode.cli.remove_skills_command") as remove: + result = runner.invoke( + app, ["skill", "remove", "--mcp", "--agents", "claude-code, gemini-cli"] + ) + + assert result.exit_code == 0, result.output + remove.assert_called_once_with(agents={"claude", "gemini"}) + class TestManagedSkillsOnLaunch: """Managed skills are delivered by download only: the launch path downloads them and never