From 4be807bbd8741acb00d80c13c71e432a3aa234b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 15:44:11 +0000 Subject: [PATCH] Studio: round 4 hardening for the Codex provider Fourth reviewer.py pass surfaced one more security finding and a handful of correctness gaps. Each is small but the env-scrub for the SDK path closes the asymmetric-fix loop opened in the previous round. * Codex SDK construction now passes an `AppServerConfig(env=...)` that overrides every non-safe-listed env key to an empty string. Upstream openai/codex/sdk/python/client.py builds the spawn env as `os.environ.copy()` then `env.update(self.config.env)`, so this scrubs HF_TOKEN / GH_TOKEN / WANDB_API_KEY / ANTHROPIC_API_KEY etc. out of the codex app-server subprocess env on the chat / parallel / synthesis paths, matching the CLI/login paths from the previous round. The helper falls back to bare `AsyncCodex()` when the SDK version does not expose AppServerConfig, with logged warning. * Install hint now names the actual upstream PyPI project, `openai-codex` (canonical), with `codex_app_server` documented as the legacy alias. The probe still accepts both import names so forward compat is preserved. * Device-auth URL regex broadened to accept upstream's current `chatgpt.com/activate` shape and any `/device|/activate|/verify` variant, not just `/codex/device`. The frontend can now open the verification page on CLI builds that print the documented ChatGPT-style URL. * `_run_codex_synthesis` now takes a `system` arg and forwards it to `thread_start(system=...)`, falling back to a prompt-prefix on older SDK revs that reject the kwarg. Previously a fan-out with "Always answer in Spanish" produced Spanish per-tab attempts but an English synthesis. * `_detect_logged_in` negative regex now also matches "Not signed in", "Please sign in" (alternative localisations / future CLI releases). Same word-boundary anchoring as before. * Frontend `CodexLoginEvent` union gains `device_code` and a `code` field. `CodexLoginButton` now renders the one-time code under the verification URL so users on a headless / remote install can copy the code without scraping the log pane. Also fixes a closure-stale bug where setError(message) was followed by a stale `error` read, losing specific backend errors; the new path keeps `lastStreamError` inside the closure. * Replaced four hardcoded `/mnt/disks/...` paths in the new regression tests with `_backend_file()` resolved from `__file__`, so the suite runs in any checkout (CI, local dev, the review worker tree). Found by the round-4 reviewer. Four new pytest cases pin the behaviour: `test_not_signed_in_wording_also_handled`, `test_device_url_accepts_generic_verification_url`, `test_synthesis_call_forwards_system_prompt`, and `test_sdk_env_scrubbed_via_appserverconfig`. 32/32 codex_provider tests pass; `tsc --noEmit` clean. --- .../core/inference/codex_availability.py | 6 +- .../backend/core/inference/codex_provider.py | 94 +++++++-- studio/backend/core/inference/providers.py | 8 +- studio/backend/tests/test_codex_provider.py | 190 ++++++++++++++++-- .../src/features/chat/api/codex-api.ts | 6 +- .../chat/components/codex-login-button.tsx | 20 +- 6 files changed, 287 insertions(+), 37 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 5b6f6c18be..f663513e41 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -226,8 +226,12 @@ async def _detect_logged_in() -> bool: # Negative prefixes win, regardless of rc. We anchor on word # boundaries so "not logged in" / "not authenticated" both match # without being fooled by the substring "logged in" inside them. + # Covers the variants seen across CLI releases and locales. negative = re.compile( - r"\b(not logged in|not authenticated|please log in|run\s+`?codex login`?)\b" + r"\b(not\s+(?:logged|signed)\s+in|" + r"not\s+authenticated|" + r"please\s+(?:log|sign)\s+in|" + r"run\s+`?codex\s+login`?)\b" ) if negative.search(combined): return False diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 903e76aec7..cf2811e4b0 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -40,6 +40,7 @@ import asyncio import importlib import importlib.util import json +import sys import time from typing import Any, AsyncGenerator, Optional @@ -71,6 +72,60 @@ class CodexUnavailableError(RuntimeError): """ +def _codex_sdk_env_override() -> dict[str, str]: + """Return an env update dict that scrubs sensitive vars from the + codex app-server subprocess env. + + The upstream openai_codex SDK's `AppServerConfig.env` is merged on + top of `os.environ.copy()` (see openai/codex/sdk/python/src/openai_codex/client.py), + so providing an empty-string mapping for every non-safe key + effectively overrides them in the spawn env. Combined with + `_codex_subprocess_env()` (used for direct CLI calls) this gives + parity between the CLI and SDK code paths: neither sees HF_TOKEN, + GH_TOKEN, WANDB_API_KEY, ANTHROPIC_API_KEY, or any other secret + that lives in the Studio parent environment. + """ + import os + + from core.inference.codex_availability import _SAFE_CODEX_ENV_KEYS + + safe = set(_SAFE_CODEX_ENV_KEYS) + return {key: "" for key in os.environ if key not in safe} + + +def _open_async_codex(async_codex_cls: Any) -> Any: + """Construct an AsyncCodex with a scrubbed env config. + + Tries `AsyncCodex(config=AppServerConfig(env=...))` first; falls + back to the bare `AsyncCodex()` form when either the SDK does not + expose AppServerConfig or its signature does not accept the env + kwarg. The fallback is a soft degradation: the app-server will see + the full Studio env, but every other Codex hardening still + applies. + """ + try: + sdk_mod = sys.modules.get("openai_codex") or sys.modules.get( + "codex_app_server" + ) + if sdk_mod is not None: + app_server_config = getattr(sdk_mod, "AppServerConfig", None) + if app_server_config is not None: + return async_codex_cls( + config = app_server_config(env = _codex_sdk_env_override()), + ) + except TypeError: + # Older SDK: AppServerConfig may not accept the env kwarg yet. + # Fall through to the bare constructor. + pass + except Exception as exc: + logger.warning( + "codex_provider.env_scrub_config_failed", + exc_type = type(exc).__name__, + error = str(exc), + ) + return async_codex_cls() + + def _import_codex() -> Any: """Return the imported Codex SDK module or raise CodexUnavailableError. @@ -88,9 +143,9 @@ def _import_codex() -> Any: return importlib.import_module(name) raise CodexUnavailableError( "Codex Python SDK is not installed on this host. " - "Install with `pip install openai-codex-app-server-sdk` " - "(import name `openai_codex`, legacy alias `codex_app_server`), " - "or use a different provider." + "Install with `pip install openai-codex` (canonical upstream " + "name, imports as `openai_codex`; legacy alias `codex_app_server` " + "is also accepted), or use a different provider." ) @@ -407,7 +462,7 @@ async def _stream_codex_single( completion_text_chars = 0 - async with async_codex_cls() as codex: + async with _open_async_codex(async_codex_cls) as codex: # ``thread_start`` accepts a model id; system prompts are # passed when supported by the SDK rev (older revs ignore the # extra kwarg). Be tolerant about kwargs that may not exist. @@ -540,7 +595,7 @@ async def _stream_codex_parallel( try: sdk = _import_codex() async_codex_cls = getattr(sdk, "AsyncCodex") - async with async_codex_cls() as codex: + async with _open_async_codex(async_codex_cls) as codex: thread_kwargs: dict[str, Any] = {"model": model} if system: thread_kwargs["system"] = system @@ -670,6 +725,7 @@ async def _stream_codex_parallel( synthesis_text = await _run_codex_synthesis( model = model, + system = system, prompt = prompt, tab_outputs = per_tab_texts, ) @@ -702,13 +758,16 @@ async def _stream_codex_parallel( async def _run_codex_synthesis( *, model: str, + system: str, prompt: str, tab_outputs: list[str], ) -> str: """Run one extra Codex call that consumes the N per-tab outputs and returns a unified synthesis. Returns the empty string on failure -- the caller already surfaced the per-tab outputs so an empty - synthesis is recoverable. + synthesis is recoverable. The Studio system prompt is forwarded to + the synthesis thread so style/role instructions like "Always answer + in Spanish" survive the fan-out. """ if not tab_outputs: return "" @@ -728,11 +787,18 @@ async def _run_codex_synthesis( try: sdk = _import_codex() async_codex_cls = getattr(sdk, "AsyncCodex") - async with async_codex_cls() as codex: + async with _open_async_codex(async_codex_cls) as codex: + thread_kwargs: dict[str, Any] = {"model": model} + if system: + thread_kwargs["system"] = system try: - thread = await codex.thread_start(model = model) + thread = await codex.thread_start(**thread_kwargs) except TypeError: - thread = await codex.thread_start() + # Older SDK rev: no `system` kwarg. Fall back to model-only + # and prepend the system prompt to the synthesis text. + thread = await codex.thread_start(model = model) + if system: + synthesis_prompt = f"{system}\n\n{synthesis_prompt}" result = await thread.run(synthesis_prompt) return ( _coerce_text(result) or getattr(result, "final_response", "") or str(result) @@ -809,11 +875,13 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: # Strip ANSI control sequences (the upstream login command wraps the # URL and code in `\x1b[34m...\x1b[0m`) before pattern matching. ansi_re = re.compile(r"\x1b\[[0-9;]*[mGKHF]") - # Anchor on the upstream URL shape: ``.../codex/device`` (optionally - # with a query string). The pattern accepts any host because some - # builds redirect via a staging host. + # Accept any plausible device-auth URL the CLI prints. Upstream has + # used `.../codex/device`, `chatgpt.com/activate`, and + # `auth.openai.com/device`; rather than guess we look for any + # https URL whose path mentions `device`, `activate`, or `verify`. url_re = re.compile( - r"https?://[^\s\x1b]+?/codex/device(?:\?[^\s\x1b]*)?", re.IGNORECASE + r"https?://[^\s\x1b]+/(?:codex/)?(?:device|activate|verify)\b[^\s\x1b]*", + re.IGNORECASE, ) # One-time-code format from upstream device_code_auth.rs: 4 chars, # dash, 4 chars. Pattern is tolerant of any uppercase alphanum. diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index 4d758f1520..f18b9c177d 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -307,10 +307,10 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { "hidden": True, "notes": ( "Dispatches chat turns through the local Codex CLI via " - "the OpenAI Codex Python SDK (pip install " - "`openai-codex-app-server-sdk`, import `openai_codex`, legacy " - "alias `codex_app_server`). Surfaced only when the CLI and " - "SDK are both installed; sign in with `codex login`." + "the OpenAI Codex Python SDK (pip install `openai-codex`, " + "imports as `openai_codex`; legacy alias `codex_app_server` " + "is accepted). Surfaced only when the CLI and SDK are both " + "installed; sign in with `codex login`." ), }, "openrouter": { diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index e5212d7b5c..57272d5ca3 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -37,6 +37,15 @@ _backend = os.path.join(os.path.dirname(__file__), "..") if _backend not in sys.path: sys.path.insert(0, _backend) +# Resolved relative to this file so the source-inspection tests work in any +# checkout location (CI, dev machines, the review worker, etc.). +_BACKEND_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "..")) + + +def _backend_file(rel: str) -> str: + """Return an absolute path inside the backend tree, regardless of cwd.""" + return os.path.join(_BACKEND_DIR, rel) + # ── Helpers ───────────────────────────────────────────────────────── @@ -529,10 +538,7 @@ class TestCodexHardenedRegressions: def test_login_status_uses_login_subcommand(self): """Upstream is `codex login status`, NOT `codex auth status`.""" - src = ( - "/mnt/disks/unslothai/ubuntu/workspace_11/unsloth_pr5724/" - "studio/backend/core/inference/codex_availability.py" - ) + src = _backend_file("core/inference/codex_availability.py") text = open(src).read() assert ( '"auth", "status"' not in text @@ -540,10 +546,7 @@ class TestCodexHardenedRegressions: assert '"login", "status"' in text def test_device_login_uses_login_subcommand(self): - src = ( - "/mnt/disks/unslothai/ubuntu/workspace_11/unsloth_pr5724/" - "studio/backend/core/inference/codex_provider.py" - ) + src = _backend_file("core/inference/codex_provider.py") text = open(src).read() assert ( '"auth", "login", "--device-auth"' not in text @@ -621,10 +624,7 @@ class TestCodexHardenedRegressions: """SSE error frame must NOT echo str(exc) verbatim (CodeQL).""" import re - src = ( - "/mnt/disks/unslothai/ubuntu/workspace_11/unsloth_pr5724/" - "studio/backend/routes/inference.py" - ) + src = _backend_file("routes/inference.py") text = open(src).read() bad = re.findall(r'f["\']Codex error:\s*\{exc\}["\']', text) assert not bad, f"raw exception in SSE: {bad}" @@ -633,10 +633,7 @@ class TestCodexHardenedRegressions: """codex.py SSE stream wrapping must also not leak str(exc).""" import re - src = ( - "/mnt/disks/unslothai/ubuntu/workspace_11/unsloth_pr5724/" - "studio/backend/routes/codex.py" - ) + src = _backend_file("routes/codex.py") text = open(src).read() for line in text.splitlines(): ls = line.strip() @@ -750,6 +747,167 @@ class TestCodexHardenedRegressions: assert "partial output" in body assert "REPLAYED" not in body + def test_not_signed_in_wording_also_handled(self): + """`Not signed in` (alternative localisation) must also be + treated as logged-out, not as positive match. + """ + import asyncio + + from core.inference import codex_availability as av + + async def _fake_run_cli(args, **kw): + return (0, "Not signed in.", "") + + orig = av._run_cli + av._run_cli = _fake_run_cli # type: ignore[assignment] + try: + assert asyncio.run(av._detect_logged_in()) is False + finally: + av._run_cli = orig # type: ignore[assignment] + + def test_device_url_accepts_generic_verification_url(self): + """The login parser must accept upstream's chatgpt.com/activate + URL as well as the canonical /codex/device shape. + """ + import re + + src = _backend_file("core/inference/codex_provider.py") + text = open(src).read() + # Find the url_re pattern literal and compile it. + m = re.search(r"url_re\s*=\s*re\.compile\(\s*\n?\s*r\"([^\"]+)\"", text) + assert m, "url_re definition not found" + pattern = re.compile(m.group(1), re.IGNORECASE) + # Upstream device URLs we expect to match. + for u in ( + "https://auth.openai.com/codex/device", + "https://chatgpt.com/activate", + "https://auth.openai.com/device/verify?code=ABCD", + ): + assert pattern.search(u), f"device URL regex missed: {u}" + + def test_synthesis_call_forwards_system_prompt(self, monkeypatch): + """`_run_codex_synthesis` must pass the system prompt so a + fan-out style instruction ("Always answer in Spanish") survives + the unification step. + """ + seen_kwargs: list[dict] = [] + seen_prompts: list[str] = [] + + class _SynThread: + async def run(self, prompt): + seen_prompts.append(prompt) + return "synth ok" + + def turn(self, prompt): + # Force buffered path via no `stream` attr. + class _T: + pass + + return _T() + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + seen_kwargs.append(kw) + return _SynThread() + + _install_fake_codex_sdk(monkeypatch, _Async) + from core.inference.codex_provider import _run_codex_synthesis + + out = asyncio.run( + _run_codex_synthesis( + model = "gpt-5.5", + system = "Always answer in Spanish.", + prompt = "What is the capital of France?", + tab_outputs = ["Paris", "Paris."], + ) + ) + # Either the kwargs carried the system prompt or it was + # prepended to the synthesis prompt as a fallback. + system_seen = ( + any("Spanish" in (kw.get("system") or "") for kw in seen_kwargs) + or any("Always answer in Spanish" in p for p in seen_prompts) + ) + assert system_seen, ( + f"system prompt dropped in synthesis. kwargs={seen_kwargs} " + f"prompts={seen_prompts}" + ) + # And the synthesis still returned the model's text. + assert "synth" in out.lower() + + def test_sdk_env_scrubbed_via_appserverconfig(self, monkeypatch): + """The SDK construction path must wire AppServerConfig(env=...) + when the SDK exposes it, so HF_TOKEN / GH_TOKEN are not leaked + to the codex app-server subprocess. + """ + monkeypatch.setenv("HF_TOKEN", "should_be_scrubbed") + monkeypatch.setenv("GH_TOKEN", "should_be_scrubbed") + monkeypatch.setenv("OPENAI_API_KEY", "ok_for_codex") + + seen_configs: list[Any] = [] + + class _FakeAppServerConfig: + def __init__(self, env=None, **kw): + self.env = env or {} + + class _Async: + def __init__(self, config=None): + seen_configs.append(config) + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _FakeThread(chunks=["ok"]) + + # Inject a fake openai_codex module exposing AppServerConfig. + import importlib.util as _iu + import types as _types + + fake_mod = _types.ModuleType("openai_codex") + fake_mod.AsyncCodex = _Async # type: ignore[attr-defined] + fake_mod.AppServerConfig = _FakeAppServerConfig # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "openai_codex", fake_mod) + real_find_spec = _iu.find_spec + monkeypatch.setattr( + "importlib.util.find_spec", + lambda n, *a, **kw: ( + _types.SimpleNamespace() + if n in ("openai_codex", "codex_app_server") + else real_find_spec(n, *a, **kw) + ), + ) + + from core.inference.codex_provider import stream_codex + + asyncio.run( + _consume_first( + stream_codex( + messages=[{"role": "user", "content": "hi"}], + model="gpt-5.5", + parallel_calls=1, + ) + ) + ) + assert seen_configs, "AsyncCodex was never instantiated" + cfg = seen_configs[0] + assert cfg is not None, "AppServerConfig was not passed to AsyncCodex" + assert "HF_TOKEN" in cfg.env and cfg.env["HF_TOKEN"] == "", ( + "HF_TOKEN not overridden to empty in SDK env" + ) + assert "GH_TOKEN" in cfg.env and cfg.env["GH_TOKEN"] == "" + # Safe-listed keys must NOT appear in the override dict (so the + # SDK keeps their os.environ values intact). + assert "OPENAI_API_KEY" not in cfg.env + def test_thread_turn_stream_path_taken(self, monkeypatch): """The canonical openai_codex API uses thread.turn(prompt).stream(); the provider must prefer that over the legacy run_streaming hook. diff --git a/studio/frontend/src/features/chat/api/codex-api.ts b/studio/frontend/src/features/chat/api/codex-api.ts index 6652bf9391..50d71cc009 100644 --- a/studio/frontend/src/features/chat/api/codex-api.ts +++ b/studio/frontend/src/features/chat/api/codex-api.ts @@ -31,8 +31,12 @@ export interface CodexStatus { } export interface CodexLoginEvent { - type: "device_url" | "log" | "error" | "done"; + // `device_code` is the one-time code the verification page asks for + // (separate from the URL); the backend extracts it from the CLI + // stdout via a dedicated regex and emits it as a structured event. + type: "device_url" | "device_code" | "log" | "error" | "done"; url?: string; + code?: string; line?: string; message?: string; ok?: boolean; diff --git a/studio/frontend/src/features/chat/components/codex-login-button.tsx b/studio/frontend/src/features/chat/components/codex-login-button.tsx index 3458a99c9e..deaacb8a71 100644 --- a/studio/frontend/src/features/chat/components/codex-login-button.tsx +++ b/studio/frontend/src/features/chat/components/codex-login-button.tsx @@ -37,6 +37,7 @@ export function CodexLoginButton({ onLoggedIn }: Props) { const [error, setError] = useState(null); const [logs, setLogs] = useState([]); const [deviceUrl, setDeviceUrl] = useState(null); + const [deviceCode, setDeviceCode] = useState(null); // Track the active stream's abort controller so a second click // (or an unmount) tears the SSE reader down cleanly. Without this // the long-running login subprocess would keep streaming into a @@ -49,9 +50,15 @@ export function CodexLoginButton({ onLoggedIn }: Props) { setError(null); setLogs([]); setDeviceUrl(null); + setDeviceCode(null); const controller = new AbortController(); abortRef.current?.abort(); abortRef.current = controller; + // Track the specific backend error inside the closure so the + // generic fallback message does not overwrite it: setError is + // async and reading `error` after `setError(event.message)` would + // still see the stale pre-stream value. + let lastStreamError: string | null = null; try { let lastOk: boolean | undefined; for await (const event of streamCodexDeviceLogin( @@ -67,9 +74,12 @@ export function CodexLoginButton({ onLoggedIn }: Props) { } catch { // Ignore -- the URL is still visible in the log. } + } else if (event.type === "device_code" && event.code) { + setDeviceCode(event.code); } else if (event.type === "log" && event.line) { setLogs((prev) => [...prev, event.line as string]); } else if (event.type === "error" && event.message) { + lastStreamError = event.message; setError(event.message); } else if (event.type === "done") { lastOk = event.ok; @@ -77,7 +87,7 @@ export function CodexLoginButton({ onLoggedIn }: Props) { } if (lastOk) { onLoggedIn?.(); - } else if (!error) { + } else if (!lastStreamError) { setError("Codex login did not complete -- see log for details."); } } catch (exc) { @@ -87,7 +97,7 @@ export function CodexLoginButton({ onLoggedIn }: Props) { } finally { setBusy(false); } - }, [busy, error, onLoggedIn]); + }, [busy, onLoggedIn]); // Abort the in-flight SSE stream on unmount so the underlying // `codex login --device-auth` subprocess does not keep streaming @@ -116,6 +126,12 @@ export function CodexLoginButton({ onLoggedIn }: Props) {

)} + {deviceCode && ( +

+ One-time code:{" "} + {deviceCode} +

+ )} {error && (

{error}

)}