diff --git a/README.md b/README.md index c8168b9d..3d61d639 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ ucode pi # Pi ucode cursor # Cursor Agent (MCP only — see below) ``` -On first launch, `ucode` will prompt for your Databricks workspace URL, authenticate, and configure that tool automatically. Subsequent launches go straight to the agent. +On first launch, `ucode` will prompt for your Databricks workspace URL, authenticate, and configure that tool automatically. Subsequent Claude and Codex launches use the generated local settings directly. Use `ucode claude --refresh` or `ucode codex --refresh` when you want to re-check Databricks and update the model/configuration. Pass flags directly to the underlying tool: @@ -342,8 +342,10 @@ The output looks like: | `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | | `ucode codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents | | `ucode codex --disable-smart-routing` | Disable routing and remove ucode's Codex routing hooks | +| `ucode codex --refresh` | Re-check Databricks, refresh models/configuration, and launch Codex | | `ucode claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents | | `ucode claude --disable-smart-routing` | Disable routing and remove ucode's Claude Code routing hooks | +| `ucode claude --refresh` | Re-check Databricks, refresh models/configuration, and launch Claude Code | | `ucode configure --skip-validate` | Write configs without sending a test message through each agent | | `ucode configure --agents claude,codex,pi --skip-unavailable` | Configure the requested agents that are available; skip the rest with a warning | | `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | @@ -367,14 +369,18 @@ The output looks like: | `ucode publish -f ` | Publish a config file exported with `ucode export` instead of the locally authored one | | `ucode publish --yes` | Publish without the confirmation prompt | +Databricks AI Tools are installed only by `ucode configure`, never by `ucode ` launches. +Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ucode configure` to +control the installation. + ## Managed Local Files `ucode` manages these files: | File | Tool | |------|------| -| `~/.codex/config.toml` | Codex | -| `~/.claude/settings.json` | Claude Code | +| `~/.codex/ucode.config.toml` (or legacy `~/.codex/config.toml`) | Codex | +| `~/.claude/ucode-settings.json` | Claude Code settings generated by ucode | | `~/.gemini/.env` | Gemini CLI | | `~/.config/opencode/opencode.json` | OpenCode | | `~/.copilot/.env` | GitHub Copilot CLI | diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 0400ba95..d2480390 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -89,11 +89,15 @@ def install_databricks_ai_tools_for_agents(tools: list[str], state: dict) -> None: - """Install Databricks AI Tools for the coding agents that support them - (gemini/pi have no ``aitools`` support and are dropped).""" + """Install Databricks AI Tools for supported agents. + + Gemini and Pi have no ``aitools`` support and are dropped. + """ if state.get("databricks_ai_tools_enabled", True) is False: return agents = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS] + if not agents: + return install_ai_tools(agents, state.get("profile")) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 04578df9..92d18f2b 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -53,9 +53,11 @@ GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" +CLAUDE_MCP_CONFIG_PATH = Path.home() / ".claude.json" # The default model is stored in Claude's default user settings, not the ucode settings. CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" +WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp" SPEC: ToolSpec = { "binary": "claude", @@ -473,6 +475,20 @@ def _register_web_search_mcp(workspace: str, search_model: str, profile: str | N return True +def _web_search_mcp_is_current(state: dict, entry: dict) -> bool: + """Return whether the desired web-search entry is already registered. + + The persisted entry acts as a cheap fingerprint, while reading Claude's config repairs a + registration removed or edited outside ucode. Avoiding the Claude CLI here matters: each + ``claude mcp`` subprocess takes roughly 0.8 seconds during a launch. + """ + if state.get(WEB_SEARCH_MCP_STATE_KEY) != entry: + return False + config = read_json_safe(CLAUDE_MCP_CONFIG_PATH) + servers = config.get("mcpServers") + return isinstance(servers, dict) and servers.get(WEB_SEARCH_MCP_NAME) == entry + + def _unregister_web_search_mcp() -> None: """Remove the web_search MCP server from all scopes. Used by revert.""" from ucode.mcp import MCP_CLEANUP_SCOPES, remove_claude_mcp_server @@ -600,7 +616,18 @@ def _compose(base: dict) -> dict: _write_managed_settings(_compose, relayed) if web_search_model: - _register_web_search_mcp(state["workspace"], web_search_model, state.get("profile")) + web_search_entry = _web_search_mcp_entry( + state["workspace"], web_search_model, state.get("profile") + ) + if not _web_search_mcp_is_current(state, web_search_entry): + # Registration runs multiple `claude mcp` subprocesses and can take several seconds. + registration_success = _register_web_search_mcp( + state["workspace"], web_search_model, state.get("profile") + ) + if registration_success: + state[WEB_SEARCH_MCP_STATE_KEY] = web_search_entry + else: + state.pop(WEB_SEARCH_MCP_STATE_KEY, None) # Persist relayed mode + proxy port so launch() wires the refresh proxy and # subscription login; cleared on a non-relayed launch. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2f34679b..4b1f555a 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -110,6 +110,21 @@ def _use_legacy_layout() -> bool: return parsed < MINIMUM_CODEX_VERSION +def has_ucode_config() -> bool: + """Return whether ucode has already written a Codex configuration.""" + if CODEX_CONFIG_PATH.exists(): + return True + if not LEGACY_CODEX_CONFIG_PATH.exists(): + return False + doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH) + profiles = doc.get("profiles") + return ( + doc.get("profile") == CODEX_PROFILE_NAME + and isinstance(profiles, dict) + and isinstance(profiles.get(CODEX_PROFILE_NAME), dict) + ) + + def _provider_block( workspace: str, databricks_profile: str | None, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c642b8cb..e5b73b24 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1533,6 +1533,7 @@ def _auto_configure_tool(tool: str) -> None: # function names via their agent module and their routing module. _ROUTING_AGENTS = {"codex": codex_agent, "claude": claude_agent} _ROUTING_MODULES = {"codex": codex_routing, "claude": claude_routing} +CAN_USE_CACHED_CONFIG_AGENTS = frozenset({"claude", "codex"}) def _reject_disabled_agent(managed: dict | None, tool: str) -> None: @@ -1611,6 +1612,10 @@ def _fetch_budget_recommendation(state: dict, managed: dict | None) -> dict | No return recommendation +def _launch_title(tool: str) -> str: + return f"Launching {TOOL_SPECS[tool]['display'].title()} with Unity Gateway" + + def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = None) -> None: """Show the workspace budget this launch spends against, when one is configured.""" agent = recommendation.get("agent") @@ -1622,7 +1627,7 @@ def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = line = recommendation_line(display_agent, recommendation.get("model"), percent) panel = render_budget_panel( recommendation, - title=f"ucode with {TOOL_SPECS[tool]['display']}", + title=_launch_title(tool), extra_lines=[line] if line else None, managed=managed, ) @@ -1723,10 +1728,49 @@ def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: _download_managed_skills(managed, state) +def _can_launch_from_cached_config( + tool: str, + state: dict, + *, + refresh: bool, + model: str | None, + explicit_provider: str | None, + enable_smart_routing_flag: bool, + workspace: str | None, + needs_auto_configure: bool, +) -> bool: + """Return whether a normal Claude/Codex launch can use its cached config.""" + if tool not in CAN_USE_CACHED_CONFIG_AGENTS: + return False + + if refresh or model or explicit_provider is not None: + return False + + smart_routing_enabled = _ROUTING_AGENTS[tool].smart_routing_enabled(state) + legacy_smart_routing_enabled = enable_smart_routing_flag or smart_routing_enabled + + # Legacy smart routing overwrites the model into ucode-settings.json and so cannot use the + # cached state. Smart routing v2 will use PTY so can use the fast path. + if legacy_smart_routing_enabled: + return False + + # If managed agent config is enabled, we cannot use the cached state in case the config changed. + if managed_agent_config_enabled(): + return False + + if not (needs_auto_configure or workspace is None): + return False + + if tool == "claude": + return claude_agent.CLAUDE_SETTINGS_PATH.exists() + return codex_agent.has_ucode_config() + + def _launch_tool( tool_name: str, ctx: typer.Context, provider: str | None = None, + refresh: bool = False, skip_preflight: bool = False, workspace: str | None = None, enable_smart_routing_flag: bool = False, @@ -1764,6 +1808,20 @@ def _launch_tool( # back to whatever `ucode configure` saved for this tool. provider = provider or get_provider_service(state, tool) routing_agent = _ROUTING_AGENTS.get(tool) + if _can_launch_from_cached_config( + tool, + state, + refresh=refresh, + model=model, + explicit_provider=explicit_provider, + enable_smart_routing_flag=enable_smart_routing_flag, + workspace=workspace, + needs_auto_configure=needs_auto_configure, + ): + print_section(_launch_title(tool)) + print_success(f"Starting {TOOL_SPECS[tool]['display']}") + launch_agent(tool, state, ctx.args) + return # Fetched before `configure_shared_state` because it decides whether this agent may launch # at all and whether the model discovery below can be skipped. # Bare `ucode` already fetched one to choose the agent; refetching would double the @@ -1947,7 +2005,7 @@ def _launch_tool( route_root_model=route_root_model, custom_model=model if tool == "claude" else None, ) - print_section(f"ucode with {TOOL_SPECS[tool]['display']}") + print_section(_launch_title(tool)) if managed is not None: print_kv("Config", "workspace-managed") if provider: @@ -2018,6 +2076,11 @@ def _launch_tool( ), ] +REFRESH_HELP = ( + "Refresh Databricks auth, gateway, models, managed config, and agent configuration before " + "launching." +) + # Ignore the workspace's managed coding-agent config for this one command, on both # `ucode configure` and the launchers. Accepted (and no-op) even when the managed-config # feature is off, so a headless launcher can always pass it. @@ -2184,6 +2247,13 @@ def codex_cmd( "before any `--` separator.", ), ] = None, + refresh: Annotated[ + bool, + typer.Option( + "--refresh", + help=REFRESH_HELP, + ), + ] = False, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, @@ -2215,6 +2285,7 @@ def codex_cmd( "codex", ctx, provider=provider, + refresh=refresh, skip_preflight=skip_preflight, workspace=workspace, enable_smart_routing_flag=enable_smart_routing_flag, @@ -2243,6 +2314,13 @@ def claude_cmd( "Pass before any `--` separator; not usable with --provider.", ), ] = None, + refresh: Annotated[ + bool, + typer.Option( + "--refresh", + help=REFRESH_HELP, + ), + ] = False, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, @@ -2285,6 +2363,7 @@ def claude_cmd( ctx, provider=provider, model=model, + refresh=refresh, skip_preflight=skip_preflight, workspace=workspace, enable_smart_routing_flag=enable_smart_routing_flag, @@ -2454,7 +2533,7 @@ def configure( typer.Option( "--enable-databricks-ai-tools/--disable-databricks-ai-tools", help="Install Databricks AI Tools (skills + plugins that teach agents to use " - "Databricks) for the configured agents. Installed by default; pass " + "Databricks) for the configured agents. Installation is configure-only; pass " "--disable-databricks-ai-tools to opt out.", ), ] = None, diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 090d1d91..832c1f51 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -573,6 +573,22 @@ def test_relayed_skips_managed_write(self, monkeypatch): class TestRegisterWebSearchMcp: + def test_skips_registration_when_entry_is_current(self, monkeypatch): + entry = claude._web_search_mcp_entry(WS, "m", "profile") + state = {claude.WEB_SEARCH_MCP_STATE_KEY: entry} + monkeypatch.setattr( + claude, + "read_json_safe", + lambda path: {"mcpServers": {claude.WEB_SEARCH_MCP_NAME: entry}}, + ) + assert claude._web_search_mcp_is_current(state, entry) is True + + def test_detects_registration_drift(self, monkeypatch): + entry = claude._web_search_mcp_entry(WS, "m", "profile") + state = {claude.WEB_SEARCH_MCP_STATE_KEY: entry} + monkeypatch.setattr(claude, "read_json_safe", lambda path: {"mcpServers": {}}) + assert claude._web_search_mcp_is_current(state, entry) is False + def test_clears_existing_then_adds(self, monkeypatch): import ucode.mcp as mcp_mod diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 54fe1aa5..9560edb2 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -25,6 +25,29 @@ def test_display(self): assert codex.SPEC["display"] == "Codex" +class TestHasUcodeConfig: + def test_detects_profile_config(self, tmp_path, monkeypatch): + config_path = tmp_path / "ucode.config.toml" + legacy_path = tmp_path / "config.toml" + legacy_path.write_text( + 'profile = "ucode"\n\n[profiles.ucode]\nmodel_provider = "ucode-databricks"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) + + assert codex.has_ucode_config() is True + + def test_ignores_unrelated_legacy_config(self, tmp_path, monkeypatch): + config_path = tmp_path / "ucode.config.toml" + legacy_path = tmp_path / "config.toml" + legacy_path.write_text('profile = "default"\n', encoding="utf-8") + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) + + assert codex.has_ucode_config() is False + + class TestRenderOverlay: def test_uses_profile_file_shape_without_legacy_profiles(self): overlay = codex.render_overlay(WS) diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 69216fac..de40affb 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -73,7 +73,7 @@ def _capture(self, monkeypatch): def test_maps_supported_tools_and_drops_others(self, monkeypatch): captured = self._capture(monkeypatch) - # gemini and pi aren't supported by `databricks aitools`, so they drop. + # Gemini and Pi aren't supported by `databricks aitools`, so they drop. install_databricks_ai_tools_for_agents( ["claude", "codex", "gemini", "pi"], {"profile": "prof"} ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4d6ace94..af085094 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -220,6 +220,13 @@ def test_codex_enable_smart_routing_is_consumed_by_ucode(self): assert mock_launch.call_args.kwargs["enable_smart_routing_flag"] is True assert mock_launch.call_args.args[1].args == [] + def test_codex_refresh_is_consumed_by_ucode(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke(app, ["codex", "--refresh"]) + + assert result.exit_code == 0, result.output + assert mock_launch.call_args.kwargs["refresh"] is True + def test_claude_enable_model_discovery_sets_ucode_env(self): with patch("ucode.cli._launch_tool") as mock_launch: result = runner.invoke(app, ["claude", "--enable-model-discovery"]) @@ -346,6 +353,12 @@ def test_model_threads_through_to_launch(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.kwargs["model"] == "cat.schema.claude-opus-5" + def test_refresh_threads_through_to_launch(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke(app, ["claude", "--refresh"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args.kwargs["refresh"] is True + def test_model_threads_to_claude_as_custom_model(self, monkeypatch): monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) with ( @@ -998,6 +1011,46 @@ def test_reverts_mcp_configs_before_clearing_state(self): class TestAutoConfigureOnFirstRun: + def test_uses_existing_claude_settings_without_preflight(self, tmp_path): + from pathlib import Path + + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text("{}", encoding="utf-8") + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state") as mock_preflight, + patch("ucode.cli.configure_tool") as mock_configure, + patch("ucode.cli.claude_agent.CLAUDE_SETTINGS_PATH", Path(settings_path)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + mock_preflight.assert_not_called() + mock_configure.assert_not_called() + mock_launch.assert_called_once() + + def test_uses_existing_codex_config_without_preflight(self): + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state") as mock_preflight, + patch("ucode.cli.configure_tool") as mock_configure, + patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), + patch("ucode.cli.install_databricks_ai_tools_for_agents") as mock_ai_tools, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["codex"]) + + assert result.exit_code == 0, result.output + mock_preflight.assert_not_called() + mock_configure.assert_not_called() + mock_ai_tools.assert_not_called() + mock_launch.assert_called_once() + def test_triggers_when_no_workspace(self): """Auto-configure runs when state has no workspace.""" empty_state = {} @@ -1073,6 +1126,69 @@ def test_skipped_when_already_configured(self): mock_auto.assert_not_called() +@pytest.mark.parametrize( + ("tool", "expected"), + [ + ("claude", "Launching Claude Code with Unity Gateway"), + ("codex", "Launching Codex with Unity Gateway"), + ], +) +def test_launch_title(tool, expected): + from ucode.cli import _launch_title + + assert _launch_title(tool) == expected + + +class TestCachedConfigPredicate: + @staticmethod + def _kwargs(**overrides): + kwargs = { + "refresh": False, + "model": None, + "explicit_provider": None, + "enable_smart_routing_flag": False, + "workspace": None, + "needs_auto_configure": False, + } + kwargs.update(overrides) + return kwargs + + def test_accepts_configured_codex_launch(self): + import ucode.cli as cli_mod + + with ( + patch("ucode.cli.managed_agent_config_enabled", return_value=False), + patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), + ): + assert ( + cli_mod._can_launch_from_cached_config("codex", MINIMAL_STATE, **self._kwargs()) + is True + ) + + @pytest.mark.parametrize( + "override", + [ + {"refresh": True}, + {"explicit_provider": "catalog.schema.provider"}, + {"enable_smart_routing_flag": True}, + {"workspace": "https://other.databricks.com"}, + ], + ) + def test_rejects_dynamic_launch_overrides(self, override): + import ucode.cli as cli_mod + + with ( + patch("ucode.cli.managed_agent_config_enabled", return_value=False), + patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), + ): + assert ( + cli_mod._can_launch_from_cached_config( + "codex", MINIMAL_STATE, **self._kwargs(**override) + ) + is False + ) + + class TestPassthroughArgs: @pytest.mark.parametrize( "tool,extra_args", @@ -2519,7 +2635,9 @@ def _patches(cfg): patch("ucode.cli._auto_configure_tool"), patch("ucode.cli.load_state", return_value=MINIMAL_STATE), patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli._can_launch_from_cached_config", return_value=False), patch("ucode.cli.configure_shared_state", cfg), + patch("ucode.cli.codex_agent.has_ucode_config", return_value=False), patch( "ucode.cli.resolve_launch_model", return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"),