Route managed Codex model discovery

scripts/run_integration.py

@@ -115,6 +115,11 @@ def arguments():
         default="gpt-5-nano",
         help="Model allowed by the OpenAI MPS selected in the configure CUJ.",
     )
+    parser.add_argument(
+        "--parent-schema",
+        default="main.ucode",
+        help="Schema containing the dedicated model-discovery Model Services.",
+    )
     parser.add_argument("--python", default=sys.executable, help="Python 3.12+ path or uv version.")
     parser.add_argument("--dependency", action="append", default=[], metavar="PACKAGE==VERSION")
     parser.add_argument("--constraints", type=Path, help="Replay a previous dependencies.txt.")
@@ -290,6 +295,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             "claude_relayed_provider": args.claude_relayed_provider,
             "codex_provider": args.codex_provider,
             "codex_provider_model": args.codex_provider_model,
+            "parent_schema": args.parent_schema,
             "dependencies": args.dependency,
             "workspace": args.workspace,
         },
@@ -530,6 +536,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 "UG_INTEGRATION_CLAUDE_OAUTH_TOKEN": oauth_token,
                 "UG_INTEGRATION_CODEX_PROVIDER": args.codex_provider,
                 "UG_INTEGRATION_CODEX_PROVIDER_MODEL": args.codex_provider_model,
+                "UG_INTEGRATION_PARENT_SCHEMA": args.parent_schema,
                 "UCODE_TEST_WORKSPACE": args.workspace or "",
                 "DATABRICKS_BEARER": bearer,
             }

src/ucode/agents/__init__.py

@@ -392,9 +392,9 @@ def configure_tool(
             state, model, provider=provider, parent_schema=parent_schema
         )
     elif tool == "claude":
-        # A Model Provider Service routes by header and pins no Databricks
-        # model, so the usual "model required" guard doesn't apply to claude.
-        if not model and not provider:
+        # A Model Provider Service or parent schema routes by header and discovers models natively,
+        # so the usual "model required" guard doesn't apply to either Claude source.
+        if not model and not provider and not parent_schema:
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         result = claude.write_tool_config(
             state,
@@ -497,12 +497,12 @@ def _availability_failure_detail(tool: str, state: dict) -> str:
     return " (" + "; ".join(parts) + ")"
 
 
-def configure_single_tool(tool: str, state: dict) -> dict:
+def configure_single_tool(tool: str, state: dict, *, parent_schema: str | None = None) -> dict:
     """Check availability, configure, and persist state for one tool only."""
-    provider = get_provider_service(state, tool)
-    # A Model Provider Service routes through the same gateway and pins no
-    # Databricks model, so the per-tool model availability check doesn't apply.
-    if not provider:
+    provider = None if parent_schema else get_provider_service(state, tool)
+    # A Model Provider Service or parent schema routes through the same gateway and pins no
+    # globally discovered Databricks model, so the availability check doesn't apply.
+    if not provider and not parent_schema:
         with spinner(f"Checking {TOOL_SPECS[tool]['display']} availability..."):
             ok = check_gateway_endpoint(state, tool)
         if not ok:
@@ -511,15 +511,19 @@ def configure_single_tool(tool: str, state: dict) -> dict:
                 f"{TOOL_SPECS[tool]['display']} is not available on this workspace.{detail}"
             )
     with managed_write_batch(_managed_settings_displays([tool])):
-        state = _configure_one(tool, state, provider)
+        state = _configure_one(tool, state, provider, parent_schema=parent_schema)
     available_tools = list(set((state.get("available_tools") or []) + [tool]))
     state["available_tools"] = available_tools
     save_state(state)
     return state
 
 
-def _configure_one(tool: str, state: dict, provider: str | None) -> dict:
+def _configure_one(
+    tool: str, state: dict, provider: str | None, *, parent_schema: str | None = None
+) -> dict:
     """Write one tool's config, routing through ``provider`` when set."""
+    if parent_schema:
+        return configure_tool(tool, state, parent_schema=parent_schema)
     if provider:
         if tool == "gemini":
             # Gemini pins a concrete target in the URL, so configure must resolve one now —
