Support ChatGPT subscription relay for ug codex

AGENTS.md

@@ -182,7 +182,7 @@ Fields live in `~/.codex/ucode.config.toml` and `/etc/codex/managed_config.toml`
 |-------|------------------------|---------------------|-------|
 | `model_provider` | Create/replace | Create/replace | Gateway provider (`Databricks`) |
 | `model` | Ignore | Create/replace | The config's default model |
-| `model_providers.Databricks` | Merge | Merge | Provider block: name, gateway base URL, wire API, `ug` auth command; other keys left alone |
+| `model_providers.Databricks` | Merge | Merge | Provider block: name, gateway base URL, wire API, `ug` auth command; other keys left alone. A relayed (ChatGPT subscription) provider sets `requires_openai_auth` instead of the auth command and skips the managed file |
 | `http_headers` | Merge | Merge | In `[model_providers.Databricks]`; merge `ug`'s routing headers by name, admin headers added under managed config. `UG_HEADERS` names not already present are added to the launch's `-c` flags only, never a file |
 | `model_catalog_json` | Create/replace | Create/replace | In `~/.codex/config.toml`; `ug`'s own catalog reference, for a static model list |
 | `otel.trace_exporter` | Ignore | Create/replace | In `~/.codex/config.toml` when the config enables tracing, for Codex desktop; another exporter and other telemetry settings left alone |

src/ucode/agents/__init__.py

