diff --git a/README.md b/README.md index f1fef58a..79dc3086 100644 --- a/README.md +++ b/README.md @@ -374,7 +374,8 @@ The output looks like: | `ug codex --model-location main.default` | Discover model services in the specified catalog and schema | | `ug claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents | | `ug claude --refresh` | Re-check Databricks, refresh models/configuration, and launch Claude Code | -| `ug claude --model-location main.default` | Discover model services in the specified catalog and schema | +| `ug claude --model-location main.default` | Add the model schema header without enabling model discovery | +| `ug claude --model-location main.default --enable-model-discovery` | Enable Claude Code model discovery in the specified catalog and schema | | `ug configure --agents claude,codex,pi` | Configure the requested agents that are available; skip the rest with a warning | | `ug configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | | `ug mcp add --location system.ai` | Register a schema's MCP servers, keeping any already configured (additive; never removes) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 79cf38e8..1593b284 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2638,7 +2638,7 @@ def claude_cmd( str | None, typer.Option( "--model-location", - help="Discover model services in `.`. Example: main.default", + help="Set the model schema header to `.`.", ), ] = None, model: Annotated[ @@ -2712,7 +2712,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: 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/README.md b/tests/README.md index 79a016b3..ee17f225 100644 --- a/tests/README.md +++ b/tests/README.md @@ -23,6 +23,10 @@ keyboard selection: nothing is selected by default, selecting Codex installs onl Codex, and submitting an empty selection installs nothing. Rendering checks cover the selected and empty checkboxes. These are local component checks, not live gateway tests. +Claude flag tests cover all combinations of `--model-location` and +`--enable-model-discovery`: the location adds the schema header, while discovery +requires a separate opt-in. These are unit/component checks, not live discovery coverage. + ## CUJ coverage matrix These are **implemented assertions**, not a claim that every version passes. diff --git a/tests/integration/README.md b/tests/integration/README.md index d9056b2c..f8b831b1 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -9,6 +9,9 @@ The existing unit tests keep their fixtures. Integration has an independent pytest configuration and uses `--confcutdir` so those fixtures cannot leak in. It is not collected by the default `uv run pytest` command. +Claude's independent `--model-location` and `--enable-model-discovery` behavior is +covered by unit/component tests; these integration journeys do not exercise it. + ## Run a specific combination Prerequisites: Python 3.12+, uv, Node/npm, and Databricks CLI >=1.0.0. The runner diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 8f484cc1..fa137a6d 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -371,12 +371,14 @@ def test_no_provider_header_without_flag(self): overlay, _ = claude.render_overlay(WS, "s4") assert "Databricks-Model-Provider-Service" not in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] - def test_parent_adds_discovery_header(self): + def test_parent_adds_header_without_enabling_discovery(self): overlay, _ = claude.render_overlay(WS, "s4", parent_schema="main.default") assert ( "Databricks-Model-Service-Parent-Schema: main.default" in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] ) + assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"] + assert claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR not in os.environ def test_bedrock_provider_pins_model_ids(self): provider_models = { diff --git a/tests/test_cli.py b/tests/test_cli.py index f12874f2..f8abafdc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -642,14 +642,25 @@ def test_claude_enable_model_discovery_sets_ucode_env(self): assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" assert mock_launch.call_args.args[1].args == [] - def test_claude_model_location_is_forwarded(self): + @pytest.mark.parametrize("model_location", [None, "main.default"]) + @pytest.mark.parametrize("enable_model_discovery", [False, True]) + def test_claude_model_location_and_discovery_are_independent( + self, model_location, enable_model_discovery + ): + args = ["claude"] + if model_location is not None: + args.extend(["--model-location", model_location]) + if enable_model_discovery: + args.append("--enable-model-discovery") with patch("ucode.cli._launch_tool") as mock_launch: - result = runner.invoke(app, ["claude", "--model-location", "main.default"]) + result = runner.invoke(app, args) assert result.exit_code == 0, result.output - assert mock_launch.call_args.kwargs["parent_schema"] == "main.default" + assert mock_launch.call_args.kwargs["parent_schema"] == model_location assert mock_launch.call_args.args[1].args == [] - assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert os.environ.get("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY") == ( + "1" if enable_model_discovery else None + ) def test_codex_model_location_is_forwarded(self): with patch("ucode.cli._launch_tool") as mock_launch: