[AIGTWY-4564] Add Codex parent schema discovery

src/ucode/agents/__init__.py

@@ -426,10 +426,13 @@ def configure_tool(
     route_root_model: str | None = None,
     custom_model: str | None = None,
     coding_agent_config_defaults: dict[str, str] | None = None,
+    parent_schema: str | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
-        result = codex.write_tool_config(state, model, provider=provider)
+        result = codex.write_tool_config(
+            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.
@@ -444,6 +447,7 @@ def configure_tool(
             route_root_model=route_root_model,
             custom_model=custom_model,
             coding_agent_config_defaults=coding_agent_config_defaults,
+            parent_schema=parent_schema,
         )
     else:
         # Every tool in this branch needs a model — including gemini under a provider,

src/ucode/agents/claude.py

@@ -24,7 +24,11 @@
     read_json_safe,
     write_json_file,
 )
-from ucode.constants import LOOPBACK_HOST
+from ucode.constants import (
+    LOOPBACK_HOST,
+    MODEL_PROVIDER_SERVICE_HEADER,
+    MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+)
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command
 from ucode.databricks import (
     build_auth_shell_command,
@@ -44,6 +48,7 @@
     reconcile_managed_file,
     revert_managed_file,
 )
+from ucode.model_service_headers import model_service_routing_headers
 from ucode.smart_routing import v2 as smart_routing_v2
 from ucode.smart_routing.claude_hooks import (
     remove_smart_routing_hooks,
@@ -166,7 +171,8 @@ def _resolve_web_search_model(state: dict) -> str | None:
     {
         "x-databricks-use-coding-agent-mode",
         "user-agent",
-        "databricks-model-provider-service",
+        MODEL_PROVIDER_SERVICE_HEADER,
+        MODEL_SERVICE_PARENT_SCHEMA_HEADER,
     }
 )
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
@@ -321,6 +327,7 @@ def render_overlay(
     relayed_base_url: str | None = None,
     route_root_model: str | None = None,
     custom_model: str | None = None,
+    parent_schema: str | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Claude settings.json.
 
@@ -356,8 +363,10 @@ def render_overlay(
         "x-databricks-use-coding-agent-mode: true",
         f"User-Agent: ucode/{ucode_version()} claude/{agent_version('claude')}",
     ]
-    if provider:
-        header_lines.append(f"Databricks-Model-Provider-Service: {provider}")
+    header_lines.extend(
+        f"{name}: {value}"
+        for name, value in model_service_routing_headers(provider, parent_schema).items()
+    )
     # Relayed: the X-Databricks-AI-Gateway-Token swap header is added per request
     # by the refresh proxy, not here — a static value would go stale mid-session.
     custom_headers = "\n".join(header_lines)
@@ -572,6 +581,7 @@ def write_tool_config(
     route_root_model: str | None = None,
     custom_model: str | None = None,
     coding_agent_config_defaults: dict[str, str] | None = None,
+    parent_schema: str | None = None,
 ) -> dict:
     backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH)
     web_search_model = _resolve_web_search_model(state)
@@ -593,6 +603,7 @@ def write_tool_config(
         relayed_base_url=relayed_base_url,
         route_root_model=route_root_model,
         custom_model=custom_model,
+        parent_schema=parent_schema,
     )
     tracing_env_vars = tracing_env(state, "claude")
     stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None

src/ucode/agents/codex.py

@@ -17,9 +17,14 @@
     ToolSpec,
     backup_existing_file,
     deep_merge_dict,
+    prune_key_paths,
     read_toml_safe,
     write_toml_file,
 )
+from ucode.constants import (
+    MODEL_PROVIDER_SERVICE_HEADER,
+    MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+)
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_token_argv
 from ucode.databricks import (
     build_auth_token_argv,
@@ -39,6 +44,7 @@
     reconcile_managed_file,
     revert_managed_file,
 )
+from ucode.model_service_headers import model_service_routing_headers
 from ucode.smart_routing import v2 as smart_routing_v2
 from ucode.smart_routing.codex_hooks import (
     remove_smart_routing_hooks,
@@ -59,6 +65,15 @@
 LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml"
 LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml"
 CODEX_MODEL_PROVIDER_NAME = "ucode-databricks"
+_MODEL_SERVICE_ROUTING_KEY_PATHS = [
+    ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers", MODEL_PROVIDER_SERVICE_HEADER],
+    [
+        "model_providers",
+        CODEX_MODEL_PROVIDER_NAME,
+        "http_headers",
+        MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+    ],
+]
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"
 MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0)
