From f914af76d0f6319127b2cbdbd5564622abbeb40e Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Wed, 9 Sep 2026 20:04:31 +0000
Subject: [PATCH 1/9] Add Claude parent schema discovery

---
 src/ucode/agents/__init__.py |  2 ++
 src/ucode/agents/claude.py   |  6 ++++++
 src/ucode/cli.py             | 18 ++++++++++++++++++
 tests/test_agent_claude.py   | 28 ++++++++++++++++++++++++++++
 tests/test_cli.py            | 13 +++++++++++++
 5 files changed, 67 insertions(+)

diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py
index 8ea3e20d..fd630147 100644
--- a/src/ucode/agents/__init__.py
+++ b/src/ucode/agents/__init__.py
@@ -426,6 +426,7 @@ 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":
@@ -444,6 +445,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,
diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 331f6690..eddf3446 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -166,6 +166,7 @@ def _resolve_web_search_model(state: dict) -> str | None:
         "x-databricks-use-coding-agent-mode",
         "user-agent",
         "databricks-model-provider-service",
+        "databricks-model-service-parent-schema",
     }
 )
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
@@ -319,6 +320,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,6 +358,8 @@ def render_overlay(
     ]
     if provider:
         header_lines.append(f"Databricks-Model-Provider-Service: {provider}")
+    elif parent_schema:
+        header_lines.append(f"Databricks-Model-Service-Parent-Schema: {parent_schema}")
     # 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)
@@ -567,6 +571,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)
@@ -587,6 +592,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
diff --git a/src/ucode/cli.py b/src/ucode/cli.py
index 9857d78a..7ea0e35c 100644
--- a/src/ucode/cli.py
+++ b/src/ucode/cli.py
@@ -1878,9 +1878,17 @@ def _launch_tool(
     managed: dict | None = None,
     recommendation: dict | None = None,
     model: str | None = None,
+    parent_schema: str | None = None,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
+        if parent_schema is not None:
+            parts = parent_schema.split(".")
+            if len(parts) != 2 or any(
+                not part or any(ch in " /" or ord(ch) < 32 or ord(ch) == 127 for ch in part)
+                for part in parts
+            ):
+                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
@@ -1970,6 +1978,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:
@@ -2064,6 +2073,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:
@@ -2387,6 +2397,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(
@@ -2450,6 +2467,7 @@ def claude_cmd(
             refresh=refresh,
             skip_preflight=skip_preflight,
             workspace_url=workspace,
+            parent_schema=parent,
         )
 
 
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index ac5ea339..658e6dfe 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -343,6 +343,25 @@ 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):
+        overlay, _ = claude.render_overlay(WS, "s4", parent_schema="main.default")
+        assert (
+            "Databricks-Model-Service-Parent-Schema: main.default"
+            in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
+        )
+
+    def test_provider_suppresses_discovery_header(self):
+        overlay, _ = claude.render_overlay(
+            WS,
+            "s4",
+            provider="main.default.anthropic",
+            parent_schema="main.default",
+        )
+        assert (
+            "Databricks-Model-Service-Parent-Schema"
+            not in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
+        )
+
     def test_bedrock_provider_pins_model_ids(self):
         provider_models = {
             "opus": "global.anthropic.claude-opus-4-8",
@@ -431,6 +450,15 @@ def test_headers_newline_delimited(self, monkeypatch):
 
 
 class TestMergeAnthropicCustomHeaders:
+    def test_removes_stale_parent_header(self):
+        existing = "X-User: keep\nDatabricks-Model-Service-Parent-Schema: main.default"
+        managed = "x-databricks-use-coding-agent-mode: true"
+
+        merged = claude._merge_anthropic_custom_headers(existing, managed)
+
+        assert "X-User: keep" in merged
+        assert "Databricks-Model-Service-Parent-Schema" not in merged
+
     def test_merges_existing_settings_with_ucode_managed_headers(self):
         headers_from_existing_settings = "\n".join(
             [
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 18637d60..1a165920 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -583,6 +583,19 @@ 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):
+        with patch("ucode.cli._launch_tool") as mock_launch:
+            result = runner.invoke(app, ["claude", "--parent", "main.default"])
+
+        assert result.exit_code == 0, result.output
+        assert mock_launch.call_args.kwargs["parent_schema"] == "main.default"
+
+    def test_invalid_parent_is_rejected(self):
+        result = runner.invoke(app, ["claude", "--parent", "main"])
+
+        assert result.exit_code == 1
+        assert "<catalog>.<schema>" in result.output
+
     def test_claude_enable_model_discovery_is_hidden_from_help(self):
         result = runner.invoke(app, ["claude", "--help"])
 

From 551dc3bea9fb0622867f0edcbea4fb683284136b Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 04:42:32 +0000
Subject: [PATCH 2/9] Address parent schema review feedback

---
 src/ucode/agents/claude.py          | 14 +++++++-------
 src/ucode/cli.py                    | 12 ++++--------
 src/ucode/constants.py              |  7 +++++++
 src/ucode/model_service_headers.py  | 17 +++++++++++++++++
 src/ucode/string_utils.py           | 14 ++++++++++++++
 tests/test_agents_init.py           | 18 ++++++++++++++++++
 tests/test_cli.py                   |  7 +------
 tests/test_model_service_headers.py | 22 ++++++++++++++++++++++
 tests/test_string_utils.py          | 28 ++++++++++++++++++++++++++++
 9 files changed, 118 insertions(+), 21 deletions(-)
 create mode 100644 src/ucode/model_service_headers.py
 create mode 100644 src/ucode/string_utils.py
 create mode 100644 tests/test_model_service_headers.py
 create mode 100644 tests/test_string_utils.py

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index f0f45225..2157b5e0 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -24,7 +24,7 @@
     read_json_safe,
     write_json_file,
 )
-from ucode.constants import LOOPBACK_HOST
+from ucode.constants import LOOPBACK_HOST, MODEL_SERVICE_ROUTING_HEADER_NAMES
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command
 from ucode.databricks import (
     build_auth_shell_command,
@@ -44,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.claude_hooks import (
     remove_smart_routing_hooks,
@@ -166,8 +167,7 @@ def _resolve_web_search_model(state: dict) -> str | None:
     {
         "x-databricks-use-coding-agent-mode",
         "user-agent",
-        "databricks-model-provider-service",
-        "databricks-model-service-parent-schema",
+        *(name.lower() for name in MODEL_SERVICE_ROUTING_HEADER_NAMES),
     }
 )
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
@@ -358,10 +358,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}")
-    elif parent_schema:
-        header_lines.append(f"Databricks-Model-Service-Parent-Schema: {parent_schema}")
+    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)
diff --git a/src/ucode/cli.py b/src/ucode/cli.py
index a2bdc759..34260ef0 100644
--- a/src/ucode/cli.py
+++ b/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,
@@ -1970,13 +1971,8 @@ def _launch_tool(
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
-        if parent_schema is not None:
-            parts = parent_schema.split(".")
-            if len(parts) != 2 or any(
-                not part or any(ch in " /" or ord(ch) < 32 or ord(ch) == 127 for ch in part)
-                for part in parts
-            ):
-                raise RuntimeError("--parent must be `<catalog>.<schema>`.")
+        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
@@ -2588,7 +2584,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(
diff --git a/src/ucode/constants.py b/src/ucode/constants.py
index f1664b17..2b21de60 100644
--- a/src/ucode/constants.py
+++ b/src/ucode/constants.py
@@ -2,3 +2,10 @@
 
 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"
+MODEL_SERVICE_ROUTING_HEADER_NAMES = (
+    MODEL_PROVIDER_SERVICE_HEADER,
+    MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+)
diff --git a/src/ucode/model_service_headers.py b/src/ucode/model_service_headers.py
new file mode 100644
index 00000000..f247bc87
--- /dev/null
+++ b/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 selects Model Services, so the modes are exclusive.
+    if provider:
+        return {MODEL_PROVIDER_SERVICE_HEADER: provider}
+    if parent_schema:
+        return {MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema}
+    return {}
diff --git a/src/ucode/string_utils.py b/src/ucode/string_utils.py
new file mode 100644
index 00000000..7d518769
--- /dev/null
+++ b/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
+    )
diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py
index 9ff7d5ee..640e9c19 100644
--- a/tests/test_agents_init.py
+++ b/tests/test_agents_init.py
@@ -387,6 +387,24 @@ def test_bedrock_returns_pinned_models(self, monkeypatch):
             "opus": "global.anthropic.claude-opus-4-8",
         }
 
+    def test_bedrock_ignores_gpt_targets(self, monkeypatch):
+        service = {
+            "provider_type": "amazon_bedrock",
+            "targets": [
+                "global.anthropic.claude-opus-4-8",
+                "openai.gpt-oss-120b-1:0",
+            ],
+        }
+        self._patch(monkeypatch, service, None)
+
+        models, error, relayed = agents_mod.resolve_provider_models(
+            "claude", self._STATE, "main.b.mixed"
+        )
+
+        assert error is None
+        assert models == {"opus": "global.anthropic.claude-opus-4-8"}
+        assert relayed is False
+
     def test_invalid_provider_returns_error(self, monkeypatch):
         self._patch(monkeypatch, None, "boom")
         models, error, relayed = agents_mod.resolve_provider_models(
diff --git a/tests/test_cli.py b/tests/test_cli.py
index ce906794..25dfcff8 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -589,12 +589,7 @@ def test_claude_parent_is_forwarded(self):
 
         assert result.exit_code == 0, result.output
         assert mock_launch.call_args.kwargs["parent_schema"] == "main.default"
-
-    def test_invalid_parent_is_rejected(self):
-        result = runner.invoke(app, ["claude", "--parent", "main"])
-
-        assert result.exit_code == 1
-        assert "<catalog>.<schema>" in result.output
+        assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1"
 
     def test_claude_enable_model_discovery_is_hidden_from_help(self):
         result = runner.invoke(app, ["claude", "--help"])
diff --git a/tests/test_model_service_headers.py b/tests/test_model_service_headers.py
new file mode 100644
index 00000000..e5c2ee87
--- /dev/null
+++ b/tests/test_model_service_headers.py
@@ -0,0 +1,22 @@
+"""Tests for model service routing headers."""
+
+from ucode.constants import MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER
+from ucode.model_service_headers import model_service_routing_headers
+
+
+def test_provider_header():
+    assert model_service_routing_headers("main.default.provider", None) == {
+        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
+    }
+
+
+def test_parent_schema_header():
+    assert model_service_routing_headers(None, "main.default") == {
+        MODEL_SERVICE_PARENT_SCHEMA_HEADER: "main.default"
+    }
+
+
+def test_provider_takes_precedence():
+    assert model_service_routing_headers("main.default.provider", "main.default") == {
+        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
+    }
diff --git a/tests/test_string_utils.py b/tests/test_string_utils.py
new file mode 100644
index 00000000..9c369df8
--- /dev/null
+++ b/tests/test_string_utils.py
@@ -0,0 +1,28 @@
+"""Tests for string validation helpers."""
+
+import pytest
+
+from ucode.string_utils import is_valid_catalog_schema
+
+
+@pytest.mark.parametrize("value", ["system.ai", "main.default", "my-catalog.my_schema"])
+def test_catalog_schema(value):
+    assert is_valid_catalog_schema(value)
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "",
+        "main",
+        "main.default.extra",
+        ".default",
+        "main.",
+        "main/development.models",
+        "main dev.models",
+        "main.\tmodels",
+        "main.\x7fmodels",
+    ],
+)
+def test_invalid_catalog_schema(value):
+    assert not is_valid_catalog_schema(value)

