From b8c7a13ed6982116c917bc4760f07126c5c4a0b4 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:59:24 +0000 Subject: [PATCH 1/3] Rename --parent to --model-location --- README.md | 2 ++ src/ucode/cli.py | 20 ++++++++++---------- tests/test_cli.py | 37 +++++++++++++++++++++++++------------ 3 files changed, 37 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 2a6d93187..ee1a09f2f 100644 --- a/README.md +++ b/README.md @@ -367,8 +367,10 @@ The output looks like: | `ug configure --profile DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | | `ug codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents | | `ug codex --refresh` | Re-check Databricks, refresh models/configuration, and launch Codex | +| `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 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 fa863fc74..be492f7ef 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2017,9 +2017,9 @@ def _launch_tool( if _child_owns_stdout(tool, ctx.args): redirect_output_to_stderr() if provider is not None and parent_schema is not None: - raise RuntimeError("--provider and --parent cannot be used together.") + raise RuntimeError("--provider and --model-location cannot be used together.") if parent_schema is not None and not is_valid_catalog_schema(parent_schema): - raise RuntimeError("--parent must be `.`.") + raise RuntimeError("--model-location must be `.`.") explicit_prompt = _has_explicit_prompt(ctx) smart_routing_enabled = smart_routing_v2.enabled() # Launchers such as isaac put their harness arguments after `--`, so the harness's own @@ -2115,7 +2115,7 @@ def _launch_tool( if managed_provider: provider = managed_provider if provider and parent_schema is not None: - raise RuntimeError("--provider and --parent cannot be used together.") + raise RuntimeError("--provider and --model-location cannot be used together.") # 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 tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider: @@ -2467,10 +2467,10 @@ def codex_cmd( "before any `--` separator.", ), ] = None, - parent: Annotated[ + model_location: Annotated[ str | None, typer.Option( - "--parent", + "--model-location", help="Discover model services in `.`. Example: main.default", ), ] = None, @@ -2535,7 +2535,7 @@ def codex_cmd( refresh=refresh, skip_preflight=skip_preflight, workspace_url=workspace, - parent_schema=parent, + parent_schema=model_location, custom_oauth=custom_oauth, ) @@ -2556,10 +2556,10 @@ def claude_cmd( "before any `--` separator.", ), ] = None, - parent: Annotated[ + model_location: Annotated[ str | None, typer.Option( - "--parent", + "--model-location", help="Discover model services in `.`. Example: main.default", ), ] = None, @@ -2634,7 +2634,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 (parent is not None and provider is None): + 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): @@ -2646,7 +2646,7 @@ def claude_cmd( refresh=refresh, skip_preflight=skip_preflight, workspace_url=workspace, - parent_schema=parent, + parent_schema=model_location, custom_oauth=custom_oauth, ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 00dbeef94..b9e17eebb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -116,6 +116,10 @@ def test_subcommand_help(self, tool): result = runner.invoke(app, [tool, "--help"]) assert result.exit_code == 0 assert "Usage:" in result.output + if tool in {"claude", "codex"}: + output = _strip_ansi(result.output) + assert "--model-location" in output + assert "--parent" not in output def test_configure_help_lists_agents_flag(self): result = runner.invoke(app, ["configure", "--help"]) @@ -636,38 +640,47 @@ 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_parent_is_forwarded(self): + def test_claude_model_location_is_forwarded(self): with patch("ucode.cli._launch_tool") as mock_launch: - result = runner.invoke(app, ["claude", "--parent", "main.default"]) + 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" - def test_codex_parent_is_forwarded(self): + def test_codex_model_location_is_forwarded(self): with patch("ucode.cli._launch_tool") as mock_launch: - result = runner.invoke(app, ["codex", "--parent", "main.default"]) + result = runner.invoke(app, ["codex", "--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 == [] - def test_codex_provider_and_parent_are_mutually_exclusive(self): + def test_codex_provider_and_model_location_are_mutually_exclusive(self): result = runner.invoke( app, - ["codex", "--provider", "main.default.provider", "--parent", "main.default"], + ["codex", "--provider", "main.default.provider", "--model-location", "main.default"], ) assert result.exit_code == 1 - assert "--provider and --parent cannot be used together" in result.output + assert "--provider and --model-location cannot be used together" in result.output - def test_claude_provider_and_parent_are_mutually_exclusive(self): + def test_claude_provider_and_model_location_are_mutually_exclusive(self): result = runner.invoke( app, - ["claude", "--provider", "main.default.provider", "--parent", "main.default"], + ["claude", "--provider", "main.default.provider", "--model-location", "main.default"], ) assert result.exit_code == 1 - assert "--provider and --parent cannot be used together" in result.output + assert "--provider and --model-location cannot be used together" in result.output + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_invalid_model_location_is_rejected(self, tool): + result = runner.invoke(app, [tool, "--model-location", "main"]) + + 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): result = runner.invoke(app, ["claude", "--help"]) @@ -1114,7 +1127,7 @@ def test_provider_sets_transient_codex_launch_marker(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_codex_launch_provider"] == "main.default.openai" - def test_parent_sets_transient_codex_launch_marker(self): + def test_model_location_sets_transient_codex_launch_marker(self): state = dict(MINIMAL_STATE) with ( patch("ucode.cli.ensure_bootstrap_dependencies"), @@ -1126,7 +1139,7 @@ def test_parent_sets_transient_codex_launch_marker(self): patch("ucode.cli._fetch_managed_config", return_value=(None, False)), patch("ucode.cli.launch_agent") as mock_launch, ): - result = runner.invoke(app, ["codex", "--parent", "main.default"]) + result = runner.invoke(app, ["codex", "--model-location", "main.default"]) assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_codex_launch_parent_schema"] == "main.default" From 32f715b656099858b40b0d1616a40f744acdca90 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:18:46 +0000 Subject: [PATCH 2/3] Decouple Claude model location from discovery opt-in --- README.md | 3 ++- src/ucode/cli.py | 8 +++++--- tests/README.md | 4 ++++ tests/integration/README.md | 3 +++ tests/test_agent_claude.py | 4 +++- tests/test_cli.py | 19 +++++++++++++++---- 6 files changed, 32 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index ee1a09f2f..2979241e1 100644 --- a/README.md +++ b/README.md @@ -370,7 +370,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 66e4b8260..73606a497 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2625,7 +2625,8 @@ def claude_cmd( str | None, typer.Option( "--model-location", - help="Discover model services in `.`. Example: main.default", + help="Set the model schema header to `.`. " + "Does not enable model discovery. Example: main.default", ), ] = None, model: Annotated[ @@ -2667,7 +2668,8 @@ def claude_cmd( typer.Option( "--enable-model-discovery", hidden=True, - help="Enable AI Gateway models in Claude Code's model picker.", + help="Enable AI Gateway models in Claude Code's model picker. " + "Combine with --model-location to scope discovery to a catalog and schema.", ), ] = False, enable_smart_routing_flag: Annotated[ @@ -2699,7 +2701,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 79a016b34..ee17f2254 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 d9056b2c4..f8b831b10 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 8f484cc1b..fa137a6db 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 e76ed3c39..7326a47ad 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: From f013720da0c4776581d483ba806ccae5c57a09ef Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:23:57 +0000 Subject: [PATCH 3/3] Shorten Claude model discovery flag help --- src/ucode/cli.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 73606a497..91cb451a0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2625,8 +2625,7 @@ def claude_cmd( str | None, typer.Option( "--model-location", - help="Set the model schema header to `.`. " - "Does not enable model discovery. Example: main.default", + help="Set the model schema header to `.`.", ), ] = None, model: Annotated[ @@ -2668,8 +2667,7 @@ def claude_cmd( typer.Option( "--enable-model-discovery", hidden=True, - help="Enable AI Gateway models in Claude Code's model picker. " - "Combine with --model-location to scope discovery to a catalog and schema.", + help="Enable AI Gateway models in Claude Code's model picker.", ), ] = False, enable_smart_routing_flag: Annotated[