Add --model launch flag for custom UC model services

src/ucode/agents/__init__.py

@@ -287,17 +287,23 @@ def configure_tool(
     model: str | None = None,
     provider: str | None = None,
     provider_models: dict[str, str] | None = None,
+    custom_model: str | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
-        result = codex.write_tool_config(state, model, provider=provider)
+        result = codex.write_tool_config(state, model, provider=provider, custom_model=custom_model)
     elif tool == "claude":
         # A Model Provider Service routes by header and pins no Databricks
         # model, so the usual "model required" guard doesn't apply to claude.
-        if not model and not provider:
+        # A custom UC model-service (`--model`) likewise satisfies the guard.
+        if not model and not provider and not custom_model:
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         result = claude.write_tool_config(
-            state, model, provider=provider, provider_models=provider_models
+            state,
+            model,
+            provider=provider,
+            provider_models=provider_models,
+            custom_model=custom_model,
         )
     else:
         # provider routing is claude/codex-only; every other tool needs a model.

src/ucode/agents/claude.py

@@ -137,6 +137,7 @@ def render_overlay(
     provider: str | None = None,
     provider_models: dict[str, str] | None = None,
     fable_enabled: bool = False,
+    custom_model: str | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Claude settings.json.
 
@@ -150,7 +151,13 @@ def render_overlay(
     understands Claude Code's own canonical model names, so no model id is
     pinned. A Bedrock-backed provider exposes different model ids (e.g.
     `us.anthropic.claude-sonnet-4-6`), passed in `provider_models` by family —
-    those get pinned via the `ANTHROPIC_DEFAULT_*_MODEL` env vars."""
+    those get pinned via the `ANTHROPIC_DEFAULT_*_MODEL` env vars.
+
+    When `custom_model` is set (from `ucode claude --model`), it pins a custom
+    UC model-service id (outside `system.ai`) as `ANTHROPIC_MODEL` — the active
+    default — *alongside* the discovered `system.ai` family aliases, so the
+    system.ai rows still populate the /model picker while the custom model
+    becomes the default. Mutually exclusive with `provider`."""
     base_url = build_tool_base_url("claude", workspace)
     # ANTHROPIC_CUSTOM_HEADERS is parsed as `key: value` pairs separated by
     # newlines (Anthropic SDK convention). Setting User-Agent here overrides
@@ -176,6 +183,14 @@ def render_overlay(
     # only one row per model. `ucode claude -- --model X` still overrides for a
     # single session via Claude Code's own --model flag.
     _ = model  # API stability; no longer pinned via env.
+    # A custom UC model-service (`--model`) pins ANTHROPIC_MODEL as the active
+    # default. This is the one case we DO set ANTHROPIC_MODEL: the family aliases
+    # below still populate the picker's system.ai rows, and the custom id becomes
+    # the resolved default on top of them. ANTHROPIC_MODEL is in
+    # CLAUDE_MANAGED_MODEL_ENV_KEYS, so it's pruned again when no `--model` is
+    # passed. Mutually exclusive with `provider` (enforced upstream in cli.py).
+    if custom_model and not provider:
+        env["ANTHROPIC_MODEL"] = custom_model
     # A Bedrock-backed provider needs its provider-side ids pinned verbatim
     # (Claude Code's canonical names aren't routable there). These come from the
     # service's targets, already de-duped to one id per family upstream.
@@ -295,6 +310,7 @@ def write_tool_config(
     model: str | None,
     provider: str | None = None,
     provider_models: dict[str, str] | None = None,
+    custom_model: str | None = None,
 ) -> dict:
     backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH)
     web_search_model = _resolve_web_search_model(state)
@@ -308,6 +324,7 @@ def write_tool_config(
         provider=provider,
         provider_models=provider_models,
         fable_enabled=bool(state.get("fable_enabled")),
+        custom_model=custom_model,
     )
     tracing_env_vars = tracing_env(state, "claude")
     stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None

src/ucode/agents/codex.py

@@ -309,12 +309,25 @@ def _parse_gpt(model: str | None) -> tuple[int, int | None, int | None, str] | N
     )
 
 
-def write_tool_config(state: dict, model: str | None = None, provider: str | None = None) -> dict:
+def write_tool_config(
+    state: dict,
+    model: str | None = None,
+    provider: str | None = None,
+    custom_model: str | None = None,
+) -> dict:
     workspace = state["workspace"]
     # With a Model Provider Service the gateway routes by header and Codex sends
     # its own canonical model name (e.g. `gpt-5`) — leave `model` unset so no
-    # Databricks endpoint id is pinned.
-    chosen_model = None if provider else _codex_model_id(model or default_model(state))
+    # Databricks endpoint id is pinned. A custom UC model-service (`--model`) is
+    # pinned VERBATIM, bypassing the `_codex_model_id` OpenAI-id rewrite — the
+    # gateway routes it by name like a `system.ai.*` id. Mutually exclusive with
+    # `provider` (enforced upstream in cli.py).
+    if provider:
+        chosen_model = None
+    elif custom_model:
+        chosen_model = custom_model
+    else:
+        chosen_model = _codex_model_id(model or default_model(state))
     databricks_profile = state.get("profile")
 
     if _use_legacy_layout():

src/ucode/cli.py

@@ -43,6 +43,7 @@
     find_profile_name_for_host,
     get_databricks_profiles,
     get_databricks_token,
+    get_model_service,
     install_databricks_cli,
     is_model_provider_feature_unavailable,
     list_profile_entries,
@@ -886,10 +887,20 @@ def _launch_tool(
     tool_name: str,
     ctx: typer.Context,
     provider: str | None = None,
+    model: str | None = None,
     skip_preflight: bool = False,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
+        # `--model` (a custom UC model-service pinned as the default) and
+        # `--provider` (routing through a Model Provider Service via a header)
+        # are two different routing modes; picking both is ambiguous.
+        if model and provider:
+            raise RuntimeError(
+                "Pass only one of --model or --provider. --model pins a custom UC "
+                "model-service as the default; --provider routes through a Model "
+                "Provider Service. Choose one routing mode."
+            )
         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
@@ -915,6 +926,17 @@ def _launch_tool(
             provider_models, error = resolve_provider_models(tool, state, provider)
             if error:
                 raise RuntimeError(error)
+        # A custom UC model-service (`--model`) pins that id as the default while
+        # KEEPING normal system.ai discovery so the family shortcuts still
+        # populate the picker/codex list. Validate existence only (compatibility
+        # is assumed per the user's explicit choice); abort with the reason on
+        # failure. Mutually exclusive with --provider (checked at the top).
+        custom_model = None
+        if model:
+            token = get_databricks_token(state["workspace"], state.get("profile"))
+            custom_model, error = get_model_service(state["workspace"], token, model)
+            if error:
+                raise RuntimeError(error)
         # Re-fetch model lists on every launch so newly-added Databricks
         # endpoints show up without a manual `ucode configure` (and so that
         # tools like pi which read multiple model bundles never run on
@@ -927,20 +949,29 @@ def _launch_tool(
             skip_model_discovery=bool(provider),
             skip_preflight=skip_preflight,
         )
-        if provider:
-            # Routing through a Model Provider Service pins no Databricks model;
-            # the agent uses its own canonical model names (header selects the
-            # provider). Skip model resolution, which would otherwise fail when
-            # the workspace has no matching Databricks models.
+        if provider or custom_model:
+            # A Model Provider Service pins no Databricks model (the header
+            # selects the provider). A custom `--model` pins its own id directly
+            # (ANTHROPIC_MODEL for claude / the `model` key for codex), so the
+            # discovery-default resolution below isn't needed — skipping it also
+            # avoids a spurious failure when the workspace has no matching
+            # Databricks models to fall back on.
             resolved_model = None
         else:
             state, resolved_model = resolve_launch_model(tool, state, None)
         state = configure_tool(
-            tool, state, resolved_model, provider=provider, provider_models=provider_models
+            tool,
+            state,
+            resolved_model,
+            provider=provider,
+            provider_models=provider_models,
+            custom_model=custom_model,
         )
         print_section(f"ucode with {TOOL_SPECS[tool]['display']}")
         if provider:
             print_kv("Provider", provider)
+        elif custom_model:
+            print_kv("Model", custom_model)
         elif resolved_model:
             print_kv("Model", resolved_model)
         if tool in ("gemini", "opencode", "copilot", "pi"):
@@ -984,10 +1015,19 @@ def codex_cmd(
             "before any `--` separator.",
         ),
     ] = None,
+    model: Annotated[
+        str | None,
+        typer.Option(
+            "--model",
+            help="Pin a custom UC model-service as the default model (outside "
+            "system.ai). Assumes the service is compatible with this agent. Pass "
+            "before any `--` separator.",
+        ),
+    ] = None,
     skip_preflight: SkipPreflightOption = False,
 ) -> None:
     """Launch Codex via Databricks."""
-    _launch_tool("codex", ctx, provider=provider, skip_preflight=skip_preflight)
+    _launch_tool("codex", ctx, provider=provider, model=model, skip_preflight=skip_preflight)
 
 
 @app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
@@ -1002,10 +1042,19 @@ def claude_cmd(
             "before any `--` separator.",
         ),
     ] = None,
+    model: Annotated[
+        str | None,
+        typer.Option(
+            "--model",
+            help="Pin a custom UC model-service as the default model (outside "
+            "system.ai). Assumes the service is compatible with this agent. Pass "
+            "before any `--` separator.",
+        ),
+    ] = None,
     skip_preflight: SkipPreflightOption = False,
 ) -> None:
     """Launch Claude Code via Databricks."""
