Configure first-run launches through the launch's --provider

src/ucode/agents/__init__.py

@@ -523,9 +523,11 @@ def _availability_failure_detail(tool: str, state: dict) -> str:
     return " (" + "; ".join(parts) + ")"
 
 
-def configure_single_tool(tool: str, state: dict, *, parent_schema: str | None = None) -> dict:
+def configure_single_tool(
+    tool: str, state: dict, *, parent_schema: str | None = None, provider: str | None = None
+) -> dict:
     """Check availability, configure, and persist state for one tool only."""
-    provider = None if parent_schema else get_provider_service(state, tool)
+    provider = None if parent_schema else provider or get_provider_service(state, tool)
     # A Model Provider Service or parent schema routes through the same gateway and pins no
     # globally discovered Databricks model, so the availability check doesn't apply.
     if not provider and not parent_schema:

src/ucode/cli.py

@@ -2443,7 +2443,9 @@ def claude_router_hook_cmd(
         sys.stdout.write(json.dumps(output))
 
 
-def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = None) -> None:
+def _auto_configure_tool(
+    tool: str, custom_oauth: CustomOAuthConfig | None = None, provider: str | None = None
+) -> None:
     """Configure a tool for launch without sending a separate validation prompt.
 
     The real agent session follows immediately; explicit configure retains the
@@ -2457,7 +2459,7 @@ def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = Non
     configure_kwargs = {"custom_oauth": custom_oauth} if custom_oauth is not None else {}
     state = configure_shared_state(workspace, profile=profile, tools=[tool], **configure_kwargs)
 
-    state = configure_single_tool(tool, state)
+    state = configure_single_tool(tool, state, provider=provider)
 
     spec = TOOL_SPECS[tool]
     console.print(
@@ -2855,10 +2857,12 @@ def _launch_tool(
             skip_cli_version_check=skip_preflight,
         )
         if needs_auto_configure:
+            # Configure through this launch's --provider: a default-gateway setup would write the
+            # managed auth keys that a relayed provider's launch then refuses to start over.
             if custom_oauth is None:
-                _auto_configure_tool(tool)
+                _auto_configure_tool(tool, provider=provider)
             else:
-                _auto_configure_tool(tool, custom_oauth=custom_oauth)
+                _auto_configure_tool(tool, custom_oauth=custom_oauth, provider=provider)
         state = ensure_provider_state(tool)
         # Remembered before the fallback below collapses the two cases: a managed config may not
         # silently override a provider the user typed on the command line (it errors instead).

tests/test_agents_init.py

@@ -206,6 +206,24 @@ def test_managed_parent_skips_global_availability_and_writes_header(self, monkey
 
         configure.assert_called_once_with("claude", state, parent_schema="main.default")
 
+    def test_explicit_relayed_provider_configures_relayed(self, monkeypatch):
+        state = {"workspace": "https://x.databricks.com"}
+        monkeypatch.setattr(
+            agents_mod,
+            "check_gateway_endpoint",
+            lambda *_a: pytest.fail("a provider must not require global model availability"),
+        )
+        monkeypatch.setattr(agents_mod, "resolve_provider_models", lambda *_a: (None, None, True))
+        configure = MagicMock(return_value=state)
+        monkeypatch.setattr(agents_mod, "configure_tool", configure)
+        monkeypatch.setattr(agents_mod, "save_state", lambda _state: None)
+
+        agents_mod.configure_single_tool("claude", state, provider="main.default.relay")
+
+        configure.assert_called_once_with(
+            "claude", state, None, provider="main.default.relay", provider_models=None, relayed=True
+        )
+
     def test_configure_selected_tools_triggers_install(self, monkeypatch):
         captured = self._stub_configure(monkeypatch)
         agents_mod.configure_selected_tools(

tests/test_cli.py

@@ -2990,7 +2990,7 @@ def test_launch_autoconfigures_without_test_prompt(self, tool, has_workspace):
             result = runner.invoke(app, [tool])
 
         assert result.exit_code == 0, result.output
-        mock_configure.assert_called_once_with(tool, configured_state)
+        mock_configure.assert_called_once_with(tool, configured_state, provider=None)
         mock_restore.assert_not_called()
         mock_launch.assert_called_once()
         assert mock_launch.call_args.args[:2] == (tool, configured_state)
@@ -3019,7 +3019,7 @@ def test_triggers_when_no_workspace(self):
             result = runner.invoke(app, ["claude"])
         assert result.exit_code == 0, result.output
         mock_bootstrap.assert_called_once_with("claude", skip_cli_version_check=False)
-        mock_auto.assert_called_once_with("claude")
+        mock_auto.assert_called_once_with("claude", provider=None)
 
     def test_triggers_when_tool_not_in_available_tools(self):
         """Auto-configure runs when workspace exists but the tool wasn't configured."""
@@ -3044,7 +3044,34 @@ def test_triggers_when_tool_not_in_available_tools(self):
             result = runner.invoke(app, ["claude"])
         assert result.exit_code == 0, result.output
         mock_bootstrap.assert_called_once_with("claude", skip_cli_version_check=False)
-        mock_auto.assert_called_once_with("claude")
+        mock_auto.assert_called_once_with("claude", provider=None)
+
+    def test_first_run_configures_through_launch_provider(self):
+        """A default-gateway first-run setup would write the managed auth keys that a relayed
+        provider's launch then refuses to start over."""
+        with (
+            _launch_policy_patches(None) as calls,
+            patch("ucode.cli.load_state", return_value={}),
+            patch("ucode.cli._auto_configure_tool") as mock_auto,
+        ):
+            calls["resolve_provider"].return_value = (None, None, True)
+            result = runner.invoke(app, ["claude", "--provider", "main.default.relay"])
+
+        assert result.exit_code == 0, result.output
+        mock_auto.assert_called_once_with("claude", provider="main.default.relay")
+        assert calls["configure"].call_args.kwargs["relayed"] is True
+
+    def test_auto_configure_forwards_provider(self):
+        with (
+            patch("ucode.cli.load_state", return_value=MINIMAL_STATE),
+            patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE),
+            patch("ucode.cli.configure_single_tool", return_value=MINIMAL_STATE) as mock_configure,
+        ):
+            cli_mod._auto_configure_tool("claude", provider="main.default.relay")
+
+        mock_configure.assert_called_once_with(
+            "claude", MINIMAL_STATE, provider="main.default.relay"
+        )
 
     def test_skipped_when_already_configured(self):
         """Auto-configure is skipped when workspace and tool are already set up."""