diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 205b68b7..1be84f91 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -302,10 +302,10 @@ def resolve_provider_models( """Validate ``provider`` for ``tool`` and return the model ids to pin. Returns ``(provider_models, error, relayed)``. ``provider_models`` is a ``{family: model_id}`` - dict re-derived from the service's live targets for a non-relayed claude service — both Bedrock - (provider-side slugs) and API-key Anthropic (canonical ids) — so the client sends exactly the ids - the MPS allows rather than Claude Code's defaults, which may not match the declared targets. It is - None when ``provider`` is None, for a relayed subscription (see below), or for a non-Claude (e.g. + dict re-derived from the service's declared targets — Bedrock (provider-side slugs), API-key + Anthropic, and relayed Anthropic that declares a curated allowlist alike — so the client uses the + ids the MPS allows rather than Claude Code's defaults. It is None when ``provider`` is None, when + the service declares no Claude targets (e.g. an ``allow_all`` relay), or for a non-Claude (e.g. codex) service. ``relayed`` is True for a credential-less Anthropic subscription relay, which the launch path wires with the relayed overlay + refresh proxy. A non-None ``error`` means the provider is invalid for the tool and the caller should not launch. @@ -322,10 +322,8 @@ def resolve_provider_models( if error or service is None: return None, error, False relayed = bool(service.get("relayed")) - # Relayed (Claude Max/Enterprise subscription) is exempt: the gateway disables - # model selection server-side for that tier, so there's nothing to reconcile. - if relayed: - return None, None, relayed + # Relayed services enforce their declared targets too, so map them like any Anthropic service + # (allow_all declares none). relayed gates auth, not model reconciliation. return map_claude_family_models(service.get("targets") or []) or None, None, relayed diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1e15e844..6bc33815 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2025,20 +2025,23 @@ def _launch_tool( # The router's per-launch pick for the root session. Codex pins it as the # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model). route_root_model = None + relayed_forward_model = None # forwarded to Claude Code's --model for a relayed provider 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. resolved_model = None - if tool == "claude" and relayed: - if model: - print_warning( - "This is a subscription-relay Model Provider Service; the gateway selects " - "the model, so --model is ignored." - ) - elif tool == "claude" and (model or provider_models): + if tool == "claude" and not relayed and (model or provider_models): route_root_model = resolve_provider_launch_model(model, provider_models or {}) + elif tool == "claude" and relayed and (model or provider_models): + # Resolve against a curated allowlist so the forwarded id is one the gateway allows; + # an allow_all relay declares none, so forward as-is. + relayed_forward_model = ( + resolve_provider_launch_model(model, provider_models, always_select=True) + if provider_models + else model + ) else: # A managed default_model is the model the admin wants sessions to start on, so it goes # in as the explicit model rather than being applied afterwards: for codex the proto has @@ -2115,6 +2118,16 @@ def _launch_tool( # per-family target pins. custom_model=model if (tool == "claude" and not provider) else None, ) + # Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`. + if ( + tool == "claude" + and provider + and relayed + and relayed_forward_model + and not forwarded_model + ): + ctx.args = ["--model", relayed_forward_model, *ctx.args] + forwarded_model = relayed_forward_model print_section(_launch_title(tool)) if managed is not None: print_kv("Config", "workspace-managed") diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7b8e442f..2e360384 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2406,16 +2406,19 @@ def map_claude_family_models(targets: list[str]) -> dict[str, str]: _CLAUDE_LAUNCH_TIER_PREFERENCE = ("opus", "sonnet", "haiku") -def resolve_provider_launch_model(model: str | None, provider_models: dict[str, str]) -> str | None: +def resolve_provider_launch_model( + model: str | None, provider_models: dict[str, str], *, always_select: bool = False +) -> str | None: """Pick the model a provider-routed Claude session starts on, or None to keep Claude Code's default. ``provider_models`` maps the Claude families a service declares to their target ids (see ``map_claude_family_models``). With an explicit ``model`` (``ucode claude --model``) the user's choice wins: a family alias resolves to that tier's declared target (erroring when the service doesn't offer it), any other value is trusted as a raw target id the service allows. Without one, - return None when the service offers opus — Claude Code's own default already works, so we avoid - setting ANTHROPIC_MODEL and the duplicate ``/model`` picker row it produces — else the most - capable tier the service does offer, so the launch doesn't dead-end on an unservable opus. + return the most capable tier the service offers — but return None if it offers opus (Claude Code's + default already resolves there via the pinned family alias), unless ``always_select`` is set. + Relayed launches set it: they forward via Claude Code's ``--model`` flag, with no pinned alias to + fall back on. """ if model: if model in ANTHROPIC_FAMILIES: @@ -2428,7 +2431,7 @@ def resolve_provider_launch_model(model: str | None, provider_models: dict[str, ) return target return model - if provider_models.get("opus"): + if not always_select and provider_models.get("opus"): return None return next( ( diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 5c287475..ac3ec98f 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -327,7 +327,8 @@ def test_anthropic_with_no_claude_targets_pins_nothing(self, monkeypatch): ) assert (models, error, relayed) == (None, None, False) - def test_relayed_anthropic_flagged(self, monkeypatch): + def test_relayed_allow_all_pins_nothing(self, monkeypatch): + # allow_all relay declares no Claude targets: nothing pinned, still flagged relayed. self._patch( monkeypatch, {"provider_type": "anthropic", "targets": [], "relayed": True}, None ) @@ -338,6 +339,24 @@ def test_relayed_anthropic_flagged(self, monkeypatch): assert models is None assert relayed is True + def test_relayed_anthropic_with_targets_pins_family(self, monkeypatch): + # A curated relay maps its declared targets by family so --model can resolve against them. + self._patch( + monkeypatch, + { + "provider_type": "anthropic", + "targets": ["claude-opus-4-8", "claude-haiku-4-5"], + "relayed": True, + }, + None, + ) + models, error, relayed = agents_mod.resolve_provider_models( + "claude", self._STATE, "main.a.relayed_ent" + ) + assert error is None + assert models == {"opus": "claude-opus-4-8", "haiku": "claude-haiku-4-5"} + assert relayed is True + def test_bedrock_returns_pinned_models(self, monkeypatch): service = { "provider_type": "amazon_bedrock", diff --git a/tests/test_cli.py b/tests/test_cli.py index 704bde48..305ecd14 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -614,7 +614,7 @@ def test_v2_model_sets_transient_launch_override(self, monkeypatch): @staticmethod def _provider_launch(monkeypatch, argv, provider_models, relayed=False): """Invoke a provider launch with model discovery/config stubbed, returning the - configure_tool mock so tests can assert what was threaded to it.""" + configure_tool and launch_agent mocks so tests can assert what was threaded to each.""" import ucode.cli as cli_mod monkeypatch.setattr(cli_mod, "ensure_bootstrap_dependencies", lambda *a, **k: None) @@ -623,19 +623,20 @@ def _provider_launch(monkeypatch, argv, provider_models, relayed=False): monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: MINIMAL_STATE) monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda s, m: None) - monkeypatch.setattr(cli_mod, "launch_agent", lambda *a, **k: None) + mock_launch = MagicMock() + monkeypatch.setattr(cli_mod, "launch_agent", mock_launch) monkeypatch.setattr( cli_mod, "resolve_provider_models", lambda t, s, p: (provider_models, None, relayed) ) mock_configure = MagicMock(return_value=MINIMAL_STATE) monkeypatch.setattr(cli_mod, "configure_tool", mock_configure) result = runner.invoke(app, argv) - return result, mock_configure + return result, mock_configure, mock_launch def test_model_and_provider_now_pin_the_launch_tier(self, monkeypatch): # --model under a provider is no longer rejected: a family alias resolves to that tier's # declared target and is threaded as route_root_model (ANTHROPIC_MODEL), not custom_model. - result, mock_configure = self._provider_launch( + result, mock_configure, _ = self._provider_launch( monkeypatch, ["claude", "--model", "haiku", "--provider", "cat.schema.svc"], {"sonnet": "claude-sonnet-5", "haiku": "claude-haiku-4-5"}, @@ -647,7 +648,7 @@ def test_model_and_provider_now_pin_the_launch_tier(self, monkeypatch): def test_provider_without_opus_auto_picks_best_servable_tier(self, monkeypatch): # No --model, and the service declares no opus target: launch on the most capable tier it # does offer (sonnet) instead of dead-ending on Claude Code's opus default. - result, mock_configure = self._provider_launch( + result, mock_configure, _ = self._provider_launch( monkeypatch, ["claude", "--provider", "cat.schema.svc"], {"sonnet": "claude-sonnet-5", "haiku": "claude-haiku-4-5"}, @@ -658,7 +659,7 @@ def test_provider_without_opus_auto_picks_best_servable_tier(self, monkeypatch): def test_provider_with_opus_keeps_claude_default(self, monkeypatch): # Opus is offered, so Claude Code's own default already works — pin nothing (no ANTHROPIC_MODEL # and no duplicate /model picker row). - result, mock_configure = self._provider_launch( + result, mock_configure, _ = self._provider_launch( monkeypatch, ["claude", "--provider", "cat.schema.svc"], {"opus": "claude-opus-4-8", "sonnet": "claude-sonnet-5"}, @@ -667,7 +668,7 @@ def test_provider_with_opus_keeps_claude_default(self, monkeypatch): assert mock_configure.call_args.kwargs["route_root_model"] is None def test_model_family_not_offered_by_provider_errors(self, monkeypatch): - result, _ = self._provider_launch( + result, _, _ = self._provider_launch( monkeypatch, ["claude", "--model", "opus", "--provider", "cat.schema.svc"], {"sonnet": "claude-sonnet-5", "haiku": "claude-haiku-4-5"}, @@ -675,17 +676,62 @@ def test_model_family_not_offered_by_provider_errors(self, monkeypatch): assert result.exit_code == 1 assert "does not offer a 'opus' model" in result.output - def test_model_ignored_for_relayed_provider(self, monkeypatch): - # A relayed (subscription) service selects the model server-side; --model can't be honored. - result, mock_configure = self._provider_launch( + def test_model_forwarded_to_claude_for_relayed_provider(self, monkeypatch): + # Relayed = a subscription: --model rides Claude Code's own flag, not gateway env. + result, mock_configure, mock_launch = self._provider_launch( monkeypatch, - ["claude", "--model", "haiku", "--provider", "cat.schema.svc"], + ["claude", "--model", "opus", "--provider", "cat.schema.svc"], None, relayed=True, ) assert result.exit_code == 0, result.output + assert mock_launch.call_args.args[2] == ["--model", "opus"] assert mock_configure.call_args.kwargs["route_root_model"] is None - assert "--model is ignored" in _strip_ansi(result.output) + assert mock_configure.call_args.kwargs["custom_model"] is None + assert "ignored" not in _strip_ansi(result.output) + + def test_relayed_provider_without_model_forwards_nothing(self, monkeypatch): + # No --model on an allow_all relay: nothing to forward. + result, _, mock_launch = self._provider_launch( + monkeypatch, + ["claude", "--provider", "cat.schema.svc"], + None, + relayed=True, + ) + assert result.exit_code == 0, result.output + assert mock_launch.call_args.args[2] == [] + + def test_relayed_allowlist_resolves_model_to_declared_target(self, monkeypatch): + # Curated relay: --model resolves to the declared id, which is what gets forwarded. + result, _, mock_launch = self._provider_launch( + monkeypatch, + ["claude", "--model", "opus", "--provider", "cat.schema.svc"], + {"opus": "claude-opus-4-8", "haiku": "claude-haiku-4-5"}, + relayed=True, + ) + assert result.exit_code == 0, result.output + assert mock_launch.call_args.args[2] == ["--model", "claude-opus-4-8"] + + def test_relayed_allowlist_auto_picks_best_tier_without_model(self, monkeypatch): + # Curated relay, no --model: forward the best allowed tier, not the (maybe forbidden) default. + result, _, mock_launch = self._provider_launch( + monkeypatch, + ["claude", "--provider", "cat.schema.svc"], + {"opus": "claude-opus-4-8", "sonnet": "claude-sonnet-5"}, + relayed=True, + ) + assert result.exit_code == 0, result.output + assert mock_launch.call_args.args[2] == ["--model", "claude-opus-4-8"] + + def test_relayed_allowlist_rejects_unavailable_family(self, monkeypatch): + result, _, _ = self._provider_launch( + monkeypatch, + ["claude", "--model", "opus", "--provider", "cat.schema.svc"], + {"sonnet": "claude-sonnet-5", "haiku": "claude-haiku-4-5"}, + relayed=True, + ) + assert result.exit_code == 1 + assert "does not offer a 'opus' model" in result.output def test_provider_sets_transient_claude_launch_marker(self): state = dict(MINIMAL_STATE) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 10e1083a..91976dfb 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -699,6 +699,13 @@ def test_raw_target_id_is_trusted(self): def test_no_models_and_no_override_is_none(self): assert db_mod.resolve_provider_launch_model(None, {}) is None + def test_always_select_picks_opus_instead_of_default(self): + # always_select: pick a tier even when opus is offered (no pinned alias to fall back on). + models = {"opus": "claude-opus-4-8", "sonnet": "claude-sonnet-5"} + assert db_mod.resolve_provider_launch_model(None, models, always_select=True) == ( + "claude-opus-4-8" + ) + class TestProviderServicePagination: """The listing is paginated; ignoring next_page_token hid services on later pages entirely."""