diff --git a/README.md b/README.md index d64d972c..68eb2c42 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,13 @@ ug codex --full-auto All agents route through Databricks AI Gateway using your workspace credentials — no API keys required. +For a scoped Claude launch, ug enables Claude Code's native gateway discovery; Claude Code owns +any model-cache update after it starts. ug does not fetch or rewrite Claude's private model cache +before launch. Set `UG_ENABLE_MODEL_DISCOVERY=0` to keep the routing header while using Claude's +native picker catalog instead. This switch does not disable normal Databricks `system.ai` model +discovery. A workspace-managed scoped source still enables discovery because administrator policy +takes precedence over the developer environment switch. + Codex uses the provider ID `Databricks` while keeping the `ucode` profile name. Re-run `ug configure --agents codex` to update existing generated configurations. This reuses history stored under the exact case-sensitive `Databricks` ID; it does not merge diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index db46a8ac..007f6001 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -11,7 +11,8 @@ import socket import subprocess import threading -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path from typing import cast @@ -25,11 +26,13 @@ write_json_file, ) from ucode.constants import ( + CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY, LOOPBACK_HOST, MCP_CLEANUP_SCOPES, MCP_USER_SCOPE, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, + scoped_model_discovery_enabled, ) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command from ucode.databricks import ( @@ -60,7 +63,13 @@ remove_smart_routing_hooks, sync_smart_routing_hooks, ) -from ucode.state import MANAGED_OVERLAY_KEY, is_tool_managed, mark_tool_managed, save_state +from ucode.state import ( + MANAGED_OVERLAY_KEY, + get_provider_service, + is_tool_managed, + mark_tool_managed, + save_state, +) from ucode.telemetry import agent_version, ug_version from ucode.tracing import tracing_env from ucode.ui import print_note, print_success, print_warning @@ -68,6 +77,7 @@ from .args import LaunchOptions, has_explicit_model_arg GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" +CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" # If set, Claude Code launches in headless mode instead of the interactive login flow. CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR = "CLAUDE_CODE_OAUTH_TOKEN" CLAUDE_CONFIG_DIR = Path.home() / ".claude" @@ -88,6 +98,35 @@ "backup_path": CLAUDE_BACKUP_PATH, } +_MISSING_ENV_VALUE = object() + + +@contextmanager +def launch_discovery_environment(enabled: bool | None) -> Iterator[None]: + """Apply Claude discovery variables only for the lifetime of one launch. + + POSIX launches replace this process, so the successful production path + never reaches restoration. Spawned, mocked, and failed launches do, and + must leave the caller's environment exactly as it was. + """ + names = (GATEWAY_MODEL_DISCOVERY_ENV_VAR, CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR) + previous = {name: os.environ.get(name, _MISSING_ENV_VALUE) for name in names} + try: + if enabled is not None: + for name in names: + if enabled: + os.environ[name] = "1" + else: + os.environ.pop(name, None) + yield + finally: + for name, value in previous.items(): + if value is _MISSING_ENV_VALUE: + os.environ.pop(name, None) + else: + os.environ[name] = cast(str, value) + + # Retained only to identify and remove state written by the legacy persisted opt-in. SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY @@ -108,6 +147,18 @@ def _minimum_version_requirement_message(version: str) -> str: ) +def _model_discovery_minimum_version_error() -> str | None: + """Return the discovery blocker after the effective launch source is known.""" + version = agent_version(SPEC["binary"]) + parsed_version = _parse_version(version) + if parsed_version is None or parsed_version >= MINIMUM_CLAUDE_VERSION: + return None + return ( + f"Model discovery requires Claude Code {MINIMUM_CLAUDE_VERSION_TEXT} or newer. " + f"Your current version is Claude Code {version}. Upgrade Claude Code and try again." + ) + + def minimum_version_error() -> str | None: if ( os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1" @@ -1424,7 +1475,14 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: a different port than the cached one. Keeps ANTHROPIC_BASE_URL (which Claude Code reads) in sync with the live proxy so requests reach it.""" state["relayed_proxy_port"] = port - save_state(state) + persisted_state = dict(state) + for key in ( + CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY, + "_claude_launch_provider", + "_claude_launch_parent_schema", + ): + persisted_state.pop(key, None) + save_state(persisted_state) settings = read_json_safe(CLAUDE_SETTINGS_PATH) env = settings.get("env") if isinstance(env, dict): @@ -1480,44 +1538,63 @@ 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 - # Smart routing needs Unix PTY support, which Windows does not provide. - if options.launch_smart_routing and os.name == "nt": + scoped_model_source = bool( + state.get("_claude_launch_provider") + or get_provider_service(state, "claude") + or state.get("_claude_launch_parent_schema") + ) + override = state.get(CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY) + scoped_discovery = ( + override if isinstance(override, bool) else scoped_model_discovery_enabled() + ) + if options.launch_smart_routing and scoped_model_source: raise RuntimeError( - "Smart routing in Claude Code is currently not supported on Windows. " - "Please use Codex or disable smart routing." - ) - if options.launch_smart_routing: - smart_routing_v2.launch_claude( - state, - tool_args, - binary=binary, - user_settings_path=CLAUDE_USER_SETTINGS_PATH, - # With no user pin, let Claude resolve its starting model from its own settings. - launch_model=options.user_pinned_model, - compose_settings=_compose_v2_settings, - launch_model_args=_launch_model_args, - model_name=_maybe_add_1m_suffix, + "Claude Code smart routing cannot be used with a Model Provider Service or model " + "location. Disable smart routing or remove the scoped model source and try again." ) - return - if workspace: - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) - settings_override = None - launch_args = list(tool_args) - if options.user_pinned_model: - os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model - settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} - launch_args = [ - *_launch_model_args(tool_args, options.user_pinned_model), - *tool_args, - ] - exec_or_spawn(_build_claude_argv(binary, launch_args, settings_override=settings_override)) + requested_discovery = os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1" + launch_discovery = scoped_discovery if scoped_model_source else requested_discovery + manage_discovery_env = launch_discovery if scoped_model_source or requested_discovery else None + + with launch_discovery_environment(manage_discovery_env): + if workspace and launch_discovery: + version_error = _model_discovery_minimum_version_error() + if version_error: + raise RuntimeError(version_error) + if state.get("claude_relayed"): + _launch_relayed(state, binary, tool_args) + return + # Smart routing needs Unix PTY support, which Windows does not provide. + if options.launch_smart_routing and os.name == "nt": + raise RuntimeError( + "Smart routing in Claude Code is currently not supported on Windows. " + "Please use Codex or disable smart routing." + ) + if options.launch_smart_routing: + smart_routing_v2.launch_claude( + state, + tool_args, + binary=binary, + user_settings_path=CLAUDE_USER_SETTINGS_PATH, + # With no user pin, let Claude resolve its starting model from its own settings. + launch_model=options.user_pinned_model, + compose_settings=_compose_v2_settings, + launch_model_args=_launch_model_args, + model_name=_maybe_add_1m_suffix, + ) + return + if workspace: + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + settings_override = None + launch_args = list(tool_args) + if options.user_pinned_model: + os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model + settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} + launch_args = [ + *_launch_model_args(tool_args, options.user_pinned_model), + *tool_args, + ] + exec_or_spawn(_build_claude_argv(binary, launch_args, settings_override=settings_override)) def validate_cmd(binary: str) -> list[str]: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 933df438..f6bc5cb7 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -43,6 +43,10 @@ from ucode.agents.codex import revert_legacy_shared_config from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import is_dry_run, restore_file, set_dry_run +from ucode.constants import ( + CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY, + scoped_model_discovery_enabled, +) from ucode.databricks import ( apply_pat_environment, build_shared_base_urls, @@ -2221,6 +2225,7 @@ def _launch_tool( ) elif not coding_agent_config_feature_disabled: print_note("No managed coding agent config found; using your own settings") + managed_provider = None if managed is not None: managed_provider = managed_provider_service(managed, tool) if explicit_provider and managed_provider and managed_provider != explicit_provider: @@ -2236,8 +2241,18 @@ def _launch_tool( provider = managed_provider if provider and parent_schema is not None: raise RuntimeError("--provider and --model-location cannot be used together.") + claude_scoped_model_source = tool == "claude" and bool(provider or parent_schema) + claude_scoped_model_discovery = _scoped_model_discovery_enabled( + managed_config_exists=managed is not None + ) # Checked after the managed config settles `provider`: an admin-set provider must trip this # guard too, or routing would be persisted as on while a provider is active. + if smart_routing_enabled and claude_scoped_model_source: + raise RuntimeError( + f"{TOOL_SPECS[tool]['display']} smart routing cannot be used with a Model " + "Provider Service or model location. Disable smart routing or remove the scoped " + "model source and try again." + ) if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider: raise RuntimeError( f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with " @@ -2378,6 +2393,10 @@ def _launch_tool( if tool == "claude": if provider: state["_claude_launch_provider"] = provider + elif parent_schema: + state["_claude_launch_parent_schema"] = parent_schema + if claude_scoped_model_source: + state[CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY] = claude_scoped_model_discovery elif tool == "codex": if provider: state["_codex_launch_provider"] = provider @@ -2567,6 +2586,17 @@ def _print_no_managed_config_guidance() -> None: ) +def _scoped_model_discovery_enabled(*, managed_config_exists: bool = False) -> bool: + """Whether an agent should refresh a provider/location-scoped catalog. + + ``UG_ENABLE_MODEL_DISCOVERY=0`` affects only an agent's scoped picker + catalog. It does not suppress the normal workspace discovery that populates + ``system.ai`` models. Managed config may force discovery because workspace + policy outranks a developer environment variable. + """ + return scoped_model_discovery_enabled(managed_config_exists=managed_config_exists) + + @app.command( "codex", cls=_PromptAwareCommand, @@ -2750,21 +2780,25 @@ def claude_cmd( claude_agent.disable_smart_routing(load_state()) print_success("Claude Code smart routing disabled; ug routing hooks removed") return - if enable_model_discovery or (model_location is not None and provider is None): - os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" - with _smart_routing_v2_flag(enable_smart_routing_flag): - with _disable_smart_routing_for_subcommand("claude", ctx): - _launch_tool( - "claude", - ctx, - provider=provider, - model=model, - refresh=refresh, - skip_preflight=skip_preflight, - workspace_url=workspace, - parent_schema=model_location, - custom_oauth=custom_oauth, - ) + requested_discovery: bool | None = None + if enable_model_discovery: + requested_discovery = True + elif provider is not None or model_location is not None: + requested_discovery = _scoped_model_discovery_enabled() + with claude_agent.launch_discovery_environment(requested_discovery): + with _smart_routing_v2_flag(enable_smart_routing_flag): + with _disable_smart_routing_for_subcommand("claude", ctx): + _launch_tool( + "claude", + ctx, + provider=provider, + model=model, + refresh=refresh, + skip_preflight=skip_preflight, + workspace_url=workspace, + parent_schema=model_location, + custom_oauth=custom_oauth, + ) @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 68035b7d..03465b45 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -11,6 +11,10 @@ # Controls provider- and location-scoped agent catalogs. MODEL_DISCOVERY_ENV_VAR = "UG_ENABLE_MODEL_DISCOVERY" +# Launch-only state handed from the CLI to Claude. A managed source sets this +# true so workspace policy can override a developer's environment variable. +CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY = "_claude_scoped_model_discovery" + def scoped_model_discovery_enabled( *, diff --git a/tests/conftest.py b/tests/conftest.py index a1a7419f..200cac71 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -60,6 +60,7 @@ def reject_privileged_write(path, _desired_text): monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write) monkeypatch.delenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", raising=False) monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False) + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) # A developer's ambient managed-config stub would otherwise short-circuit every fetch in the suite. monkeypatch.delenv("UCODE_MANAGED_CONFIG_STUB", raising=False) # The model-services listing is memoized for the life of the process, so without this a cached diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 8f484cc1..ec0952a5 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1438,20 +1438,27 @@ def boom(name, entry, scope=claude.MCP_USER_SCOPE): class TestClaudeLaunch: def test_gateway_discovery_enabled_for_relayed_provider(self, monkeypatch): - calls: list[tuple[dict, str, list[str]]] = [] + calls: list[tuple[dict, str, list[str], str | None]] = [] 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)), + lambda state, binary, tool_args: calls.append( + ( + state, + binary, + tool_args, + os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR), + ) + ), ) 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"])] + assert claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ + assert calls == [(state, "claude", ["--debug"], "1")] def test_relayed_launch_uses_refresh_proxy(self, monkeypatch): calls: list[tuple] = [] @@ -1607,41 +1614,237 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ model_name=claude._maybe_add_1m_suffix, ) + @pytest.mark.parametrize( + "scope_state", + [ + {"_claude_launch_provider": "main.default.anthropic"}, + {"_claude_launch_parent_schema": "main.default"}, + ], + ) + def test_direct_scoped_launch_rejects_smart_routing_before_v2(self, monkeypatch, scope_state): + launch_v2 = Mock() + monkeypatch.setattr(v2, "launch_claude", launch_v2) + + with pytest.raises(RuntimeError, match="cannot be used.*model location"): + claude.launch( + {"workspace": WS, **scope_state}, + [], + options=LaunchOptions(launch_smart_routing=True), + ) + + launch_v2.assert_not_called() + + def test_saved_provider_enforces_discovery_minimum_at_launch(self, monkeypatch): + launch = Mock() + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247") + monkeypatch.setattr(claude, "exec_or_spawn", launch) + + with pytest.raises(RuntimeError, match="requires Claude Code 2.1.248.*Upgrade"): + claude.launch( + { + "workspace": WS, + "provider_services": {"claude": "main.default.saved"}, + }, + [], + options=LaunchOptions(), + ) + + launch.assert_not_called() + + def test_managed_provider_enforces_discovery_minimum_despite_disable(self, monkeypatch): + launch = Mock() + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247") + monkeypatch.setattr(claude, "exec_or_spawn", launch) + + with pytest.raises(RuntimeError, match="requires Claude Code 2.1.248.*Upgrade"): + claude.launch( + { + "workspace": WS, + "_claude_launch_provider": "main.default.managed", + "_claude_scoped_model_discovery": True, + }, + [], + options=LaunchOptions(), + ) + + launch.assert_not_called() + def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): calls: list[list[str]] = [] + discovery_env: list[str | None] = [] monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_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)) + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda argv: ( + calls.append(argv), + discovery_env.append(os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR)), + ), + ) claude.launch({"workspace": WS, "profile": "test"}, ["--debug"], options=LaunchOptions()) assert os.environ["OAUTH_TOKEN"] == "token" - assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert discovery_env == ["1"] + assert claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] def test_gateway_discovery_enabled_under_provider(self, monkeypatch): calls: list[list[str]] = [] + discovery_env: list[str | None] = [] monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) - monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) + monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) 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)) + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda argv: ( + calls.append(argv), + discovery_env.append(os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR)), + ), + ) claude.launch( { "workspace": WS, "profile": "test", - "_claude_launch_provider": "main.default.anthropic", + "provider_services": {"claude": "main.default.anthropic"}, }, ["--debug"], options=LaunchOptions(), ) - assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert discovery_env == ["1"] + assert claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] + def test_direct_scoped_launch_honors_public_disable_and_restores_environment(self, monkeypatch): + child_env: list[tuple[str | None, str | None]] = [] + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "caller-internal") + monkeypatch.setenv(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR, "caller-claude") + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda _argv: child_env.append( + ( + os.environ.get(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR), + os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR), + ) + ), + ) + + claude.launch( + {"workspace": WS, "_claude_launch_parent_schema": "main.default"}, + [], + options=LaunchOptions(), + ) + + assert child_env == [(None, None)] + assert os.environ[claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR] == "caller-internal" + assert os.environ[claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR] == "caller-claude" + + def test_managed_scoped_override_forces_direct_discovery(self, monkeypatch): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) + monkeypatch.delenv(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) + child_env: list[str | None] = [] + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda _argv: child_env.append( + os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR) + ), + ) + + claude.launch( + { + "workspace": WS, + "_claude_launch_provider": "main.default.anthropic", + "_claude_scoped_model_discovery": True, + }, + [], + options=LaunchOptions(), + ) + + assert child_env == ["1"] + assert claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ + assert claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ + + def test_failed_launch_restores_discovery_environment(self, monkeypatch): + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "caller-internal") + monkeypatch.setenv(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR, "caller-claude") + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr( + claude, "exec_or_spawn", Mock(side_effect=RuntimeError("launch failed")) + ) + + with pytest.raises(RuntimeError, match="launch failed"): + claude.launch( + {"workspace": WS, "_claude_launch_parent_schema": "main.default"}, + [], + options=LaunchOptions(), + ) + + assert os.environ[claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR] == "caller-internal" + assert os.environ[claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR] == "caller-claude" + + def test_scoped_launch_preserves_native_gateway_cache(self, tmp_path, monkeypatch): + config_dir = tmp_path / "claude-config" + cache_path = config_dir / "cache" / "gateway-models.json" + cache_path.parent.mkdir(parents=True) + sentinel = b'{"sentinel":"native Claude cache"}\n' + cache_path.write_bytes(sentinel) + sentinel_mtime_ns = 1_700_000_000_123_456_789 + os.utime(cache_path, ns=(sentinel_mtime_ns, sentinel_mtime_ns)) + launch_events: list[str] = [] + launch_env: list[tuple[str | None, str | None]] = [] + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(config_dir)) + monkeypatch.setattr( + claude, + "get_databricks_token", + lambda *_args: launch_events.append("token") or "token", + ) + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda _argv: ( + launch_events.append("launch"), + launch_env.append( + ( + os.environ.get(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR), + os.environ.get(claude.CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR), + ) + ), + ), + ) + + claude.launch( + { + "workspace": WS, + "_claude_launch_parent_schema": "main.default", + "_claude_scoped_model_discovery": True, + }, + [], + options=LaunchOptions(), + ) + + assert launch_env == [("1", "1")] + assert launch_events == ["token", "launch"] + assert cache_path.read_bytes() == sentinel + assert cache_path.stat().st_mtime_ns == sentinel_mtime_ns + assert not (cache_path.parent / ".gateway-models.lock").exists() + class TestWriteToolConfigPrunesStaleModelEnv: """Stale ucode-managed model env keys (ANTHROPIC_MODEL, etc.) from earlier diff --git a/tests/test_cli.py b/tests/test_cli.py index bff007c9..34936f51 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -519,6 +519,24 @@ def capture(_tool, ctx, **_kwargs): assert result.exit_code == 0, result.output assert captured == [("1", ["fix the parser"])] + @pytest.mark.parametrize( + "source_args", + [ + ["--provider", "main.default.provider"], + ["--model-location", "main.default"], + ], + ) + def test_claude_scoped_source_rejects_smart_routing(self, source_args): + patches = _patch_launch("claude") + with contextlib.ExitStack() as stack: + for item in patches: + stack.enter_context(item) + result = runner.invoke(app, ["claude", "--enable-smart-routing", *source_args]) + + assert result.exit_code == 1 + assert "smart routing cannot be used" in result.output + assert "Disable smart routing or remove the scoped model source" in result.output + @pytest.mark.parametrize( ("args", "forwarded", "has_separator"), [ @@ -635,21 +653,34 @@ def test_codex_forwarded_model_is_not_printed_in_launch_summary( assert mock_launch.call_args.args[2] == forwarded_args def test_claude_enable_model_discovery_sets_ucode_env(self): - with patch("ucode.cli._launch_tool") as mock_launch: + launch_env: list[str | None] = [] + with patch( + "ucode.cli._launch_tool", + side_effect=lambda *_args, **_kwargs: launch_env.append( + os.environ.get("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY") + ), + ) as mock_launch: result = runner.invoke(app, ["claude", "--enable-model-discovery"]) assert result.exit_code == 0, result.output - assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert launch_env == ["1"] + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ assert mock_launch.call_args.args[1].args == [] def test_claude_model_location_is_forwarded(self): - with patch("ucode.cli._launch_tool") as mock_launch: + launch_env: list[str | None] = [] + with patch( + "ucode.cli._launch_tool", + side_effect=lambda *_args, **_kwargs: launch_env.append( + os.environ.get("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY") + ), + ) as mock_launch: result = runner.invoke(app, ["claude", "--model-location", "main.default"]) assert result.exit_code == 0, result.output assert mock_launch.call_args.kwargs["parent_schema"] == "main.default" - assert mock_launch.call_args.args[1].args == [] - assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert launch_env == ["1"] + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ def test_codex_model_location_is_forwarded(self): with patch("ucode.cli._launch_tool") as mock_launch: @@ -1108,7 +1139,57 @@ def test_provider_sets_transient_claude_launch_marker(self): result = runner.invoke(app, ["claude", "--provider", "main.default.anthropic"]) assert result.exit_code == 0, result.output - assert mock_launch.call_args.args[1]["_claude_launch_provider"] == "main.default.anthropic" + launch_state = mock_launch.call_args.args[1] + assert launch_state["_claude_launch_provider"] == "main.default.anthropic" + assert launch_state["_claude_scoped_model_discovery"] is True + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ + + def test_saved_claude_provider_version_error_is_actionable_at_cli(self, monkeypatch): + state = { + **MINIMAL_STATE, + "provider_services": {"claude": "main.default.saved"}, + } + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) + monkeypatch.setattr(cli_mod.claude_agent, "agent_version", lambda _binary: "2.1.247") + + def launch_direct(_tool, launch_state, tool_args, *, options): + cli_mod.claude_agent.launch(launch_state, tool_args, options=options) + + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent", side_effect=launch_direct), + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 1 + assert "Model discovery requires Claude Code 2.1.248 or newer" in result.output + assert "Upgrade Claude Code and try again" in result.output + + def test_model_location_sets_transient_claude_launch_marker(self): + state = dict(MINIMAL_STATE) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.resolve_launch_model", return_value=(state, "system.ai.claude")), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["claude", "--model-location", "main.default"]) + + assert result.exit_code == 0, result.output + launch_state = mock_launch.call_args.args[1] + assert launch_state["_claude_launch_parent_schema"] == "main.default" + assert launch_state["_claude_scoped_model_discovery"] is True + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ def test_provider_sets_transient_codex_launch_marker(self): state = dict(MINIMAL_STATE) @@ -1144,6 +1225,55 @@ def test_model_location_sets_transient_codex_launch_marker(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_codex_launch_parent_schema"] == "main.default" + @pytest.mark.parametrize( + ("tool", "source_args", "expected_marker"), + [ + ( + "claude", + ["--provider", "main.default.anthropic"], + ("_claude_launch_provider", "main.default.anthropic"), + ), + ( + "claude", + ["--model-location", "main.default"], + ("_claude_launch_parent_schema", "main.default"), + ), + ], + ) + def test_claude_discovery_disable_keeps_scope_but_suppresses_native_discovery( + self, monkeypatch, tool, source_args, expected_marker + ): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + state = dict(MINIMAL_STATE) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state) as mock_shared, + patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)), + patch("ucode.cli.resolve_launch_model", return_value=(state, "system.ai.model")), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, [tool, *source_args]) + + assert result.exit_code == 0, result.output + launch_state = mock_launch.call_args.args[1] + assert launch_state[expected_marker[0]] == expected_marker[1] + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ + assert launch_state["_claude_scoped_model_discovery"] is False + if source_args[0] == "--model-location": + # The policy flag suppresses only the agent-native scoped catalog; + # ordinary system.ai discovery still refreshes for this launch. + assert mock_shared.call_args.kwargs["skip_model_discovery"] is False + + def test_managed_config_overrides_developer_discovery_disable(self, monkeypatch): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + + assert cli_mod._scoped_model_discovery_enabled() is False + assert cli_mod._scoped_model_discovery_enabled(managed_config_exists=True) is True + class TestGeminiProviderLaunch: @staticmethod