From 8ce76d8a1928d200db8ff6a562d3db7e995230c8 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:59:16 +0000 Subject: [PATCH] Enable Claude model discovery for provider launches --- src/ucode/cli.py | 2 +- tests/test_cli.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c25aecba..e4bfd870 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2778,7 +2778,7 @@ 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): + 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): diff --git a/tests/test_cli.py b/tests/test_cli.py index 36ed8ab9..90e410fa 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -651,6 +651,15 @@ def test_claude_model_location_is_forwarded(self): assert mock_launch.call_args.args[1].args == [] 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"]) + + 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 os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + def test_codex_model_location_is_forwarded(self): with patch("ucode.cli._launch_tool") as mock_launch: result = runner.invoke(app, ["codex", "--model-location", "main.default"])