diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c8d71188..cf7c6541 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2209,6 +2209,14 @@ def claude_cmd( skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = 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( @@ -2233,6 +2241,8 @@ def claude_cmd( claude_agent.disable_smart_routing(load_state()) print_success("Claude Code smart routing disabled; ucode routing hooks removed") return + if enable_model_discovery: + os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" _launch_tool( "claude", ctx, diff --git a/tests/test_cli.py b/tests/test_cli.py index 99b99dcc..bacfe050 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -220,6 +220,20 @@ def test_codex_enable_smart_routing_is_consumed_by_ucode(self): assert mock_launch.call_args.kwargs["enable_smart_routing_flag"] is True assert mock_launch.call_args.args[1].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"]) + + 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 == [] + + def test_claude_enable_model_discovery_is_hidden_from_help(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),