From 373ad4124af79023bf81a81ea1e1094835409eee Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 28 Oct 2026 16:23:57 +0000 Subject: [PATCH 1/3] Support ChatGPT subscription relay for ug codex Relayed OpenAI Model Provider Services ({"openai": {"relayed": {}}}) are now usable by codex. A relayed launch keeps Codex's own ChatGPT login (requires_openai_auth, no ug auth command) and runs a loopback proxy that adds the Databricks token in X-Databricks-AI-Gateway-Token on the AI Gateway /codex/ routes, mirroring relayed Claude. Co-authored-by: Isaac --- AGENTS.md | 2 +- src/ucode/agents/__init__.py | 2 +- src/ucode/agents/codex.py | 204 ++++++++++++++++++++++++++++------- src/ucode/databricks.py | 25 +++-- src/ucode/gateway_proxy.py | 9 +- tests/test_agent_codex.py | 112 +++++++++++++++++++ tests/test_databricks.py | 21 ++++ tests/test_gateway_proxy.py | 58 ++++++++++ 8 files changed, 383 insertions(+), 50 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 374dcf7bd..95047ce7d 100644 --- a/AGENTS.md +++ b/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 | diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 1c8964d00..ae2b7c4be 100644 --- a/src/ucode/agents/__init__.py +++ b/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, diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index d813403f4..b543da6c2 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -7,6 +7,7 @@ import os import plistlib import re +import shutil import signal import socket import sys @@ -14,7 +15,7 @@ import threading from collections.abc import Callable from pathlib import Path -from subprocess import CalledProcessError, list2cmdline +from subprocess import CalledProcessError, TimeoutExpired, list2cmdline from xml.sax.saxutils import escape import tomlkit @@ -94,7 +95,7 @@ from ucode.state import get_provider_service, is_tool_managed, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_user_agent, ug_version from ucode.ug_headers import ug_headers -from ucode.ui import print_warning_err +from ucode.ui import print_note, print_success, print_warning_err from .args import LaunchOptions from .codex_catalog import prepare_codex_catalog, validate_codex_catalog @@ -109,10 +110,12 @@ LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "Databricks" LEGACY_CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" -# ug owns the whole provider http_headers table, so it is pruned before each merge and rewritten -# from render_overlay — dropping stale routing and admin headers that deep_merge cannot delete. +# ug owns the provider http_headers table and its auth/requires_openai_auth mode, so they are +# pruned before each merge and rewritten from render_overlay — dropping stale routing and admin +# headers or a previous relayed/non-relayed mode that deep_merge cannot delete. _PROVIDER_HTTP_HEADERS_KEY_PATHS = [ - ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], + ["model_providers", CODEX_MODEL_PROVIDER_NAME, key] + for key in ("http_headers", "auth", "requires_openai_auth") ] # Windows requirements.toml replaces the whole provider entry, so launch-scoped headers set with # `-c` would be dropped; Codex reads them from these variables instead. @@ -220,6 +223,7 @@ def _provider_block( parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, managed_http_headers: dict[str, str] | None = None, + relayed: bool = False, ) -> dict: if custom_oauth: auth_argv = build_custom_auth_token_argv(workspace, custom_oauth) @@ -228,7 +232,8 @@ def _provider_block( if not use_pat: # Desktop helpers have no terminal for interactive sign-in. auth_argv.append("--browser-login") - base_url = build_tool_base_url("codex", workspace) + # A relayed launch points base_url at its loopback proxy's bound port. + base_url = f"http://{LOOPBACK_HOST}/v1" if relayed else build_tool_base_url("codex", workspace) http_headers = { "User-Agent": ucode_user_agent(ug_version(), f"codex/{agent_version('codex')}"), } @@ -242,11 +247,18 @@ def _provider_block( if request_tags: http_headers[AI_GATEWAY_REQUEST_TAGS_HEADER] = request_tags _apply_managed_headers(http_headers, managed_http_headers) - return { + block = { "name": "Databricks AI Gateway", "base_url": base_url, "wire_api": "responses", "http_headers": http_headers, + } + if relayed: + # Codex sends its own ChatGPT login; a provider `auth` command can't be combined with it. + block["requires_openai_auth"] = True + return block + return { + **block, # Run the `ug auth-token` executable directly (not via `sh -c`) so the # helper works on Windows, where there is no POSIX shell (issue #116). "auth": { @@ -267,6 +279,7 @@ def render_overlay( parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, managed_http_headers: dict[str, str] | None = None, + relayed: bool = False, ) -> dict: overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: @@ -280,6 +293,7 @@ def render_overlay( parent_schema=parent_schema, custom_oauth=custom_oauth, managed_http_headers=managed_http_headers, + relayed=relayed, ), } return overlay @@ -438,6 +452,7 @@ def write_tool_config( model: str | None = None, provider: str | None = None, parent_schema: str | None = None, + relayed: bool = False, ) -> dict: workspace = state["workspace"] # Leave model selection to Codex. The gateway still receives the configured @@ -450,6 +465,11 @@ def write_tool_config( static_models = static_models if isinstance(static_models, list) and static_models else None if _use_legacy_layout(): + if relayed: + raise RuntimeError( + "A ChatGPT subscription relay needs Codex " + f"{LEGACY_LAYOUT_CODEX_VERSION_TEXT} or newer. Upgrade Codex and retry." + ) if static_models and not provider: raise RuntimeError( "This Codex version cannot use the managed static model catalog. " @@ -516,6 +536,7 @@ def write_tool_config( parent_schema=parent_schema, custom_oauth=state.get("custom_oauth"), managed_http_headers=state.get("codex_http_headers"), + relayed=relayed, ) def compose(base: dict, *, include_catalog: bool = True) -> dict: @@ -550,7 +571,14 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: enabled=False, ) write_toml_file(CODEX_CONFIG_PATH, doc) - _reconcile_managed_config(state, lambda base: compose(base, include_catalog=False)) + _reconcile_managed_config( + state, lambda base: compose(base, include_catalog=False), relayed=relayed + ) + # Persisted so launch() wires the relay proxy and ChatGPT login; cleared on a non-relayed run. + if relayed: + state["codex_relayed"] = True + else: + state.pop("codex_relayed", None) state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) # The desktop proxy reads authentication from the saved workspace. @@ -604,7 +632,10 @@ def managed_config_is_current(state: dict) -> bool: path = codex_managed_config_path() if path is None: return True - required_scope = "managed" if managed_writes_allowed() else None + if state.get("codex_relayed"): + required_scope = "relay-compatible" + else: + required_scope = "managed" if managed_writes_allowed() else None return managed_file_is_verified(state, "codex", path, required_scope=required_scope) @@ -893,7 +924,30 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: return False -def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> None: +def _managed_relayed_conflicts(path: Path) -> list[str]: + """Return managed settings that would override a ChatGPT relay launch's ``-c`` provider.""" + text = read_managed_file(path) + if text is None: + return [] + try: + doc = _parse_managed_config(text) + except RuntimeError as exc: + raise RuntimeError( + f"Cannot safely inspect Codex managed settings at {path}: {exc}. Repair the file or " + "contact your administrator." + ) from exc + conflicts = ["model_provider"] if "model_provider" in doc else [] + provider = (doc.get("model_providers") or {}).get(CODEX_MODEL_PROVIDER_NAME) or {} + return conflicts + [ + f"model_providers.{CODEX_MODEL_PROVIDER_NAME}.{key}" + for key in ("base_url", "auth") + if key in provider + ] + + +def _reconcile_managed_config( + state: dict, compose: Callable[[dict], dict], relayed: bool = False +) -> None: """Reconcile Codex's highest-precedence config while preserving unrelated policy.""" path = codex_managed_config_path() if path is not None and path.name == CODEX_REQUIREMENTS_FILE: @@ -912,6 +966,18 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N f"Refusing to use Codex managed settings through symlink {path}. Replace it with a " "regular file or contact your administrator." ) + if relayed: + # The relay's loopback proxy only runs during `ug codex`, so skip the machine-wide file. + conflicts = _managed_relayed_conflicts(path) + if conflicts: + raise RuntimeError( + "Codex ChatGPT relay cannot start because enterprise managed settings define " + f"{', '.join(conflicts)} at {path}. Ask your administrator to remove those " + "entries or use standard Databricks authentication. If ucode previously created " + "them, run `ug revert` first." + ) + mark_managed_file_verified(state, "codex", path, scope="relay-compatible") + return current_text = read_managed_file(path) try: existing = _parse_managed_config(current_text) if current_text is not None else {} @@ -1417,36 +1483,72 @@ def _launch_codex_with_otel_proxy( base_argv: list[str], tool_args: list[str], workspace: str, + *, + otel_tracing: bool = True, + relay: tuple | None = None, ) -> None: - """Run the loopback OTLP refresh proxy for the session, with Codex as a child. + """Run the loopback proxies for the session (OTLP refresh and/or the started ChatGPT + ``relay``), with Codex as a child. - The proxy must outlive the launch, so Codex runs as a child rather than - exec-replacing this process; mirrors Claude's relayed launch. The proxy binds an - OS-assigned port and tears everything down when Codex exits (or fails to spawn). + The proxies must outlive the launch, so Codex runs as a child rather than + exec-replacing this process; mirrors Claude's relayed launch. Each binds an + OS-assigned port and everything is torn down when Codex exits (or fails to spawn). """ - # Only tracing launches need the proxy (and its `httpx` import). - from ucode import gateway_proxy - - server, cache, client = gateway_proxy.start_otel_proxy( - workspace, _otel_token_provider(state, workspace) - ) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - endpoint = f"http://{LOOPBACK_HOST}:{server.server_address[1]}/v1/traces" - otel_args = codex_config_args(_otel_proxy_overlay(endpoint)) - proc = subprocess_cross_os.popen([*base_argv, *otel_args, *tool_args]) + proxies = [relay] if relay else [] + otel_args: list[str] = [] try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + if otel_tracing: + # Only tracing launches need the proxy (and its `httpx` import). + from ucode import gateway_proxy + + otel = gateway_proxy.start_otel_proxy(workspace, _otel_token_provider(state, workspace)) + proxies.append(otel) + threading.Thread(target=otel[0].serve_forever, daemon=True).start() + endpoint = f"http://{LOOPBACK_HOST}:{otel[0].server_address[1]}/v1/traces" + otel_args = codex_config_args(_otel_proxy_overlay(endpoint)) + proc = subprocess_cross_os.popen([*base_argv, *otel_args, *tool_args]) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() finally: - cache.stop() - server.shutdown() - client.close() + for server, cache, client in proxies: + cache.stop() + server.shutdown() + client.close() raise SystemExit(returncode) +def _has_chatgpt_login(binary: str) -> bool: + """True when Codex holds a ChatGPT login (not an API key) that a relayed launch will use.""" + try: + result = subprocess_cross_os.run( + [shutil.which(binary) or binary, "login", "status"], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, TimeoutExpired): + return False + return result.returncode == 0 and "ChatGPT" in f"{result.stdout}{result.stderr}" + + +def _ensure_chatgpt_login(binary: str) -> None: + """Run Codex's browser sign-in if it has no ChatGPT login. ug never sees the token.""" + if _has_chatgpt_login(binary): + return + print_note("Opening browser to sign in with your ChatGPT subscription...") + try: + subprocess_cross_os.run([shutil.which(binary) or binary, "login"], check=True, timeout=300) + except CalledProcessError as exc: + raise RuntimeError("`codex login` failed.") from exc + except TimeoutExpired as exc: + raise RuntimeError("`codex login` timed out.") from exc + print_success("ChatGPT subscription authenticated") + + def _run_codex( state: dict, base_argv: list[str], @@ -1454,13 +1556,16 @@ def _run_codex( *, otel_tracing: bool, workspace: str | None, + relay: tuple | None = None, ) -> None: - """Launch Codex — via the loopback proxy when OTLP tracing is on, else exec-replace.""" + """Launch Codex — via loopback proxies when OTLP tracing or a relay is on, else exec-replace.""" if tool_args[:1] == ["update"]: # exec replaces ug, so reattach only on a later validated refresh. detach_app_model_catalog() - if otel_tracing and workspace: - _launch_codex_with_otel_proxy(state, base_argv, tool_args, workspace) + if (otel_tracing or relay) and workspace: + _launch_codex_with_otel_proxy( + state, base_argv, tool_args, workspace, otel_tracing=otel_tracing, relay=relay + ) else: exec_or_spawn([*base_argv, *tool_args]) @@ -1471,7 +1576,13 @@ def launch( *, options: LaunchOptions, ) -> None: + relayed = bool(state.get("codex_relayed")) if options.launch_smart_routing: + if relayed: + raise RuntimeError( + "Smart routing can't be used with a ChatGPT subscription relay. Re-run with " + "`--disable-smart-routing` or pick another provider." + ) _launch_smart_routing(state, tool_args) return clear_model_preferences(state) @@ -1500,7 +1611,8 @@ def launch( otel_tracing = bool(workspace and state.get("codex_otel_tracing")) if workspace: token = _launch_token(state, workspace) - os.environ["OAUTH_TOKEN"] = token + if not relayed: # Codex authenticates with its ChatGPT login, not this token. + os.environ["OAUTH_TOKEN"] = token if _use_legacy_layout(): print_warning_err( f"Codex {agent_version(binary)} is outdated. Upgrade Codex to " @@ -1532,7 +1644,7 @@ def launch( updating = tool_args[:1] == ["update"] if updating and _is_ucode_catalog_reference(profile_doc.get("model_catalog_json")): profile_doc.pop("model_catalog_json") - if workspace and token and (provider or parent_schema) and not updating: + if workspace and token and (provider or parent_schema) and not updating and not relayed: try: if provider is not None: catalog_source = CodexCatalogSource.PROVIDER @@ -1571,12 +1683,30 @@ def launch( slugs = catalog_slugs(catalog) if slugs: profile_doc["model"] = slugs[0] + relay = None + if relayed and workspace and not updating: + _ensure_chatgpt_login(binary) + # Only relayed launches need the proxy (and its `httpx` import). + from ucode import gateway_proxy + + relay = gateway_proxy.start_relay_proxy( + workspace, + lambda force: _launch_token(state, workspace, force_refresh=force), + 0, + codex=True, + ) + threading.Thread(target=relay[0].serve_forever, daemon=True).start() + port = relay[0].server_address[1] + profile_doc["model_providers"][CODEX_MODEL_PROVIDER_NAME]["base_url"] = ( + f"http://{LOOPBACK_HOST}:{port}/v1" + ) _run_codex( state, [binary, *codex_config_args(profile_doc)], tool_args, otel_tracing=otel_tracing, workspace=workspace, + relay=relay, ) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 281242029..8f12f38de 100644 --- a/src/ucode/databricks.py +++ b/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 diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 6e8c43f3a..c359907f9 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -640,15 +640,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, ) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 7f3288f35..1edca207d 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -1641,6 +1641,68 @@ 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) + monkeypatch.setattr(codex, "_has_chatgpt_login", lambda binary: True) + 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(), + ) + + 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.return_value = Mock(returncode=1, stdout="", stderr="Not logged in") + codex._ensure_chatgpt_login("codex") + assert run.call_args.args[0][1:] == ["login"] + + 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 +1823,56 @@ 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) + + 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" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 522618729..35b443c96 100644 --- a/tests/test_databricks.py +++ b/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) diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index fdeb1a2c6..7563c05ca 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -608,3 +608,61 @@ 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_codex_relay_upstream_and_plain_handler(self, monkeypatch): + captured = {} + + def fake_start(workspace, token_provider, port, **kwargs): + captured.update(kwargs) + return None + + monkeypatch.setattr(gateway_proxy, "_start_proxy", fake_start) + gateway_proxy.start_relay_proxy("https://w", lambda _f: "t", 0, codex=True) + assert captured["upstream_path"] == "ai-gateway/codex/" + assert captured["handler_type"] is gateway_proxy._ProxyHandler + assert captured["token_header"] == gateway_proxy.AI_GATEWAY_TOKEN_HEADER + + def test_codex_request_path_joins_upstream_once(self, monkeypatch): + class _StubCache: + def run_refresher(self): + return None + + def stop(self): + return None + + monkeypatch.setattr(gateway_proxy, "TokenCache", lambda *_a, **_k: _StubCache()) + server, _cache, client = gateway_proxy.start_relay_proxy( + "https://w", lambda _f: "t", 0, codex=True + ) + try: + # The handler requests the client's path minus its leading slash. + url = client.build_request("POST", "/v1/responses".lstrip("/")).url + assert str(url) == "https://w/ai-gateway/codex/v1/responses" + finally: + server.server_close() + client.close() + + def test_passes_chatgpt_headers_and_swaps_databricks_token(self): + client = _FakeClient([_FakeResp(200, b"ok")]) + h = object.__new__(gateway_proxy._ProxyHandler) + h.client, h.cache = client, _FakeCache() + h.token_header = gateway_proxy.AI_GATEWAY_TOKEN_HEADER + 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 = "/v1/responses", "POST" + h.wfile = _Collect() + h.request_version, h.requestline, h._headers_buffer = "HTTP/1.1", "POST /", [] + 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" From 8232935d23d933f663e79b5aafcc449064b8b507 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 29 Oct 2026 20:13:36 +0000 Subject: [PATCH 2/3] Don't request compressed bodies the relay client didn't ask for httpx adds Accept-Encoding: gzip, deflate by default, so the gateway compressed responses that the proxy relays raw. Codex sends no Accept-Encoding and can't decode gzip, so it printed error bodies as binary. Ask for identity unless the client set its own encoding. Co-Authored-By: Claude Opus 5.5 --- src/ucode/gateway_proxy.py | 4 ++++ tests/test_gateway_proxy.py | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index c359907f9..10808e32b 100644 --- a/src/ucode/gateway_proxy.py +++ b/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 diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 7563c05ca..c5f0d2817 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -43,6 +43,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"}) From a4af9a1eead59fdfcd29ba14857fe83e97623d60 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 29 Oct 2026 22:09:16 +0000 Subject: [PATCH 3/3] Address relay review: Windows provider conflict, tighter tests Treat any managed model_providers.Databricks entry in Windows requirements.toml as a relay conflict, since that file replaces the whole provider entry. Fold the Codex relay proxy tests into one that runs through the bound handler, and model login success/failure realistically. Co-Authored-By: Claude Opus 5.5 --- src/ucode/agents/codex.py | 3 +++ tests/test_agent_codex.py | 17 +++++++++++-- tests/test_gateway_proxy.py | 51 ++++++++++--------------------------- 3 files changed, 32 insertions(+), 39 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index b543da6c2..56648c48f 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -938,6 +938,9 @@ def _managed_relayed_conflicts(path: Path) -> list[str]: ) from exc conflicts = ["model_provider"] if "model_provider" in doc else [] provider = (doc.get("model_providers") or {}).get(CODEX_MODEL_PROVIDER_NAME) or {} + if provider and path.name == CODEX_REQUIREMENTS_FILE: + # requirements.toml replaces the whole provider entry, dropping the relay's `-c` keys. + return [*conflicts, f"model_providers.{CODEX_MODEL_PROVIDER_NAME}"] return conflicts + [ f"model_providers.{CODEX_MODEL_PROVIDER_NAME}.{key}" for key in ("base_url", "auth") diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 1edca207d..ddf523365 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -1658,7 +1658,8 @@ def start_relay_proxy(workspace, token_provider, port, *, codex=False): monkeypatch.setattr(gateway_proxy, "start_relay_proxy", start_relay_proxy) monkeypatch.setattr(codex.subprocess_cross_os, "popen", popen) - monkeypatch.setattr(codex, "_has_chatgpt_login", lambda binary: True) + 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")) ) @@ -1675,6 +1676,7 @@ def start_relay_proxy(workspace, token_provider, port, *, codex=False): 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] @@ -1690,9 +1692,12 @@ def test_relayed_launch_logs_in_when_not_chatgpt(self, tmp_path, monkeypatch): 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.return_value = Mock(returncode=1, stdout="", stderr="Not logged in") + 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) @@ -1862,6 +1867,14 @@ def test_relayed_rejects_managed_provider_override(self, tmp_path, monkeypatch): 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 = [] diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index c5f0d2817..bee91a88c 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -9,6 +9,7 @@ import socket import threading import time +from unittest.mock import Mock import httpx import pytest @@ -620,44 +621,22 @@ def test_relayed_model_keeps_oauth_passthrough(self): class TestCodexRelayProxy: - def test_codex_relay_upstream_and_plain_handler(self, monkeypatch): - captured = {} - - def fake_start(workspace, token_provider, port, **kwargs): - captured.update(kwargs) - return None - - monkeypatch.setattr(gateway_proxy, "_start_proxy", fake_start) - gateway_proxy.start_relay_proxy("https://w", lambda _f: "t", 0, codex=True) - assert captured["upstream_path"] == "ai-gateway/codex/" - assert captured["handler_type"] is gateway_proxy._ProxyHandler - assert captured["token_header"] == gateway_proxy.AI_GATEWAY_TOKEN_HEADER - - def test_codex_request_path_joins_upstream_once(self, monkeypatch): - class _StubCache: - def run_refresher(self): - return None - - def stop(self): - return None - - monkeypatch.setattr(gateway_proxy, "TokenCache", lambda *_a, **_k: _StubCache()) - server, _cache, client = gateway_proxy.start_relay_proxy( + 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 ) - try: - # The handler requests the client's path minus its leading slash. - url = client.build_request("POST", "/v1/responses".lstrip("/")).url - assert str(url) == "https://w/ai-gateway/codex/v1/responses" - finally: - server.server_close() - client.close() + 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" - def test_passes_chatgpt_headers_and_swaps_databricks_token(self): client = _FakeClient([_FakeResp(200, b"ok")]) - h = object.__new__(gateway_proxy._ProxyHandler) - h.client, h.cache = client, _FakeCache() - h.token_header = gateway_proxy.AI_GATEWAY_TOKEN_HEADER + 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", @@ -666,9 +645,7 @@ def test_passes_chatgpt_headers_and_swaps_databricks_token(self): "Content-Length": str(len(body)), } h.rfile = io.BytesIO(body) - h.path, h.command = "/v1/responses", "POST" - h.wfile = _Collect() - h.request_version, h.requestline, h._headers_buffer = "HTTP/1.1", "POST /", [] + 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"