diff --git a/AGENTS.md b/AGENTS.md index 95047ce7d..933ea33d0 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.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/README.md b/README.md index 6959df57f..b101928bc 100644 --- a/README.md +++ b/README.md @@ -510,6 +510,14 @@ 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 + +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 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..da1163a95 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. @@ -427,8 +429,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 +584,10 @@ 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")) + 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 @@ -796,19 +801,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 +835,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 +843,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 +887,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,21 +915,22 @@ 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: 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( @@ -912,7 +940,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 +952,85 @@ 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: + """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). 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(): + return False + try: + doc = _read_app_config() + providers = doc.get("model_providers") + 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: + # 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() + return isinstance(current, dict) + 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() + 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, 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"] = 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_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 + 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..6a106554e 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,14 @@ 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: + # 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, 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. server.serve_forever() except (OSError, RuntimeError) as exc: @@ -4491,6 +4498,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/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 ddf523365..b90ab0087 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2345,6 +2345,111 @@ 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() + # 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()) + 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"]["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 + 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"]["DatabricksRelay"] == 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() + + # 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() + + # 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) + 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") + monkeypatch.setattr(codex, "_explicit_configure", True) + + 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..61a5578bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3837,6 +3837,30 @@ 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, + # 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, + ["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, "fallback_on_bind_error": False} + relay[0].serve_forever.assert_called_once() + otel[0].serve_forever.assert_called_once() + class TestDoctorCommand: def test_invokes_doctor(self):