-    _launch_tool("claude", ctx, provider=provider, skip_preflight=skip_preflight)
+    _launch_tool("claude", ctx, provider=provider, model=model, skip_preflight=skip_preflight)
 
 
 @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})

src/ucode/databricks.py

@@ -26,7 +26,7 @@
 from typing import Literal, cast, overload
 from urllib import error as urllib_error
 from urllib import request as urllib_request
-from urllib.parse import urlencode, urlparse
+from urllib.parse import quote, urlencode, urlparse
 
 from databricks.sql.exc import ServerOperationError
 
@@ -1183,19 +1183,29 @@ def list_model_services(
 ) -> tuple[list[str], str | None]:
     """List all `system.ai.*` model ids via the UC model-services API.
 
-    Pages through ``/api/2.1/unity-catalog/model-services`` (metastore scope)
-    with a bounded ``page_size`` (the endpoint 499s without one) and returns the
-    de-duplicated, sorted list of ``system.ai.<model-name>`` ids. Returns
-    (ids, reason); reason is None on success, otherwise it describes why the
-    list is empty (HTTP/network error or no services).
+    Pages through ``/api/2.1/unity-catalog/model-services`` scoped to the
+    ``system.ai`` schema (``parent=schemas/system.ai``) with a bounded
+    ``page_size`` (the endpoint 499s without one) and returns the
+    de-duplicated, sorted list of ``system.ai.<model-name>`` ids. The
+    ``system.ai.`` client-side filter is kept as a safety net in case the
+    ``parent`` scope isn't honored on some workspaces. Returns (ids, reason);
+    reason is None on success, otherwise it describes why the list is empty
+    (HTTP/network error or no services).
     """
     hostname = workspace_hostname(workspace)
     ids: list[str] = []
     page_token: str | None = None
     seen_tokens: set[str] = set()
     last_reason: str | None = None
     for _ in range(max_pages):
