From 0d3cde4a16cb5468ba2859ae1d288729641cf894 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 28 Oct 2026 16:31:31 +0000 Subject: [PATCH 1/5] Serve the ChatGPT subscription relay to the Codex app For a relayed OpenAI provider, ug's per-user login service (the Codex desktop OTLP proxy service) also runs the Codex relay proxy on a stable loopback port, and ~/.codex/config.toml points the Codex app's Databricks provider at it with requires_openai_auth and the MPS header. Codex's own ChatGPT login supplies Authorization. Only ug's own block is removed on switch-back or revert; a user-defined provider is left alone with a warning. Co-authored-by: Isaac --- AGENTS.md | 1 + README.md | 7 +++ src/ucode/agents/codex.py | 118 +++++++++++++++++++++++++++++++++----- src/ucode/cli.py | 14 +++-- tests/test_agent_codex.py | 80 ++++++++++++++++++++++++++ tests/test_cli.py | 21 +++++++ 6 files changed, 222 insertions(+), 19 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 95047ce7d..d8ae981fb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -184,6 +184,7 @@ Fields live in `~/.codex/ucode.config.toml` and `/etc/codex/managed_config.toml` | `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. 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_provider`, `model_providers.Databricks` (app) | Ignore | Ignore | In `~/.codex/config.toml`, only for a relayed (ChatGPT subscription) provider, for Codex desktop: a loopback relay block with the MPS header served by the login service; a user-defined provider is left alone and ug removes only its own block | | `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 | | `mcp_servers` | Ignore | Merge | Managed file; add/update the config's MCP server entries, other entries left alone | diff --git a/README.md b/README.md index 6959df57f..821d8acca 100644 --- a/README.md +++ b/README.md @@ -510,6 +510,13 @@ This policy is separate from agent-file backups and is not removed by `ug revert See [managed workspace policy](docs/managed-workspace-policy.md) for permissions and trade-offs. Managed policy headers must not contain secrets: the policy is readable by all local users. +### Codex desktop ChatGPT relay + +For a relayed OpenAI provider, `ug configure --agents codex --provider ` also points the Codex +app's `~/.codex/config.toml` at a loopback relay served by the same login service below (plain +`codex` then uses it too). Restart the app and sign in with ChatGPT. A provider you already set is +left alone; `ug revert` removes ug's entries. + ### Codex desktop tracing With workspace tracing enabled, run `ug configure --agents codex` on the machine diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 56648c48f..5f96c6fde 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -427,8 +427,8 @@ def revert_legacy_shared_config() -> bool: # Launch-scoped catalogs from _model_catalog_path; only the reverted ug profile names them. for path in CODEX_MODEL_CATALOG_PATH.parent.glob(f"{CODEX_MODEL_CATALOG_PATH.stem}-*.json"): path.unlink(missing_ok=True) - desktop_tracing_changed = _sync_desktop_tracing(None) - return legacy_changed or app_catalog_changed or desktop_tracing_changed + desktop_changed = _sync_desktop_tracing(None) | _sync_desktop_relay({}, None, None) + return legacy_changed or app_catalog_changed or desktop_changed def configured_paths(state: dict) -> list[str]: @@ -582,7 +582,9 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) # The desktop proxy reads authentication from the saved workspace. - _sync_desktop_tracing(workspace if state.get("codex_otel_tracing") else None) + tracing = bool(state.get("codex_otel_tracing")) + _sync_desktop_tracing(workspace if tracing else None) + _sync_desktop_relay(state, workspace, provider if relayed else None, tracing=tracing) return state @@ -796,19 +798,31 @@ def _owned_otel_service(path: Path) -> bool: return path.is_file() and _CODEX_OTEL_UNIT_MARKER in _read_otel_service(path) -def _ensure_otel_service(workspace: str, port: int) -> None: +def _ensure_otel_service(workspace: str, port: int, relay_port: int | None = None) -> None: path = _otel_service_path() if path.is_symlink(): raise RuntimeError(f"Refusing to replace a symlinked service file: {path}") if path.exists() and not _owned_otel_service(path): raise RuntimeError(f"Refusing to replace an unmanaged service file: {path}") + # Also serve the ChatGPT relay for the Codex app; before --port so its anchor still matches. + relay = ["--relay-port", str(relay_port)] if relay_port else [] if current_os() == OS.WINDOWS: # pythonw runs the proxy without a console window that closing would kill. # shortcut: the task gets the logon environment, not this shell's DATABRICKS_* overrides; # pass them through if a Windows user needs them. pythonw = Path(sys.executable).with_name("pythonw.exe") python = str(pythonw if pythonw.is_file() else sys.executable) - argv = [python, "-m", "ucode.cli", "otel-proxy", "--host", workspace, "--port", str(port)] + argv = [ + python, + "-m", + "ucode.cli", + "otel-proxy", + "--host", + workspace, + *relay, + "--port", + str(port), + ] _ensure_scheduled_task(path, argv) return ug_path = ug_binary() @@ -818,7 +832,7 @@ def _ensure_otel_service(workspace: str, port: int) -> None: ) if current_os() == OS.MACOS: _ensure_launch_agent( - path, [ug_path, "otel-proxy", "--host", workspace, "--port", str(port)] + path, [ug_path, "otel-proxy", "--host", workspace, *relay, "--port", str(port)] ) return lines = [ @@ -826,7 +840,7 @@ def _ensure_otel_service(workspace: str, port: int) -> None: "[Service]", # systemd expands `$` in ExecStart arguments but never in the executable path. f"ExecStart={_systemd_quote(ug_path)} otel-proxy --host " - f"{_systemd_quote(workspace, exec_start=True)} --port {port}", + f"{_systemd_quote(workspace, exec_start=True)} {' '.join([*relay, '--port', str(port)])}", "Restart=on-failure", "RestartSec=1s", ] @@ -870,6 +884,16 @@ def _remove_otel_service() -> None: _run_systemctl(["daemon-reload"]) +def _service_port(flag: str) -> int | None: + """The ``--{flag}`` port recorded in ug's own login service unit, if there is one.""" + unit = _otel_service_path() + text = _read_otel_service(unit) if _owned_otel_service(unit) else "" + match = re.search(rf" --{flag} (\d+)(?= --|$)", text, re.MULTILINE) + # A plist puts the port in its own ; task XML ends it with . + match = match or re.search(rf"--{flag}\D+(\d+)", text) + return int(match[1]) if match else None + + def _desktop_trace_exporter(port: int) -> dict: return _otel_proxy_overlay(f"http://{LOOPBACK_HOST}:{port}/v1/traces")["otel"]["trace_exporter"] @@ -888,15 +912,12 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: if not isinstance(otel, dict): raise RuntimeError(f"otel in {_legacy_config_path()} is not a table") # ug's own unit records the port, so only an exporter at that port is ug's. - unit = _otel_service_path() - unit_text = _read_otel_service(unit) if _owned_otel_service(unit) else "" - match = re.search(r" --port (\d+)$", unit_text, re.MULTILINE) - # A plist puts the port in its own ; task XML ends it with . - match = match or re.search(r"--port\D+(\d+)", unit_text) - port = int(match[1]) if match else None + port = _service_port("port") ours = port is not None and otel.get("trace_exporter") == _desktop_trace_exporter(port) + relay_port = _service_port("relay-port") if workspace is None: - _remove_otel_service() + if not relay_port: # the Codex app's relay still needs the service + _remove_otel_service() if ours: del otel["trace_exporter"] if not otel: @@ -912,7 +933,7 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: return False port = port or _allocate_otel_port() # Before the exporter, so a missing systemd user manager leaves Codex settings untouched. - _ensure_otel_service(workspace, port) + _ensure_otel_service(workspace, port, relay_port) if ours: return False otel["trace_exporter"] = _desktop_trace_exporter(port) @@ -924,6 +945,73 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: return False +def _sync_desktop_relay( + state: dict, workspace: str | None, provider: str | None, *, tracing: bool = False +) -> bool: + """Point the Codex app at ug's login service as a ChatGPT relay for ``provider``; None removes. + + It goes in the Codex App settings (``ug codex`` sessions' own ``-c`` provider outranks it and + keeps a per-session proxy). Only a provider at the service's recorded relay port is ug's. + Returns whether that file changed; setup problems only warn. + """ + if current_os() == OS.OTHER or is_dry_run(): + return False + try: + doc = _read_app_config() + providers = doc.get("model_providers") + current = providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, dict) else None + relay_port = _service_port("relay-port") + base_url = f"http://{LOOPBACK_HOST}:{relay_port}/v1" + ours = ( + isinstance(current, dict) + and relay_port is not None + and current.get("base_url") == base_url + ) + if provider is None or workspace is None: + if tracing and workspace and relay_port: + _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port()) + elif not tracing: + _remove_otel_service() + if ours and isinstance(providers, dict): + del providers[CODEX_MODEL_PROVIDER_NAME] + if not providers: + del doc["model_providers"] + if doc.get("model_provider") == CODEX_MODEL_PROVIDER_NAME: + del doc["model_provider"] + write_toml_file(_legacy_config_path(), doc) + return ours + if (current is not None and not ours) or doc.get( + "model_provider", CODEX_MODEL_PROVIDER_NAME + ) != (CODEX_MODEL_PROVIDER_NAME): + print_warning_err( + f"{_legacy_config_path()} already sets a model provider; skipped the Codex app " + "ChatGPT relay." + ) + return False + relay_port = relay_port or _allocate_otel_port() + base_url = f"http://{LOOPBACK_HOST}:{relay_port}/v1" + # Before the config, so a missing systemd user manager leaves Codex settings untouched. + _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port(), relay_port) + block = _provider_block( + workspace, + state.get("profile"), + provider=provider, + managed_http_headers=state.get("codex_http_headers"), + relayed=True, + ) + block["base_url"] = base_url + if doc.get("model_provider") == CODEX_MODEL_PROVIDER_NAME and current == block: + return False + doc["model_provider"] = CODEX_MODEL_PROVIDER_NAME + doc.setdefault("model_providers", {})[CODEX_MODEL_PROVIDER_NAME] = block + write_toml_file(_legacy_config_path(), doc) + print_warning_err("Restart the Codex app and sign in with ChatGPT to use the relay.") + return True + except (OSError, RuntimeError, ValueError) as exc: + print_warning_err(f"Could not update the Codex app ChatGPT relay: {exc}") + return False + + 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) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2d32b7010..0cb546518 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -7,6 +7,7 @@ import os import shutil import subprocess +import threading from collections.abc import Iterator from contextlib import contextmanager, redirect_stdout from enum import StrEnum @@ -2604,8 +2605,9 @@ def auth_token_cmd( def otel_proxy_cmd( host: Annotated[str, typer.Option("--host")], port: Annotated[int, typer.Option("--port", min=1, max=65535)], + relay_port: Annotated[int | None, typer.Option("--relay-port", min=1, max=65535)] = None, ) -> None: - """Run the same OTLP proxy as a Codex CLI session, for desktop.""" + """Run the same OTLP proxy (and, with --relay-port, ChatGPT relay) as Codex CLI, for desktop.""" from ucode import gateway_proxy try: @@ -2613,9 +2615,13 @@ def otel_proxy_cmd( if not state: raise RuntimeError(f"No saved configuration for {host}. Run `ug configure` first.") apply_pat_environment(state) - server, _cache, _client = gateway_proxy.start_otel_proxy( - host, codex_agent._otel_token_provider(state, host), port=port - ) + token_provider = codex_agent._otel_token_provider(state, host) + server, _cache, _client = gateway_proxy.start_otel_proxy(host, token_provider, port=port) + if relay_port: + relay, _cache, _client = gateway_proxy.start_relay_proxy( + host, token_provider, relay_port, codex=True + ) + threading.Thread(target=relay.serve_forever, daemon=True).start() # systemd stops this with SIGTERM, which exits without cleanup, and nothing needs it. server.serve_forever() except (OSError, RuntimeError) as exc: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index ddf523365..06fea26f9 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2345,6 +2345,86 @@ def test_desktop_tracing_lifecycle(self, tmp_path, monkeypatch): assert "otel" not in read_toml_safe(app_path) assert not unit_path.exists() + @pytest.mark.skipif(sys.platform != "linux", reason="Linux systemd configuration") + def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): + managed_path, _ = self._sudo_counting_env(tmp_path, monkeypatch) + app_path = codex._legacy_config_path() + app_path.parent.mkdir(parents=True, exist_ok=True) + unit_path = codex._otel_service_path() + ports = iter(range(43198, 43298)) + monkeypatch.setattr(codex, "_allocate_otel_port", lambda: next(ports)) + monkeypatch.setattr(codex, "ug_binary", lambda: "/opt/ug") + monkeypatch.setattr(codex.subprocess_cross_os, "run", Mock()) + relayed = {"workspace": WS} + + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + provider = read_toml_safe(app_path)["model_providers"]["Databricks"] + assert read_toml_safe(app_path)["model_provider"] == "Databricks" + assert provider["base_url"] == "http://127.0.0.1:43198/v1" + assert provider["requires_openai_auth"] is True + assert "auth" not in provider + assert provider["http_headers"][codex.MODEL_PROVIDER_SERVICE_HEADER] == "cat.s.mps" + assert f'--host "{WS}" --relay-port 43198 --port 43199' in unit_path.read_text() + + # The ports are stable across reconfigures, and tracing shares the one service. + codex.write_tool_config( + {**relayed, "codex_otel_tracing": True}, provider="cat.s.mps", relayed=True + ) + assert read_toml_safe(app_path)["model_providers"]["Databricks"] == provider + assert "--relay-port 43198 --port 43199" in unit_path.read_text() + assert read_toml_safe(app_path)["otel"]["trace_exporter"] + + # Dropping relay keeps the service for tracing and removes only ug's provider entries. + codex.write_tool_config({**relayed, "codex_otel_tracing": True}) + doc = read_toml_safe(app_path) + assert "model_provider" not in doc and "model_providers" not in doc + assert "--relay-port" not in unit_path.read_text() + assert "--port 43199" in unit_path.read_text() + + managed_path.unlink(missing_ok=True) # a non-relayed run owns it; relay refuses that file + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + codex.write_tool_config(dict(relayed)) + assert not unit_path.exists() + assert "model_providers" not in read_toml_safe(app_path) + + managed_path.unlink(missing_ok=True) + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + assert codex.revert_legacy_shared_config() is True + assert "model_providers" not in read_toml_safe(app_path) + assert not unit_path.exists() + + @pytest.mark.skipif(sys.platform != "linux", reason="Linux systemd configuration") + def test_desktop_relay_preserves_user_provider_and_dry_run(self, tmp_path, monkeypatch, capsys): + self._sudo_counting_env(tmp_path, monkeypatch) + app_path = codex._legacy_config_path() + app_path.parent.mkdir(parents=True, exist_ok=True) + user = 'model_provider = "mine"\n[model_providers.mine]\nbase_url = "https://x"\n' + app_path.write_text(user) + monkeypatch.setattr(codex.subprocess_cross_os, "run", Mock()) + monkeypatch.setattr(codex, "ug_binary", lambda: "/opt/ug") + + codex.write_tool_config({"workspace": WS}, provider="cat.s.mps", relayed=True) + assert app_path.read_text() == user + assert "model provider" in capsys.readouterr().err + assert not codex._otel_service_path().exists() + + app_path.write_text("") + monkeypatch.setattr(codex, "is_dry_run", lambda: True) + assert codex._sync_desktop_relay({}, WS, "cat.s.mps") is False + assert app_path.read_text() == "" + + def test_desktop_relay_service_argv_includes_relay_port(self, tmp_path, monkeypatch): + self._sudo_counting_env(tmp_path, monkeypatch) + ensure = Mock() + monkeypatch.setattr(codex, "_ensure_launch_agent", ensure) + monkeypatch.setattr(codex, "_ensure_scheduled_task", ensure) + monkeypatch.setattr(codex, "ug_binary", lambda: str(tmp_path / "ug")) + for target in (managed_files.OS.MACOS, managed_files.OS.WINDOWS): + monkeypatch.setattr(codex, "current_os", lambda target=target: target) + codex._ensure_otel_service(WS, 43199, 43198) + argv = ensure.call_args.args[1] + assert argv[argv.index("--host") + 2 :] == ["--relay-port", "43198", "--port", "43199"] + def test_desktop_tracing_launch_agent_lifecycle(self, tmp_path, monkeypatch): self._sudo_counting_env(tmp_path, monkeypatch) app_path = codex._legacy_config_path() diff --git a/tests/test_cli.py b/tests/test_cli.py index 4c4c90ba0..04a73e71d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3837,6 +3837,27 @@ def test_reuses_cli_proxy_and_auth_for_the_configured_workspace(self): token.assert_called_once_with(workspace, "saved-profile", force_refresh=True) server.serve_forever.assert_called_once() + def test_relay_port_also_starts_the_codex_relay_proxy(self): + workspace = "https://configured.example" + real_save_state({"workspace": workspace, "profile": "saved-profile"}) + otel, relay = ( + (MagicMock(), MagicMock(), MagicMock()), + (MagicMock(), MagicMock(), MagicMock()), + ) + with ( + patch("ucode.gateway_proxy.start_otel_proxy", return_value=otel), + patch("ucode.gateway_proxy.start_relay_proxy", return_value=relay) as start_relay, + ): + result = runner.invoke( + app, + ["otel-proxy", "--host", workspace, "--relay-port", "43198", "--port", "43199"], + ) + assert result.exit_code == 0, result.output + assert start_relay.call_args.args[0] == workspace + assert start_relay.call_args.args[2] == 43198 + assert start_relay.call_args.kwargs == {"codex": True} + otel[0].serve_forever.assert_called_once() + class TestDoctorCommand: def test_invokes_doctor(self): From ce04d8f1e03d0ff97865b2a9d5665df9136b4b5d Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 29 Oct 2026 18:51:08 +0000 Subject: [PATCH 2/5] Fix the Codex app relay README command Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 821d8acca..78611e52c 100644 --- a/README.md +++ b/README.md @@ -512,9 +512,9 @@ Managed policy headers must not contain secrets: the policy is readable by all l ### Codex desktop ChatGPT relay -For a relayed OpenAI provider, `ug configure --agents codex --provider ` also points the Codex -app's `~/.codex/config.toml` at a loopback relay served by the same login service below (plain -`codex` then uses it too). Restart the app and sign in with ChatGPT. A provider you already set is +When Codex uses a relayed OpenAI provider (picked in `ug configure`, or `ug codex --provider `), +`ug` also points the Codex app's `~/.codex/config.toml` at a loopback relay served by the same login +service below (plain `codex` then uses it too). Restart the app and sign in with ChatGPT. A provider you already set is left alone; `ug revert` removes ug's entries. ### Codex desktop tracing From 928357fb6e80ad84a7d79f6512be78fc169a0c02 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 30 Oct 2026 14:38:11 +0000 Subject: [PATCH 3/5] Set up the Codex app relay only from ug configure, like Claude Desktop A launch's --provider no longer rewrites ~/.codex/config.toml or the login service; only an explicit ug configure sets up or removes the app's ChatGPT relay. Co-authored-by: Isaac --- AGENTS.md | 2 +- README.md | 9 +++++---- src/ucode/agents/codex.py | 13 ++++++++++++- src/ucode/cli.py | 1 + tests/test_agent_codex.py | 5 +++++ 5 files changed, 24 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index d8ae981fb..f7c9bbd62 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -184,7 +184,7 @@ Fields live in `~/.codex/ucode.config.toml` and `/etc/codex/managed_config.toml` | `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. 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_provider`, `model_providers.Databricks` (app) | Ignore | Ignore | In `~/.codex/config.toml`, only for a relayed (ChatGPT subscription) provider, for Codex desktop: a loopback relay block with the MPS header served by the login service; a user-defined provider is left alone and ug removes only its own block | +| `model_provider`, `model_providers.Databricks` (app) | Ignore | Ignore | In `~/.codex/config.toml`, only for a relayed (ChatGPT subscription) provider, for Codex desktop, and only by `ug configure` (like Claude Desktop): a loopback relay block with the MPS header served by the login service; a user-defined provider is left alone and ug removes only its own block | | `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 | | `mcp_servers` | Ignore | Merge | Managed file; add/update the config's MCP server entries, other entries left alone | diff --git a/README.md b/README.md index 78611e52c..b101928bc 100644 --- a/README.md +++ b/README.md @@ -512,10 +512,11 @@ Managed policy headers must not contain secrets: the policy is readable by all l ### Codex desktop ChatGPT relay -When Codex uses a relayed OpenAI provider (picked in `ug configure`, or `ug codex --provider `), -`ug` also points the Codex app's `~/.codex/config.toml` at a loopback relay served by the same login -service below (plain `codex` then uses it too). Restart the app and sign in with ChatGPT. A provider you already set is -left alone; `ug revert` removes ug's entries. +When `ug configure` sets up Codex with a relayed OpenAI provider, it also points the Codex app's +`~/.codex/config.toml` at a loopback relay served by the same login service below (plain `codex` +then uses it too). Like Claude Desktop, `ug codex --provider` doesn't change the app. Restart the +app and sign in with ChatGPT. A provider you already set is left alone; `ug revert` removes ug's +entries. ### Codex desktop tracing diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 5f96c6fde..428154aa5 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -584,7 +584,8 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: # The desktop proxy reads authentication from the saved workspace. tracing = bool(state.get("codex_otel_tracing")) _sync_desktop_tracing(workspace if tracing else None) - _sync_desktop_relay(state, workspace, provider if relayed else None, tracing=tracing) + if _explicit_configure: # like Claude Desktop, only `ug configure` changes the app + _sync_desktop_relay(state, workspace, provider if relayed else None, tracing=tracing) return state @@ -945,6 +946,16 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: return False +_explicit_configure = False + + +def request_explicit_configure(requested: bool = True) -> None: + """Mark this run as an explicit ``ug configure`` (``False`` clears it), the only run that sets + up or removes the Codex app's ChatGPT relay; a launch's ``--provider`` never changes the app.""" + global _explicit_configure + _explicit_configure = requested + + def _sync_desktop_relay( state: dict, workspace: str | None, provider: str | None, *, tracing: bool = False ) -> bool: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0cb546518..7269def9d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -4497,6 +4497,7 @@ def configure( scope = "user" if user_scope else "global" set_config_scope(scope) claude_agent.request_explicit_configure() + codex_agent.request_explicit_configure() custom_oauth = _custom_oauth_config(client_id, redirect_url, scopes) if custom_oauth is not None and use_pat: raise RuntimeError("--client-id cannot be combined with --use-pat.") diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 06fea26f9..eb8f3e26a 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2357,6 +2357,10 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): monkeypatch.setattr(codex.subprocess_cross_os, "run", Mock()) relayed = {"workspace": WS} + # Like Claude Desktop, a launch's --provider never changes the app; only `ug configure` does. + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + assert not app_path.exists() and not unit_path.exists() + monkeypatch.setattr(codex, "_explicit_configure", True) codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) provider = read_toml_safe(app_path)["model_providers"]["Databricks"] assert read_toml_safe(app_path)["model_provider"] == "Databricks" @@ -2402,6 +2406,7 @@ def test_desktop_relay_preserves_user_provider_and_dry_run(self, tmp_path, monke app_path.write_text(user) monkeypatch.setattr(codex.subprocess_cross_os, "run", Mock()) monkeypatch.setattr(codex, "ug_binary", lambda: "/opt/ug") + monkeypatch.setattr(codex, "_explicit_configure", True) codex.write_tool_config({"workspace": WS}, provider="cat.s.mps", relayed=True) assert app_path.read_text() == user From 8b518ef967d867aeb4bd2373ddf49aa6e456516a Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 30 Oct 2026 15:28:07 +0000 Subject: [PATCH 4/5] Give the Codex app relay its own provider name The app merges /etc/codex/managed_config.toml over ~/.codex/config.toml key by key, so a later non-relayed run's managed Databricks auth merged into the relay entry and Codex refused to start (auth plus requires_openai_auth). The relay now uses DatabricksRelay, which nothing merges into, and ug owns it by name, like Claude Desktop's reserved profile, so configure and revert can replace or remove it even after the service's port record is lost. Co-authored-by: Isaac --- AGENTS.md | 2 +- src/ucode/agents/codex.py | 48 +++++++++++++++++++-------------------- tests/test_agent_codex.py | 15 +++++++++--- 3 files changed, 37 insertions(+), 28 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f7c9bbd62..933ea33d0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -184,7 +184,7 @@ Fields live in `~/.codex/ucode.config.toml` and `/etc/codex/managed_config.toml` | `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. 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_provider`, `model_providers.Databricks` (app) | Ignore | Ignore | In `~/.codex/config.toml`, only for a relayed (ChatGPT subscription) provider, for Codex desktop, and only by `ug configure` (like Claude Desktop): a loopback relay block with the MPS header served by the login service; a user-defined provider is left alone and ug removes only its own block | +| `model_provider`, `model_providers.DatabricksRelay` (app) | Ignore | Ignore | In `~/.codex/config.toml`, only for a relayed (ChatGPT subscription) provider, for Codex desktop, and only by `ug configure` (like Claude Desktop): a loopback relay block with the MPS header served by the login service. Its own name keeps a managed `Databricks` provider from merging into it, and marks it as ug's to replace or remove; a user-defined `model_provider` is left alone | | `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 | | `mcp_servers` | Ignore | Merge | Managed file; add/update the config's MCP server entries, other entries left alone | diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 428154aa5..b8a7a3562 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -110,6 +110,8 @@ LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "Databricks" LEGACY_CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" +# The Codex app relay's own provider in ~/.codex/config.toml, apart from the managed `Databricks`. +CODEX_APP_RELAY_PROVIDER_NAME = "DatabricksRelay" # 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. @@ -962,7 +964,8 @@ def _sync_desktop_relay( """Point the Codex app at ug's login service as a ChatGPT relay for ``provider``; None removes. It goes in the Codex App settings (``ug codex`` sessions' own ``-c`` provider outranks it and - keeps a per-session proxy). Only a provider at the service's recorded relay port is ug's. + keeps a per-session proxy). Its own provider name is ug's, like Claude Desktop's reserved + profile, so a managed ``Databricks`` provider never merges into it. Returns whether that file changed; setup problems only warn. """ if current_os() == OS.OTHER or is_dry_run(): @@ -970,37 +973,34 @@ def _sync_desktop_relay( try: doc = _read_app_config() providers = doc.get("model_providers") - current = providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, dict) else None - relay_port = _service_port("relay-port") - base_url = f"http://{LOOPBACK_HOST}:{relay_port}/v1" - ours = ( - isinstance(current, dict) - and relay_port is not None - and current.get("base_url") == base_url + current = ( + providers.get(CODEX_APP_RELAY_PROVIDER_NAME) if isinstance(providers, dict) else None ) + relay_port = _service_port("relay-port") if provider is None or workspace is None: if tracing and workspace and relay_port: _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port()) elif not tracing: _remove_otel_service() - if ours and isinstance(providers, dict): - del providers[CODEX_MODEL_PROVIDER_NAME] - if not providers: - del doc["model_providers"] - if doc.get("model_provider") == CODEX_MODEL_PROVIDER_NAME: - del doc["model_provider"] - write_toml_file(_legacy_config_path(), doc) - return ours - if (current is not None and not ours) or doc.get( - "model_provider", CODEX_MODEL_PROVIDER_NAME - ) != (CODEX_MODEL_PROVIDER_NAME): + if current is None or not isinstance(providers, dict): + return False + del providers[CODEX_APP_RELAY_PROVIDER_NAME] + if not providers: + del doc["model_providers"] + if doc.get("model_provider") == CODEX_APP_RELAY_PROVIDER_NAME: + del doc["model_provider"] + write_toml_file(_legacy_config_path(), doc) + return True + if ( + doc.get("model_provider", CODEX_APP_RELAY_PROVIDER_NAME) + != CODEX_APP_RELAY_PROVIDER_NAME + ): print_warning_err( f"{_legacy_config_path()} already sets a model provider; skipped the Codex app " "ChatGPT relay." ) return False relay_port = relay_port or _allocate_otel_port() - base_url = f"http://{LOOPBACK_HOST}:{relay_port}/v1" # Before the config, so a missing systemd user manager leaves Codex settings untouched. _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port(), relay_port) block = _provider_block( @@ -1010,11 +1010,11 @@ def _sync_desktop_relay( managed_http_headers=state.get("codex_http_headers"), relayed=True, ) - block["base_url"] = base_url - if doc.get("model_provider") == CODEX_MODEL_PROVIDER_NAME and current == block: + block["base_url"] = f"http://{LOOPBACK_HOST}:{relay_port}/v1" + if doc.get("model_provider") == CODEX_APP_RELAY_PROVIDER_NAME and current == block: return False - doc["model_provider"] = CODEX_MODEL_PROVIDER_NAME - doc.setdefault("model_providers", {})[CODEX_MODEL_PROVIDER_NAME] = block + doc["model_provider"] = CODEX_APP_RELAY_PROVIDER_NAME + doc.setdefault("model_providers", {})[CODEX_APP_RELAY_PROVIDER_NAME] = block write_toml_file(_legacy_config_path(), doc) print_warning_err("Restart the Codex app and sign in with ChatGPT to use the relay.") return True diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index eb8f3e26a..4c5e078f0 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2362,8 +2362,8 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): assert not app_path.exists() and not unit_path.exists() monkeypatch.setattr(codex, "_explicit_configure", True) codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) - provider = read_toml_safe(app_path)["model_providers"]["Databricks"] - assert read_toml_safe(app_path)["model_provider"] == "Databricks" + provider = read_toml_safe(app_path)["model_providers"]["DatabricksRelay"] + assert read_toml_safe(app_path)["model_provider"] == "DatabricksRelay" assert provider["base_url"] == "http://127.0.0.1:43198/v1" assert provider["requires_openai_auth"] is True assert "auth" not in provider @@ -2374,7 +2374,7 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): codex.write_tool_config( {**relayed, "codex_otel_tracing": True}, provider="cat.s.mps", relayed=True ) - assert read_toml_safe(app_path)["model_providers"]["Databricks"] == provider + assert read_toml_safe(app_path)["model_providers"]["DatabricksRelay"] == provider assert "--relay-port 43198 --port 43199" in unit_path.read_text() assert read_toml_safe(app_path)["otel"]["trace_exporter"] @@ -2397,6 +2397,15 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): assert "model_providers" not in read_toml_safe(app_path) assert not unit_path.exists() + # The name marks the entry as ug's, so it is replaced even after its service is lost. + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + unit_path.unlink() + codex.write_tool_config(dict(relayed), provider="cat.s.mps", relayed=True) + relay = read_toml_safe(app_path)["model_providers"]["DatabricksRelay"] + port = relay["base_url"].removeprefix("http://127.0.0.1:").removesuffix("/v1") + assert relay["base_url"] != provider["base_url"] + assert f"--relay-port {port} " in unit_path.read_text() + @pytest.mark.skipif(sys.platform != "linux", reason="Linux systemd configuration") def test_desktop_relay_preserves_user_provider_and_dry_run(self, tmp_path, monkeypatch, capsys): self._sudo_counting_env(tmp_path, monkeypatch) From c8d77b66ccea53f3393ef24225b508b352204890 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 30 Oct 2026 17:42:28 +0000 Subject: [PATCH 5/5] Address app relay review: keep the relay on its workspace and port - A traced launch in another workspace no longer moves the login service (and so the app's relay, whose provider headers name the configured MPS) to that workspace; configure syncs the relay before tracing so it can still move it. - The relay proxy no longer falls back to a random port: the app is pinned to the configured one, so the service fails and restarts instead. - Retry a repeated port allocation so the relay and OTLP proxies never share a port. - Removal writes the app config before touching the service, so a failed write leaves the relay served. Co-authored-by: Isaac --- src/ucode/agents/codex.py | 30 +++++++++++++++++++----------- src/ucode/cli.py | 3 ++- src/ucode/gateway_proxy.py | 2 ++ tests/test_agent_codex.py | 13 ++++++++++++- tests/test_cli.py | 5 ++++- 5 files changed, 39 insertions(+), 14 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index b8a7a3562..da1163a95 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -585,9 +585,9 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: save_state(state) # The desktop proxy reads authentication from the saved workspace. tracing = bool(state.get("codex_otel_tracing")) - _sync_desktop_tracing(workspace if tracing else None) if _explicit_configure: # like Claude Desktop, only `ug configure` changes the app _sync_desktop_relay(state, workspace, provider if relayed else None, tracing=tracing) + _sync_desktop_tracing(workspace if tracing else None) return state @@ -927,6 +927,10 @@ def _sync_desktop_tracing(workspace: str | None) -> bool: del doc["otel"] write_toml_file(_legacy_config_path(), doc) return ours + # The app's relay stays on the workspace `ug configure` set it up for; a launch elsewhere + # must not move the service (and the app's MPS headers) to another workspace. + if relay_port and workspace not in _read_otel_service(_otel_service_path()): + return False # shortcut: a managed-file exporter still outranks this one, check it if admins set one. if "trace_exporter" in otel and not ours: print_warning_err( @@ -978,19 +982,19 @@ def _sync_desktop_relay( ) relay_port = _service_port("relay-port") if provider is None or workspace is None: + # Config first, so a failed write leaves the app's relay still served. + if isinstance(current, dict) and isinstance(providers, dict): + del providers[CODEX_APP_RELAY_PROVIDER_NAME] + if not providers: + del doc["model_providers"] + if doc.get("model_provider") == CODEX_APP_RELAY_PROVIDER_NAME: + del doc["model_provider"] + write_toml_file(_legacy_config_path(), doc) if tracing and workspace and relay_port: _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port()) elif not tracing: _remove_otel_service() - if current is None or not isinstance(providers, dict): - return False - del providers[CODEX_APP_RELAY_PROVIDER_NAME] - if not providers: - del doc["model_providers"] - if doc.get("model_provider") == CODEX_APP_RELAY_PROVIDER_NAME: - del doc["model_provider"] - write_toml_file(_legacy_config_path(), doc) - return True + return isinstance(current, dict) if ( doc.get("model_provider", CODEX_APP_RELAY_PROVIDER_NAME) != CODEX_APP_RELAY_PROVIDER_NAME @@ -1001,8 +1005,12 @@ def _sync_desktop_relay( ) return False relay_port = relay_port or _allocate_otel_port() + port = _service_port("port") + # Each allocation releases its port, so a second one can repeat the first; retry it. + while not port or port == relay_port: + port = _allocate_otel_port() # Before the config, so a missing systemd user manager leaves Codex settings untouched. - _ensure_otel_service(workspace, _service_port("port") or _allocate_otel_port(), relay_port) + _ensure_otel_service(workspace, port, relay_port) block = _provider_block( workspace, state.get("profile"), diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 7269def9d..6a106554e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2618,8 +2618,9 @@ def otel_proxy_cmd( token_provider = codex_agent._otel_token_provider(state, host) server, _cache, _client = gateway_proxy.start_otel_proxy(host, token_provider, port=port) if relay_port: + # The Codex app is pinned to this port, so fail (and restart) rather than move. relay, _cache, _client = gateway_proxy.start_relay_proxy( - host, token_provider, relay_port, codex=True + host, token_provider, relay_port, codex=True, fallback_on_bind_error=False ) threading.Thread(target=relay.serve_forever, daemon=True).start() # systemd stops this with SIGTERM, which exits without cleanup, and nothing needs it. diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 10808e32b..95592cd10 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -646,6 +646,7 @@ def start_relay_proxy( port: int, *, codex: bool = False, + fallback_on_bind_error: bool = True, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the subscription relay proxy: Claude's, or Codex's with ``codex`` (plain forwarding, since the Anthropic hybrid reroute doesn't apply).""" @@ -657,6 +658,7 @@ def start_relay_proxy( token_header=AI_GATEWAY_TOKEN_HEADER, handler_type=_ProxyHandler if codex else _RelayProxyHandler, force_refresh_near_expiry=False, + fallback_on_bind_error=fallback_on_bind_error, ) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 4c5e078f0..b90ab0087 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2351,7 +2351,8 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): app_path = codex._legacy_config_path() app_path.parent.mkdir(parents=True, exist_ok=True) unit_path = codex._otel_service_path() - ports = iter(range(43198, 43298)) + # A repeated allocation is retried, so the service never gets the same port twice. + ports = iter([43198, 43198, *range(43199, 43298)]) monkeypatch.setattr(codex, "_allocate_otel_port", lambda: next(ports)) monkeypatch.setattr(codex, "ug_binary", lambda: "/opt/ug") monkeypatch.setattr(codex.subprocess_cross_os, "run", Mock()) @@ -2406,6 +2407,16 @@ def test_desktop_relay_lifecycle(self, tmp_path, monkeypatch): assert relay["base_url"] != provider["base_url"] assert f"--relay-port {port} " in unit_path.read_text() + # A traced launch in another workspace leaves the service, and so the relay, on this one. + monkeypatch.setattr(codex, "_explicit_configure", False) + codex.write_tool_config({"workspace": "https://other.example", "codex_otel_tracing": True}) + assert f'--host "{WS}" --relay-port {port} ' in unit_path.read_text() + + # Removal writes the config first, so a failed write keeps the relay served. + monkeypatch.setattr(codex, "write_toml_file", Mock(side_effect=OSError("denied"))) + assert codex._sync_desktop_relay({}, WS, None) is False + assert f"--relay-port {port} " in unit_path.read_text() + @pytest.mark.skipif(sys.platform != "linux", reason="Linux systemd configuration") def test_desktop_relay_preserves_user_provider_and_dry_run(self, tmp_path, monkeypatch, capsys): self._sudo_counting_env(tmp_path, monkeypatch) diff --git a/tests/test_cli.py b/tests/test_cli.py index 04a73e71d..61a5578bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3847,6 +3847,8 @@ def test_relay_port_also_starts_the_codex_relay_proxy(self): with ( patch("ucode.gateway_proxy.start_otel_proxy", return_value=otel), patch("ucode.gateway_proxy.start_relay_proxy", return_value=relay) as start_relay, + # Run the relay thread's target inline so its serve_forever call is deterministic. + patch("ucode.cli.threading.Thread", lambda target, daemon: MagicMock(start=target)), ): result = runner.invoke( app, @@ -3855,7 +3857,8 @@ def test_relay_port_also_starts_the_codex_relay_proxy(self): assert result.exit_code == 0, result.output assert start_relay.call_args.args[0] == workspace assert start_relay.call_args.args[2] == 43198 - assert start_relay.call_args.kwargs == {"codex": True} + assert start_relay.call_args.kwargs == {"codex": True, "fallback_on_bind_error": False} + relay[0].serve_forever.assert_called_once() otel[0].serve_forever.assert_called_once()