diff --git a/README.md b/README.md index 9658dbc86..c11648b3b 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,8 @@ to discovery or change ug's saved default. `ug copilot` lists every Claude and GPT model you can use on the workspace in Copilot's `/model` picker, and you can switch between them mid-session. Claude models use the Anthropic Messages API, GPT-5 and newer use Responses, and older -GPT models use Chat Completions. Gemini and open-source models aren't offered. +GPT models use Chat Completions. Gemini and open-source models aren't offered, except under +`--model-location` (below). `ug copilot --model claude-opus-5-5` (or `system.ai.claude-opus-5-5`, or Copilot's own `dbx-anthropic/system.ai.claude-opus-5-5`) picks the model a session starts on; an unknown or unsupported model fails before Copilot starts and lists the @@ -95,15 +96,17 @@ Bedrock service, and the chat models of an OpenAI, Azure OpenAI or Microsoft Fou allows every model, `/model` lists the models the gateway reports for it, which the provider account may not serve; name the one to start on with `--model`. Services that relay a Claude subscription work only with Claude Code. Interactive `ug configure` offers the matching services under **External Models**, and a managed workspace -config can set one for Copilot. +config can set one for Copilot. `ug copilot --model-location .` uses the model +services of one schema instead: services named like Claude or GPT models use their usual API, and +the rest go through the gateway's Anthropic endpoint, as for Claude Code. 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. Interactive `ug configure` also offers a **Custom Model Location** (`.`) for -Claude Code and Codex alongside Databricks Hosted and External Models. The choice is saved -with the rest of the configuration, so later `ug claude` / `ug codex` launches use it; an +Claude Code, Codex and Copilot alongside Databricks Hosted and External Models. The choice is saved +with the rest of the configuration, so later `ug claude` / `ug codex` / `ug copilot` launches use it; an explicit `--provider` or `--model-location` still overrides it for that launch. `ug codex` validates discovered models with the installed Codex binary and publishes diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index a53e89526..ebef7eb88 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -16,7 +16,7 @@ from pathlib import Path from ucode.config_io import ToolSpec -from ucode.constants import AGENT_CLAUDE, AGENT_CODEX +from ucode.constants import AGENT_CLAUDE, AGENT_CODEX, AGENT_COPILOT from ucode.databricks import ( AnthropicModelCatalog, get_databricks_token, @@ -80,6 +80,8 @@ DEFAULT_TOOL = "codex" BUNDLE_VERSION = 1 _MANAGED_SETTINGS_TOOLS = {"claude", "codex"} +# Agents that can take their models from a Unity Catalog schema (`--model-location`). +MODEL_LOCATION_TOOLS = (AGENT_CLAUDE, AGENT_CODEX, AGENT_COPILOT) _NATIVE_UPGRADE_COMMANDS = { "claude": ["claude", "upgrade"], "codex": ["codex", "update"], @@ -483,8 +485,9 @@ def configure_tool( picker_catalog=picker_catalog, ) elif tool == "copilot": - # Under a model provider service, Copilot starts on one of the service's models, chosen at launch. - if not model and not provider: + # Under a model provider service or a model location, Copilot starts on one of its models, + # chosen at launch. + if not model and not provider and not parent_schema: raise RuntimeError(f"A {tool} model must be selected before configuration.") result = copilot.write_tool_config(state, model) else: @@ -605,10 +608,10 @@ def configure_single_tool( def _persisted_model_location(tool: str, state: dict, provider: str | None) -> str | None: """The model location ``ug configure`` saved for ``tool``, unless a provider is in play. - Only claude and codex discover models by parent schema; an explicit or persisted provider - keeps the source it names. + Only ``MODEL_LOCATION_TOOLS`` discover models by parent schema; an explicit or persisted + provider keeps the source it names. """ - if tool not in (AGENT_CLAUDE, AGENT_CODEX) or provider or get_provider_service(state, tool): + if tool not in MODEL_LOCATION_TOOLS or provider or get_provider_service(state, tool): return None return get_model_location(state, tool) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 3d091088c..47080feba 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -12,7 +12,9 @@ `ug copilot --provider` registers a Model Provider Service's models instead: its Claude models on the Anthropic endpoint, its other chat models on the gateway's native OpenAI route, with every -registry provider sending the service's routing header. +registry provider sending the service's routing header. `ug copilot --model-location` registers +the model services of one Unity Catalog schema: those named like Claude or GPT models on their +usual wire, the rest on the Anthropic endpoint like Claude Code, which the gateway translates. Copilot holds its provider credential for the whole session, so every provider points at a per-session loopback proxy (`gateway_proxy.start_gateway_auth_proxy`): Copilot authenticates to @@ -66,6 +68,7 @@ bare_model_name, build_copilot_base_url, build_tool_base_url, + list_anthropic_model_catalog, list_provider_service_models, model_version_sort_key, resolve_provider_service, @@ -75,6 +78,7 @@ acquire_exclusive_file_lock, release_file_lock, ) +from ucode.smart_routing.routing import unwrap_anthropic_gateway_model from ucode.state import is_tool_managed, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_user_agent, ug_version from ucode.ui import print_warning @@ -117,6 +121,8 @@ CATALOG_STATE_KEY = "copilot_catalog" # The Model Provider Service this launch routes through; set by the launch command, never saved. LAUNCH_PROVIDER_SERVICE_KEY = "_copilot_launch_provider_service" +# The Unity Catalog schema whose model services this launch uses; never saved. +LAUNCH_MODEL_LOCATION_KEY = "_copilot_launch_model_location" # The wire APIs Copilot speaks; `responses` and `completions` are Copilot's `wireApi` values. WIRE_ANTHROPIC = "anthropic" @@ -379,6 +385,27 @@ def provider_service_catalog( return catalog +def location_catalog(state: dict, schema: str) -> list[CopilotModel]: + """The model services of Unity Catalog schema ``schema`` that this launch registers. + + The listing names some services by gateway alias (``anthropic-aigw--``); Copilot + sends the full name, which the gateway routes directly.""" + workspace = state["workspace"] + listing = list_anthropic_model_catalog( + workspace, launch_token(state, workspace), parent_schema=schema + ) + models = dict.fromkeys(unwrap_anthropic_gateway_model(model) for model in listing.model_ids) + catalog = [ + CopilotModel(model, classify_copilot_model(model) or WIRE_ANTHROPIC) + for model in models + if bare_model_name(model) not in _UNSUPPORTED_MODELS + ] + if not catalog: + reason = f": {listing.error_msg}" if listing.error_msg else "" + raise RuntimeError(f"Found no model services to use in {schema}{reason}.") + return sorted(catalog, key=_catalog_sort_key) + + def catalog_from_listings( entries: Iterable[ModelServiceEntry] | None, fallback: Iterable[str] ) -> list[CopilotModel] | None: @@ -567,6 +594,7 @@ def render_providers_config( session_secret: str, *, provider_service: str | None = None, + model_location: str | None = None, ) -> dict: """Copilot's providers.json for ``catalog``: one provider per wire, all on the session proxy. @@ -613,8 +641,12 @@ def render_providers_config( if catalog_id is not None: entry["modelId"] = catalog_id limits = copilot_model_limits(model.id) if model.wire == WIRE_ANTHROPIC else None - if limits is not None and provider_service and sum(limits) > _CLAUDE_200K_WINDOW: - # A Model Provider Service may not serve the 1M window. + if ( + limits is not None + and (provider_service or model_location) + and sum(limits) > _CLAUDE_200K_WINDOW + ): + # A Model Provider Service, or a location's model service, may not serve the 1M window. limits = _CLAUDE_200K_LIMITS if limits is not None: entry["maxPromptTokens"], entry["maxOutputTokens"] = limits @@ -852,14 +884,16 @@ def _wait_for_exit(proc: _ChildProcess) -> int: def _session_catalog( - state: dict, tool_args: list[str], provider_service: str | None + state: dict, tool_args: list[str], provider_service: str | None, model_location: str | None ) -> list[CopilotModel]: - """The models this launch registers: the Model Provider Service's, else the Databricks - catalog.""" + """The models this launch registers: the Model Provider Service's, the model location's, + else the Databricks catalog.""" if provider_service: return provider_service_catalog( state, provider_service, explicit_model_arg_value(tool_args) ) + if model_location: + return location_catalog(state, model_location) catalog = copilot_catalog(state) if not catalog: raise RuntimeError( @@ -873,11 +907,15 @@ def _session_catalog( def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: provider_service = state.pop(LAUNCH_PROVIDER_SERVICE_KEY, None) - catalog = _session_catalog(state, tool_args, provider_service) + model_location = state.pop(LAUNCH_MODEL_LOCATION_KEY, None) + catalog = _session_catalog(state, tool_args, provider_service, model_location) # Resolve before the proxy starts so a bad model fails without starting Copilot. tool_args, selected = rewrite_model_args(tool_args, catalog) start_model = selected or launch_model( - state, catalog, requested=options.launch_model, managed_list=not provider_service + state, + catalog, + requested=options.launch_model, + managed_list=not (provider_service or model_location), ) scrub_stored_credentials() @@ -901,7 +939,11 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None try: proxy_base_url = f"http://{LOOPBACK_HOST}:{server.server_address[1]}" config = render_providers_config( - catalog, proxy_base_url, session_secret, provider_service=provider_service + catalog, + proxy_base_url, + session_secret, + provider_service=provider_service, + model_location=model_location, ) wire_override = os.environ.get("COPILOT_PROVIDER_WIRE_MODEL") if wire_override: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0dc1cc36c..ec468e3a0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -22,6 +22,7 @@ from ucode import custom_oauth from ucode.agents import ( + MODEL_LOCATION_TOOLS, TOOL_SPECS, LaunchOptions, check_gateway_endpoint, @@ -53,6 +54,7 @@ from ucode.constants import ( AGENT_CLAUDE, AGENT_CODEX, + AGENT_COPILOT, ISAAC_SMART_ROUTER_FEEDBACK_ENV_VAR, SMART_ROUTING_ENV_KEYS, UG_HEADERS_ENV_VAR, @@ -882,7 +884,7 @@ def _setup_single_agent( ) -> dict: """Log in, discover the agent's models, and write its config, then print the summary. - With ``apply_managed`` the workspace's managed config governs claude/codex (the + With ``apply_managed`` the workspace's managed config governs claude/codex/copilot (the ``ug configure --agent`` path: admin-enabled applies the admin config, a disabled agent is rejected). ``ug agents add`` passes ``apply_managed=False`` to set up a self-managed agent standalone, independent of the managed config. @@ -902,16 +904,17 @@ def _setup_single_agent( if apply_managed and tool in ("claude", "codex"): # Explicit configure always reads the workspace policy fresh, so it can install it. managed, _ = _refresh_configured_policy(state) + elif apply_managed: + # Other agents only read the policy; they never install it. + managed, _ = _read_policy_for_other_agent(state) + if apply_managed: refuse_user_scope_over_managed_config(state.get("workspace"), managed) + if apply_managed and tool in MODEL_LOCATION_TOOLS: _reject_disabled_agent(managed, tool) if managed is not None: state = resolve_state(managed, state, tool) if not managed_provider_service(managed, tool): parent_schema = managed_unity_catalog_location(managed, tool) - elif apply_managed: - # Other agents only read the policy; they never install it. - managed, _ = _read_policy_for_other_agent(state) - refuse_user_scope_over_managed_config(state.get("workspace"), managed) state = configure_single_tool(tool, state, parent_schema=parent_schema) if apply_managed and not managed_enabled_tools(managed or {}): state = _record_unmanaged_configured_agents(state, [tool]) @@ -924,9 +927,9 @@ def _setup_single_agent( skills_configured = _configure_managed_skills(managed) if resolved_mcps is not None and skills_configured: _sync_claude_cowork_mcp_servers(state, managed_mcp_servers=resolved_mcps) - install_databricks_ai_tools_for_agents( - [tool], state, force_refresh=tool not in ("claude", "codex") - ) + # Only the Claude and Codex branch above reads the managed config fresh. + managed_refreshed = apply_managed and tool in ("claude", "codex") + install_databricks_ai_tools_for_agents([tool], state, force_refresh=not managed_refreshed) spec = TOOL_SPECS[tool] provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state) console.print( @@ -969,9 +972,9 @@ def _maybe_select_provider_service(tool: str, state: dict) -> dict: (``system.ai``), a Model Provider Service, or a custom Unity Catalog model location — and persist (or clear) the choice. - No-op for tools other than claude/codex/gemini/copilot (gemini and copilot have no custom - location). Falls back to Databricks when no matching provider services are found or the - listing fails and no custom location is offered. + No-op for tools other than claude/codex/gemini/copilot (gemini has no custom location). Falls + back to Databricks when no matching provider services are found or the listing fails and no + custom location is offered. """ if tool not in ("claude", "codex", "gemini", "copilot"): return state @@ -989,7 +992,7 @@ def _use_databricks() -> dict: token = get_databricks_token(state["workspace"], state.get("profile")) with spinner("Checking for model provider services..."): names, reason = list_tool_provider_services(tool, state["workspace"], token) - offers_location = tool in ("claude", "codex") + offers_location = tool in MODEL_LOCATION_TOOLS if reason is not None: # Most workspaces don't have the feature enabled — that's the common case, # so skip the provider choice silently. Only surface unexpected failures. @@ -1133,7 +1136,7 @@ def _configure_workspace_command( resolved = resolve_state(managed, state, tool_name) parent_schema = ( managed_unity_catalog_location(managed, tool_name) - if tool_name in ("claude", "codex") + if tool_name in MODEL_LOCATION_TOOLS and not managed_provider_service(managed, tool_name) else None ) @@ -3368,7 +3371,7 @@ def _launch_tool( parent_schema is None and explicit_provider is None and not persisted_provider - and tool in {"claude", "codex"} + and tool in MODEL_LOCATION_TOOLS ): parent_schema = get_model_location(state, tool) # An explicit --model-location displaces a persisted provider, as --provider displaces a @@ -3410,7 +3413,7 @@ def _launch_tool( managed_provider = managed_provider_service(managed or {}, tool) managed_parent_schema = ( managed_unity_catalog_location(managed or {}, tool) - if tool in {"claude", "codex"} and not managed_provider + if tool in MODEL_LOCATION_TOOLS and not managed_provider else None ) if managed_provider: @@ -3575,16 +3578,24 @@ def _launch_tool( route_root_model = None managed_model = None relayed_forward_model = None # forwarded to Claude Code's --model for a relayed provider - if provider or managed_parent_schema or picker_catalog: + if ( + provider + or managed_parent_schema + or picker_catalog + or (tool == AGENT_COPILOT and parent_schema) + ): # Routing through a Model Provider Service pins no Databricks model; # scoped UC discovery likewise lets the agent select from the parent schema. Skip model # resolution, which would otherwise fail when global discovery found no models. resolved_model = None managed_source_model = ( managed_launch_model(managed or {}, recommendation, tool) - if tool == "claude" and (managed_provider or managed_parent_schema) + if tool in (AGENT_CLAUDE, AGENT_COPILOT) + and (managed_provider or managed_parent_schema) else None ) + if tool == AGENT_COPILOT: + managed_model = managed_source_model if tool == "claude" and managed_parent_schema: # Native discovery supplies the catalog, but the managed policy still controls # which model Claude starts on. @@ -3723,16 +3734,19 @@ def _launch_tool( _print_budget_panel(recommendation, tool, managed) # The managed config's MCP servers and skills are both applied at `ug configure`, not here, # so the launch hot path makes no per-launch discovery calls for them. - if tool == "claude": + if tool == AGENT_CLAUDE: if provider: state["_claude_launch_provider"] = provider - elif tool == "codex": + elif tool == AGENT_CODEX: if provider: state["_codex_launch_provider"] = provider elif parent_schema: state["_codex_launch_parent_schema"] = parent_schema - elif tool == "copilot" and provider: - state[copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY] = provider + elif tool == AGENT_COPILOT: + if provider: + state[copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY] = provider + elif parent_schema: + state[copilot_agent.LAUNCH_MODEL_LOCATION_KEY] = parent_schema launch_options = _launch_options( tool, ctx.args, @@ -4205,6 +4219,13 @@ def copilot_cmd( "(..); `/model` then lists that service's models.", ), ] = None, + model_location: Annotated[ + str | None, + typer.Option( + "--model-location", + help="Discover model services in `.`. Example: main.default", + ), + ] = None, skip_preflight: SkipPreflightOption = False, workspace: WorkspaceOption = None, ) -> None: @@ -4215,6 +4236,7 @@ def copilot_cmd( provider=provider, skip_preflight=skip_preflight, workspace_url=workspace, + parent_schema=model_location, ) diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 6764ba84b..8ded72215 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -7,6 +7,7 @@ AGENT_CLAUDE = "claude" AGENT_CODEX = "codex" +AGENT_COPILOT = "copilot" ENABLE_SMART_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_V2" ENABLE_SUBAGENT_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_SUBAGENT_ONLY" diff --git a/tests/integration/README.md b/tests/integration/README.md index bcf4b24b7..5b7b6d8a4 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -272,6 +272,8 @@ and then a GPT task chosen with Copilot's `dbx-responses/` form, an unkno fails before Copilot starts, and the TUI switches from Claude to GPT with `/model` mid-session. Model evidence comes from Copilot's own `~/.copilot/session-state/*/events.jsonl` (`utils/agents/copilot.py`). +`test_ug_copilot_model_location.py` launches Copilot from fresh state with `--model-location` on the +model-discovery schema and completes a file task on its Claude-compatible model service. `test_ug_copilot_providers.py` runs Copilot through the MPS fixtures below: a file task on the Anthropic service's model, on the OpenAI, Azure OpenAI and Microsoft Foundry services' `gpt-5-nano`, and on a named model of the Bedrock allow-all service; launching on the relayed service fails before Copilot starts. @@ -517,6 +519,7 @@ test_ug_copilot_headless.py # Copilot via the session proxy: Claude test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert, upgrade test_ug_copilot_models.py # Copilot model registry: Claude then GPT, unknown model, TUI /model switch test_ug_copilot_providers.py # Copilot via Model Provider Services: Anthropic, OpenAI, Azure OpenAI, Foundry, Bedrock allow-all; relayed refused +test_ug_copilot_model_location.py # Copilot on a Unity Catalog model location: fresh --model-location file task test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding @@ -723,7 +726,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **88 live cases** (including 12 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds sixteen live cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds seventeen live cases. One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **44 `managed_fixture` cases** (twelve of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -949,7 +952,7 @@ A second non-blocking job, **Managed OpenCode self-managed** (`managed_fixture a in `test_ug_agents_self_managed.py`), installs Claude, Codex and OpenCode and runs against the managed e2e workspace (`E2E_ADMIN_WORKSPACE`) with the injected `managed_workspace_default` config, which enables Claude/Codex but not OpenCode. It is likewise outside `cujs`. -A non-blocking **Copilot** job (`live and copilot`, fifteen cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, sixteen cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's diff --git a/tests/integration/test_ug_copilot_model_location.py b/tests/integration/test_ug_copilot_model_location.py new file mode 100644 index 000000000..005985488 --- /dev/null +++ b/tests/integration/test_ug_copilot_model_location.py @@ -0,0 +1,46 @@ +"""CUJs for running GitHub Copilot CLI on a Unity Catalog model location via ug.""" + +import pytest +from utils.constants import COPILOT_HEADLESS_FLAGS +from utils.evidence import ( + FileTask, + assert_completed_task_model, + assert_no_session_registry_left, + copilot_selected_models, +) + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + + +def test_ug_copilot_runs_a_task_on_a_model_location( + live_session, workspace, parent_schema, claude_parent_model +): + """Scenario: from fresh state, run a headless task with `--model-location` set to the + model-discovery schema and `--model` naming its Claude-compatible model service. + + Expected: ug sets up the workspace, Copilot starts on that service through `dbx-anthropic`, + the completed answer contains the fixture value and was sent as the service's full name, and + no session registry is left. + """ + session = live_session + task = FileTask(session) + session.record("copilot-model-location.json", {"location": parent_schema}) + + result = session.run( + "copilot", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--model", + claude_parent_model, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, claude_parent_model) + assert f"dbx-anthropic/{claude_parent_model}" in copilot_selected_models(session) + assert_no_session_registry_left(session) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index ec9eb7cc2..b1ac9a0f0 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -12,7 +12,7 @@ from ucode import custom_oauth from ucode.agents import copilot -from ucode.databricks import ModelServiceEntry +from ucode.databricks import AnthropicModelCatalog, ModelServiceEntry from ucode.telemetry import client_platform WS = "https://example.databricks.com" @@ -555,6 +555,48 @@ def test_invalid_service_raises_its_error(self, monkeypatch): copilot.provider_service_catalog({"workspace": WS}, MPS, None) +def _stub_location(monkeypatch, model_ids, error=None): + monkeypatch.setattr(copilot, "launch_token", lambda state, workspace, **_: "dbx-token") + seen = {} + + def anthropic_catalog(workspace, token, *, parent_schema): + seen["parent_schema"] = parent_schema + return AnthropicModelCatalog( + model_ids=list(model_ids), model_id_to_display_name={}, error_msg=error + ) + + monkeypatch.setattr(copilot, "list_anthropic_model_catalog", anthropic_catalog) + return seen + + +class TestLocationCatalog: + def test_registers_the_schemas_services_by_full_name(self, monkeypatch): + seen = _stub_location( + monkeypatch, + [ + "anthropic-aigw-1a2b3c4d-main.ucode.claude_haiku", + "main.ucode.claude-sonnet-4-6", + "anthropic-aigw-0f0f0f0f-main.ucode.gpt-5-6-luna", + "main.ucode.claude-opus-4-1", + ], + ) + + models = copilot.location_catalog({"workspace": WS}, "main.ucode") + + assert [(m.id, m.wire) for m in models] == [ + ("main.ucode.claude-sonnet-4-6", "anthropic"), + ("main.ucode.claude_haiku", "anthropic"), + ("main.ucode.gpt-5-6-luna", "responses"), + ] + assert seen["parent_schema"] == "main.ucode" + + def test_an_empty_schema_says_so(self, monkeypatch): + _stub_location(monkeypatch, [], error="AI Gateway returned no Anthropic model ids") + + with pytest.raises(RuntimeError, match="no model services to use in main.ucode"): + copilot.location_catalog({"workspace": WS}, "main.ucode") + + class TestCopilotModelLimits: @pytest.mark.parametrize( ("model", "limits"), @@ -1203,6 +1245,26 @@ def test_provider_service_launch_registers_its_models(self, tmp_path, monkeypatc assert seen["env"]["COPILOT_MODEL"] == "dbx-anthropic/claude-sonnet-5" assert copilot.LAUNCH_PROVIDER_SERVICE_KEY not in state + def test_model_location_launch_registers_the_schemas_services(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + seen = _stub_launch(monkeypatch) + _stub_location(monkeypatch, ["main.ucode.claude-sonnet-4-6", "main.ucode.claude_haiku"]) + state = dict(STATE, copilot_static_models=[GPT]) + state[copilot.LAUNCH_MODEL_LOCATION_KEY] = "main.ucode" + + with pytest.raises(SystemExit): + copilot.launch(state, [], options=copilot.LaunchOptions()) + + registry = seen["registry"] + assert [m["id"] for m in registry["models"]] == [ + "main.ucode.claude-sonnet-4-6", + "main.ucode.claude_haiku", + ] + assert "headers" not in registry["providers"][0] + assert registry["models"][0]["maxPromptTokens"] == 168_000 + assert seen["env"]["COPILOT_MODEL"] == "dbx-anthropic/main.ucode.claude-sonnet-4-6" + assert copilot.LAUNCH_MODEL_LOCATION_KEY not in state + def test_provider_service_launch_rejects_a_model_it_does_not_serve(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) seen = _stub_launch(monkeypatch) diff --git a/tests/test_cli.py b/tests/test_cli.py index bd237ac5d..3d275460e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1186,39 +1186,57 @@ def test_codex_model_location_is_forwarded(self): assert mock_launch.call_args.kwargs["parent_schema"] == "main.default" assert mock_launch.call_args.args[1].args == [] - def test_codex_provider_and_model_location_are_mutually_exclusive(self): - with _launch_policy_patches(None): + @pytest.mark.parametrize("tool", ["claude", "codex", "copilot"]) + def test_provider_and_model_location_are_mutually_exclusive(self, tool): + with _launch_policy_patches(None) as calls: + calls["state"]["available_tools"] = [tool] result = runner.invoke( app, - [ - "codex", - "--provider", - "main.default.provider", - "--model-location", - "main.default", - ], + [tool, "--provider", "main.default.provider", "--model-location", "main.default"], ) assert result.exit_code == 1 assert "--provider and --model-location cannot be used together" in result.output - def test_claude_provider_and_model_location_are_mutually_exclusive(self): - with _launch_policy_patches(None): - result = runner.invoke( - app, - [ - "claude", - "--provider", - "main.default.provider", - "--model-location", - "main.default", - ], - ) + def test_copilot_launch_uses_saved_model_location(self): + with _launch_policy_patches(None) as calls: + calls["state"]["model_locations"] = {"copilot": "main.saved"} + result = runner.invoke(app, ["copilot"]) - assert result.exit_code == 1 - assert "--provider and --model-location cannot be used together" in result.output + assert result.exit_code == 0, result.output + launch_state = calls["launch"].call_args.args[1] + assert launch_state[copilot_agent.LAUNCH_MODEL_LOCATION_KEY] == "main.saved" - @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_copilot_explicit_provider_overrides_saved_location(self): + with _launch_policy_patches(None) as calls: + calls["state"]["available_tools"] = ["copilot"] + calls["state"]["model_locations"] = {"copilot": "main.saved"} + result = runner.invoke(app, ["copilot", "--provider", "main.default.openai"]) + + assert result.exit_code == 0, result.output + launch_state = calls["launch"].call_args.args[1] + assert launch_state[copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY] == "main.default.openai" + assert copilot_agent.LAUNCH_MODEL_LOCATION_KEY not in launch_state + + def test_copilot_explicit_model_location_overrides_saved_provider(self): + with _launch_policy_patches(None, persisted_provider="main.default.openai") as calls: + result = runner.invoke(app, ["copilot", "--model-location", "main.flag"]) + + assert result.exit_code == 0, result.output + launch_state = calls["launch"].call_args.args[1] + assert launch_state[copilot_agent.LAUNCH_MODEL_LOCATION_KEY] == "main.flag" + assert copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY not in launch_state + + def test_copilot_saved_model_location_is_ignored_under_managed_config(self): + managed = {"enabled_agents": {"copilot": {"model_config": {}}}} + with _launch_policy_patches(managed) as calls: + calls["state"]["model_locations"] = {"copilot": "main.saved"} + result = runner.invoke(app, ["copilot"]) + + assert result.exit_code == 0, result.output + assert copilot_agent.LAUNCH_MODEL_LOCATION_KEY not in calls["launch"].call_args.args[1] + + @pytest.mark.parametrize("tool", ["claude", "codex", "copilot"]) def test_invalid_model_location_is_rejected(self, tool): with _launch_policy_patches(None): result = runner.invoke(app, [tool, "--model-location", "main"]) @@ -2188,6 +2206,24 @@ def test_model_location_sets_transient_codex_launch_marker(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_codex_launch_parent_schema"] == "main.default" + def test_model_location_sets_transient_copilot_launch_marker(self): + state = dict(MINIMAL_STATE, available_tools=["copilot"]) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["copilot", "--model-location", "main.default"]) + + assert result.exit_code == 0, result.output + assert mock_configure.call_args.args[2] is None + launch_state = mock_launch.call_args.args[1] + assert launch_state[copilot_agent.LAUNCH_MODEL_LOCATION_KEY] == "main.default" + class TestGeminiProviderLaunch: @staticmethod @@ -3764,6 +3800,51 @@ def test_managed_default_is_the_launch_model(self): assert result.exit_code == 0, result.output assert calls["launch"].call_args.kwargs["options"].launch_model == "system.ai.gpt-5-4" + @pytest.mark.parametrize( + "model_source", + [ + {"unity_catalog_location": "main.models"}, + {"model_provider_service": "main.default.mps"}, + ], + ) + def test_managed_default_is_the_launch_model_under_a_managed_source(self, model_source): + managed = { + "enabled_agents": { + "copilot": {"model_config": {"default_model": "gpt-5-4", **model_source}} + } + } + with _launch_policy_patches(managed) as calls: + result = runner.invoke(app, ["copilot"]) + + assert result.exit_code == 0, result.output + assert calls["launch"].call_args.kwargs["options"].launch_model == "gpt-5-4" + + @pytest.mark.parametrize( + ("model_config", "expected_provider", "expected_location"), + [ + ( + {"model_provider_service": "main.default.managed-mps"}, + "main.default.managed-mps", + None, + ), + ({"unity_catalog_location": "main.managed"}, None, "main.managed"), + ], + ids=["mps", "uc-parent"], + ) + def test_managed_source_overrides_saved_provider( + self, model_config, expected_provider, expected_location + ): + managed = {"enabled_agents": {"copilot": {"model_config": model_config}}} + with _launch_policy_patches(managed, persisted_provider="main.default.developer") as calls: + calls["state"]["model_locations"] = {"copilot": "main.saved"} + result = runner.invoke(app, ["copilot"]) + + assert result.exit_code == 0, result.output + assert calls["shared"].call_args.kwargs["skip_model_discovery"] is True + launch_state = calls["launch"].call_args.args[1] + assert launch_state.get(copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY) == expected_provider + assert launch_state.get(copilot_agent.LAUNCH_MODEL_LOCATION_KEY) == expected_location + def test_a_managed_default_still_discovers_the_catalog(self): assert not cli_mod._skip_model_discovery( self.MANAGED, "copilot", provider=None, managed_parent_schema=None, family_defaults=None @@ -4728,6 +4809,17 @@ def test_picker_offers_location_when_provider_feature_unavailable(self, monkeypa assert asked == [["databricks", "location"]] assert cli_mod.get_model_location(result, "codex") == "main.models" + def test_picker_offers_location_for_copilot(self, monkeypatch): + state, _saved, asked = self._picker( + monkeypatch, "copilot", services=(["a.b.mps"], None), answers=["location"] + ) + + result = cli_mod._maybe_select_provider_service("copilot", state) + + assert asked == [["databricks", "mps", "location"]] + assert cli_mod.get_model_location(result, "copilot") == "main.models" + assert cli_mod.get_provider_service(result, "copilot") is None + def test_picker_reprompts_until_location_is_valid(self, monkeypatch): state, _saved, _asked = self._picker( monkeypatch, @@ -4988,6 +5080,86 @@ def test_single_codex_agent_passes_managed_parent_directly(self, monkeypatch): configure.assert_called_once_with("codex", state, parent_schema="main.models") install_ai_tools.assert_called_once_with(["codex"], state, force_refresh=False) + def test_managed_copilot_parent_is_passed_to_generic_configure(self, monkeypatch): + state = { + **MINIMAL_STATE, + "available_tools": [], + "copilot_models": [], + "provider_services": {"copilot": "main.default.developer"}, + } + managed = { + "enabled_agents": { + "copilot": {"model_config": {"unity_catalog_location": "main.models"}} + } + } + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s, **_k: (managed, False)) + monkeypatch.setattr( + cli_mod, + "check_gateway_endpoint", + lambda *_a: pytest.fail("managed model sources do not require global models"), + ) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda *_a: None) + configured: list[dict] = [] + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda resolved, tools, **kwargs: configured.append(kwargs) or resolved, + ) + + assert cli_mod.configure_workspace_command(workspaces=[("https://w.com", None)]) == 0 + + assert configured[0]["parent_schemas"] == {"copilot": "main.models"} + + def test_single_copilot_agent_passes_managed_parent_directly(self, monkeypatch): + state = { + **MINIMAL_STATE, + "provider_services": {"copilot": "main.default.developer"}, + } + managed = { + "enabled_agents": { + "copilot": {"model_config": {"unity_catalog_location": "main.models"}} + } + } + monkeypatch.setattr(cli_mod, "install_databricks_cli", MagicMock()) + monkeypatch.setattr(cli_mod, "install_tool_binary", MagicMock()) + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + refresh = MagicMock(return_value=(managed, False)) + monkeypatch.setattr(cli_mod, "refresh_managed_config", refresh) + configure = MagicMock(return_value=state) + monkeypatch.setattr(cli_mod, "configure_single_tool", configure) + install_ai_tools = MagicMock() + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", install_ai_tools) + + result = runner.invoke( + app, ["configure", "--agent", "copilot", "--workspace", "https://w.com"] + ) + + assert result.exit_code == 0, result.output + assert "(Provider: Databricks)" in _strip_ansi(result.output) + refresh.assert_called_once_with(state, force_refresh=False) + configure.assert_called_once_with("copilot", state, parent_schema="main.models") + install_ai_tools.assert_called_once_with(["copilot"], state, force_refresh=True) + + @pytest.mark.parametrize("apply_managed", [True, False]) + def test_single_copilot_agent_refreshes_ai_tools_since_it_only_reads_the_policy( + self, monkeypatch, apply_managed + ): + state = dict(MINIMAL_STATE) + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda *a, **k: (None, False)) + monkeypatch.setattr(cli_mod, "configure_single_tool", lambda *a, **k: state) + install_ai_tools = MagicMock() + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", install_ai_tools) + + cli_mod._setup_single_agent( + [("https://w.com", None)], "copilot", apply_managed=apply_managed + ) + + install_ai_tools.assert_called_once_with(["copilot"], state, force_refresh=True) + def test_managed_config_fails_when_no_enabled_agent_is_available(self, monkeypatch): import ucode.cli as cli_mod