@@ -466,7 +466,7 @@ def configure_tool(
     result: dict | tuple[dict, str]
     if tool == AGENT_CODEX:
         result = codex.write_tool_config(
-            state, model, provider=provider, parent_schema=parent_schema
+            state, model, provider=provider, parent_schema=parent_schema, relayed=relayed
         )
     elif tool == AGENT_CLAUDE:
         # A Model Provider Service or parent schema routes by header and discovers models natively,

src/ucode/databricks.py

@@ -2939,11 +2939,13 @@ def _provider_service_entry(raw_service: object) -> dict | None:
             api_types = entry.get("native_api_types")
             if isinstance(api_types, list):
                 target_api_types[model_id] = tuple(t for t in api_types if isinstance(t, str))
-    # Relayed = credential-less Anthropic (subscription relay). Only whether
+    # Relayed = credential-less Anthropic or OpenAI (subscription relay). Only whether
     # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
     # server-side, so both launch identically.
-    anthropic_cfg = config.get("anthropic")
-    relayed = isinstance(anthropic_cfg, dict) and "relayed" in anthropic_cfg
+    relayed = any(
+        isinstance(cfg, dict) and "relayed" in cfg
+        for cfg in (config.get("anthropic"), config.get("openai"))
+    )
     raw_type = config.get("provider_type")
     return {
         "name": full_name,
@@ -3038,7 +3040,7 @@ def list_tool_provider_services(
 def _provider_service_unusable_reason(tool: str, service: dict, service_name: str) -> str | None:
     """Why ``tool`` can't route through ``service``, or None when it can.
 
-    Beyond the provider-type match, only claude can use a relayed (subscription) service, and
+    Beyond the provider-type match, only claude (Anthropic) or codex (OpenAI) can use a relayed (subscription) service, and
     claude and copilot can use a Bedrock service only if it has a Claude target — otherwise there's
     no routable model id to pin — or is ``allow_all_targets``, in which case the model id comes
     from elsewhere (the managed config's default or ``--model``). (Anthropic services use canonical
@@ -3051,11 +3053,18 @@ def _provider_service_unusable_reason(tool: str, service: dict, service_name: st
             f"Model provider service '{service_name}' is a '{provider_type}' provider, "
             f"which {tool} can't route to (supported: {supported})."
         )
-    if tool != "claude" and service.get("relayed"):
-        return (
-            f"Model provider service '{service_name}' relays a Claude subscription, which only "
-            f"Claude Code can use. Pick a Model Provider Service with an API key for {tool}."
+    if service.get("relayed"):
+        relay_tool, subscription, display = (
+            ("codex", "ChatGPT", "Codex")
+            if provider_type == "openai"
+            else ("claude", "Claude", "Claude Code")
         )
+        if tool != relay_tool:
+            return (
+                f"Model provider service '{service_name}' relays a {subscription} subscription, "
+                f"which only {display} can use. Pick a Model Provider Service with an API key "
+                f"for {tool}."
+            )
     if (
         tool in _CLAUDE_ONLY_BEDROCK_TOOLS
         and provider_type in BEDROCK_PROVIDER_TYPES

src/ucode/gateway_proxy.py

@@ -215,6 +215,10 @@ def forwarded_request_headers(
         key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward
     }
     headers[token_header] = f"Bearer {token}"
+    # Don't let httpx's default `Accept-Encoding: gzip, deflate` ask for a compressed body the
+    # client never requested: the response is relayed raw, and Codex can't decode gzip.
+    if not any(key.lower() == "accept-encoding" for key in headers):
+        headers["Accept-Encoding"] = "identity"
     return headers
 
 
@@ -640,15 +644,18 @@ def start_relay_proxy(
     workspace: str,
     token_provider: Callable[[bool], str],
     port: int,
+    *,
+    codex: bool = False,
 ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]:
-    """Start the Claude subscription relay proxy."""
+    """Start the subscription relay proxy: Claude's, or Codex's with ``codex`` (plain forwarding,
+    since the Anthropic hybrid reroute doesn't apply)."""
     return _start_proxy(
         workspace,
         token_provider,
         port,
-        upstream_path="ai-gateway/anthropic/",
+        upstream_path="ai-gateway/codex/" if codex else "ai-gateway/anthropic/",
         token_header=AI_GATEWAY_TOKEN_HEADER,
-        handler_type=_RelayProxyHandler,
+        handler_type=_ProxyHandler if codex else _RelayProxyHandler,
         force_refresh_near_expiry=False,
     )
 

tests/test_agent_codex.py

@@ -1641,6 +1641,73 @@ def start_otel_proxy(workspace, token_provider):
         assert "Authorization" not in otel  # no credential in argv; the proxy injects it
         assert argv[-2:] == ["exec", "hi"]
 
+    def test_relayed_launch_checks_login_and_routes_through_relay_proxy(
+        self, tmp_path, monkeypatch
+    ):
+        self._patch(tmp_path, monkeypatch)
+        server = Mock(server_address=("127.0.0.1", 54321))
+        cache, client = Mock(), Mock()
+        process = Mock()
+        process.wait.return_value = 0
+        popen = Mock(return_value=process)
+        started = {}
+
+        def start_relay_proxy(workspace, token_provider, port, *, codex=False):
+            started.update(workspace=workspace, token=token_provider(False), codex=codex)
+            return server, cache, client
+
+        monkeypatch.setattr(gateway_proxy, "start_relay_proxy", start_relay_proxy)
+        monkeypatch.setattr(codex.subprocess_cross_os, "popen", popen)
+        login_status = Mock(return_value=True)
+        monkeypatch.setattr(codex, "_has_chatgpt_login", login_status)
+        monkeypatch.setattr(
+            codex, "_fetch_codex_model_catalog", Mock(side_effect=AssertionError("no catalog"))
+        )
+        monkeypatch.delenv("OAUTH_TOKEN", raising=False)
+
+        with pytest.raises(SystemExit):
+            codex.launch(
+                {
+                    "workspace": WS,
+                    "codex_relayed": True,
+                    "_codex_launch_provider": "cat.s.mps",
+                },
+                ["exec", "hi"],
+                options=LaunchOptions(),
+            )
+
+        login_status.assert_called_once()
+        assert started == {"workspace": WS, "token": "tok", "codex": True}
+        assert "OAUTH_TOKEN" not in os.environ
+        argv = popen.call_args.args[0]
+        assert any("http://127.0.0.1:54321/v1" in arg for arg in argv)
+        assert argv[-2:] == ["exec", "hi"]
+        cache.stop.assert_called_once_with()
+        server.shutdown.assert_called_once_with()
+
+    def test_relayed_launch_logs_in_when_not_chatgpt(self, tmp_path, monkeypatch):
+        self._patch(tmp_path, monkeypatch)
+        run = Mock(return_value=Mock(returncode=0, stdout="", stderr="Logged in using an API key"))
+        monkeypatch.setattr(codex.subprocess_cross_os, "run", run)
+        assert codex._has_chatgpt_login("codex") is False
+        run.return_value = Mock(returncode=0, stdout="", stderr="Logged in using ChatGPT")
+        assert codex._has_chatgpt_login("codex") is True
+        run.side_effect = [Mock(returncode=1, stdout="", stderr="Not logged in"), Mock()]
+        codex._ensure_chatgpt_login("codex")
+        assert run.call_args.args[0][1:] == ["login"]
+        run.side_effect = [Mock(returncode=1, stdout="", stderr=""), CalledProcessError(1, "x")]
+        with pytest.raises(RuntimeError, match="`codex login` failed"):
+            codex._ensure_chatgpt_login("codex")
+
+    def test_relayed_launch_rejects_smart_routing(self, tmp_path, monkeypatch):
+        self._patch(tmp_path, monkeypatch)
+        with pytest.raises(RuntimeError, match="Smart routing can't be used"):
+            codex.launch(
+                {"workspace": WS, "codex_relayed": True},
+                [],
+                options=LaunchOptions(launch_smart_routing=True),
+            )
+
     def test_no_otel_config_when_tracing_disabled(self, tmp_path, monkeypatch):
         launches = self._patch(tmp_path, monkeypatch)
 
@@ -1761,6 +1828,64 @@ def test_writes_managed_config_by_default(self, tmp_path, monkeypatch):
         assert "Databricks" in doc["model_providers"]
         assert read_toml_safe(config_path)["model_provider"] == "Databricks"
 
+    def test_relayed_writes_proxyless_provider_and_skips_managed_file(self, tmp_path, monkeypatch):
+        config_path, managed_path = self._patch(tmp_path, monkeypatch)
+        config_path.parent.mkdir(parents=True)
+        config_path.write_text(
+            '[model_providers.Databricks]\nauth = { command = "old" }\n', encoding="utf-8"
+        )
+        marked = []
+        monkeypatch.setattr(
+            codex, "mark_managed_file_verified", lambda *a, **k: marked.append(k.get("scope"))
+        )
+        state = {"workspace": WS}
+
+        codex.write_tool_config(state, provider="cat.s.mps", relayed=True)
+
+        provider = read_toml_safe(config_path)["model_providers"]["Databricks"]
+        assert provider["requires_openai_auth"] is True
+        assert "auth" not in provider
+        assert not managed_path.exists()
+        assert marked == ["relay-compatible"]
+        assert state["codex_relayed"] is True
+
+        marked.clear()
+        codex.write_tool_config(state, provider="cat.s.mps")
+        provider = read_toml_safe(config_path)["model_providers"]["Databricks"]
+        assert "requires_openai_auth" not in provider
+        assert "auth" in provider
+        assert "codex_relayed" not in state
+
+    def test_relayed_rejects_managed_provider_override(self, tmp_path, monkeypatch):
+        _, managed_path = self._patch(tmp_path, monkeypatch)
+        managed_path.parent.mkdir(parents=True)
+        managed_path.write_text(
+            'model_provider = "Databricks"\n[model_providers.Databricks]\nbase_url = "x"\n',
+            encoding="utf-8",
+        )
+
+        with pytest.raises(RuntimeError, match=r"model_provider, .*base_url.*ug revert"):
+            codex.write_tool_config({"workspace": WS}, provider="cat.s.mps", relayed=True)
+
+        # requirements.toml replaces the whole entry, so even a name-only one would drop the relay.
+        requirements = tmp_path / "ProgramData" / "requirements.toml"
+        requirements.parent.mkdir(parents=True)
+        requirements.write_text('[model_providers.Databricks]\nname = "x"\n', encoding="utf-8")
+        monkeypatch.setattr(codex, "codex_managed_config_path", lambda: requirements)
+        with pytest.raises(RuntimeError, match=r"define model_providers\.Databricks at"):
+            codex.write_tool_config({"workspace": WS}, provider="cat.s.mps", relayed=True)
+
+    def test_relayed_state_requires_relay_compatible_scope(self, tmp_path, monkeypatch):
+        self._patch(tmp_path, monkeypatch)
+        scopes = []
+        monkeypatch.setattr(
+            codex,
+            "managed_file_is_verified",
+            lambda state, tool, path, required_scope=None: scopes.append(required_scope) or True,
+        )
+        codex.managed_config_is_current({"codex_relayed": True})
+        assert scopes == ["relay-compatible"]
+
     def test_windows_writes_requirements_shaped_config(self, tmp_path, monkeypatch):
         config_path, _ = self._patch(tmp_path, monkeypatch)
         managed_path = tmp_path / "ProgramData" / "requirements.toml"

tests/test_databricks.py

@@ -888,6 +888,19 @@ def test_flags_relayed_anthropic(self, monkeypatch):
         assert by_name["main.schema1.claude-max-svc"]["relayed"] is True
         assert by_name["main.schema1.anthropic-svc"]["relayed"] is False
 
+    def test_flags_relayed_openai(self):
+        entry = db_mod._provider_service_entry(
+            {
+                "name": "model-provider-services/main.s.chatgpt",
+                "config": {
+                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI",
+                    "openai": {"relayed": {}},
+                },
+            }
+        )
+        assert entry["relayed"] is True
+        assert entry["provider_type"] == "openai"
+
     def test_extracts_targets(self, monkeypatch):
         monkeypatch.setattr(
             db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
@@ -1526,6 +1539,14 @@ def test_copilot_skips_other_provider_types(self, provider_type):
         service = {"provider_type": provider_type, "targets": [], "allow_all_targets": True}
         assert not db_mod.service_usable_for_tool("copilot", service)
 
+    def test_openai_relay_is_codex_only(self):
+        service = {"provider_type": "openai", "targets": [], "relayed": True}
+        assert db_mod.service_usable_for_tool("codex", service)
+        assert not db_mod.service_usable_for_tool("copilot", service)
+        assert "ChatGPT" in db_mod._provider_service_unusable_reason("copilot", service, "svc")
+        anthropic = {"provider_type": "anthropic", "targets": [], "relayed": True}
+        assert not db_mod.service_usable_for_tool("codex", anthropic)
+
     def test_copilot_skips_a_subscription_relay(self):
         service = {"provider_type": "anthropic", "targets": [], "relayed": True}
         assert not db_mod.service_usable_for_tool("copilot", service)

tests/test_gateway_proxy.py

@@ -9,6 +9,7 @@
 import socket
 import threading
 import time
+from unittest.mock import Mock
 
 import httpx
 import pytest
@@ -43,6 +44,15 @@ def test_passes_authorization_through_untouched(self):
         out = gateway_proxy.forwarded_request_headers(handler, "dbx-token")
         assert out["Authorization"] == "Bearer anthropic-oauth"
 
+    def test_requests_identity_encoding_unless_client_asks(self):
+        # Codex sends no Accept-Encoding and can't decode gzip; httpx would otherwise ask for it.
+        out = gateway_proxy.forwarded_request_headers(_FakeHandler({}), "t")
+        assert out["Accept-Encoding"] == "identity"
+        handler = _FakeHandler({"accept-encoding": "gzip, br"})
+        out = gateway_proxy.forwarded_request_headers(handler, "t")
+        assert out["accept-encoding"] == "gzip, br"
+        assert "Accept-Encoding" not in out
+
     def test_overwrites_client_supplied_swap_header(self):
         # A stale settings.json value must not survive; the proxy replaces it.
         handler = _FakeHandler({"X-Databricks-AI-Gateway-Token": "Bearer stale"})
@@ -608,3 +618,37 @@ def test_relayed_model_keeps_oauth_passthrough(self):
         assert sent["Authorization"] == "Bearer anthropic-oauth"
         assert sent[gateway_proxy.AI_GATEWAY_TOKEN_HEADER] == "Bearer tok1"
         assert sent["Databricks-Model-Provider-Service"] == "cat.s.relayed_mps"
+
+
+class TestCodexRelayProxy:
+    def test_relays_chatgpt_login_to_codex_upstream(self, monkeypatch):
+        cache = Mock(token="tok1")
+        monkeypatch.setattr(gateway_proxy, "TokenCache", lambda *_a, **_k: cache)
+        server, _cache, upstream = gateway_proxy.start_relay_proxy(
+            "https://w", lambda _f: "t", 0, codex=True
+        )
+        server.server_close()
+        upstream.close()
+        # The handler requests the client's path minus its leading slash.
+        url = upstream.build_request("POST", "v1/responses").url
+        assert str(url) == "https://w/ai-gateway/codex/v1/responses"
+
+        client = _FakeClient([_FakeResp(200, b"ok")])
+        h = object.__new__(server.RequestHandlerClass)
+        h.client = client  # the bound class keeps the real cache, token header, and handler
+        h.wfile, h.request_version, h._headers_buffer = _Collect(), "HTTP/1.1", []
+        body = b'{"model": "claude-opus-4-1"}'  # not rerouted: no Anthropic hybrid logic
+        h.headers = {
+            "Authorization": "Bearer chatgpt-token",
+            "ChatGPT-Account-ID": "acct",
+            "Databricks-Model-Provider-Service": "cat.s.mps",
+            "Content-Length": str(len(body)),
+        }
+        h.rfile = io.BytesIO(body)
+        h.path, h.command, h.requestline = "/v1/responses", "POST", "POST /v1/responses HTTP/1.1"
+        h._handle()
+        sent = client.sent_headers[0]
+        assert sent["Authorization"] == "Bearer chatgpt-token"
+        assert sent["ChatGPT-Account-ID"] == "acct"
+        assert sent[gateway_proxy.AI_GATEWAY_TOKEN_HEADER] == "Bearer tok1"
+        assert sent["Databricks-Model-Provider-Service"] == "cat.s.mps"