-        params: dict[str, str] = {"page_size": str(page_size)}
+        # Scope the listing to the system.ai schema so the server returns only
+        # foundation models instead of every schema's services — much less to
+        # page through. The client-side `system.ai.` filter below stays as a
+        # safety net.
+        params: dict[str, str] = {
+            "page_size": str(page_size),
+            "parent": "schemas/system.ai",
+        }
         if page_token:
             params["page_token"] = page_token
         url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}"
@@ -1225,6 +1235,85 @@ def list_model_services(
     return [], last_reason or "model-services listing returned no models"
 
 
+def get_model_service(workspace: str, token: str, full_name: str) -> tuple[str | None, str | None]:
+    """Look up a single UC model-service by its ``<catalog>.<schema>.<name>``.
+
+    Returns ``(routable_id, None)`` when the service exists, where
+    ``routable_id`` is the ``name`` with the ``model-services/`` prefix stripped
+    (i.e. the ``<catalog>.<schema>.<name>`` the gateway routes by). Returns
+    ``(None, reason)`` with an actionable message when the service can't be
+    found or the API call fails.
+
+    Validates existence only — it does NOT check the provider type or model
+    family, since a custom ``--model`` service is assumed compatible with the
+    agent per the caller's explicit choice.
+
+    The single-resource endpoint (``GET .../model-services/<full_name>``) is
+    tried first; if that path isn't served (404-style) we fall back to a
+    schema-scoped listing filtered to the exact name, so the lookup works
+    regardless of which shape the workspace's API exposes.
+    """
+    full_name = (full_name or "").strip()
+    if not full_name:
+        return None, "No model-service name was provided."
+    hostname = workspace_hostname(workspace)
+
+    # 1. Try the singular GET first.
+    encoded = quote(full_name, safe="")
+    url = f"https://{hostname}/api/2.1/unity-catalog/model-services/{encoded}"
+    payload, reason = _http_get_json(url, token, timeout=30)
+    if isinstance(payload, dict):
+        routable = _model_service_routable_id(payload)
+        if routable:
+            return routable, None
+        # A 200 without a usable name is unexpected; fall through to listing.
+    not_found = bool(reason) and "HTTP 404" in reason
+
+    # 2. Fall back to a schema-scoped listing filtered to the exact name.
+    parts = full_name.split(".")
+    if len(parts) >= 2:
+        schema_parent = ".".join(parts[:2])  # <catalog>.<schema>
+        params = {
+            "page_size": str(_MODEL_SERVICES_PAGE_SIZE),
+            "parent": f"schemas/{schema_parent}",
+        }
+        list_url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}"
+        list_payload, list_reason = _get_model_services_page(list_url, token)
+        if isinstance(list_payload, dict):
+            for service in list_payload.get("model_services", []):
+                if not isinstance(service, dict):
+                    continue
+                routable = _model_service_routable_id(service)
+                if routable == full_name:
+                    return routable, None
+            return None, (
+                f"Model-service '{full_name}' was not found in schema "
+                f"'{schema_parent}'. Pass a <catalog>.<schema>.<name> that exists "
+                f"on this workspace."
+            )
+        # Listing itself failed; prefer the more specific failure to report.
+        reason = list_reason or reason
+
+    if not_found or (reason and "HTTP 404" in reason):
+        return None, (
+            f"Model-service '{full_name}' was not found. Pass a "
+            f"<catalog>.<schema>.<name> that exists on this workspace."
+        )
+    return None, f"Could not look up model-service '{full_name}': {reason}"
+
+
+def _model_service_routable_id(service: dict) -> str | None:
+    """Return the ``<catalog>.<schema>.<name>`` routable id from a model-service
+    entry, stripping the ``model-services/`` prefix. None if there's no name."""
+    name = service.get("name")
+    if not isinstance(name, str):
+        return None
+    name = name.strip()
+    if name.startswith(_MODEL_SERVICE_NAME_PREFIX):
+        name = name[len(_MODEL_SERVICE_NAME_PREFIX) :]
+    return name or None
+
+
 def discover_model_services(
     workspace: str, token: str
 ) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]:

