copilot: run ug copilot on a Unity Catalog model location

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 or Azure OpenAI service. For a
 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 <catalog>.<schema>` 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** (`<catalog>.<schema>`) 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

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"],
@@ -469,8 +471,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:
@@ -591,10 +594,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)
 

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
@@ -65,6 +67,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,
@@ -74,6 +77,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
@@ -116,6 +120,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"
@@ -373,6 +379,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-<hash>-<name>``); 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:
@@ -561,6 +588,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.
 
@@ -607,8 +635,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
@@ -846,14 +878,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(
@@ -867,11 +901,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()
 
@@ -895,7 +933,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:

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,
@@ -52,6 +53,7 @@
 from ucode.constants import (
     AGENT_CLAUDE,
     AGENT_CODEX,
+    AGENT_COPILOT,
     ISAAC_SMART_ROUTER_FEEDBACK_ENV_VAR,
     SMART_ROUTING_ENV_KEYS,
 )
@@ -810,7 +812,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.
@@ -827,7 +829,7 @@ def _setup_single_agent(
     state = states[0]
     parent_schema = None
     managed = None
-    if apply_managed and tool in ("claude", "codex"):
+    if apply_managed and tool in MODEL_LOCATION_TOOLS:
         managed, _ = refresh_managed_config(state, force_refresh=True)
         _reject_disabled_agent(managed, tool)
         if managed is not None:
@@ -848,9 +850,8 @@ 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")
-    )
+    managed_refreshed = apply_managed and tool in MODEL_LOCATION_TOOLS
+    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(
@@ -893,9 +894,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
@@ -913,7 +914,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.
@@ -1056,7 +1057,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
             )
@@ -3291,7 +3292,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
@@ -3333,7 +3334,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:
@@ -3498,16 +3499,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.
@@ -3646,16 +3655,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,
@@ -4114,6 +4126,13 @@ def copilot_cmd(
             "(<catalog>.<schema>.<name>); `/model` then lists that service's models.",
         ),
     ] = None,
+    model_location: Annotated[
+        str | None,
+        typer.Option(
+            "--model-location",
+            help="Discover model services in `<catalog>.<schema>`. Example: main.default",
+        ),
+    ] = None,
     skip_preflight: SkipPreflightOption = False,
     workspace: WorkspaceOption = None,
 ) -> None:
@@ -4124,6 +4143,7 @@ def copilot_cmd(
         provider=provider,
         skip_preflight=skip_preflight,
         workspace_url=workspace,
+        parent_schema=model_location,
     )
 
 

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"