[andy-xu-db]

🥞 Stacked PR

Use this link to review incremental changes.

• stack/codex-relay-cli [Files changed] ← this PR
  • stack/codex-relay-desktop [Files changed]

---

Summary

ug codex --provider <mps> now works with a relayed OpenAI Model Provider Service ({"openai": {"relayed": {}}}, created by the "ChatGPT/Codex Subscription" option in databricks-eng/universe#2891795). Codex runs on the developer's own ChatGPT subscription through AI Gateway, the same way relayed Claude Code uses a Claude subscription.

• Provider services: a relayed OpenAI service is usable only by codex, and a relayed Anthropic service only by claude, as before.
• Auth: the relayed provider block sets requires_openai_auth = true and has no ug auth-token command, so Codex sends its own ChatGPT login (Authorization + ChatGPT-Account-ID). ug never reads that token. Codex rejects an auth command combined with requires_openai_auth, so a relayed block can't keep the command.
• Proxy: the launch runs a loopback proxy (start_relay_proxy(..., codex=True), upstream ai-gateway/codex/, the plain handler). It adds the Databricks token in X-Databricks-AI-Gateway-Token on the /codex/v1/responses and /codex/v1/models swap routes (databricks-eng/universe#2853663, #2854855). Codex's -c provider points at the bound port, and the proxy can run alongside the OTLP proxy for one child.
• Encoding: the proxy asks the gateway for Accept-Encoding: identity when the client sent none. Otherwise httpx's default gzip request reaches the gateway, and the relayed raw body comes back compressed. Codex can't decode gzip and would print error bodies as binary.
• Sign-in: before launch, codex login status must report a ChatGPT login (an API key doesn't count); otherwise ug runs codex login.
• Managed file: a relayed configure never writes /etc/codex/managed_config.toml (or Windows requirements.toml). If that file defines model_provider or the Databricks provider's base_url/auth, which would beat the -c flags, configure fails and asks you to run ug revert or contact your admin, matching relayed Claude.
• Skipped for relay: the provider catalog fetch and OAUTH_TOKEN, since ug has no ChatGPT token and Codex lists models through the proxy. Smart routing with a relayed provider is refused.

The Codex desktop app half is in the next PR in this stack.

Testing

• New unit tests in tests/test_databricks.py, tests/test_gateway_proxy.py (including the real upstream join /v1/responses → {ws}/ai-gateway/codex/v1/responses) and tests/test_agent_codex.py (relayed overlay, stale auth pruned, managed conflicts and the relay-compatible scope, launch login and proxy wiring, no catalog fetch, smart-routing refusal).
• pytest tests/test_databricks.py tests/test_gateway_proxy.py tests/test_agent_codex.py tests/test_cli.py tests/test_agent_claude.py tests/test_agents_init.py with the ENABLE_SMART_ROUTING_*/SMART_ROUTER_NAME/UG_AUTO_UPGRADE_AGENT env vars unset: 1927 passed. just lint clean.
• Not yet run end to end against a real ChatGPT login and a relayed MPS. The staging gateway has enableOpenAiRelayedAuth on, but this machine has no ChatGPT login and its managed Codex file pins the default provider.

This pull request and its description were written by Isaac.

[andy-xu-db]

/ug-review

[github-actions]

Polly AI Review

Blocking issues

None substantiated.

Security vulnerabilities

None substantiated in the changed code. Native Codex authentication behavior was not exercised.

Non-blocking notes

• src/ucode/agents/codex.py:943: can we verify name-only Windows provider entries? The documented whole-entry replacement may defeat relay overrides, but this case is unverified.
• src/ucode/agents/codex.py:1483: can we extend the launch test to enable relay and OTLP together, including startup failure, and assert both endpoints and cleanup?
• tests/test_agent_codex.py:1686: can we model a successful login separately from failed status, cover login failure, and assert that the launch actually checks login?
• tests/test_gateway_proxy.py:614: can we fold the constructor assertions into test_codex_request_path_joins_upstream_once and remove the extra constructor mock?

Approach

The existing relay handler, token cache, and child-process lifecycle are appropriate building blocks. No broader refactor is needed. Managed-conflict detection should be checked against Windows’ documented provider-replacement behavior.

Summary

Diff size (computed): 8 files; +383 / -50 lines (433 changed text lines); 0 binary files.

Category	Files	Added	Deleted	Share of changed text lines
Tests and test support	3	191	0	44.1%
Documentation	1	1	1	0.5%
Dependencies and lockfiles	0	0	0	0.0%
Implementation / other	4	191	49	55.4%

Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements.

Changes

The change enables ug codex to use a relayed OpenAI provider with the developer’s ChatGPT subscription. It adds provider eligibility, subscription-login checks, relay configuration, proxy lifecycle handling, and managed-file conflict detection.

The implementation, tests, and ownership documentation support that outcome. No dependencies, lockfiles, generated files, or shared test fixtures change. The size is proportionate; consolidating the two proxy-construction scenarios would remove duplicate setup without reducing coverage.

Tests

Test-by-test assessment

All changed tests are additions. No existing tests or assertions are removed. The table also assesses the new inline setup.

Test / case	Behavior protected	Layer	Needed?	Action / rationale
tests/test_agent_codex.py::TestCodexLaunch::test_relayed_launch_checks_login_and_routes_through_relay_proxy	Relay URL and token wiring, argument forwarding, skipped catalog/OAUTH_TOKEN assignment, cleanup	Component	keep	Distinct launch contract. Its inline start_relay_proxy replacement captures explicit inputs without new shared scaffolding. Add a login spy; the current stub does not establish that login checking occurs.
tests/test_agent_codex.py::TestCodexLaunch::test_relayed_launch_logs_in_when_not_chatgpt	API-key status is rejected, ChatGPT status accepted, missing subscription login triggers login	Unit	keep	Appropriate subprocess boundary. Give status and login separate responses: returning code 1 from the mocked login despite check=True does not model a successful subprocess. Add a login-failure case.
tests/test_agent_codex.py::TestCodexLaunch::test_relayed_launch_rejects_smart_routing	Relay and smart routing cannot launch together	Unit	keep	Directly protects the launch guard. Existing CLI validation also covers the public entry point.
tests/test_agent_codex.py::TestCodexManagedConfig::test_relayed_writes_proxyless_provider_and_skips_managed_file	Stale auth removal, subscription-auth flag, no managed-file creation, relay state, return to nonrelay auth	Component	keep	Covers persisted configuration and a useful reverse transition. It does not exercise switching from a complete existing nonrelay managed configuration.
tests/test_agent_codex.py::TestCodexManagedConfig::test_relayed_rejects_managed_provider_override	Conflicting managed configuration produces recovery guidance	Component	keep	Protects the intended refusal contract. Extend this scenario once Windows name-only provider behavior is established.
tests/test_agent_codex.py::TestCodexManagedConfig::test_relayed_state_requires_relay_compatible_scope	Relay verification requests the dedicated compatibility scope	Unit	keep	Complements the writer’s scope assertion; checks the verification contract, not native configuration precedence.
tests/test_databricks.py::TestListModelProviderServices::test_flags_relayed_openai	OpenAI relay metadata is parsed correctly	Unit	keep	Direct parser coverage with minimal setup.
tests/test_databricks.py::TestServiceUsableForTool::test_openai_relay_is_codex_only	OpenAI relay is accepted for Codex and rejected for another tool; Anthropic relay is rejected for Codex	Unit	keep	Tests eligibility at the smallest reliable layer. Existing Anthropic/Claude cases cover the complementary dialect.
tests/test_gateway_proxy.py::TestCodexRelayProxy::test_codex_relay_upstream_and_plain_handler	Codex upstream, plain-handler selection, Databricks token header	Unit	consolidate	Move these assertions into the following real-construction scenario. Its inline fake_start constructor capture can then be removed.
tests/test_gateway_proxy.py::TestCodexRelayProxy::test_codex_request_path_joins_upstream_once	The actual HTTPX base URL joins /v1/responses without duplicating the prefix	Component	consolidate	Retain this scenario and absorb the constructor assertions. Inline _StubCache.run_refresher and .stop narrowly isolate construction from background token work.
tests/test_gateway_proxy.py::TestCodexRelayProxy::test_passes_chatgpt_headers_and_swaps_databricks_token	ChatGPT authorization/account headers survive while the Databricks token uses its separate header	Component	keep	Protects the combined credential boundary. It does not establish native Codex behavior or deployed gateway acceptance.

The most useful missing regression is a combined relay-plus-OTLP launch, including failure after the relay starts. Extend the existing launch component scenario rather than adding a separate E2E suite. The login test also needs realistic success and failure responses.

Existing shared proxy tests cover refresh, retries, streaming, rejection, and loopback transport; duplicating that matrix is unnecessary. A live check would uniquely establish native Codex configuration precedence and deployed subscription authentication.

No tests were executed during this review. The PR reports 1927 passing tests across the six named test files and a clean just lint; those results were not independently verified. Live ChatGPT/MPS verification remains outstanding.

For manual verification in an isolated profile, run ug codex --provider <relayed-mps> -- exec "Reply with RELAY_OK". Confirm successful inference and unchanged managed-file contents, then repeat with OTLP enabled.

Scope

No unrelated changes were found. Provider parsing, eligibility, configuration ownership, login handling, proxy lifecycle, tests, and documentation all contribute to CLI subscription relay support. Desktop implementation remains outside this diff.

---

_{Automated review by Polly · workflow run}