diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ee0d4d013..68e3f1d00 100644 --- a/src/ucode/agents/claude.py +++ b/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 = [ diff --git a/src/ucode/cli.py b/src/ucode/cli.py index cfbe67eea..656abd746 100644 --- a/src/ucode/cli.py +++ b/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 " - "`..`). Pinned via ANTHROPIC_MODEL so the gateway " - "resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. " + "`..`). 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.", ), diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d079cfd2d..376f762eb 100644 --- a/src/ucode/databricks.py +++ b/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} diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index ebed95c38..e9bb3436a 100644 --- a/tests/test_agent_claude.py +++ b/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()) diff --git a/tests/test_cli.py b/tests/test_cli.py index 06c6f9c9c..7e5124b8b 100644 --- a/tests/test_cli.py +++ b/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"}} diff --git a/tests/test_databricks.py b/tests/test_databricks.py index ca619be8b..d75c1d467 100644 --- a/tests/test_databricks.py +++ b/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}