From a2754bd8f0e9da5dcbfb704bc0858d30686e6ff6 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Mon, 31 Aug 2026 23:45:35 +0000 Subject: [PATCH 1/2] MPS: forward --model to Claude Code for a relayed provider A relayed (Claude Max/Team/Enterprise subscription) Model Provider Service is a subscription relay: the AI Gateway passes the client's requested model through to Anthropic and, for Team/Enterprise, validates it against the service's allowlist (only personal Max skips that) -- it never selects a model server-side. Real Enterprise relays are allow_all_targets with no declared Claude targets, so there is nothing for ucode to resolve --model against. So for a relayed provider, forward --model to Claude Code's own --model flag -- exactly what 'ucode claude --provider -- --model X' already does, now automatic. Claude Code selects the model natively and the relay honors it. Non-relayed providers (API-key / Bedrock) keep the env-pinning path unchanged. Verified end-to-end against a relayed Enterprise MPS: '-- --model opus' launches on Opus via the subscription. Co-authored-by: Isaac --- src/ucode/cli.py | 13 ++++++------- tests/test_cli.py | 38 ++++++++++++++++++++++++++------------ 2 files changed, 32 insertions(+), 19 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1e15e844..5fba6e69 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2031,13 +2031,8 @@ def _launch_tool( # 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): + # Relayed services forward --model to Claude Code's own flag at launch (below), not env. + if tool == "claude" and not relayed and (model or provider_models): route_root_model = resolve_provider_launch_model(model, provider_models or {}) else: # A managed default_model is the model the admin wants sessions to start on, so it goes @@ -2115,6 +2110,10 @@ def _launch_tool( # per-family target pins. custom_model=model if (tool == "claude" and not provider) else None, ) + # Relayed = a Claude subscription: forward --model to Claude Code's own flag, like `-- --model X`. + if tool == "claude" and provider and relayed and model and not forwarded_model: + ctx.args = ["--model", model, *ctx.args] + forwarded_model = model print_section(_launch_title(tool)) if managed is not None: print_kv("Config", "workspace-managed") diff --git a/tests/test_cli.py b/tests/test_cli.py index 704bde48..c54480c0 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,30 @@ 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 a relayed launch: 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_provider_sets_transient_claude_launch_marker(self): state = dict(MINIMAL_STATE) From d0e701fe1af33718dc38e2abf7db5d827e1b42cf Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 1 Sep 2026 15:02:00 +0000 Subject: [PATCH 2/2] MPS: resolve relayed --model against a curated target allowlist Follow-up to #427 (which forwards --model to Claude Code for a relayed provider). That covers allow_all relays; a relayed Team/Enterprise service that declares a curated model allowlist still needs its --model reconciled against those targets, since the AI Gateway enforces them -- forwarding a bare alias or Claude Code's subscription default can 403. For a relayed provider, resolve --model (and the bare-launch default) against the declared targets, then forward the resolved id via Claude Code's own --model flag: --model opus -> the declared opus target; no --model -> the best allowed tier; an allow_all relay declares nothing -> forward as-is (unchanged). Stops exempting relayed in resolve_provider_models; adds resolve_provider_launch_model always_select (relayed has no pinned family alias to fall back on). Unit-tested; the curated-allowlist path is not e2e-verified (real relays we've seen are allow_all). One assumption to confirm live: Claude Code's --model accepts a full canonical id (e.g. claude-opus-4-8), not only opus/sonnet/haiku aliases. Co-authored-by: Isaac --- src/ucode/agents/__init__.py | 14 ++++++-------- src/ucode/cli.py | 24 +++++++++++++++++++----- src/ucode/databricks.py | 13 ++++++++----- tests/test_agents_init.py | 21 ++++++++++++++++++++- tests/test_cli.py | 34 +++++++++++++++++++++++++++++++++- tests/test_databricks.py | 7 +++++++ 6 files changed, 93 insertions(+), 20 deletions(-) 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 5fba6e69..6bc33815 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2025,15 +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 - # Relayed services forward --model to Claude Code's own flag at launch (below), not env. 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 @@ -2110,10 +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 --model to Claude Code's own flag, like `-- --model X`. - if tool == "claude" and provider and relayed and model and not forwarded_model: - ctx.args = ["--model", model, *ctx.args] - forwarded_model = model + # 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 c54480c0..305ecd14 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -691,7 +691,7 @@ def test_model_forwarded_to_claude_for_relayed_provider(self, monkeypatch): assert "ignored" not in _strip_ansi(result.output) def test_relayed_provider_without_model_forwards_nothing(self, monkeypatch): - # No --model on a relayed launch: nothing to forward. + # No --model on an allow_all relay: nothing to forward. result, _, mock_launch = self._provider_launch( monkeypatch, ["claude", "--provider", "cat.schema.svc"], @@ -701,6 +701,38 @@ def test_relayed_provider_without_model_forwards_nothing(self, monkeypatch): 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) with ( 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."""