@@ -542,7 +546,11 @@ def _configure_one(tool: str, state: dict, provider: str | None) -> dict:
 
 
 def configure_selected_tools(
-    state: dict, tools: list[str], *, install_ai_tools: bool = True
+    state: dict,
+    tools: list[str],
+    *,
+    install_ai_tools: bool = True,
+    parent_schemas: dict[str, str] | None = None,
 ) -> dict:
     """Configure the given tools. Caller is responsible for ensuring each tool
     is available on the workspace.
@@ -553,7 +561,9 @@ def configure_selected_tools(
     """
     with managed_write_batch(_managed_settings_displays(tools)):
         for tool in tools:
-            state = _configure_one(tool, state, get_provider_service(state, tool))
+            parent_schema = (parent_schemas or {}).get(tool)
+            provider = None if parent_schema else get_provider_service(state, tool)
+            state = _configure_one(tool, state, provider, parent_schema=parent_schema)
 
     existing = state.get("available_tools") or []
     state["available_tools"] = sorted(set(existing) | set(tools))

src/ucode/agents/claude.py

@@ -447,7 +447,7 @@ def render_overlay(
     # provider and Claude Code's own canonical model names are sent verbatim —
     # pinning a Databricks model id here would mislabel the picker and isn't
     # routable.
-    elif claude_models and not provider:
+    elif claude_models and not provider and not parent_schema:
         # Picker rows show the raw routable id (e.g. "system.ai.claude-opus-4-8[1m]")
         # so users can see which gateway-routable model is behind each shortcut.
         # We deliberately don't set the `_NAME` companion env vars — the raw id
@@ -486,7 +486,7 @@ def render_overlay(
         overlay["permissions"] = {"deny": ["WebSearch"]}
         keys.append(["permissions", "deny"])
 
-    if static_models and not provider and not relayed:
+    if static_models and not provider and not parent_schema and not relayed:
         overlay["availableModels"] = list(static_models)
         overlay["enforceAvailableModels"] = True
         overlay["modelPicker"] = {
@@ -865,6 +865,14 @@ def write_tool_config(
     # revert would restore that snapshot instead of deleting the file.
     if not is_tool_managed(state, "claude"):
         backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH)
+    previous_keys = ((state.get("managed_configs") or {}).get("claude") or {}).get("keys", [])
+    # Native discovery must not inherit UG's prior static allow-list. Untracked picker values
+    # remain user/admin-owned and are preserved by the merge below.
+    stale_picker_keys = (
+        [key for key in CLAUDE_MANAGED_PICKER_KEYS if [key] in previous_keys]
+        if provider or parent_schema
+        else []
+    )
     web_search_model = _resolve_web_search_model(state)
     # Relayed inference points at a local refresh proxy; its loopback base URL is
     # recorded in state so launch starts the proxy on the matching port.
@@ -889,7 +897,8 @@ def write_tool_config(
     )
     managed_file_keys = list(managed_keys)
     for path in (
-        [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS]
+        [[key] for key in stale_picker_keys]
+        + [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS]
         + [["env", key] for key in CLAUDE_CONDITIONAL_ENV_KEYS]
         + [["env", key] for key in CLAUDE_REMOVED_ENV_KEYS]
         + [["env", key] for key in CLAUDE_OTEL_TRACE_ENV_KEYS]
@@ -953,6 +962,8 @@ def _compose(
                 else:
                     target_env[key] = selected_default_model
         merged = deep_merge_dict(base, overlay_for_merge)
+        for key in stale_picker_keys:
+            merged.pop(key, None)
         overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY]
         merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers(
             existing_custom_headers, overlay_custom_headers
@@ -1013,7 +1024,7 @@ def _compose(
         state,
         lambda base: _compose(
             base,
-            enforce_model_default_hierarchy=provider is None,
+            enforce_model_default_hierarchy=provider is None and parent_schema is None,
             managed_settings_snapshots=managed_snapshots,
         ),
         managed_file_keys,

src/ucode/agents/codex.py

@@ -917,17 +917,22 @@ def launch(
     binary = SPEC["binary"]
     workspace = state.get("workspace")
     launch_provider = state.get("_codex_launch_provider")
-    provider = (
+    transient_provider = (
         launch_provider.strip()
         if isinstance(launch_provider, str) and launch_provider.strip()
-        else get_provider_service(state, "codex")
+        else None
     )
     launch_parent_schema = state.get("_codex_launch_parent_schema")
     parent_schema = (
         launch_parent_schema.strip()
         if isinstance(launch_parent_schema, str) and launch_parent_schema.strip()
         else None
     )
+    # Launch-scoped admin routing wins over persisted developer configuration. A transient provider
+    # is most specific; otherwise a transient UC parent must suppress a saved provider.
+    provider = transient_provider or (
+        None if parent_schema else get_provider_service(state, "codex")
+    )
     if workspace and (provider or parent_schema):
         _reject_managed_model_catalog()
     token = None

src/ucode/cli.py

@@ -91,6 +91,7 @@
     managed_provider_family_models,
     managed_provider_service,
     managed_supplies_models,
+    managed_unity_catalog_location,
     managed_unservable_models,
     recommended_agent,
     resolve_state,
@@ -775,16 +776,25 @@ def configure_workspace_command(
             clear_custom_oauth=custom_oauth is None,
         )
         state = states[0]
-        state = configure_single_tool(tool, state)
-        # No managed refresh precedes this branch, so read fresh here: `ug configure` never decides
-        # from a stale cache.
-        install_databricks_ai_tools_for_agents([tool], state, force_refresh=True)
+        parent_schema = None
+        if tool in ("claude", "codex"):
+            managed, _ = refresh_managed_config(state, force_refresh=True)
+            _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)
+        state = configure_single_tool(tool, state, parent_schema=parent_schema)
+        install_databricks_ai_tools_for_agents(
+            [tool], state, force_refresh=tool not in ("claude", "codex")
+        )
         spec = TOOL_SPECS[tool]
+        provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state)
         console.print(
             Panel(
                 f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]\n"
                 f"[bold]{spec['display']}:[/bold] [green]configured[/green] "
-                f"[dim](Provider: {_provider_summary(tool, state)})[/dim]",
+                f"[dim](Provider: {provider_summary})[/dim]",
                 title="Configuration Complete",
                 style="green",
                 expand=False,
@@ -812,13 +822,25 @@ def configure_workspace_command(
     if managed is not None and managed_tools:
         configured_tools: list[str] = []
         for tool_name in managed_tools:
-            if check_gateway_endpoint(state, tool_name):
+            resolved = resolve_state(managed, state, tool_name)
+            parent_schema = (
+                managed_unity_catalog_location(managed, tool_name)
+                if tool_name in ("claude", "codex")
+                and not managed_provider_service(managed, tool_name)
+                else None
+            )
+            if (
+                get_provider_service(resolved, tool_name)
+                or parent_schema
+                or check_gateway_endpoint(resolved, tool_name)
+            ):
                 if not install_tool_binary(tool_name, strict=False):
                     continue
                 configured = configure_selected_tools(
-                    resolve_state(managed, state, tool_name),
+                    resolved,
                     [tool_name],
                     install_ai_tools=not is_dry_run(),
+                    parent_schemas={tool_name: parent_schema} if parent_schema else None,
                 )
                 # Each iteration resolves from `state` and persists a copy, so carry the
                 # accumulated available_tools forward — otherwise the last agent's save drops
@@ -2337,6 +2359,22 @@ def _launch_tool(
             raise RuntimeError("--model-location must be `<catalog>.<schema>`.")
         # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast.
         _reject_disabled_agent(managed, 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
+            else None
+        )
+        if managed_provider:
+            provider = managed_provider
+            parent_schema = None
+        elif managed_parent_schema:
+            # Managed UC discovery supersedes a developer's persisted provider without
+            # rewriting it; the admin's location exists only for this launch.
+            provider = None
+            parent_schema = managed_parent_schema
+        if tool == "claude" and (managed_provider or managed_parent_schema):
+            os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1"
         # The environment switch remains a developer override; managed config is the workspace
         # policy equivalent and must take effect before launch options are computed.
         managed_smart_routing_enabled = _managed_smart_routing_enabled(managed, tool)
@@ -2353,7 +2391,9 @@ def _launch_tool(
             state["workspace"],
             profile=state.get("profile"),
             tools=[tool],
-            skip_model_discovery=bool(provider) or managed_models_known,
+            skip_model_discovery=(
+                bool(provider) or bool(managed_parent_schema) or managed_models_known
+            ),
             skip_preflight=skip_preflight,
             **configure_kwargs,
         )
@@ -2374,10 +2414,6 @@ def _launch_tool(
                 )
         elif not coding_agent_config_feature_disabled:
             print_note("No managed coding agent config found; using your own settings")
-        if managed is not None:
-            managed_provider = managed_provider_service(managed, tool)
-            if managed_provider:
-                provider = managed_provider
         if provider and parent_schema is not None:
             raise RuntimeError("--provider and --model-location cannot be used together.")
         # Checked after the managed config settles `provider`: an admin-set provider must trip this
@@ -2428,13 +2464,12 @@ 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:
+        if provider or managed_parent_schema:
             # Routing through a Model Provider Service pins no Databricks model;
-            # the agent uses its own canonical model names (header selects the
-            # provider). Skip model resolution, which would otherwise fail when
-            # the workspace has no matching Databricks models.
+            # managed 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
-            if tool == "claude" and (model or provider_models):
+            if provider and tool == "claude" and (model or provider_models):
                 if relayed:
                     # Resolve against a curated allowlist so the forwarded id is one the gateway
                     # allows; an allow_all relay declares none, so forward as-is.
@@ -2445,7 +2480,7 @@ def _launch_tool(
                     )
                 else:
                     route_root_model = resolve_provider_launch_model(model, provider_models or {})
-            if tool == "gemini":
+            if provider and tool == "gemini":
                 # Gemini is the exception: the request still names a concrete model
                 # in the URL, so pin one of the service's targets (--model or default).
                 resolved_model, gemini_error = resolve_gemini_provider_model(state, provider, model)

src/ucode/managed_resolve.py

@@ -190,6 +190,11 @@ def managed_provider_service(managed: dict, tool: str) -> str | None:
     return _str(_agent_model_config(managed, tool).get("model_provider_service"))
 
 
+def managed_unity_catalog_location(managed: dict, tool: str) -> str | None:
+    """Return only the Unity Catalog model location managed config specifies for ``tool``."""
+    return _str(_agent_model_config(managed, tool).get("unity_catalog_location"))
+
+
 def managed_static_models(managed: dict, tool: str) -> list[str] | None:
     """The explicit model allow-list (``model_config.model_services``) the config sets for ``tool``.
 

tests/AGENTS.md

@@ -30,9 +30,10 @@ tests. Keep work scoped to the behavior requested by the user.
    exception is the `managed_fixture` marker: it uses the built-in
    `UCODE_MANAGED_CONFIG_STUB` hook to inject the admin CodingAgentConfig INPUT so the
    real `ug configure` path can be exercised across config shapes the live workspace does
-   not publish. The gateway, agent binaries, ug internals, and ug state stay real; the
-   config fetch/wire contract stays covered by the un-stubbed `managed` tests; and the
-   hook must never be used to disable validation or conceal a failure.
+   not publish and one workspace fetch can be replayed across isolated cases. The gateway,
+   agent binaries, ug internals, and ug state stay real; the config fetch/wire contract stays
+   covered by the un-stubbed `managed` tests; and the hook must never be used to disable
+   validation or conceal a failure.
 5. **Real responses and binaries.** Pin requested ug and agent versions. Never
    substitute a missing binary/service. Reuse explicit e2e workspace/auth settings;
    never pick a developer's Databricks profile automatically.