@@ -147,6 +162,7 @@ def _provider_block(
     databricks_profile: str | None,
     use_pat: bool = False,
     provider: str | None = None,
+    parent_schema: str | None = None,
     custom_oauth: CustomOAuthConfig | None = None,
 ) -> dict:
     if custom_oauth:
@@ -157,10 +173,7 @@ def _provider_block(
     http_headers = {
         "User-Agent": f"ucode/{ucode_version()} codex/{agent_version('codex')}",
     }
-    # Route to an external Model Provider Service; the gateway selects the
-    # provider from this header on every request.
-    if provider:
-        http_headers["Databricks-Model-Provider-Service"] = provider
+    http_headers.update(model_service_routing_headers(provider, parent_schema))
     return {
         "name": "Databricks AI Gateway",
         "base_url": base_url,
@@ -183,14 +196,20 @@ def render_overlay(
     databricks_profile: str | None = None,
     use_pat: bool = False,
     provider: str | None = None,
+    parent_schema: str | None = None,
     custom_oauth: CustomOAuthConfig | None = None,
 ) -> dict:
     overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME}
     if model:
         overlay["model"] = model
     overlay["model_providers"] = {
         CODEX_MODEL_PROVIDER_NAME: _provider_block(
-            workspace, databricks_profile, use_pat, provider, custom_oauth
+            workspace,
+            databricks_profile,
+            use_pat=use_pat,
+            provider=provider,
+            parent_schema=parent_schema,
+            custom_oauth=custom_oauth,
         ),
     }
     return overlay
@@ -202,6 +221,7 @@ def render_legacy_overlay(
     databricks_profile: str | None = None,
     use_pat: bool = False,
     provider: str | None = None,
+    parent_schema: str | None = None,
     custom_oauth: CustomOAuthConfig | None = None,
 ) -> dict:
     """Overlay for Codex CLI < 0.134.0, which only reads `~/.codex/config.toml`.
@@ -217,7 +237,12 @@ def render_legacy_overlay(
         "profiles": {CODEX_PROFILE_NAME: profile_block},
         "model_providers": {
             CODEX_MODEL_PROVIDER_NAME: _provider_block(
-                workspace, databricks_profile, use_pat, provider, custom_oauth
+                workspace,
+                databricks_profile,
+                use_pat=use_pat,
+                provider=provider,
+                parent_schema=parent_schema,
+                custom_oauth=custom_oauth,
             ),
         },
     }
@@ -311,7 +336,12 @@ def revert_legacy_shared_config() -> bool:
     return _strip_legacy_ucode_entries(_legacy_config_path())
 
 
-def write_tool_config(state: dict, model: str | None = None, provider: str | None = None) -> dict:
+def write_tool_config(
+    state: dict,
+    model: str | None = None,
+    provider: str | None = None,
+    parent_schema: str | None = None,
+) -> dict:
     workspace = state["workspace"]
     # Leave model selection to Codex. The gateway still receives the configured
     # provider and authentication settings, while Codex uses its own default.
@@ -332,9 +362,11 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
             databricks_profile,
             use_pat=bool(state.get("use_pat")),
             provider=provider,
+            parent_schema=parent_schema,
             custom_oauth=state.get("custom_oauth"),
         )
         doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH)
+        prune_key_paths(doc, _MODEL_SERVICE_ROUTING_KEY_PATHS)
         deep_merge_dict(doc, overlay)
         # deep_merge can't drop keys, so clear model preferences from an earlier run.
         profiles = doc.get("profiles")
@@ -358,10 +390,12 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
         databricks_profile,
         use_pat=bool(state.get("use_pat")),
         provider=provider,
+        parent_schema=parent_schema,
         custom_oauth=state.get("custom_oauth"),
     )
 
     def compose(base: dict) -> dict:
+        prune_key_paths(base, _MODEL_SERVICE_ROUTING_KEY_PATHS)
         deep_merge_dict(base, copy.deepcopy(overlay))
         # deep_merge can't drop keys, so clear model preferences from an earlier run.
         if chosen_model is None:

src/ucode/cli.py

@@ -128,6 +128,7 @@
     set_current_workspace,
     set_provider_service,
 )
+from ucode.string_utils import is_valid_catalog_schema
 from ucode.tracing import configure_tracing_command
 from ucode.ui import (
     console,
@@ -1965,10 +1966,13 @@ def _launch_tool(
     managed: dict | None = None,
     recommendation: dict | None = None,
     model: str | None = None,
+    parent_schema: str | None = None,
     custom_oauth: CustomOAuthConfig | None = None,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
+        if parent_schema is not None and not is_valid_catalog_schema(parent_schema):
+            raise RuntimeError("--parent must be `<catalog>.<schema>`.")
         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
@@ -2063,6 +2067,7 @@ def _launch_tool(
                 )
             if managed_provider:
                 provider = managed_provider
+        effective_parent_schema = None if provider else parent_schema
         # 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:
@@ -2157,6 +2162,7 @@ def _launch_tool(
             # Claude's explicit model is launch-scoped and is passed through LaunchOptions below.
             custom_model=None,
             coding_agent_config_defaults=coding_agent_config_defaults,
+            parent_schema=effective_parent_schema,
         )
         # Relayed = a Claude subscription: forward --model to Claude Code's own flag, like `-- --model X`.
         if tool == "claude" and provider and relayed and model and not forwarded_model:
@@ -2420,6 +2426,13 @@ def codex_cmd(
             "before any `--` separator.",
         ),
     ] = None,
+    parent: Annotated[
+        str | None,
+        typer.Option(
+            "--parent",
+            help="Discover model services in `<catalog>.<schema>`.",
+        ),
+    ] = None,
     refresh: Annotated[
         bool,
         typer.Option(
@@ -2480,6 +2493,7 @@ def codex_cmd(
             refresh=refresh,
             skip_preflight=skip_preflight,
             workspace_url=workspace,
+            parent_schema=parent,
             custom_oauth=custom_oauth,
         )
 
@@ -2500,6 +2514,13 @@ def claude_cmd(
             "before any `--` separator.",
         ),
     ] = None,
+    parent: Annotated[
+        str | None,
+        typer.Option(
+            "--parent",
+            help="Discover model services in `<catalog>.<schema>`.",
+        ),
+    ] = None,
     model: Annotated[
         str | None,
         typer.Option(
@@ -2571,7 +2592,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:
+    if enable_model_discovery or (parent 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):
         _launch_tool(
@@ -2582,6 +2603,7 @@ def claude_cmd(
             refresh=refresh,
             skip_preflight=skip_preflight,
             workspace_url=workspace,
+            parent_schema=parent,
             custom_oauth=custom_oauth,
         )
 

src/ucode/constants.py

@@ -2,3 +2,6 @@
 
 LOCALHOST = "localhost"
 LOOPBACK_HOST = "127.0.0.1"
+
+MODEL_PROVIDER_SERVICE_HEADER = "databricks-model-provider-service"
+MODEL_SERVICE_PARENT_SCHEMA_HEADER = "databricks-model-service-parent-schema"

src/ucode/model_service_headers.py

@@ -0,0 +1,17 @@
+"""Headers used to route model service requests."""
+
+from __future__ import annotations
+
+from ucode.constants import MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER
+
+
+def model_service_routing_headers(
+    provider: str | None,
+    parent_schema: str | None,
+) -> dict[str, str]:
+    # A provider selects one MPS; a parent discovers Model Services.
+    if provider:
+        return {MODEL_PROVIDER_SERVICE_HEADER: provider}
+    if parent_schema:
+        return {MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema}
+    return {}

src/ucode/string_utils.py

@@ -0,0 +1,14 @@
+"""Shared string validation helpers."""
+
+from __future__ import annotations
+
+
+def is_valid_catalog_schema(value: str) -> bool:
+    """Return whether value is a safe ``<catalog>.<schema>`` reference."""
+    parts = value.split(".")
+    return len(parts) == 2 and all(
+        part
+        and part.isprintable()
+        and not any(character.isspace() or character == "/" for character in part)
+        for part in parts
+    )