From 98854964a5aec4434fb5bd0e740aabdaca868804 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 31 Aug 2026 22:09:35 +0000 Subject: [PATCH 1/2] Use cached launch path for explicit workspace --- src/ucode/cli.py | 25 ++++++++++++++----------- tests/test_cli.py | 27 ++++++++++++++++++++++++--- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 56fe3a0e..1e15e844 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1820,8 +1820,7 @@ def _can_launch_from_cached_config( model: str | None, explicit_provider: str | None, enable_smart_routing_flag: bool, - workspace: str | None, - needs_auto_configure: bool, + workspace_url: str | None, ) -> bool: """Return whether a normal Claude/Codex launch can use its cached config.""" if tool not in CAN_USE_CACHED_CONFIG_AGENTS: @@ -1846,7 +1845,12 @@ def _can_launch_from_cached_config( if managed_agent_config_enabled(): return False - if not (needs_auto_configure or workspace is None): + # `_launch_tool` selects an explicit workspace before loading state. A matching workspace here + # therefore means its cached state was selected (or it was just auto-configured) and is safe to + # launch. Keep rejecting a mismatched state rather than launching against the wrong workspace. + if workspace_url is not None and state.get("workspace") != normalize_workspace_url( + workspace_url + ): return False if tool == "claude": @@ -1860,7 +1864,7 @@ def _launch_tool( provider: str | None = None, refresh: bool = False, skip_preflight: bool = False, - workspace: str | None = None, + workspace_url: str | None = None, enable_smart_routing_flag: bool = False, managed: dict | None = None, recommendation: dict | None = None, @@ -1880,8 +1884,8 @@ def _launch_tool( # An explicit --workspace targets that workspace for this launch (and # auto-configures it if unseen), so `ucode claude --provider ... --workspace ...` # works without a prior `ucode configure`. - if workspace: - set_current_workspace(normalize_workspace_url(workspace)) + if workspace_url: + set_current_workspace(normalize_workspace_url(workspace_url)) existing = load_state() # Workspaces configured with --use-pat export the profile's PAT as # DATABRICKS_BEARER up front so every auth check below (and the @@ -1908,8 +1912,7 @@ def _launch_tool( model=model, explicit_provider=explicit_provider, enable_smart_routing_flag=enable_smart_routing_flag, - workspace=workspace, - needs_auto_configure=needs_auto_configure, + workspace_url=workspace_url, ): print_section(_launch_title(tool)) if forwarded_model: @@ -2325,7 +2328,7 @@ def _launch_managed_default( tool, ctx, skip_preflight=skip_preflight, - workspace=workspace, + workspace_url=workspace, managed=managed, recommendation=recommendation, ) @@ -2401,7 +2404,7 @@ def codex_cmd( provider=provider, refresh=refresh, skip_preflight=skip_preflight, - workspace=workspace, + workspace_url=workspace, enable_smart_routing_flag=enable_smart_routing_flag, ) @@ -2480,7 +2483,7 @@ def claude_cmd( model=model, refresh=refresh, skip_preflight=skip_preflight, - workspace=workspace, + workspace_url=workspace, enable_smart_routing_flag=enable_smart_routing_flag, ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 65b830b6..bffe17a0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1334,6 +1334,28 @@ def test_uses_existing_codex_config_without_preflight(self): mock_ai_tools.assert_not_called() mock_launch.assert_called_once() + def test_workspace_flag_uses_cached_codex_config_without_preflight(self): + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.set_current_workspace") as mock_set_workspace, + 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.smart_routing_v2.enabled", return_value=False), + patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke( + app, ["codex", "--workspace", "https://example.databricks.com/"] + ) + + assert result.exit_code == 0, result.output + mock_set_workspace.assert_called_once_with("https://example.databricks.com") + mock_preflight.assert_not_called() + mock_configure.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 = {} @@ -1434,8 +1456,7 @@ def _kwargs(**overrides): "model": None, "explicit_provider": None, "enable_smart_routing_flag": False, - "workspace": None, - "needs_auto_configure": False, + "workspace_url": None, } kwargs.update(overrides) return kwargs @@ -1500,7 +1521,7 @@ def test_rejects_codex_v2_launch_with_incomplete_model_cache( {"refresh": True}, {"explicit_provider": "catalog.schema.provider"}, {"enable_smart_routing_flag": True}, - {"workspace": "https://other.databricks.com"}, + {"workspace_url": "https://other.databricks.com"}, ], ) def test_rejects_dynamic_launch_overrides(self, override): From 89d94c811a42040e9e92cec9c66ca9c0e37d3b1e Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 31 Aug 2026 22:48:39 +0000 Subject: [PATCH 2/2] ruff --- tests/test_cli.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index bffe17a0..704bde48 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1346,9 +1346,7 @@ def test_workspace_flag_uses_cached_codex_config_without_preflight(self): patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), patch("ucode.cli.launch_agent") as mock_launch, ): - result = runner.invoke( - app, ["codex", "--workspace", "https://example.databricks.com/"] - ) + result = runner.invoke(app, ["codex", "--workspace", "https://example.databricks.com/"]) assert result.exit_code == 0, result.output mock_set_workspace.assert_called_once_with("https://example.databricks.com")