From a1c5bd08aec34eec37af0d2db0ec6869abae8a98 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:47:08 +0000 Subject: [PATCH 1/3] Enable Claude model discovery by default for unmanaged launches --- README.md | 5 ++ src/ucode/cli.py | 15 ++--- tests/README.md | 2 +- tests/integration/README.md | 5 ++ tests/integration/test_ug_configure_claude.py | 17 +++++- tests/test_cli.py | 59 ++++++++++++++----- 6 files changed, 74 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 9ae0bb692..94fcacf80 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,11 @@ On first launch of a model-backed agent, `ug` prompts for a Databricks workspace, authenticates, and writes local agent config. Later launches reuse the saved workspace and credentials. +Without a managed workspace config, `ug claude` automatically discovers gateway +models for Claude Code's `/model` picker. Discovery defaults to `system.ai` when +no provider or model location is selected. Use `--provider` or `--model-location` +to select another model source; managed workspace configs control their own sources. + ## Configure ```bash diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 404790756..593d8bfd1 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2600,7 +2600,10 @@ def _launch_tool( # rewriting it; the admin's location exists only for this launch. provider = None parent_schema = managed_parent_schema - if tool == "claude" and (managed_provider or managed_parent_schema): + # Unmanaged Claude launches discover gateway models automatically; with no + # provider or parent header the gateway defaults to system.ai. Managed + # configs opt into discovery by selecting an MPS or Unity Catalog location. + if tool == "claude" and (managed is None or managed_provider or managed_parent_schema): os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" # The environment switch remains a developer override; managed config is the workspace # policy equivalent and must take effect before launch options are computed. @@ -3151,14 +3154,6 @@ def claude_cmd( str | None, typer.Option("--scopes", hidden=True, help="Comma-separated custom OAuth scopes."), ] = None, - enable_model_discovery: Annotated[ - bool, - typer.Option( - "--enable-model-discovery", - hidden=True, - help="Enable AI Gateway models in Claude Code's model picker.", - ), - ] = False, enable_smart_routing_flag: Annotated[ bool, typer.Option( @@ -3188,8 +3183,6 @@ 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 or provider is not 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( diff --git a/tests/README.md b/tests/README.md index c6be4ec3d..f73dd721e 100644 --- a/tests/README.md +++ b/tests/README.md @@ -33,7 +33,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | Test | User action | Expected evidence | | --- | --- | --- | -| `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Claude TUI and read a file | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; normal exit; reopen with working keyboard input | +| `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; launch plain `ug claude`, read a file, and open `/model` | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; native discovery caches `system.ai` models and the picker shows a discovered model without an opt-in flag; normal exit; reopen with working keyboard input | | `test_ug_configure_claude_anthropic_mps` | Select Anthropic MPS in the real configure picker; launch Claude | Saved provider in status; completed TUI file task; normal exit | | `test_ug_configure_codex_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Codex TUI and read a file | Generated helper invokes `ug` with clean token stdout; completed assistant answer contains the file value; normal exit and reopen | | `test_ug_configure_codex_openai_mps` | Select OpenAI MPS in the real configure picker; launch Codex | Saved provider in status; completed TUI file task; normal exit | diff --git a/tests/integration/README.md b/tests/integration/README.md index a1b13314a..1923d49e3 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -80,6 +80,11 @@ bearer override and drive their real local web-search MCP handshake/tool listing These assert protocol stdout without stripping ANSI escapes and make no workspace requests. The live Hosted configure journeys additionally execute the actual `ug` auth helper written into each agent's configuration before completing a real TUI task. +The Claude journey also opens `/model` after a plain `ug claude` launch, requires +its native gateway cache to contain `system.ai` models, and checks that a discovered +model appears in the picker. No managed config, provider, model location, discovery +flag, or inherited discovery environment variable enables this path. This runs with +the pinned Claude version (currently 2.1.268 in CI). ## Test layout and format diff --git a/tests/integration/test_ug_configure_claude.py b/tests/integration/test_ug_configure_claude.py index 8f3c29b42..54c5a8b62 100644 --- a/tests/integration/test_ug_configure_claude.py +++ b/tests/integration/test_ug_configure_claude.py @@ -13,10 +13,12 @@ @pytest.mark.smoke def test_ug_configure_claude_databricks(live_session, workspace): - """Scenario: configure Claude with Databricks Hosted and use its TUI. + """Scenario: configure Databricks Hosted, launch plain ug claude, and open /model. - Expected: configure succeeds, Claude returns a file value through the real - gateway, exits normally, and can reopen the configuration ug created. + Expected: without managed config, a provider, a model location, or a discovery + flag, Claude caches system.ai models from the gateway and shows a discovered + model in its real picker. It returns a file value through the real gateway, + exits normally, and can reopen the configuration ug created. The generated auth helper uses ug and prints only the supplied bearer. Optional AI Tools are disabled; the selected agent version is kept pinned. """ @@ -49,6 +51,15 @@ def test_ug_configure_claude_databricks(live_session, workspace): tui.boot() tui.submit(task.prompt) tui.wait_for_task(task) + screen = tui.open_model_picker() + cache = json.loads((session.home / ".claude/cache/gateway-models.json").read_text()) + assert cache["baseUrl"] == workspace.rstrip("/") + "/ai-gateway/anthropic", cache + models = cache["models"] + assert models and all(model["id"].startswith("system.ai.") for model in models), cache + assert any( + model["id"] in screen or (model.get("display_name") and model["display_name"] in screen) + for model in models + ), screen tui.exit_normally() task.assert_completed(session, "claude") session.assert_not_routed() diff --git a/tests/test_cli.py b/tests/test_cli.py index 2a0824a7e..e22654637 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -729,30 +729,41 @@ def test_codex_forwarded_model_is_not_printed_in_launch_summary( assert "Model: system.ai.gpt-5-6-luna" not in output 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: - result = runner.invoke(app, ["claude", "--enable-model-discovery"]) + @pytest.mark.parametrize("tool", ["claude", "codex", "gemini"]) + def test_unmanaged_launch_enables_discovery_only_for_claude(self, tool): + with _launch_policy_patches(None) as calls: + result = runner.invoke(app, [tool]) assert result.exit_code == 0, result.output - assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" - assert mock_launch.call_args.args[1].args == [] + assert os.environ.get("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY") == ( + "1" if tool == "claude" else None + ) + assert calls["configure"].call_args.kwargs["provider"] is None + assert calls["configure"].call_args.kwargs["parent_schema"] is None + assert calls["launch"].call_args.args[2] == [] def test_claude_model_location_is_forwarded(self): - with patch("ucode.cli._launch_tool") as mock_launch: + with _launch_policy_patches(None) as calls: 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 calls["configure"].call_args.kwargs["parent_schema"] == "main.default" + assert calls["launch"].call_args.args[2] == [] assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" - def test_claude_provider_enables_model_discovery(self): - with patch("ucode.cli._launch_tool") as mock_launch: - result = runner.invoke(app, ["claude", "--provider", "main.default.anthropic"]) + @pytest.mark.parametrize("persisted", [False, True], ids=["explicit", "saved"]) + def test_claude_provider_enables_model_discovery(self, persisted): + provider = "main.default.anthropic" + with _launch_policy_patches( + None, persisted_provider=provider if persisted else None + ) as calls: + result = runner.invoke( + app, ["claude", *([] if persisted else ["--provider", provider])] + ) assert result.exit_code == 0, result.output - assert mock_launch.call_args.kwargs["provider"] == "main.default.anthropic" - assert mock_launch.call_args.args[1].args == [] + assert calls["configure"].call_args.kwargs["provider"] == provider + assert calls["launch"].call_args.args[2] == [] assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" def test_codex_model_location_is_forwarded(self): @@ -803,7 +814,7 @@ def test_invalid_model_location_is_rejected(self, tool): assert result.exit_code == 1 assert "--model-location must be `.`." in _strip_ansi(result.output) - def test_claude_enable_model_discovery_is_hidden_from_help(self): + def test_claude_help_does_not_offer_model_discovery_flag(self): result = runner.invoke(app, ["claude", "--help"]) assert result.exit_code == 0, result.output @@ -1027,6 +1038,26 @@ def test_persisted_provider_is_not_mistaken_for_an_explicit_option(self): class TestManagedClaudeModelDiscovery: + @pytest.mark.parametrize( + "managed", + [ + {}, + { + "enabled_agents": { + "claude": {"model_config": {"model_services": ["system.ai.claude-sonnet-5"]}} + } + }, + ], + ids=["empty", "static-models"], + ) + def test_managed_config_without_discovery_source_does_not_enable_discovery(self, managed): + with _launch_policy_patches(managed) as calls: + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + calls["launch"].assert_called_once() + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ + MPS_CONFIG = { "enabled_agents": { "claude": { From 75aa2bddd69aa14305c2ca4dfaba08ff131eafe7 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:50:12 +0000 Subject: [PATCH 2/3] Trim redundant Claude discovery test cases --- tests/test_cli.py | 52 +++++++++++++---------------------------------- 1 file changed, 14 insertions(+), 38 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index e22654637..5bb6be074 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -729,18 +729,13 @@ def test_codex_forwarded_model_is_not_printed_in_launch_summary( assert "Model: system.ai.gpt-5-6-luna" not in output assert mock_launch.call_args.args[2] == forwarded_args - @pytest.mark.parametrize("tool", ["claude", "codex", "gemini"]) - def test_unmanaged_launch_enables_discovery_only_for_claude(self, tool): + def test_unmanaged_claude_launch_enables_model_discovery(self): with _launch_policy_patches(None) as calls: - result = runner.invoke(app, [tool]) + result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output - assert os.environ.get("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY") == ( - "1" if tool == "claude" else None - ) - assert calls["configure"].call_args.kwargs["provider"] is None - assert calls["configure"].call_args.kwargs["parent_schema"] is None - assert calls["launch"].call_args.args[2] == [] + assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + calls["launch"].assert_called_once() def test_claude_model_location_is_forwarded(self): with _launch_policy_patches(None) as calls: @@ -751,18 +746,12 @@ def test_claude_model_location_is_forwarded(self): assert calls["launch"].call_args.args[2] == [] assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" - @pytest.mark.parametrize("persisted", [False, True], ids=["explicit", "saved"]) - def test_claude_provider_enables_model_discovery(self, persisted): - provider = "main.default.anthropic" - with _launch_policy_patches( - None, persisted_provider=provider if persisted else None - ) as calls: - result = runner.invoke( - app, ["claude", *([] if persisted else ["--provider", provider])] - ) + def test_claude_provider_enables_model_discovery(self): + with _launch_policy_patches(None) as calls: + result = runner.invoke(app, ["claude", "--provider", "main.default.anthropic"]) assert result.exit_code == 0, result.output - assert calls["configure"].call_args.kwargs["provider"] == provider + assert calls["configure"].call_args.kwargs["provider"] == "main.default.anthropic" assert calls["launch"].call_args.args[2] == [] assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" @@ -814,12 +803,6 @@ def test_invalid_model_location_is_rejected(self, tool): assert result.exit_code == 1 assert "--model-location must be `.`." in _strip_ansi(result.output) - def test_claude_help_does_not_offer_model_discovery_flag(self): - result = runner.invoke(app, ["claude", "--help"]) - - assert result.exit_code == 0, result.output - assert "--enable-model-discovery" not in result.output - def test_codex_disable_removes_hooks_without_launching(self): with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), @@ -1038,19 +1021,12 @@ def test_persisted_provider_is_not_mistaken_for_an_explicit_option(self): class TestManagedClaudeModelDiscovery: - @pytest.mark.parametrize( - "managed", - [ - {}, - { - "enabled_agents": { - "claude": {"model_config": {"model_services": ["system.ai.claude-sonnet-5"]}} - } - }, - ], - ids=["empty", "static-models"], - ) - def test_managed_config_without_discovery_source_does_not_enable_discovery(self, managed): + def test_managed_static_models_do_not_enable_discovery(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"model_services": ["system.ai.claude-sonnet-5"]}} + } + } with _launch_policy_patches(managed) as calls: result = runner.invoke(app, ["claude"]) From b1d9f318b5ed527cb287bafc8aa583b743a1ae6a Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:09:28 +0000 Subject: [PATCH 3/3] Allow additional models in Claude discovery cache --- tests/integration/test_ug_configure_claude.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_ug_configure_claude.py b/tests/integration/test_ug_configure_claude.py index 54c5a8b62..9b2ebb882 100644 --- a/tests/integration/test_ug_configure_claude.py +++ b/tests/integration/test_ug_configure_claude.py @@ -55,10 +55,11 @@ def test_ug_configure_claude_databricks(live_session, workspace): cache = json.loads((session.home / ".claude/cache/gateway-models.json").read_text()) assert cache["baseUrl"] == workspace.rstrip("/") + "/ai-gateway/anthropic", cache models = cache["models"] - assert models and all(model["id"].startswith("system.ai.") for model in models), cache + system_ai_models = [model for model in models if model["id"].startswith("system.ai.")] + assert system_ai_models, cache assert any( model["id"] in screen or (model.get("display_name") and model["display_name"] in screen) - for model in models + for model in system_ai_models ), screen tui.exit_normally() task.assert_completed(session, "claude")