feat: add --model flag to ucode copilot

src/ucode/agents/__init__.py

@@ -366,8 +366,15 @@ def configure_tool(
     return result
 
 
-def launch(tool: str, state: dict, tool_args: list[str]) -> None:
-    _MODULES[tool].launch(state, tool_args)
+def launch(tool: str, state: dict, tool_args: list[str], model: str | None = None) -> None:
+    # Only copilot's launch loop re-derives its model on every token refresh
+    # (default_model(state)), so only it needs an explicit override threaded
+    # through past that point; other tools already bake --model into the
+    # config file written by configure_tool.
+    if tool == "copilot":
+        copilot.launch(state, tool_args, model_override=model)
+    else:
+        _MODULES[tool].launch(state, tool_args)
 
 
 def check_gateway_endpoint(state: dict, tool: str) -> bool:
@@ -412,8 +419,15 @@ def _availability_failure_detail(tool: str, state: dict) -> str:
     return " (" + "; ".join(parts) + ")"
 
 
-def configure_single_tool(tool: str, state: dict) -> dict:
-    """Check availability, configure, and persist state for one tool only."""
+def configure_single_tool(tool: str, state: dict, explicit_model: str | None = None) -> dict:
+    """Check availability, configure, and persist state for one tool only.
+
+    ``explicit_model`` is a caller-requested model (e.g. `ucode copilot --model X` on a
+    first-time configure) that must win over the automatic default pick, so the model written
+    to the config and smoke-tested is X — otherwise a bad automatic pick can fail validation
+    before X is ever tried. The availability check above is model-independent
+    (``check_gateway_endpoint``), so it still gates on discovery finding any model at all.
+    """
     provider = get_provider_service(state, tool)
     # A Model Provider Service routes through the same gateway and pins no
     # Databricks model, so the per-tool model availability check doesn't apply.
@@ -425,14 +439,16 @@ def configure_single_tool(tool: str, state: dict) -> dict:
             raise RuntimeError(
                 f"{TOOL_SPECS[tool]['display']} is not available on this workspace.{detail}"
             )
-    state = _configure_one(tool, state, provider)
+    state = _configure_one(tool, state, provider, explicit_model=explicit_model)
     available_tools = list(set((state.get("available_tools") or []) + [tool]))
     state["available_tools"] = available_tools
     save_state(state)
     return state
 
 
-def _configure_one(tool: str, state: dict, provider: str | None) -> dict:
+def _configure_one(
+    tool: str, state: dict, provider: str | None, explicit_model: str | None = None
+) -> dict:
     """Write one tool's config, routing through ``provider`` when set."""
     if provider:
         provider_models, error, relayed = resolve_provider_models(tool, state, provider)
@@ -443,7 +459,7 @@ def _configure_one(tool: str, state: dict, provider: str | None) -> dict:
         )
     if tool == "codex":
         return configure_tool("codex", state)
-    state, model = resolve_launch_model(tool, state, None)
+    state, model = resolve_launch_model(tool, state, explicit_model)
     return configure_tool(tool, state, model)
 
 
@@ -504,8 +520,14 @@ def ensure_provider_state(tool: str) -> dict:
     return state
 
 
-def validate_tool(tool: str) -> tuple[bool, str]:
-    """Invoke a tool with a simple prompt to verify it works. Returns (ok, error_msg)."""
+def validate_tool(tool: str, model: str | None = None) -> tuple[bool, str]:
+    """Invoke a tool with a simple prompt to verify it works. Returns (ok, error_msg).
+
+    ``model`` is an explicit --model request (currently only honored for copilot): without it,
+    validation would smoke-test the automatic default pick instead of the model the caller
+    actually asked for, so a bad automatic pick could fail validation before the requested
+    model is ever tried.
+    """
     spec = TOOL_SPECS[tool]
     binary = spec["binary"]
     module = _MODULES[tool]
@@ -517,7 +539,13 @@ def validate_tool(tool: str) -> tuple[bool, str]:
     env = None
     if hasattr(module, "validate_env"):
         try:
