From b07173e8bb55bfff1574779cdecff7c33df339d3 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:50:03 +0000 Subject: [PATCH 1/2] [AIGTWY-4565] Enable Claude MPS model discovery --- src/ucode/agents/claude.py | 15 ++------------- tests/test_agent_claude.py | 26 ++++++++++++++++++++++---- 2 files changed, 24 insertions(+), 17 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a84a2911..7838c6f2 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -49,7 +49,7 @@ remove_smart_routing_hooks, sync_smart_routing_hooks, ) -from ucode.state import MANAGED_OVERLAY_KEY, get_provider_service, mark_tool_managed, save_state +from ucode.state import MANAGED_OVERLAY_KEY, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version from ucode.tracing import tracing_env from ucode.ui import print_note, print_success, print_warning @@ -1150,13 +1150,6 @@ def _original_launch_model(state: dict) -> str | None: return default_model(state) -def _has_provider_launch(state: dict) -> bool: - transient = state.get("_claude_launch_provider") - return (isinstance(transient, str) and bool(transient.strip())) or bool( - get_provider_service(state, "claude") - ) - - def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[str]: if not launch_model or has_explicit_model_arg(tool_args): return [] @@ -1329,11 +1322,7 @@ def launch( model_name=_maybe_add_1m_suffix, ) return - if ( - workspace - and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1" - and not _has_provider_launch(state) - ): + if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": # Discovery is launch-scoped. Pass it in the process environment rather # than persisting it in Claude's private or OS-managed settings. os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 11c50978..c3e99bc2 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -205,10 +205,7 @@ def test_smart_routing_does_not_persist_gateway_model_discovery(self, monkeypatc overlay, _ = claude.render_overlay(WS, "s4") assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"] - def test_gateway_model_discovery_skipped_under_provider(self, monkeypatch): - # A Model Provider Service routes every request to the external provider, - # so a discovered gateway endpoint id would reach a provider that can't - # resolve it — discovery must be off in that mode. + def test_gateway_model_discovery_not_persisted_under_provider(self, monkeypatch): monkeypatch.setenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", "1") overlay, _ = claude.render_overlay(WS, "s4", provider="main.x.claude-svc") assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"] @@ -1234,6 +1231,27 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] + def test_gateway_discovery_enabled_under_provider(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv(v2.ENV_VAR, raising=False) + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + { + "workspace": WS, + "profile": "test", + "_claude_launch_provider": "main.default.anthropic", + }, + ["--debug"], + options=LaunchOptions(), + ) + + assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] + class TestWriteToolConfigPrunesStaleModelEnv: """Stale ucode-managed model env keys (ANTHROPIC_MODEL, etc.) from earlier From 09edc66e85ab9835958820856c606cc2f40fa31c Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:04:41 +0000 Subject: [PATCH 2/2] [AIGTWY-4565] Fix Claude MPS discovery launch config --- src/ucode/agents/claude.py | 10 +++---- tests/test_agent_claude.py | 54 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 7838c6f2..c6194cd6 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -701,7 +701,7 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: _reconcile_managed_settings( state, - lambda base: _compose(base, enforce_model_default_hierarchy=True), + lambda base: _compose(base, enforce_model_default_hierarchy=provider is None), managed_file_keys, relayed, ) @@ -1301,6 +1301,10 @@ def launch( ) -> None: binary = SPEC["binary"] workspace = state.get("workspace") + if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": + # Discovery is launch-scoped. Pass it in the process environment rather + # than persisting it in Claude's private or OS-managed settings. + os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return @@ -1322,10 +1326,6 @@ def launch( model_name=_maybe_add_1m_suffix, ) return - if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": - # Discovery is launch-scoped. Pass it in the process environment rather - # than persisting it in Claude's private or OS-managed settings. - os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) if options.claude_launch_model: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index c3e99bc2..3222acfe 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -872,6 +872,44 @@ def test_managed_file_applies_model_default_precedence(self, monkeypatch): "haiku": "system.ai.claude-haiku-5", # Ucode default took priority. } + def test_managed_file_omits_workspace_defaults_for_provider(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": {"ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8"} + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = { + "workspace": WS, + "claude_models": { + "opus": "system.ai.claude-opus-4-8", + "haiku": "system.ai.claude-haiku-4-6", + }, + } + + claude.write_tool_config(state, None, provider="main.default.anthropic") + + env = json.loads(managed_writes[0][1])["env"] + assert not set(claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.values()) & env.keys() + + def test_managed_file_keeps_provider_model_pins(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + state = {"workspace": WS, "claude_models": {"opus": "system.ai.claude-opus-4-8"}} + + claude.write_tool_config( + state, + None, + provider="main.default.bedrock", + provider_models={"opus": "us.anthropic.claude-opus-4-6"}, + ) + + env = json.loads(managed_writes[0][1])["env"] + assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "us.anthropic.claude-opus-4-6" + def test_managed_file_removes_fable_default_when_fable_is_disabled(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( monkeypatch, @@ -1064,6 +1102,22 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: + def test_gateway_discovery_enabled_for_relayed_provider(self, monkeypatch): + calls: list[tuple[dict, str, list[str]]] = [] + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") + monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False) + monkeypatch.setattr( + claude, + "_launch_relayed", + lambda state, binary, tool_args: calls.append((state, binary, tool_args)), + ) + state = {"workspace": WS, "claude_relayed": True} + + claude.launch(state, ["--debug"], options=LaunchOptions()) + + assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert calls == [(state, "claude", ["--debug"])] + def test_relayed_launch_uses_refresh_proxy(self, monkeypatch): calls: list[tuple] = []