From 2a50e7aef840e5be6dd4177514d5cfcaafc95fe8 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 15:48:12 +0000
Subject: [PATCH 3/9] Inline model routing header selection

---
 src/ucode/agents/claude.py          | 17 +++++++++++------
 src/ucode/model_service_headers.py  | 17 -----------------
 tests/test_model_service_headers.py | 22 ----------------------
 3 files changed, 11 insertions(+), 45 deletions(-)
 delete mode 100644 src/ucode/model_service_headers.py
 delete mode 100644 tests/test_model_service_headers.py

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 2157b5e0..22d77e61 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -24,7 +24,12 @@
     read_json_safe,
     write_json_file,
 )
-from ucode.constants import LOOPBACK_HOST, MODEL_SERVICE_ROUTING_HEADER_NAMES
+from ucode.constants import (
+    LOOPBACK_HOST,
+    MODEL_PROVIDER_SERVICE_HEADER,
+    MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+    MODEL_SERVICE_ROUTING_HEADER_NAMES,
+)
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command
 from ucode.databricks import (
     build_auth_shell_command,
@@ -44,7 +49,6 @@
     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,
@@ -358,10 +362,11 @@ def render_overlay(
         "x-databricks-use-coding-agent-mode: true",
         f"User-Agent: ucode/{ucode_version()} claude/{agent_version('claude')}",
     ]
-    header_lines.extend(
-        f"{name}: {value}"
-        for name, value in model_service_routing_headers(provider, parent_schema).items()
-    )
+    # A provider selects one MPS; a parent discovers Model Services, so they are exclusive.
+    if provider:
+        header_lines.append(f"{MODEL_PROVIDER_SERVICE_HEADER}: {provider}")
+    elif parent_schema:
+        header_lines.append(f"{MODEL_SERVICE_PARENT_SCHEMA_HEADER}: {parent_schema}")
     # 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)
diff --git a/src/ucode/model_service_headers.py b/src/ucode/model_service_headers.py
deleted file mode 100644
index f247bc87..00000000
--- a/src/ucode/model_service_headers.py
+++ /dev/null
@@ -1,17 +0,0 @@
-"""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 selects Model Services, so the modes are exclusive.
-    if provider:
-        return {MODEL_PROVIDER_SERVICE_HEADER: provider}
-    if parent_schema:
-        return {MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema}
-    return {}
diff --git a/tests/test_model_service_headers.py b/tests/test_model_service_headers.py
deleted file mode 100644
index e5c2ee87..00000000
--- a/tests/test_model_service_headers.py
+++ /dev/null
@@ -1,22 +0,0 @@
-"""Tests for model service routing headers."""
-
-from ucode.constants import MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER
-from ucode.model_service_headers import model_service_routing_headers
-
-
-def test_provider_header():
-    assert model_service_routing_headers("main.default.provider", None) == {
-        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
-    }
-
-
-def test_parent_schema_header():
-    assert model_service_routing_headers(None, "main.default") == {
-        MODEL_SERVICE_PARENT_SCHEMA_HEADER: "main.default"
-    }
-
-
-def test_provider_takes_precedence():
-    assert model_service_routing_headers("main.default.provider", "main.default") == {
-        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
-    }

From e308bea35ef17ad1c1c6b915f907e6bcec858aa4 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 15:52:08 +0000
Subject: [PATCH 4/9] Simplify model routing header constants

---
 src/ucode/agents/claude.py |  4 ++--
 src/ucode/constants.py     |  8 ++------
 tests/test_agent_claude.py | 14 +++++++-------
 3 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 22d77e61..e619a328 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -28,7 +28,6 @@
     LOOPBACK_HOST,
     MODEL_PROVIDER_SERVICE_HEADER,
     MODEL_SERVICE_PARENT_SCHEMA_HEADER,
-    MODEL_SERVICE_ROUTING_HEADER_NAMES,
 )
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command
 from ucode.databricks import (
@@ -171,7 +170,8 @@ def _resolve_web_search_model(state: dict) -> str | None:
     {
         "x-databricks-use-coding-agent-mode",
         "user-agent",
-        *(name.lower() for name in MODEL_SERVICE_ROUTING_HEADER_NAMES),
+        MODEL_PROVIDER_SERVICE_HEADER,
+        MODEL_SERVICE_PARENT_SCHEMA_HEADER,
     }
 )
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
diff --git a/src/ucode/constants.py b/src/ucode/constants.py
index 2b21de60..0c7cd80c 100644
--- a/src/ucode/constants.py
+++ b/src/ucode/constants.py
@@ -3,9 +3,5 @@
 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"
-MODEL_SERVICE_ROUTING_HEADER_NAMES = (
-    MODEL_PROVIDER_SERVICE_HEADER,
-    MODEL_SERVICE_PARENT_SCHEMA_HEADER,
-)
+MODEL_PROVIDER_SERVICE_HEADER = "databricks-model-provider-service"
+MODEL_SERVICE_PARENT_SCHEMA_HEADER = "databricks-model-service-parent-schema"
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index 5b763fa2..2a96aa81 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -288,7 +288,7 @@ def test_relayed_sends_mps_header_but_not_swap_token(self):
             relayed_base_url="http://127.0.0.1:9",
         )
         headers = overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
-        assert "Databricks-Model-Provider-Service: c.s.mps" in headers
+        assert "databricks-model-provider-service: c.s.mps" in headers
         assert "X-Databricks-AI-Gateway-Token" not in headers
 
     def test_model_overrides_when_all_provided(self):
@@ -349,7 +349,7 @@ def test_fable_not_pinned_under_provider(self):
     def test_provider_adds_routing_header(self):
         overlay, _ = claude.render_overlay(WS, "s4", provider="main.aarushi.aarushi-claude")
         assert (
-            "Databricks-Model-Provider-Service: main.aarushi.aarushi-claude"
+            "databricks-model-provider-service: main.aarushi.aarushi-claude"
             in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
         )
 
@@ -369,12 +369,12 @@ def test_provider_skips_model_pinning(self):
 
     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"]
+        assert "databricks-model-provider-service" not in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
 
     def test_parent_adds_discovery_header(self):
         overlay, _ = claude.render_overlay(WS, "s4", parent_schema="main.default")
         assert (
-            "Databricks-Model-Service-Parent-Schema: main.default"
+            "databricks-model-service-parent-schema: main.default"
             in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
         )
 
@@ -386,7 +386,7 @@ def test_provider_suppresses_discovery_header(self):
             parent_schema="main.default",
         )
         assert (
-            "Databricks-Model-Service-Parent-Schema"
+            "databricks-model-service-parent-schema"
             not in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
         )
 
@@ -409,7 +409,7 @@ def test_bedrock_provider_pins_model_ids(self):
         # Bedrock ids are pinned verbatim — no `[1m]` suffix mangling.
         assert "[1m]" not in env["ANTHROPIC_DEFAULT_OPUS_MODEL"]
         assert (
-            "Databricks-Model-Provider-Service: main.bob.bedrock-svc"
+            "databricks-model-provider-service: main.bob.bedrock-svc"
             in env["ANTHROPIC_CUSTOM_HEADERS"]
         )
 
@@ -426,7 +426,7 @@ def test_non_relayed_provider_pins_tier_via_anthropic_model(self):
         env = overlay["env"]
         assert env["ANTHROPIC_MODEL"] == "claude-haiku-4-5"
         assert (
-            "Databricks-Model-Provider-Service: main.mcao.anthropic-mps"
+            "databricks-model-provider-service: main.mcao.anthropic-mps"
             in (env["ANTHROPIC_CUSTOM_HEADERS"])
         )
         assert "apiKeyHelper" in overlay

From 74a278d18348046079ced9d3c34d4e0276725c90 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 17:03:00 +0000
Subject: [PATCH 5/9] Share model service routing headers

---
 src/ucode/agents/claude.py          | 10 +++++-----
 src/ucode/model_service_headers.py  | 17 +++++++++++++++++
 tests/test_model_service_headers.py | 22 ++++++++++++++++++++++
 3 files changed, 44 insertions(+), 5 deletions(-)
 create mode 100644 src/ucode/model_service_headers.py
 create mode 100644 tests/test_model_service_headers.py

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index e619a328..c83b161e 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -48,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,
@@ -362,11 +363,10 @@ def render_overlay(
         "x-databricks-use-coding-agent-mode: true",
         f"User-Agent: ucode/{ucode_version()} claude/{agent_version('claude')}",
     ]
-    # A provider selects one MPS; a parent discovers Model Services, so they are exclusive.
-    if provider:
-        header_lines.append(f"{MODEL_PROVIDER_SERVICE_HEADER}: {provider}")
-    elif parent_schema:
-        header_lines.append(f"{MODEL_SERVICE_PARENT_SCHEMA_HEADER}: {parent_schema}")
+    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)
diff --git a/src/ucode/model_service_headers.py b/src/ucode/model_service_headers.py
new file mode 100644
index 00000000..7b54ec59
--- /dev/null
+++ b/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 {}
diff --git a/tests/test_model_service_headers.py b/tests/test_model_service_headers.py
new file mode 100644
index 00000000..e5c2ee87
--- /dev/null
+++ b/tests/test_model_service_headers.py
@@ -0,0 +1,22 @@
+"""Tests for model service routing headers."""
+
+from ucode.constants import MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER
+from ucode.model_service_headers import model_service_routing_headers
+
+
+def test_provider_header():
+    assert model_service_routing_headers("main.default.provider", None) == {
+        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
+    }
+
+
+def test_parent_schema_header():
+    assert model_service_routing_headers(None, "main.default") == {
+        MODEL_SERVICE_PARENT_SCHEMA_HEADER: "main.default"
+    }
+
+
+def test_provider_takes_precedence():
+    assert model_service_routing_headers("main.default.provider", "main.default") == {
+        MODEL_PROVIDER_SERVICE_HEADER: "main.default.provider"
+    }

From ce55c11328c3e3e078873a1fcf7e663d6b2eb20a Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 17:04:31 +0000
Subject: [PATCH 6/9] Add Codex parent schema discovery

---
 src/ucode/agents/__init__.py |  4 ++-
 src/ucode/agents/codex.py    | 61 +++++++++++++++++++++++++++++++----
 src/ucode/cli.py             |  8 +++++
 tests/test_agent_codex.py    | 62 ++++++++++++++++++++++++++++++++++--
 tests/test_cli.py            |  7 ++++
 5 files changed, 131 insertions(+), 11 deletions(-)

diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py
index fd630147..403fde73 100644
--- a/src/ucode/agents/__init__.py
+++ b/src/ucode/agents/__init__.py
@@ -430,7 +430,9 @@ def configure_tool(
 ) -> 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.
diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py
index f0f10429..3a6927ca 100644
--- a/src/ucode/agents/codex.py
+++ b/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,28 @@
 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,
+    ],
+    # Remove keys written before routing headers were normalized.
+    [
+        "model_providers",
+        CODEX_MODEL_PROVIDER_NAME,
+        "http_headers",
+        "Databricks-Model-Provider-Service",
+    ],
+    [
+        "model_providers",
+        CODEX_MODEL_PROVIDER_NAME,
+        "http_headers",
+        "Databricks-Model-Service-Parent-Schema",
+    ],
+]
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"
 MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0)
@@ -147,6 +175,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 +186,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,6 +209,7 @@ 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}
@@ -190,7 +217,12 @@ def render_overlay(
         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 +234,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 +250,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 +349,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 +375,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 +403,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:
diff --git a/src/ucode/cli.py b/src/ucode/cli.py
index 34260ef0..905ee2c5 100644
--- a/src/ucode/cli.py
+++ b/src/ucode/cli.py
@@ -2426,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(
@@ -2486,6 +2493,7 @@ def codex_cmd(
             refresh=refresh,
             skip_preflight=skip_preflight,
             workspace_url=workspace,
+            parent_schema=parent,
             custom_oauth=custom_oauth,
         )
 
diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py
index 4f4af970..a711c17d 100644
--- a/tests/test_agent_codex.py
+++ b/tests/test_agent_codex.py
@@ -131,7 +131,7 @@ def test_auth_refresh_interval(self):
     def test_provider_adds_routing_header(self):
         overlay = codex.render_overlay(WS, provider="main.aarushi.aarushi-openai")
         headers = overlay["model_providers"]["ucode-databricks"]["http_headers"]
-        assert headers["Databricks-Model-Provider-Service"] == "main.aarushi.aarushi-openai"
+        assert headers["databricks-model-provider-service"] == "main.aarushi.aarushi-openai"
 
     def test_provider_omits_model(self):
         overlay = codex.render_overlay(WS, model=None, provider="main.aarushi.aarushi-openai")
@@ -140,7 +140,19 @@ def test_provider_omits_model(self):
     def test_no_provider_header_without_flag(self):
         overlay = codex.render_overlay(WS)
         headers = overlay["model_providers"]["ucode-databricks"]["http_headers"]
-        assert "Databricks-Model-Provider-Service" not in headers
+        assert "databricks-model-provider-service" not in headers
+
+    def test_parent_adds_discovery_header(self):
+        overlay = codex.render_overlay(WS, parent_schema="main.default")
+        headers = overlay["model_providers"]["ucode-databricks"]["http_headers"]
+        assert headers["databricks-model-service-parent-schema"] == "main.default"
+
+    def test_provider_suppresses_discovery_header(self):
+        overlay = codex.render_overlay(
+            WS, provider="main.default.openai", parent_schema="main.default"
+        )
+        headers = overlay["model_providers"]["ucode-databricks"]["http_headers"]
+        assert "databricks-model-service-parent-schema" not in headers
 
 
 class TestRenderOverlayUserAgent:
@@ -224,7 +236,51 @@ def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatc
         doc = read_toml_safe(config_path)
         assert "model" not in doc
         headers = doc["model_providers"]["ucode-databricks"]["http_headers"]
-        assert headers["Databricks-Model-Provider-Service"] == "main.aarushi.aarushi-openai"
+        assert headers["databricks-model-provider-service"] == "main.aarushi.aarushi-openai"
+
+    def test_replaces_stale_routing_headers(self, tmp_path, monkeypatch):
+        config_path = tmp_path / ".codex" / "ucode.config.toml"
+        monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path)
+        monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml")
+        monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0")
+        monkeypatch.setattr(codex, "save_state", lambda state: None)
+        state = {"workspace": WS, "codex_models": []}
+
+        codex.write_tool_config(state, provider="main.default.openai")
+        codex.write_tool_config(state, parent_schema="main.default")
+
+        headers = read_toml_safe(config_path)["model_providers"]["ucode-databricks"]["http_headers"]
+        assert headers["databricks-model-service-parent-schema"] == "main.default"
+        assert "databricks-model-provider-service" not in headers
+
+        codex.write_tool_config(state)
+
+        headers = read_toml_safe(config_path)["model_providers"]["ucode-databricks"]["http_headers"]
+        assert "databricks-model-service-parent-schema" not in headers
+        assert "databricks-model-provider-service" not in headers
+
+    def test_legacy_replaces_stale_routing_headers(self, tmp_path, monkeypatch):
+        config_dir = tmp_path / ".codex"
+        legacy_path = config_dir / "config.toml"
+        monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_dir / "ucode.config.toml")
+        monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml")
+        monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path)
+        monkeypatch.setattr(codex, "LEGACY_CODEX_BACKUP_PATH", tmp_path / "legacy-backup.toml")
+        monkeypatch.setattr(codex, "agent_version", lambda binary: "0.133.0")
+        monkeypatch.setattr(codex, "save_state", lambda state: None)
+        state = {"workspace": WS, "codex_models": []}
+
+        codex.write_tool_config(state, provider="main.default.openai")
+        codex.write_tool_config(state, parent_schema="main.default")
+
+        headers = read_toml_safe(legacy_path)["model_providers"]["ucode-databricks"]["http_headers"]
+        assert headers["databricks-model-service-parent-schema"] == "main.default"
+        assert "databricks-model-provider-service" not in headers
+
+        codex.write_tool_config(state)
+        headers = read_toml_safe(legacy_path)["model_providers"]["ucode-databricks"]["http_headers"]
+        assert "databricks-model-service-parent-schema" not in headers
+        assert "databricks-model-provider-service" not in headers
 
     def test_clears_profile_model_preferences_before_launch(self, tmp_path, monkeypatch):
         config_path = tmp_path / ".codex" / "ucode.config.toml"
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 25dfcff8..37cff4c8 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -591,6 +591,13 @@ def test_claude_parent_is_forwarded(self):
         assert mock_launch.call_args.kwargs["parent_schema"] == "main.default"
         assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1"
 
+    def test_codex_parent_is_forwarded(self):
+        with patch("ucode.cli._launch_tool") as mock_launch:
+            result = runner.invoke(app, ["codex", "--parent", "main.default"])
+
+        assert result.exit_code == 0, result.output
+        assert mock_launch.call_args.kwargs["parent_schema"] == "main.default"
+
     def test_claude_enable_model_discovery_is_hidden_from_help(self):
         result = runner.invoke(app, ["claude", "--help"])
 

From 7c90344404cfdc7bc94ca6ba300244a2b327fe20 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 20:41:05 +0000
Subject: [PATCH 7/9] Support Responses-capable Bedrock services in Codex

---
 src/ucode/databricks.py     | 54 ++++++++++++++++++++++---------------
 src/ucode/managed_setup.py  |  3 +--
 tests/test_databricks.py    | 50 +++++++++++++++++++++++++++++++---
 tests/test_managed_setup.py |  3 ++-
 4 files changed, 83 insertions(+), 27 deletions(-)

diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index 49dba37f..a3d7b40e 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -2097,14 +2097,11 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 
 
 # Maps the gateway routing dialect a coding tool speaks to the Model Provider
-# Service `provider_type`s it can be backed by. claude speaks Anthropic's API,
-# which both the `anthropic` and `amazon_bedrock` provider types serve (Bedrock
-# just exposes different model ids); codex speaks OpenAI's; gemini speaks
-# Google's, served by a Gemini Enterprise provider. Tags are the short form
-# produced by `_provider_type_tag` (e.g. `amazon_bedrock`).
+# Service `provider_type`s it can be backed by. Bedrock exposes both Anthropic
+# and OpenAI dialects, so target capabilities decide whether it is usable.
 _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = {
     "claude": ("anthropic", "amazon_bedrock"),
-    "codex": ("openai",),
+    "codex": ("openai", "amazon_bedrock"),
     "gemini": ("gemini_enterprise",),
 }
 
@@ -2112,6 +2109,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 # `us.anthropic.claude-sonnet-4-6`) instead of the agent's canonical model
 # names, so ucode must pin them explicitly.
 BEDROCK_PROVIDER_TYPES: tuple[str, ...] = ("amazon_bedrock",)
+CODEX_NATIVE_API_TYPE = "openai/v1/responses"
 
 
 def tool_supports_provider_type(tool: str, provider_type: str) -> bool:
@@ -2141,7 +2139,7 @@ def list_model_provider_services(
 
     Returns ``(services, reason)`` where each service is
     ``{"name": "<catalog>.<schema>.<service>", "provider_type": "anthropic"|...,
-    "targets": [model_id, ...], "allow_all_targets": bool, "relayed": bool}``.
+    "targets": [model_id, ...], "target_api_types": {model_id: [api_type, ...]}, ...}``.
     ``targets`` is the provider-side model ids the service exposes (used to pin
     Bedrock model names). ``relayed`` is True for a credential-less Anthropic
     service (Claude Max/Team/Enterprise subscription relay). A non-None
@@ -2225,13 +2223,21 @@ def _provider_service_entry(raw_service: object) -> dict | None:
     raw_config = service.get("config")
     config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {}
     targets: list[str] = []
+    target_api_types: dict[str, list[str]] = {}
     raw_targets = config.get("targets")
     for target in raw_targets if isinstance(raw_targets, list) else []:
         if not isinstance(target, dict):
             continue
-        model_id = cast("dict[str, object]", target).get("model")
+        target_config = cast("dict[str, object]", target)
+        model_id = target_config.get("model")
         if isinstance(model_id, str) and model_id:
             targets.append(model_id)
+            raw_api_types = target_config.get("native_api_types")
+            target_api_types[model_id] = [
+                api_type
+                for api_type in (raw_api_types if isinstance(raw_api_types, list) else [])
+                if isinstance(api_type, str)
+            ]
     # Relayed = credential-less Anthropic (subscription relay). Only whether
     # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
     # server-side, so both launch identically.
@@ -2242,6 +2248,7 @@ def _provider_service_entry(raw_service: object) -> dict | None:
         "name": full_name,
         "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None),
         "targets": targets,
+        "target_api_types": target_api_types,
         "allow_all_targets": bool(config.get("allow_all_targets")),
         "relayed": relayed,
     }
@@ -2330,16 +2337,20 @@ def list_tool_provider_services(
 def service_usable_for_tool(tool: str, service: dict) -> bool:
     """True when ``tool`` can actually route through ``service``.
 
-    Beyond the provider-type match, a Bedrock service is only usable for claude
-    if it exposes at least one Claude model in its targets — otherwise there's no
-    routable model id to pin. (Anthropic services use canonical names, so any
-    match is usable.)
+    Bedrock services must expose a target compatible with the tool's dialect.
     """
     provider_type = service.get("provider_type", "")
     if not tool_supports_provider_type(tool, provider_type):
         return False
     if provider_type in BEDROCK_PROVIDER_TYPES:
-        return bool(map_claude_family_models(service.get("targets") or []))
+        if tool == "claude":
+            return bool(map_claude_family_models(service.get("targets") or []))
+        if tool == "codex":
+            raw_api_types = service.get("target_api_types")
+            return isinstance(raw_api_types, dict) and any(
+                isinstance(api_types, list) and CODEX_NATIVE_API_TYPE in api_types
+                for api_types in raw_api_types.values()
+            )
     return True
 
 
@@ -2366,9 +2377,7 @@ def resolve_provider_service(
         # fetched directly. Only when that 404s is it really absent.
         match, get_reason = get_model_provider_service(service_name, workspace, token)
         if match is None:
-            usable = [
-                s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"])
-            ]
+            usable = [s["name"] for s in services if service_usable_for_tool(tool, s)]
             suffix = f" Available for {tool}: {', '.join(usable)}." if usable else ""
             detail = f" ({get_reason})" if get_reason and "404" not in get_reason else ""
             return None, f"Model provider service '{service_name}' was not found.{detail}{suffix}"
@@ -2379,12 +2388,15 @@ def resolve_provider_service(
             f"Model provider service '{service_name}' is a '{provider_type}' provider, "
             f"which {tool} can't route to (supported: {supported})."
         )
-    if provider_type in BEDROCK_PROVIDER_TYPES and not map_claude_family_models(
-        match.get("targets") or []
-    ):
+    if provider_type in BEDROCK_PROVIDER_TYPES and not service_usable_for_tool(tool, match):
+        if tool == "claude":
+            return None, (
+                f"Model provider service '{service_name}' exposes no Claude models — "
+                f"add Claude targets to it or pick a different service."
+            )
         return None, (
-            f"Model provider service '{service_name}' exposes no Claude models — "
-            f"add Claude targets to it or pick a different service."
+            f"Model provider service '{service_name}' exposes no targets supporting "
+            f"{CODEX_NATIVE_API_TYPE}."
         )
     return match, None
 
diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py
index 3c52f909..d3d1fc8e 100644
--- a/src/ucode/managed_setup.py
+++ b/src/ucode/managed_setup.py
@@ -84,8 +84,7 @@ def supports_provider_service(tool: str, provider_type: str) -> bool:
     """True when ``tool`` can route through a ``provider_type`` Model Provider Service.
 
     Thin pass-through to :func:`ucode.databricks.tool_supports_provider_type` so the wizard has one
-    obvious place to ask. Only claude (anthropic / amazon_bedrock) and codex (openai) have MPS
-    support today; the other harnesses are Databricks-hosted only.
+    obvious place to ask. Claude and Codex also accept compatible Bedrock services.
     """
     return tool_supports_provider_type(tool, provider_type)
 
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index 20301da3..919c1aca 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -562,7 +562,29 @@ class TestListModelProviderServices:
                             "model": "us.anthropic.claude-sonnet-4-6",
                             "native_api_types": ["anthropic/v1/messages"],
                         },
-                        {"model": "global.anthropic.claude-opus-4-8"},
+                        {
+                            "model": "global.anthropic.claude-opus-4-8",
+                            "native_api_types": ["anthropic/v1/messages"],
+                        },
+                    ],
+                },
+            },
+            {
+                "name": "model-provider-services/main.schema2.bedrock-mixed-svc",
+                "config": {
+                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_AMAZON_BEDROCK",
+                    "targets": [
+                        {
+                            "model": "openai.gpt-oss-20b-1:0",
+                            "native_api_types": [
+                                "openai/v1/chat/completions",
+                                "openai/v1/responses",
+                            ],
+                        },
+                        {
+                            "model": "anthropic.claude-haiku-4-5-20251001-v1:0",
+                            "native_api_types": ["anthropic/v1/messages"],
+                        },
                     ],
                 },
             },
@@ -586,6 +608,7 @@ def test_strips_prefix_and_tags_provider_type(self, monkeypatch):
             "name": "main.schema1.anthropic-svc",
             "provider_type": "anthropic",
             "targets": [],
+            "target_api_types": {},
             "allow_all_targets": False,
             "relayed": False,
         }
@@ -614,6 +637,10 @@ def test_extracts_targets(self, monkeypatch):
             "us.anthropic.claude-sonnet-4-6",
             "global.anthropic.claude-opus-4-8",
         ]
+        assert bedrock["target_api_types"] == {
+            "us.anthropic.claude-sonnet-4-6": ["anthropic/v1/messages"],
+            "global.anthropic.claude-opus-4-8": ["anthropic/v1/messages"],
+        }
 
     def test_returns_reason_on_failure(self, monkeypatch):
         monkeypatch.setattr(
@@ -635,15 +662,16 @@ def test_claude_includes_anthropic_and_usable_bedrock(self, monkeypatch):
         assert names == [
             "main.schema1.anthropic-svc",
             "main.schema1.claude-max-svc",
+            "main.schema2.bedrock-mixed-svc",
             "main.schema2.bedrock-svc",
         ]
 
-    def test_codex_filters_to_openai(self, monkeypatch):
+    def test_codex_includes_bedrock_with_responses_targets(self, monkeypatch):
         monkeypatch.setattr(
             db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
         )
         names, _ = db_mod.list_tool_provider_services("codex", WS, "token")
-        assert names == ["main.schema1.openai-svc"]
+        assert names == ["main.schema1.openai-svc", "main.schema2.bedrock-mixed-svc"]
 
 
 class TestMapClaudeFamilyModels:
@@ -889,6 +917,22 @@ def test_bedrock_with_claude_ok(self, monkeypatch):
         assert error is None
         assert service["provider_type"] == "amazon_bedrock"
 
+    def test_bedrock_with_responses_ok_for_codex(self, monkeypatch):
+        self._patch(monkeypatch)
+        service, error = db_mod.resolve_provider_service(
+            "codex", "main.schema2.bedrock-mixed-svc", WS, "token"
+        )
+        assert error is None
+        assert service["provider_type"] == "amazon_bedrock"
+
+    def test_bedrock_without_responses_rejected_for_codex(self, monkeypatch):
+        self._patch(monkeypatch)
+        service, error = db_mod.resolve_provider_service(
+            "codex", "main.schema2.bedrock-svc", WS, "token"
+        )
+        assert service is None
+        assert db_mod.CODEX_NATIVE_API_TYPE in error
+
     def test_wrong_type_rejected(self, monkeypatch):
         self._patch(monkeypatch)
         service, error = db_mod.resolve_provider_service(
diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py
index 140bc3d9..c8a66066 100644
--- a/tests/test_managed_setup.py
+++ b/tests/test_managed_setup.py
@@ -356,8 +356,9 @@ def test_claude_supports_anthropic_and_bedrock(self):
         assert supports_provider_service("claude", "anthropic")
         assert supports_provider_service("claude", "amazon_bedrock")
 
-    def test_codex_supports_openai(self):
+    def test_codex_supports_openai_and_bedrock(self):
         assert supports_provider_service("codex", "openai")
+        assert supports_provider_service("codex", "amazon_bedrock")
 
     def test_claude_does_not_support_openai(self):
         assert not supports_provider_service("claude", "openai")

From 1d4b1b659be83f8708ea3bbd6b765e3974b8a172 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 21:21:56 +0000
Subject: [PATCH 8/9] Revert Codex Bedrock MPS support

---
 src/ucode/databricks.py     | 54 +++++++++++++++----------------------
 src/ucode/managed_setup.py  |  3 ++-
 tests/test_databricks.py    | 50 +++-------------------------------
 tests/test_managed_setup.py |  3 +--
 4 files changed, 27 insertions(+), 83 deletions(-)

diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index a3d7b40e..49dba37f 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -2097,11 +2097,14 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 
 
 # Maps the gateway routing dialect a coding tool speaks to the Model Provider
-# Service `provider_type`s it can be backed by. Bedrock exposes both Anthropic
-# and OpenAI dialects, so target capabilities decide whether it is usable.
+# Service `provider_type`s it can be backed by. claude speaks Anthropic's API,
+# which both the `anthropic` and `amazon_bedrock` provider types serve (Bedrock
+# just exposes different model ids); codex speaks OpenAI's; gemini speaks
+# Google's, served by a Gemini Enterprise provider. Tags are the short form
+# produced by `_provider_type_tag` (e.g. `amazon_bedrock`).
 _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = {
     "claude": ("anthropic", "amazon_bedrock"),
-    "codex": ("openai", "amazon_bedrock"),
+    "codex": ("openai",),
     "gemini": ("gemini_enterprise",),
 }
 
@@ -2109,7 +2112,6 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 # `us.anthropic.claude-sonnet-4-6`) instead of the agent's canonical model
 # names, so ucode must pin them explicitly.
 BEDROCK_PROVIDER_TYPES: tuple[str, ...] = ("amazon_bedrock",)
-CODEX_NATIVE_API_TYPE = "openai/v1/responses"
 
 
 def tool_supports_provider_type(tool: str, provider_type: str) -> bool:
@@ -2139,7 +2141,7 @@ def list_model_provider_services(
 
     Returns ``(services, reason)`` where each service is
     ``{"name": "<catalog>.<schema>.<service>", "provider_type": "anthropic"|...,
-    "targets": [model_id, ...], "target_api_types": {model_id: [api_type, ...]}, ...}``.
+    "targets": [model_id, ...], "allow_all_targets": bool, "relayed": bool}``.
     ``targets`` is the provider-side model ids the service exposes (used to pin
     Bedrock model names). ``relayed`` is True for a credential-less Anthropic
     service (Claude Max/Team/Enterprise subscription relay). A non-None
@@ -2223,21 +2225,13 @@ def _provider_service_entry(raw_service: object) -> dict | None:
     raw_config = service.get("config")
     config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {}
     targets: list[str] = []
-    target_api_types: dict[str, list[str]] = {}
     raw_targets = config.get("targets")
     for target in raw_targets if isinstance(raw_targets, list) else []:
         if not isinstance(target, dict):
             continue
-        target_config = cast("dict[str, object]", target)
-        model_id = target_config.get("model")
+        model_id = cast("dict[str, object]", target).get("model")
         if isinstance(model_id, str) and model_id:
             targets.append(model_id)
-            raw_api_types = target_config.get("native_api_types")
-            target_api_types[model_id] = [
-                api_type
-                for api_type in (raw_api_types if isinstance(raw_api_types, list) else [])
-                if isinstance(api_type, str)
-            ]
     # Relayed = credential-less Anthropic (subscription relay). Only whether
     # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
     # server-side, so both launch identically.
@@ -2248,7 +2242,6 @@ def _provider_service_entry(raw_service: object) -> dict | None:
         "name": full_name,
         "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None),
         "targets": targets,
-        "target_api_types": target_api_types,
         "allow_all_targets": bool(config.get("allow_all_targets")),
         "relayed": relayed,
     }
@@ -2337,20 +2330,16 @@ def list_tool_provider_services(
 def service_usable_for_tool(tool: str, service: dict) -> bool:
     """True when ``tool`` can actually route through ``service``.
 
-    Bedrock services must expose a target compatible with the tool's dialect.
+    Beyond the provider-type match, a Bedrock service is only usable for claude
+    if it exposes at least one Claude model in its targets — otherwise there's no
+    routable model id to pin. (Anthropic services use canonical names, so any
+    match is usable.)
     """
     provider_type = service.get("provider_type", "")
     if not tool_supports_provider_type(tool, provider_type):
         return False
     if provider_type in BEDROCK_PROVIDER_TYPES:
-        if tool == "claude":
-            return bool(map_claude_family_models(service.get("targets") or []))
-        if tool == "codex":
-            raw_api_types = service.get("target_api_types")
-            return isinstance(raw_api_types, dict) and any(
-                isinstance(api_types, list) and CODEX_NATIVE_API_TYPE in api_types
-                for api_types in raw_api_types.values()
-            )
+        return bool(map_claude_family_models(service.get("targets") or []))
     return True
 
 
@@ -2377,7 +2366,9 @@ def resolve_provider_service(
         # fetched directly. Only when that 404s is it really absent.
         match, get_reason = get_model_provider_service(service_name, workspace, token)
         if match is None:
-            usable = [s["name"] for s in services if service_usable_for_tool(tool, s)]
+            usable = [
+                s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"])
+            ]
             suffix = f" Available for {tool}: {', '.join(usable)}." if usable else ""
             detail = f" ({get_reason})" if get_reason and "404" not in get_reason else ""
             return None, f"Model provider service '{service_name}' was not found.{detail}{suffix}"
@@ -2388,15 +2379,12 @@ def resolve_provider_service(
             f"Model provider service '{service_name}' is a '{provider_type}' provider, "
             f"which {tool} can't route to (supported: {supported})."
         )
-    if provider_type in BEDROCK_PROVIDER_TYPES and not service_usable_for_tool(tool, match):
-        if tool == "claude":
-            return None, (
-                f"Model provider service '{service_name}' exposes no Claude models — "
-                f"add Claude targets to it or pick a different service."
-            )
+    if provider_type in BEDROCK_PROVIDER_TYPES and not map_claude_family_models(
+        match.get("targets") or []
+    ):
         return None, (
-            f"Model provider service '{service_name}' exposes no targets supporting "
-            f"{CODEX_NATIVE_API_TYPE}."
+            f"Model provider service '{service_name}' exposes no Claude models — "
+            f"add Claude targets to it or pick a different service."
         )
     return match, None
 
diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py
index d3d1fc8e..3c52f909 100644
--- a/src/ucode/managed_setup.py
+++ b/src/ucode/managed_setup.py
@@ -84,7 +84,8 @@ def supports_provider_service(tool: str, provider_type: str) -> bool:
     """True when ``tool`` can route through a ``provider_type`` Model Provider Service.
 
     Thin pass-through to :func:`ucode.databricks.tool_supports_provider_type` so the wizard has one
-    obvious place to ask. Claude and Codex also accept compatible Bedrock services.
+    obvious place to ask. Only claude (anthropic / amazon_bedrock) and codex (openai) have MPS
+    support today; the other harnesses are Databricks-hosted only.
     """
     return tool_supports_provider_type(tool, provider_type)
 
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index 919c1aca..20301da3 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -562,29 +562,7 @@ class TestListModelProviderServices:
                             "model": "us.anthropic.claude-sonnet-4-6",
                             "native_api_types": ["anthropic/v1/messages"],
                         },
-                        {
-                            "model": "global.anthropic.claude-opus-4-8",
-                            "native_api_types": ["anthropic/v1/messages"],
-                        },
-                    ],
-                },
-            },
-            {
-                "name": "model-provider-services/main.schema2.bedrock-mixed-svc",
-                "config": {
-                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_AMAZON_BEDROCK",
-                    "targets": [
-                        {
-                            "model": "openai.gpt-oss-20b-1:0",
-                            "native_api_types": [
-                                "openai/v1/chat/completions",
-                                "openai/v1/responses",
-                            ],
-                        },
-                        {
-                            "model": "anthropic.claude-haiku-4-5-20251001-v1:0",
-                            "native_api_types": ["anthropic/v1/messages"],
-                        },
+                        {"model": "global.anthropic.claude-opus-4-8"},
                     ],
                 },
             },
@@ -608,7 +586,6 @@ def test_strips_prefix_and_tags_provider_type(self, monkeypatch):
             "name": "main.schema1.anthropic-svc",
             "provider_type": "anthropic",
             "targets": [],
-            "target_api_types": {},
             "allow_all_targets": False,
             "relayed": False,
         }
@@ -637,10 +614,6 @@ def test_extracts_targets(self, monkeypatch):
             "us.anthropic.claude-sonnet-4-6",
             "global.anthropic.claude-opus-4-8",
         ]
-        assert bedrock["target_api_types"] == {
-            "us.anthropic.claude-sonnet-4-6": ["anthropic/v1/messages"],
-            "global.anthropic.claude-opus-4-8": ["anthropic/v1/messages"],
-        }
 
     def test_returns_reason_on_failure(self, monkeypatch):
         monkeypatch.setattr(
@@ -662,16 +635,15 @@ def test_claude_includes_anthropic_and_usable_bedrock(self, monkeypatch):
         assert names == [
             "main.schema1.anthropic-svc",
             "main.schema1.claude-max-svc",
-            "main.schema2.bedrock-mixed-svc",
             "main.schema2.bedrock-svc",
         ]
 
-    def test_codex_includes_bedrock_with_responses_targets(self, monkeypatch):
+    def test_codex_filters_to_openai(self, monkeypatch):
         monkeypatch.setattr(
             db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
         )
         names, _ = db_mod.list_tool_provider_services("codex", WS, "token")
-        assert names == ["main.schema1.openai-svc", "main.schema2.bedrock-mixed-svc"]
+        assert names == ["main.schema1.openai-svc"]
 
 
 class TestMapClaudeFamilyModels:
@@ -917,22 +889,6 @@ def test_bedrock_with_claude_ok(self, monkeypatch):
         assert error is None
         assert service["provider_type"] == "amazon_bedrock"
 
-    def test_bedrock_with_responses_ok_for_codex(self, monkeypatch):
-        self._patch(monkeypatch)
-        service, error = db_mod.resolve_provider_service(
-            "codex", "main.schema2.bedrock-mixed-svc", WS, "token"
-        )
-        assert error is None
-        assert service["provider_type"] == "amazon_bedrock"
-
-    def test_bedrock_without_responses_rejected_for_codex(self, monkeypatch):
-        self._patch(monkeypatch)
-        service, error = db_mod.resolve_provider_service(
-            "codex", "main.schema2.bedrock-svc", WS, "token"
-        )
-        assert service is None
-        assert db_mod.CODEX_NATIVE_API_TYPE in error
-
     def test_wrong_type_rejected(self, monkeypatch):
         self._patch(monkeypatch)
         service, error = db_mod.resolve_provider_service(
diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py
index c8a66066..140bc3d9 100644
--- a/tests/test_managed_setup.py
+++ b/tests/test_managed_setup.py
@@ -356,9 +356,8 @@ def test_claude_supports_anthropic_and_bedrock(self):
         assert supports_provider_service("claude", "anthropic")
         assert supports_provider_service("claude", "amazon_bedrock")
 
-    def test_codex_supports_openai_and_bedrock(self):
+    def test_codex_supports_openai(self):
         assert supports_provider_service("codex", "openai")
-        assert supports_provider_service("codex", "amazon_bedrock")
 
     def test_claude_does_not_support_openai(self):
         assert not supports_provider_service("claude", "openai")

From eb476bcb8e7861e5ce0a0ee8889a56c96cab1951 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 11 Sep 2026 21:59:52 +0000
Subject: [PATCH 9/9] Remove unused Codex header cleanup paths

---
 src/ucode/agents/codex.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py
index 3a6927ca..deb9c131 100644
--- a/src/ucode/agents/codex.py
+++ b/src/ucode/agents/codex.py
@@ -73,19 +73,6 @@
         "http_headers",
         MODEL_SERVICE_PARENT_SCHEMA_HEADER,
     ],
-    # Remove keys written before routing headers were normalized.
-    [
-        "model_providers",
-        CODEX_MODEL_PROVIDER_NAME,
-        "http_headers",
-        "Databricks-Model-Provider-Service",
-    ],
-    [
-        "model_providers",
-        CODEX_MODEL_PROVIDER_NAME,
-        "http_headers",
-        "Databricks-Model-Service-Parent-Schema",
-    ],
 ]
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"