tests/test_agent_claude.py

@@ -113,6 +113,31 @@ def test_model_overrides_not_set_when_no_models(self):
         env = overlay["env"]
         assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in env
 
+    def test_custom_model_sets_anthropic_model_and_keeps_family_aliases(self):
+        # `--model` pins the custom UC id as ANTHROPIC_MODEL (the active default)
+        # AND retains the discovered system.ai family aliases so those picker
+        # rows still populate.
+        models = {
+            "opus": "system.ai.claude-opus-4-8",
+            "sonnet": "system.ai.claude-sonnet-4-6",
+        }
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models=models, custom_model="cat.schema.my-model"
+        )
+        env = overlay["env"]
+        assert env["ANTHROPIC_MODEL"] == "cat.schema.my-model"
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+        assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.claude-sonnet-4-6[1m]"
+
+    def test_custom_model_ignored_under_provider(self):
+        # `--model` and `--provider` are mutually exclusive upstream; if both
+        # somehow reach the overlay, the provider header wins and ANTHROPIC_MODEL
+        # is not pinned (the custom id isn't routable through the header path).
+        overlay, _ = claude.render_overlay(
+            WS, "s4", provider="cat.schema.prov", custom_model="cat.schema.my-model"
+        )
+        assert "ANTHROPIC_MODEL" not in overlay["env"]
+
     def test_fable_not_pinned_by_default(self):
         # Fable is opt-in: even when the workspace advertises it, the env var is
         # absent unless fable_enabled is passed.