diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index bbb78bef64..5b6f6c18be 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -60,6 +60,50 @@ _DEFAULT_SUPPORTED_MODELS: tuple[str, ...] = ( # an internal alpha may publish under it. _SDK_MODULE_NAMES: tuple[str, ...] = ("openai_codex", "codex_app_server") +# Safe-list of environment variables forwarded to the codex subprocess. +# Studio's parent env contains secrets (HF_TOKEN, GH_TOKEN, WANDB_API_KEY, +# OPENAI key for non-codex providers, etc.); a malicious or shimmed codex +# binary earlier on PATH would receive all of them via plain os.environ +# inheritance. We pass only what codex actually needs: PATH for spawning +# its own helpers, HOME / USER for auth config lookup, the Windows / +# macOS equivalents, the codex-specific CODEX_HOME override, and the +# OPENAI_API_KEY that codex's own ``--with-api-key`` flow expects. +_SAFE_CODEX_ENV_KEYS: tuple[str, ...] = ( + "PATH", + "HOME", + "USER", + "USERNAME", + "SHELL", + "LANG", + "LC_ALL", + "TMPDIR", + "TEMP", + "TMP", + "SYSTEMROOT", + "WINDIR", + "APPDATA", + "LOCALAPPDATA", + "PROGRAMDATA", + "CODEX_HOME", + "OPENAI_API_KEY", + "OPENAI_BASE_URL", +) + + +def _codex_subprocess_env() -> dict[str, str]: + """Return a scrubbed env mapping for codex subprocess spawning. + + Forwards only keys from `_SAFE_CODEX_ENV_KEYS` that are actually set + in the parent environment, so secrets from other providers never + reach the codex CLI. + """ + env: dict[str, str] = {} + for key in _SAFE_CODEX_ENV_KEYS: + value = os.environ.get(key) + if value is not None: + env[key] = value + return env + def _which_codex() -> Optional[str]: """Return absolute path to the ``codex`` CLI, or None if missing. @@ -119,7 +163,7 @@ async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, *args, stdout = asyncio.subprocess.PIPE, stderr = asyncio.subprocess.PIPE, - env = os.environ.copy(), + env = _codex_subprocess_env(), ) except FileNotFoundError: return -1, "", "codex binary not on PATH" diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 6461dfbdf0..903e76aec7 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -319,10 +319,15 @@ async def _stream_thread_run( path. Used when neither streaming helper resolves and as the final fallback. - On any streaming exception we log and fall through to the buffered - path so a partially-broken streaming helper does not take the - whole turn down. + Cross-turn side-effect protection: once any chunk has been emitted + via a streaming helper, we never fall through to the buffered + ``thread.run(prompt)`` path -- a partial-stream failure would + otherwise re-execute the same Codex turn and duplicate side + effects (file writes, shell commands, etc.). The buffered path + runs only when streaming helpers produced zero output. """ + emitted_any = False + # 1. Canonical: thread.turn(prompt).stream() turn_factory = getattr(thread, "turn", None) if turn_factory is not None: @@ -338,6 +343,7 @@ async def _stream_thread_run( async for event in stream_obj: text = _coerce_text(getattr(event, "payload", event)) if text: + emitted_any = True yield text return except Exception as exc: @@ -345,7 +351,13 @@ async def _stream_thread_run( "codex_provider.turn_stream_failed_fallback", exc_type = type(exc).__name__, error = str(exc), + emitted_any = emitted_any, ) + if emitted_any: + # The Codex turn already ran far enough to emit text; + # do not re-execute via run() or run_streaming() -- the + # side-effects (commands / writes) would replay. + return # 2. Legacy: thread.run_streaming(prompt) run_streaming = getattr(thread, "run_streaming", None) @@ -357,6 +369,7 @@ async def _stream_thread_run( async for event in stream_obj: text = _coerce_text(event) if text: + emitted_any = True yield text return except Exception as exc: @@ -364,9 +377,14 @@ async def _stream_thread_run( "codex_provider.run_streaming_failed_fallback", exc_type = type(exc).__name__, error = str(exc), + emitted_any = emitted_any, ) + if emitted_any: + return # 3. Buffered fallback: await the full TurnResult, emit one chunk. + # Only reached when no streaming helper emitted anything, so this + # is the first (and only) execution of the turn. result = await thread.run(prompt) text = _coerce_text(result) or getattr(result, "final_response", "") or str(result) if text: @@ -763,9 +781,13 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: # SIGTERM the whole group on cancel without sending it to ourselves. # On Windows, ``creationflags=CREATE_NEW_PROCESS_GROUP`` (0x200) gives # an equivalent isolation for ``proc.send_signal(signal.CTRL_BREAK_EVENT)``. + # Env is scrubbed to the codex safe-list (see codex_availability) so a + # shimmed `codex` on PATH does not inherit other provider secrets. + from core.inference.codex_availability import _codex_subprocess_env spawn_kwargs: dict[str, Any] = { "stdout": asyncio.subprocess.PIPE, "stderr": asyncio.subprocess.STDOUT, + "env": _codex_subprocess_env(), } if os.name == "posix": spawn_kwargs["start_new_session"] = True diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index bfa44176f8..e5212d7b5c 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -673,6 +673,83 @@ class TestCodexHardenedRegressions: ), "raw exception text leaked into codex_tab_error SSE frame" assert "Codex tab failed" in body or "exception_type" in body + def test_codex_subprocess_env_scrubbed(self, monkeypatch): + """The codex subprocess env must not include other-provider secrets.""" + from core.inference.codex_availability import _codex_subprocess_env + + monkeypatch.setenv("HF_TOKEN", "hf_should_not_leak") + monkeypatch.setenv("GH_TOKEN", "gh_should_not_leak") + monkeypatch.setenv("WANDB_API_KEY", "wandb_should_not_leak") + monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic_should_not_leak") + monkeypatch.setenv("OPENAI_API_KEY", "openai_codex_uses_this") + monkeypatch.setenv("CODEX_HOME", "/custom/.codex") + monkeypatch.setenv("PATH", "/usr/bin") + + env = _codex_subprocess_env() + for secret in ( + "HF_TOKEN", + "GH_TOKEN", + "WANDB_API_KEY", + "ANTHROPIC_API_KEY", + ): + assert secret not in env, f"{secret} leaked into codex env" + # Codex-relevant keys must be preserved. + assert env.get("OPENAI_API_KEY") == "openai_codex_uses_this" + assert env.get("CODEX_HOME") == "/custom/.codex" + assert env.get("PATH") == "/usr/bin" + + def test_partial_stream_failure_does_not_replay_turn(self, monkeypatch): + """If turn.stream() fails after emitting some text, the buffered + run() fallback must NOT fire -- replaying would duplicate side + effects (file writes, shell commands). + """ + run_calls = {"n": 0} + + class _PartialStreamTurn: + async def stream(self): + yield {"text": "partial output "} + raise RuntimeError("network glitch mid-stream") + + class _ThreadPartialFail: + def turn(self, prompt): + return _PartialStreamTurn() + + async def run(self, prompt): + run_calls["n"] += 1 + return "REPLAYED -- BAD" + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _ThreadPartialFail() + + _install_fake_codex_sdk(monkeypatch, _Async) + from core.inference.codex_provider import stream_codex + + chunks: list[str] = [] + + async def _collect(): + async for c in stream_codex( + messages=[{"role": "user", "content": "hi"}], + model="gpt-5.5", + parallel_calls=1, + ): + chunks.append(c) + + asyncio.run(_collect()) + assert run_calls["n"] == 0, ( + "buffered run() fired after partial stream emission -- " + "would replay side effects" + ) + body = "".join(chunks) + assert "partial output" in body + assert "REPLAYED" not in body + 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/components/codex-login-button.tsx b/studio/frontend/src/features/chat/components/codex-login-button.tsx index 7f7737cc28..3458a99c9e 100644 --- a/studio/frontend/src/features/chat/components/codex-login-button.tsx +++ b/studio/frontend/src/features/chat/components/codex-login-button.tsx @@ -19,7 +19,7 @@ * probe and flip back into the "ready" state automatically. */ -import { useCallback, useRef, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { streamCodexDeviceLogin, @@ -89,6 +89,15 @@ export function CodexLoginButton({ onLoggedIn }: Props) { } }, [busy, error, onLoggedIn]); + // Abort the in-flight SSE stream on unmount so the underlying + // `codex login --device-auth` subprocess does not keep streaming + // (and consuming a device-auth session) after the dialog closes. + useEffect(() => { + return () => { + abortRef.current?.abort(); + }; + }, []); + return (