MPS: resolve relayed --model against a curated target allowlist

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
 
 

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")

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(
         (

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",

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,25 +668,70 @@ 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"},
         )
         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)

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."""