[AIGTWY-4876] Support managed CUJ1 setup

src/ucode/agents/claude.py

@@ -210,6 +210,8 @@ def _otel_trace_env(workspace: str) -> dict[str, str]:
     "sonnet": "ANTHROPIC_DEFAULT_SONNET_MODEL",
     "haiku": "ANTHROPIC_DEFAULT_HAIKU_MODEL",
 }
+CLAUDE_CUSTOM_MODEL_FAMILIES = ("opus", "sonnet", "haiku")
+CLAUDE_CUSTOM_MODEL_SELECTOR = "opus"
 # Launch-scoped feature flags that ucode may write into Claude settings. These
 # must be removed again when the corresponding launch flag is absent.
 CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",)
@@ -1864,6 +1866,15 @@ def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[s
     return ["--model", launch_model]
 
 
+def _launch_custom_model_settings(model: str) -> dict:
+    """Pin a Databricks model through launch-scoped Claude family aliases."""
+    return {
+        "env": {
+            CLAUDE_DEFAULT_MODEL_ENV_KEYS[family]: model for family in CLAUDE_CUSTOM_MODEL_FAMILIES
+        },
+    }
+
+
 def _resolve_picker_model_id(model: str, settings_env: dict) -> str:
     """Resolve configured aliases and context suffixes for comparisons only."""
     model = re.sub(r"\[(?:1m|200k)\]$", "", model)
@@ -1875,6 +1886,24 @@ def _resolve_picker_model_id(model: str, settings_env: dict) -> str:
     return re.sub(r"\[(?:1m|200k)\]$", "", model)
 
 
+def _is_managed_launch_model(state: dict, model: str) -> bool:
+    """Return whether *model* is present in Claude's managed model catalog."""
+    picker_models = state.get("_claude_launch_picker_models")
+    if not isinstance(picker_models, list) or not picker_models:
+        picker_models = state.get("claude_static_models")
+    if not isinstance(picker_models, list) or not picker_models:
+        return False
+
+    settings_env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env")
+    settings_env = settings_env if isinstance(settings_env, dict) else {}
+    resolved_model = _resolve_picker_model_id(model, settings_env)
+    return any(
+        isinstance(picker_model, str)
+        and _resolve_picker_model_id(picker_model, settings_env) == resolved_model
+        for picker_model in picker_models
+    )
+
+
 def _resolve_launch_binary(binary: str) -> str:
     """Resolve Claude's native executable without sending arguments through a batch shim."""
     if os.name != "nt":
@@ -2100,7 +2129,18 @@ def launch(
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
     settings_override = None
     launch_args = list(tool_args)
-    if options.user_pinned_model:
+    launch_custom_model = state.get("_claude_launch_custom_model")
+    if isinstance(launch_custom_model, str) and launch_custom_model:
+        if _is_managed_launch_model(state, launch_custom_model):
+            launch_args = [
+                *_launch_model_args(tool_args, launch_custom_model),
+                *tool_args,
+            ]
+        else:
+            # Claude rejects raw model ids outside its catalog.
+            os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR
+            settings_override = _launch_custom_model_settings(launch_custom_model)
+    elif options.user_pinned_model:
         os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model
         settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}}
         launch_args = [

src/ucode/cli.py

@@ -2851,6 +2851,8 @@ def _launch_tool(
             # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim.
             if model and tool != "claude":
                 resolved_model = model
+        if model and tool == "claude" and not provider:
+            route_root_model = None
         if coding_agent_config_defaults and not state.get("claude_static_models") and not relayed:
             picker_catalog = claude_agent.default_model_picker_catalog(
                 coding_agent_config_defaults,
@@ -2867,7 +2869,6 @@ def _launch_tool(
             picker_catalog=picker_catalog,
             relayed=relayed,
             route_root_model=route_root_model,
-            # Claude's explicit model is launch-scoped and is passed through LaunchOptions below.
             custom_model=None,
             coding_agent_config_defaults=coding_agent_config_defaults,
             parent_schema=parent_schema,
@@ -2876,7 +2877,7 @@ def _launch_tool(
             # Claude re-adds an out-of-catalog saved model to /model even when built-ins are
             # replaced. Keep the catalog launch-scoped and leave the user's settings alone.
             state["_claude_launch_picker_models"] = picker_catalog.model_ids
-            if managed is None and (explicit_provider or parent_schema):
+            if managed is None and (explicit_provider or parent_schema) and not model:
                 # The permanent Default row should also resolve within the selected catalog.
                 state["_claude_launch_default_model"] = (
                     claude_agent.default_model(
@@ -2917,6 +2918,9 @@ def _launch_tool(
         if tool == "claude":
             if provider:
                 state["_claude_launch_provider"] = provider
+            elif model:
+                # Set after configure_tool so the selection stays launch-scoped.
+                state["_claude_launch_custom_model"] = model
         elif tool == "codex":
             if provider:
                 state["_codex_launch_provider"] = provider
@@ -3211,8 +3215,8 @@ def claude_cmd(
         typer.Option(
             "--model",
             help="Launch on a specific Databricks model id (e.g. a UC "
-            "`<catalog>.<schema>.<name>`). Pinned via ANTHROPIC_MODEL so the gateway "
-            "resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. "
+            "`<catalog>.<schema>.<name>`). Pinned via launch-scoped Claude family aliases so the "
+            "gateway resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. "
             "With --provider, pass a family (opus/sonnet/haiku) or a target the service allows to "
             "start on that tier instead of Claude Code's opus default. Pass before any `--` separator.",
         ),

src/ucode/databricks.py

@@ -676,6 +676,14 @@ def workspace_hostname(workspace: str) -> str:
     return parsed.hostname
 
 
+def workspace_origin(workspace: str) -> str:
+    """Return the workspace scheme and authority, preserving an explicit port."""
+    parsed = urlparse(normalize_workspace_url(workspace))
+    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+        raise RuntimeError(f"Unable to derive origin from workspace URL: {workspace}")
+    return f"{parsed.scheme}://{parsed.netloc}"
+
+
 def _parse_databricks_cli_version(output: str) -> tuple[int, int, int] | None:
     # Example output: "Databricks CLI v0.299.2"
     match = re.search(r"v?(\d+)\.(\d+)\.(\d+)", output)
@@ -1836,7 +1844,7 @@ def list_model_services(
         if cached is not None:
             return list(cached), None
 
-    hostname = workspace_hostname(workspace)
+    origin = workspace_origin(workspace)
     ids: list[str] = []
     page_token: str | None = None
     seen_tokens: set[str] = set()
@@ -1848,7 +1856,7 @@ def list_model_services(
         }
         if page_token:
             params["page_token"] = page_token
-        url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}"
+        url = f"{origin}/api/2.1/unity-catalog/model-services?{urlencode(params)}"
         payload, reason = _get_model_services_page(url, token)
         if payload is None:
             # Surface the failure only if we have nothing yet; a mid-pagination
@@ -2038,8 +2046,7 @@ def _managed_config_user_agent() -> str:
 
 def fetch_managed_coding_agent_configs(workspace: str, token: str) -> tuple[list[dict], str | None]:
     """List the workspace's managed CodingAgentConfig(s) via the AI Gateway."""
-    hostname = workspace_hostname(workspace)
-    url = f"https://{hostname}{_CODING_AGENT_CONFIGS_API_PATH}"
+    url = f"{workspace_origin(workspace)}{_CODING_AGENT_CONFIGS_API_PATH}"
     payload, reason = _http_get_json(
         url, token, timeout=30, headers={"User-Agent": _managed_config_user_agent()}
     )
@@ -2062,8 +2069,7 @@ def fetch_model_recommendation(workspace: str, token: str) -> tuple[dict, str |
     The request takes no parameters: the server matches the caller's live spend against the managed
     config's budget tiers and resolves the agent first, then that agent's model.
     """
-    hostname = workspace_hostname(workspace)
-    url = f"https://{hostname}{_CODING_AGENT_CONFIGS_API_PATH}:recommendModel"
+    url = f"{workspace_origin(workspace)}{_CODING_AGENT_CONFIGS_API_PATH}:recommendModel"
     payload, reason = _http_post_json(url, token, {}, timeout=30)
     if reason is not None:
         return {}, reason
@@ -2086,8 +2092,8 @@ def fetch_external_model_prices(workspace: str, token: str) -> tuple[list[dict],
     Returns ``(models, reason)`` with each model the raw API entry; ``reason`` is non-None on failure
     (callers omit cost rather than fail).
     """
-    hostname = workspace_hostname(workspace)
-    base_url = f"https://{hostname}{_EXTERNAL_PROVIDER_MODELS_API_PATH}"
+    origin = workspace_origin(workspace)
+    base_url = f"{origin}{_EXTERNAL_PROVIDER_MODELS_API_PATH}"
     models: list[dict] = []
     page_token: str | None = None
     seen_tokens: set[str] = set()
@@ -2405,7 +2411,7 @@ def list_model_provider_services(
             # reach the next.
             return [dict(service) for service in cached], None
 
-    origin = normalize_workspace_url(workspace)
+    origin = workspace_origin(workspace)
     services: list[dict] = []
     page_token: str | None = None
     seen_tokens: set[str] = set()
@@ -2489,7 +2495,7 @@ def get_model_provider_service(
     server-side filter) makes a service that plainly exists look absent. Addressing it directly
     removes that whole class of false negative.
     """
-    origin = normalize_workspace_url(workspace)
+    origin = workspace_origin(workspace)
     url = f"{origin}/api/2.1/unity-catalog/model-provider-services/{service_name}"
     payload, reason = _http_get_json(url, token, timeout=30)
     if payload is None:
@@ -2975,7 +2981,7 @@ def _get_anthropic_models_json(
     parent_schema: str | None = None,
     provider: str | None = None,
 ) -> tuple[dict | list | None, str | None]:
-    origin = normalize_workspace_url(workspace)
+    origin = workspace_origin(workspace)
     headers = None
     if provider is not None:
         headers = {MODEL_PROVIDER_SERVICE_HEADER: provider}
@@ -3217,7 +3223,7 @@ class GatewayProbe(NamedTuple):
 
 
 def _probe_model_services(workspace: str, token: str) -> GatewayProbe:
-    base = f"{normalize_workspace_url(workspace)}/api/2.1/unity-catalog/model-services"
+    base = f"{workspace_origin(workspace)}/api/2.1/unity-catalog/model-services"
     page_token: str | None = None
     for page in range(_MODEL_SERVICE_PROBE_MAX_PAGES):
         params: dict[str, object] = {"page_size": _MODEL_SERVICE_PROBE_PAGE_SIZE}

tests/test_agent_claude.py

@@ -2636,6 +2636,77 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch):
         settings = json.loads(calls[0][2])
         assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model"
 
+    def test_launch_custom_model_uses_family_aliases_without_native_model(
+        self, monkeypatch, tmp_path
+    ):
+        calls: list[list[str]] = []
+        monkeypatch.setenv("ANTHROPIC_MODEL", "stale-model")
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token")
+        monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv))
+        settings_path = tmp_path / "ucode-settings.json"
+        saved_model = "system.ai.claude-haiku-4-5"
+        settings_path.write_text(
+            json.dumps({"model": saved_model, "env": {"USER_SETTING": "keep"}})
+        )
+        monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path)
+        original_settings = settings_path.read_bytes()
+
+        custom_model = "system.ai.claude-opus-4-8"
+        claude.launch(
+            {"workspace": WS, "_claude_launch_custom_model": custom_model},
+            ["--debug"],
+            options=LaunchOptions(user_pinned_model=custom_model),
+        )
+
+        assert os.environ["ANTHROPIC_MODEL"] == claude.CLAUDE_CUSTOM_MODEL_SELECTOR
+        assert calls[0][0:2] == ["claude", "--settings"]
+        settings = json.loads(calls[0][2])
+        assert settings["model"] == saved_model
+        for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES:
+            key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family]
+            assert settings["env"][key] == custom_model
+        assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in settings["env"]
+        assert "--model" not in calls[0]
+        assert calls[0][-1] == "--debug"
+        assert settings_path.read_bytes() == original_settings
+
+    def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch):
+        calls: list[list[str]] = []
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token")
+        monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv))
+
+        claude.launch(
+            {"workspace": WS, "_claude_launch_custom_model": "system.ai.claude-opus-4-8"},
+            ["--model", "claude-sonnet-5"],
+            options=LaunchOptions(user_pinned_model="system.ai.claude-opus-4-8"),
+        )
+
+        assert calls[0][-2:] == ["--model", "claude-sonnet-5"]
+
+    def test_launch_managed_custom_model_uses_native_model_over_stale_settings(
+        self, monkeypatch, tmp_path
+    ):
+        calls: list[list[str]] = []
+        settings_path = tmp_path / "settings.json"
+        settings_path.write_text(json.dumps({"model": "system.ai.claude-haiku-4-5"}))
+        monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path)
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token")
+        monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv))
+
+        custom_model = "system.ai.claude-opus-4-8"
+        claude.launch(
+            {
+                "workspace": WS,
+                "_claude_launch_custom_model": custom_model,
+                "claude_static_models": ["system.ai.claude-opus-4-8", "system.ai.claude-haiku-4-5"],
+            },
+            [],
+            options=LaunchOptions(user_pinned_model=custom_model),
+        )
+
+        assert calls[0][1:3] == ["--settings", str(settings_path)]
+        assert calls[0][-2:] == ["--model", custom_model]
+
     def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch):
         monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False)
         monkeypatch.setattr(v2, "launch_claude", Mock())

tests/test_cli.py

@@ -814,6 +814,11 @@ def test_claude_model_location_preserves_explicit_model(self):
         assert calls["launch"].call_args.kwargs["options"].user_pinned_model == (
             "main.default.claude-opus-5"
         )
+        assert "_claude_launch_default_model" not in calls["launch"].call_args.args[1]
+        assert (
+            calls["launch"].call_args.args[1]["_claude_launch_custom_model"]
+            == "main.default.claude-opus-5"
+        )
         assert calls["launch"].call_args.args[2] == []
 
     @pytest.mark.parametrize(
@@ -1505,13 +1510,32 @@ def test_model_is_launch_scoped_for_claude(self, monkeypatch):
         ):
             result = runner.invoke(app, ["claude", "--model", "cat.schema.claude-opus-5"])
         assert result.exit_code == 0, result.output
-        # The model is passed through invocation-scoped LaunchOptions, not persisted in settings.
         assert mock_configure.call_args.kwargs["custom_model"] is None
         assert mock_configure.call_args.kwargs["route_root_model"] is None
+        assert (
+            mock_launch.call_args.args[1]["_claude_launch_custom_model"]
+            == "cat.schema.claude-opus-5"
+        )
         assert (
             mock_launch.call_args.kwargs["options"].user_pinned_model == "cat.schema.claude-opus-5"
         )
 
+    def test_explicit_model_clears_managed_default_route_root(self):
+        managed = {
+            "enabled_agents": {
+                "claude": {"model_config": {"default_model": "system.ai.claude-sonnet-5"}}
+            }
+        }
+        with _launch_policy_patches(managed) as calls:
+            result = runner.invoke(app, ["claude", "--model", "system.ai.claude-opus-4-8"])
+
+        assert result.exit_code == 0, result.output
+        assert calls["configure"].call_args.kwargs["route_root_model"] is None
+        assert (
+            calls["launch"].call_args.args[1]["_claude_launch_custom_model"]
+            == "system.ai.claude-opus-4-8"
+        )
+
     def test_v2_model_sets_transient_launch_override(self, monkeypatch):
         monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1")
         state = {**MINIMAL_STATE, "claude_models": {"opus": "system.ai.claude-opus-4-8"}}

tests/test_databricks.py

@@ -49,6 +49,7 @@
     resolve_current_budget_spend,
     upgrade_databricks_cli,
     workspace_hostname,
+    workspace_origin,
 )
 
 WS = "https://example.databricks.com"
@@ -192,6 +193,9 @@ def test_invalid_url_raises(self):
         with pytest.raises((RuntimeError, ValueError)):
             workspace_hostname("")
 
+    def test_origin_preserves_explicit_scheme_and_port(self):
+        assert workspace_origin("http://127.0.0.1:54321/path") == "http://127.0.0.1:54321"
+
 
 class _FakeResponseWithHeaders(_FakeResponse):
     def __init__(self, payload: dict, headers: dict):
@@ -657,6 +661,19 @@ def fake_get(url, token, timeout=10):
         # foundation models rather than walking the whole metastore.
         assert all("parent=schemas%2Fsystem.ai" in u for u in urls)
 
+    def test_preserves_http_origin_and_port_for_local_gateway_recorder(self, monkeypatch):
+        urls = []
+
+        def fake_page(url, token):
+            urls.append(url)
+            return {"model_services": [_model_service("system.ai.gpt-5")]}, None
+
+        monkeypatch.setattr(db_mod, "_get_model_services_page", fake_page)
+
+        db_mod.list_model_services("http://127.0.0.1:54321", "token", use_cache=False)
+
+        assert urls[0].startswith("http://127.0.0.1:54321/api/")
+
     def test_retries_page_before_giving_up(self, monkeypatch):
         payload = {"model_services": [_model_service("system.ai.gpt-5")]}
         calls = {"n": 0}