-            env = module.validate_env(load_state())
+            # `copilot.validate_env` is the only variant accepting `model_override`; reference
+            # it directly rather than through the union-typed `module` so the extra kwarg
+            # type-checks (ty can't narrow `module`'s type from the `tool == "copilot"` check).
+            if tool == "copilot":
+                env = copilot.validate_env(load_state(), model_override=model)
+            else:
+                env = module.validate_env(load_state())
         except RuntimeError:
             env = None
     try:

src/ucode/agents/copilot.py

@@ -168,30 +168,34 @@ def write_tool_config(
     return state, token
 
 
-def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> tuple[str, str]:
-    model = default_model(state)
+def _refresh_token_once(
+    state: dict, *, force_refresh: bool = False, model_override: str | None = None
+) -> tuple[str, str]:
+    model = model_override or default_model(state)
     if not model:
         raise RuntimeError("No Copilot model is available on this workspace.")
     _, token = write_tool_config(state, model, force_refresh=force_refresh)
     return model, token
 
 
-def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
+def _refresh_forever(
+    state: dict, stop_event: threading.Event, model_override: str | None = None
+) -> None:
     while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
         try:
-            _refresh_token_once(state, force_refresh=True)
+            _refresh_token_once(state, force_refresh=True, model_override=model_override)
         except RuntimeError:
             continue
 
 
-def launch(state: dict, tool_args: list[str]) -> None:
-    model, token = _refresh_token_once(state)
+def launch(state: dict, tool_args: list[str], model_override: str | None = None) -> None:
+    model, token = _refresh_token_once(state, model_override=model_override)
     env = build_runtime_env(state["workspace"], model, token)
 
     stop_event = threading.Event()
     refresher = threading.Thread(
         target=_refresh_forever,
-        args=(state, stop_event),
+        args=(state, stop_event, model_override),
         daemon=True,
     )
     refresher.start()
@@ -225,12 +229,12 @@ def mcp_config_args() -> list[str]:
     return ["--additional-mcp-config", f"@{COPILOT_MCP_CONFIG_PATH}"]
 
 
-def validate_env(state: dict) -> dict[str, str]:
+def validate_env(state: dict, model_override: str | None = None) -> dict[str, str]:
     """Inject BYOK env vars for the validation subprocess (Copilot doesn't auto-load .env)."""
     workspace = state.get("workspace")
     if not workspace:
         raise RuntimeError("No workspace configured.")
-    model = default_model(state)
+    model = model_override or default_model(state)
     if not model:
         raise RuntimeError("No Copilot model is available on this workspace.")
     token = get_databricks_token(workspace, state.get("profile"))

src/ucode/cli.py

@@ -1469,16 +1469,22 @@ def claude_router_hook_cmd(
         sys.stdout.write(json.dumps(output))
 
 
-def _auto_configure_tool(tool: str) -> None:
-    """First-time setup for a single tool — mirrors configure_workspace_command."""
+def _auto_configure_tool(tool: str, model: str | None = None) -> None:
+    """First-time setup for a single tool — mirrors configure_workspace_command.
+
+    ``model`` is an explicit --model request (currently only threaded for copilot): without it,
+    a first-time `ucode copilot --model X` would auto-configure and validate against the
+    automatic sonnet/opus/haiku/codex default instead of X, so a bad automatic pick could fail
+    validation and abort the launch before X is ever tried.
+    """
     existing = load_state()
     workspace = existing.get("workspace")
     profile = existing.get("profile")
     if not workspace:
         workspace, profile = _prompt_for_configuration(tool)
     state = configure_shared_state(workspace, profile=profile, tools=[tool])
 
-    state = configure_single_tool(tool, state)
+    state = configure_single_tool(tool, state, explicit_model=model)
 
     spec = TOOL_SPECS[tool]
     console.print(
@@ -1493,7 +1499,7 @@ def _auto_configure_tool(tool: str) -> None:
     )
 
     with spinner(f"Validating {spec['display']}..."):
-        ok, err = validate_tool(tool)
+        ok, err = validate_tool(tool, model=model)
     if ok:
         print_success(f"{spec['display']} is working")
     else:
@@ -1733,7 +1739,7 @@ def _launch_tool(
         )
         ensure_bootstrap_dependencies(tool, update_existing=needs_auto_configure)
         if needs_auto_configure:
-            _auto_configure_tool(tool)
+            _auto_configure_tool(tool, model=model if tool == "copilot" else None)
         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).
@@ -1958,7 +1964,11 @@ def _launch_tool(
             _register_managed_mcp_servers(managed, tool, state)
             _apply_managed_skills(managed, tool, state)
         print_success(f"Starting {TOOL_SPECS[tool]['display']}")
-        launch_agent(tool, state, ctx.args)
+        # Pass the already-settled resolved_model (which absorbs --model, a managed config's
+        # default, and any budget recommendation), not the raw --model flag — otherwise a launch
+        # without an explicit --model would have copilot.launch() call default_model(state) fresh
+        # and silently rewrite whatever configure_tool just wrote, and every refresh after it.
+        launch_agent(tool, state, ctx.args, model=resolved_model if tool == "copilot" else None)
     except RuntimeError as exc:
         print_err(str(exc))
         raise typer.Exit(1) from None
@@ -2271,12 +2281,21 @@ def opencode_cmd(
 @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
 def copilot_cmd(
     ctx: typer.Context,
+    model: Annotated[
+        str | None,
+        typer.Option(
+            "--model",
+            help="Launch on a specific Databricks model id (e.g. a UC "
+            "`<catalog>.<schema>.<name>`). Outranks the automatic sonnet/opus/haiku/codex "
+            "pick and stays pinned across ucode's automatic token refreshes.",
+        ),
+    ] = None,
     skip_preflight: SkipPreflightOption = False,
     skip_managed_config: SkipManagedConfigOption = False,
 ) -> None:
     """Launch GitHub Copilot CLI via Databricks."""
     _disable_managed_config_if_requested(skip_managed_config)
-    _launch_tool("copilot", ctx, skip_preflight=skip_preflight)
+    _launch_tool("copilot", ctx, model=model, skip_preflight=skip_preflight)
 
 
 @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})

tests/test_agent_copilot.py

@@ -3,6 +3,9 @@
 from __future__ import annotations
 
 import json
+import threading
+
+import pytest
 
 from ucode.agents import copilot
 
@@ -198,6 +201,57 @@ def test_includes_required_vars(self):
             assert key in copilot.MANAGED_KEYS
 
 
+class TestRefreshTokenOnceModelOverride:
+    def test_model_override_wins_over_default_model(self, monkeypatch):
+        seen: dict = {}
+
+        def fake_write_tool_config(state, model, force_refresh=False):
+            seen["model"] = model
+            return state, "tok"
+
+        monkeypatch.setattr(copilot, "write_tool_config", fake_write_tool_config)
+        state = {"claude_models": {"sonnet": "discovered"}}
+
+        model, token = copilot._refresh_token_once(state, model_override="explicit-model")
+
+        assert model == "explicit-model"
+        assert token == "tok"
+        assert seen["model"] == "explicit-model"
+
+    def test_falls_back_to_default_model_without_override(self, monkeypatch):
+        monkeypatch.setattr(
+            copilot, "write_tool_config", lambda state, model, force_refresh=False: (state, "tok")
+        )
+        state = {"claude_models": {"sonnet": "discovered"}}
+
+        model, _ = copilot._refresh_token_once(state)
+
+        assert model == "discovered"
+
+    def test_raises_when_no_override_and_no_default(self):
+        with pytest.raises(RuntimeError, match="No Copilot model is available"):
+            copilot._refresh_token_once({})
+
+
+class TestRefreshForeverModelOverride:
+    def test_forwards_model_override_to_each_refresh(self, monkeypatch):
+        calls: list[str | None] = []
+
+        def fake_refresh_token_once(state, *, force_refresh=False, model_override=None):
+            calls.append(model_override)
+            if len(calls) >= 2:
+                stop_event.set()
+            return model_override, "tok"
+
+        monkeypatch.setattr(copilot, "_refresh_token_once", fake_refresh_token_once)
+        monkeypatch.setattr(copilot, "TOKEN_REFRESH_INTERVAL_SECONDS", 0)
+
+        stop_event = threading.Event()
+        copilot._refresh_forever({}, stop_event, model_override="pinned-model")
+
+        assert calls == ["pinned-model", "pinned-model"]
+
+
 class TestValidateCmd:
     def test_starts_with_binary(self):
         cmd = copilot.validate_cmd("copilot")