From 00892fa044055261717cd410d6007ff9ad12246e Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Thu, 27 Aug 2026 20:27:14 -0400 Subject: [PATCH 1/2] Speed up Claude and Codex launches --- README.md | 12 +++- src/ucode/agents/__init__.py | 8 ++- src/ucode/agents/claude.py | 27 ++++++++- src/ucode/agents/codex.py | 15 +++++ src/ucode/cli.py | 68 +++++++++++++++++++++- tests/test_agent_claude.py | 16 ++++++ tests/test_agent_codex.py | 23 ++++++++ tests/test_agents_init.py | 2 +- tests/test_cli.py | 106 +++++++++++++++++++++++++++++++++++ 9 files changed, 269 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index f26ea6ba..65c8ad61 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: @@ -285,8 +285,10 @@ their next ucode run. | `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 | @@ -307,14 +309,18 @@ their next ucode run. | `ucode apply` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | | `ucode apply --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 0785342f..0ad76595 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -90,11 +90,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 5fad5225..d8f4090c 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -42,7 +42,9 @@ CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" +CLAUDE_MCP_CONFIG_PATH = Path.home() / ".claude.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" +WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp" SPEC: ToolSpec = { "binary": "claude", @@ -435,6 +437,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 @@ -562,7 +578,16 @@ 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): + if _register_web_search_mcp(state["workspace"], web_search_model, state.get("profile")): + state[WEB_SEARCH_MCP_STATE_KEY] = web_search_entry + else: + 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 2b7415e2..36603882 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -111,6 +111,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 910af654..bf6bc234 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1597,10 +1597,43 @@ def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: _download_managed_skills(managed, state) +def _can_use_local_agent_config( + tool: str, + state: dict, + *, + refresh: bool, + model: str | None, + explicit_provider: str | None, + enable_smart_routing_flag: bool, + skip_preflight: bool, + workspace: str | None, + needs_auto_configure: bool, +) -> bool: + """Return whether a normal Claude/Codex launch can trust local config.""" + if tool not in ("claude", "codex"): + return False + if refresh or model or explicit_provider is not None: + return False + if enable_smart_routing_flag or skip_preflight: + return False + if managed_agent_config_enabled(): + return False + if not (needs_auto_configure or workspace is None): + return False + + routing_agent = _ROUTING_AGENTS.get(tool) + if routing_agent is not None and routing_agent.smart_routing_enabled(state): + 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, @@ -1638,6 +1671,21 @@ 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_use_local_agent_config( + tool, + state, + refresh=refresh, + model=model, + explicit_provider=explicit_provider, + enable_smart_routing_flag=enable_smart_routing_flag, + skip_preflight=skip_preflight, + workspace=workspace, + needs_auto_configure=needs_auto_configure, + ): + print_section(f"ucode with {TOOL_SPECS[tool]['display']}") + 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 @@ -2043,6 +2091,14 @@ def codex_cmd( "before any `--` separator.", ), ] = None, + refresh: Annotated[ + bool, + typer.Option( + "--refresh", + help="Refresh Databricks auth, gateway, models, managed config, and Codex configuration " + "before launching.", + ), + ] = False, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, @@ -2074,6 +2130,7 @@ def codex_cmd( "codex", ctx, provider=provider, + refresh=refresh, skip_preflight=skip_preflight, workspace=workspace, enable_smart_routing_flag=enable_smart_routing_flag, @@ -2102,6 +2159,14 @@ def claude_cmd( "Pass before any `--` separator; not usable with --provider.", ), ] = None, + refresh: Annotated[ + bool, + typer.Option( + "--refresh", + help="Refresh Databricks auth, gateway, models, managed config, and Claude settings " + "before launching.", + ), + ] = False, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, @@ -2134,6 +2199,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, @@ -2303,7 +2369,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 b45b6303..5b4c15ff 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -546,6 +546,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 f346af90..0618fab1 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 ad45a9a7..8f68a734 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_codex_disable_removes_hooks_without_launching(self): with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), @@ -291,6 +298,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): with ( patch("ucode.cli.ensure_bootstrap_dependencies"), @@ -862,6 +875,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 = {} @@ -937,6 +990,58 @@ def test_skipped_when_already_configured(self): mock_auto.assert_not_called() +class TestLocalAgentConfigPredicate: + @staticmethod + def _kwargs(**overrides): + kwargs = { + "refresh": False, + "model": None, + "explicit_provider": None, + "enable_smart_routing_flag": False, + "skip_preflight": 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_use_local_agent_config("codex", MINIMAL_STATE, **self._kwargs()) + is True + ) + + @pytest.mark.parametrize( + "override", + [ + {"refresh": True}, + {"explicit_provider": "catalog.schema.provider"}, + {"enable_smart_routing_flag": True}, + {"skip_preflight": 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_use_local_agent_config( + "codex", MINIMAL_STATE, **self._kwargs(**override) + ) + is False + ) + + class TestPassthroughArgs: @pytest.mark.parametrize( "tool,extra_args", @@ -2384,6 +2489,7 @@ def _patches(cfg): 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", 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"), From 3491e5ec33592bf728e48af2a7c45051ab2d7072 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 28 Aug 2026 18:16:11 +0000 Subject: [PATCH 2/2] address comments and clean up launch config logic --- src/ucode/agents/claude.py | 8 +++++--- src/ucode/cli.py | 42 +++++++++++++++++++++++--------------- tests/test_cli.py | 22 +++++++++++++++----- 3 files changed, 47 insertions(+), 25 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index d8f4090c..59da5bae 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -582,10 +582,12 @@ def _compose(base: dict) -> dict: state["workspace"], web_search_model, state.get("profile") ) if not _web_search_mcp_is_current(state, web_search_entry): - if _register_web_search_mcp(state["workspace"], web_search_model, state.get("profile")): + # 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[WEB_SEARCH_MCP_STATE_KEY] = web_search_entry else: state.pop(WEB_SEARCH_MCP_STATE_KEY, None) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index bf6bc234..ba87fed7 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1407,6 +1407,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: @@ -1485,6 +1486,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") @@ -1496,7 +1501,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, ) @@ -1597,7 +1602,7 @@ def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: _download_managed_skills(managed, state) -def _can_use_local_agent_config( +def _can_launch_from_cached_config( tool: str, state: dict, *, @@ -1605,25 +1610,26 @@ def _can_use_local_agent_config( model: str | None, explicit_provider: str | None, enable_smart_routing_flag: bool, - skip_preflight: bool, workspace: str | None, needs_auto_configure: bool, ) -> bool: - """Return whether a normal Claude/Codex launch can trust local config.""" - if tool not in ("claude", "codex"): + """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 - if enable_smart_routing_flag or skip_preflight: + 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 - routing_agent = _ROUTING_AGENTS.get(tool) - if routing_agent is not None and routing_agent.smart_routing_enabled(state): - return False if tool == "claude": return claude_agent.CLAUDE_SETTINGS_PATH.exists() return codex_agent.has_ucode_config() @@ -1671,18 +1677,17 @@ 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_use_local_agent_config( + if _can_launch_from_cached_config( tool, state, refresh=refresh, model=model, explicit_provider=explicit_provider, enable_smart_routing_flag=enable_smart_routing_flag, - skip_preflight=skip_preflight, workspace=workspace, needs_auto_configure=needs_auto_configure, ): - print_section(f"ucode with {TOOL_SPECS[tool]['display']}") + print_section(_launch_title(tool)) print_success(f"Starting {TOOL_SPECS[tool]['display']}") launch_agent(tool, state, ctx.args) return @@ -1864,7 +1869,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: @@ -1925,6 +1930,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. @@ -2095,8 +2105,7 @@ def codex_cmd( bool, typer.Option( "--refresh", - help="Refresh Databricks auth, gateway, models, managed config, and Codex configuration " - "before launching.", + help=REFRESH_HELP, ), ] = False, skip_preflight: SkipPreflightOption = False, @@ -2163,8 +2172,7 @@ def claude_cmd( bool, typer.Option( "--refresh", - help="Refresh Databricks auth, gateway, models, managed config, and Claude settings " - "before launching.", + help=REFRESH_HELP, ), ] = False, skip_preflight: SkipPreflightOption = False, diff --git a/tests/test_cli.py b/tests/test_cli.py index 8f68a734..7011a69a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -990,7 +990,20 @@ def test_skipped_when_already_configured(self): mock_auto.assert_not_called() -class TestLocalAgentConfigPredicate: +@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 = { @@ -998,7 +1011,6 @@ def _kwargs(**overrides): "model": None, "explicit_provider": None, "enable_smart_routing_flag": False, - "skip_preflight": False, "workspace": None, "needs_auto_configure": False, } @@ -1013,7 +1025,7 @@ def test_accepts_configured_codex_launch(self): patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), ): assert ( - cli_mod._can_use_local_agent_config("codex", MINIMAL_STATE, **self._kwargs()) + cli_mod._can_launch_from_cached_config("codex", MINIMAL_STATE, **self._kwargs()) is True ) @@ -1023,7 +1035,6 @@ def test_accepts_configured_codex_launch(self): {"refresh": True}, {"explicit_provider": "catalog.schema.provider"}, {"enable_smart_routing_flag": True}, - {"skip_preflight": True}, {"workspace": "https://other.databricks.com"}, ], ) @@ -1035,7 +1046,7 @@ def test_rejects_dynamic_launch_overrides(self, override): patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), ): assert ( - cli_mod._can_use_local_agent_config( + cli_mod._can_launch_from_cached_config( "codex", MINIMAL_STATE, **self._kwargs(**override) ) is False @@ -2488,6 +2499,7 @@ 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(