From cbc3c43655a42fcd5d73400d2c9e8a52593c4603 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 16:46:00 +0000 Subject: [PATCH 01/40] Studio: add Codex SDK as a chat provider with parallel-calls fan-out Wires the OpenAI Codex CLI / Python SDK (codex_app_server) into Studio as a new chat provider type. Hosts that don't have the CLI or the SDK installed never see the entry; on logged-out hosts the provider config dialog renders a device-auth Sign-in button that surfaces the verification URL and streams CLI progress back over SSE. Backend - new core/inference/codex_availability.py probes the CLI + SDK and reports {installed, logged_in, version, supported_models}; it never imports codex_app_server at module top level so the rest of the backend keeps starting cleanly on hosts that don't have the SDK. - new core/inference/codex_provider.py wraps AsyncCodex and translates Codex events into OpenAI chat-completion chunks. Supports the thread.run_streaming path with a non-streaming fallback for older SDK revs. - parallel_calls > 1 fans the turn out across N tasks (capped at 20) via asyncio.gather and emits codex_tab_open / codex_tab_chunk / codex_tab_close tool-events per attempt plus a final codex_gather synthesis event. A separate standalone Codex call produces the unified answer. - new routes/codex.py exposes GET /api/codex/status and POST /api/codex/login. The login route shells out to codex auth login --device-auth and streams events; the first event carries the verification URL so the frontend can window.open it. - ChatCompletionRequest gains a parallel_calls field bounded [1, 20] by pydantic. The codex registry entry stays hidden by default; the /api/codex/status probe is the authoritative gate. - routes/inference.py dispatches provider_type=codex through the local CLI/SDK pipeline instead of the standard HTTP client, with graceful error surfacing for CodexUnavailableError. Frontend - new api/codex-api.ts exposes fetchCodexStatus() and an async generator streamCodexDeviceLogin() that drives the SSE stream and yields parsed events. - new components/codex-parallel-tabs.tsx renders the tabbed parallel- calls UI with a Synthesis tab highlighted once the codex_gather event arrives. Pure reducer keeps the state transitions unit- testable. - new components/codex-login-button.tsx posts to /api/codex/login, opens the verification URL in a new tab via window.open, and shows the streamed CLI log as it lands. - external-providers.ts exports CODEX_PROVIDER_TYPE, CODEX_MAX_PARALLEL_CALLS, isCodexProviderType, and clampCodexParallelCalls. Codex is marked text-only so the composer hides image-attach affordances when selected. Tests - tests/test_codex_provider.py (14 cases) covers the availability probe across the four install / login states, the streaming + parallel-calls translation against a fake codex_app_server module injected into sys.modules, the [1, 20] pydantic clamp, the CodexUnavailableError surfacing path, and the parallel_calls=1 single-call shape (no tab tool-events). --- .../core/inference/codex_availability.py | 221 ++++++ .../backend/core/inference/codex_provider.py | 669 ++++++++++++++++++ studio/backend/core/inference/providers.py | 35 + studio/backend/main.py | 5 + studio/backend/models/inference.py | 14 + studio/backend/routes/__init__.py | 2 + studio/backend/routes/codex.py | 92 +++ studio/backend/routes/inference.py | 76 +- studio/backend/tests/test_codex_provider.py | 499 +++++++++++++ .../src/features/chat/api/codex-api.ts | 146 ++++ .../chat/components/codex-login-button.tsx | 120 ++++ .../chat/components/codex-parallel-tabs.tsx | 227 ++++++ .../src/features/chat/external-providers.ts | 34 + 13 files changed, 2138 insertions(+), 2 deletions(-) create mode 100644 studio/backend/core/inference/codex_availability.py create mode 100644 studio/backend/core/inference/codex_provider.py create mode 100644 studio/backend/routes/codex.py create mode 100644 studio/backend/tests/test_codex_provider.py create mode 100644 studio/frontend/src/features/chat/api/codex-api.ts create mode 100644 studio/frontend/src/features/chat/components/codex-login-button.tsx create mode 100644 studio/frontend/src/features/chat/components/codex-parallel-tabs.tsx diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py new file mode 100644 index 0000000000..c7488a0ed0 --- /dev/null +++ b/studio/backend/core/inference/codex_availability.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Codex CLI / SDK availability probe. + +This module never imports `codex_app_server` at module top level. The +SDK is optional and not pinned in pyproject.toml -- if it's installed +locally, we use it; if it isn't, the probe simply returns +``installed=False`` and the provider stays hidden in the frontend. + +The frontend calls ``GET /api/codex/status`` at startup to decide +whether to surface the "codex" entry in the provider picker. Three +states matter: + +* ``installed=False`` -- either the CLI is missing OR the SDK + (``codex_app_server``) is not importable. The picker hides the + entry entirely. +* ``installed=True, logged_in=False`` -- everything resolves on the + Python side but ``codex auth status`` (or equivalent) reports no + active credentials. The provider config dialog shows a + ``Sign in to Codex`` button instead of the regular API-key field. +* ``installed=True, logged_in=True`` -- ready to use; the picker + shows the regular model dropdown. + +Detection is best-effort and cheap: we shell out to ``which codex`` +plus ``codex --version`` for the CLI and use ``importlib.util.find_spec`` +for the SDK. No long-running CLI commands are invoked here so the +status endpoint is safe to poll on every page load. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import shutil +from typing import Any, Optional + +import structlog + +logger = structlog.get_logger(__name__) + + +# Default catalog of models surfaced in the picker when the CLI is +# present but doesn't advertise a list. The SDK accepts arbitrary model +# ids; this is purely a sensible default. +_DEFAULT_SUPPORTED_MODELS: tuple[str, ...] = ( + "gpt-5.4", + "gpt-5.4-mini", + "gpt-5.5", + "o3", +) + + +def _which_codex() -> Optional[str]: + """Return absolute path to the ``codex`` CLI, or None if missing. + + Uses :func:`shutil.which` so the lookup honours ``PATH`` exactly + the way the user's shell would. Returns ``None`` on any failure + so callers can treat "missing" and "broken probe" the same way. + """ + try: + return shutil.which("codex") + except Exception as exc: + # shutil.which itself is documented as raising only on + # genuinely unusual conditions, but a hardened wrapper costs + # nothing and keeps the status endpoint from 500'ing. + logger.warning("codex_availability.which_failed", error = str(exc)) + return None + + +def _sdk_importable() -> bool: + """True iff ``codex_app_server`` is importable in this interpreter. + + We deliberately use :func:`importlib.util.find_spec` instead of an + ``import codex_app_server`` so the import never actually runs -- + that keeps the cost negligible and avoids the SDK's own side + effects (which include reaching out to the CLI subprocess for an + RPC ping) during a simple availability check. + """ + try: + return importlib.util.find_spec("codex_app_server") is not None + except Exception as exc: + logger.warning("codex_availability.find_spec_failed", error = str(exc)) + return False + + +async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, str]: + """Run a short ``codex`` CLI command and return (rc, stdout, stderr). + + The probe uses 4s as the wall-clock cap because ``codex --version`` + and ``codex auth status`` both return in well under a second on a + healthy install. A longer probe would block the + ``/api/codex/status`` route -- and that route fires on every chat + page load, so a tight cap matters. + """ + try: + proc = await asyncio.create_subprocess_exec( + "codex", + *args, + stdout = asyncio.subprocess.PIPE, + stderr = asyncio.subprocess.PIPE, + env = os.environ.copy(), + ) + except FileNotFoundError: + return -1, "", "codex binary not on PATH" + except Exception as exc: + logger.warning( + "codex_availability.spawn_failed", + args = args, + error = str(exc), + ) + return -1, "", str(exc) + + try: + stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout = timeout) + except asyncio.TimeoutError: + proc.kill() + try: + await proc.wait() + except Exception: + pass + return -1, "", f"codex {' '.join(args)} timed out after {timeout:.1f}s" + + return ( + proc.returncode if proc.returncode is not None else -1, + stdout_b.decode("utf-8", errors = "replace").strip(), + stderr_b.decode("utf-8", errors = "replace").strip(), + ) + + +async def _detect_version() -> Optional[str]: + rc, stdout, stderr = await _run_cli(["--version"]) + if rc != 0: + return None + # ``codex --version`` prints something like "codex-cli 0.133.0". + # Surface the whole line so the UI can show the exact build the + # user has installed -- it's useful when troubleshooting. + text = stdout or stderr + return text.split("\n", 1)[0].strip() if text else None + + +async def _detect_logged_in() -> bool: + """Best-effort: assume a non-zero ``codex auth status`` rc means + "not logged in". The CLI surface has shifted over releases (some + versions print to stderr, some to stdout, some use "Logged in as + ..." vs "Not authenticated"). The return code is the most stable + signal we have, so we lean on it and fall back to substring + inspection only when the rc itself is ambiguous (rc=0 but no + output, or rc=-1 from a timeout / crash). + """ + rc, stdout, stderr = await _run_cli(["auth", "status"]) + combined = f"{stdout}\n{stderr}".lower() + if rc == 0: + # Some 0.x releases exit 0 even when the user is logged out. + # If we got any output, look for the obvious "not logged in" + # signals; if there's nothing on either pipe at all, treat + # rc=0 as authenticated (the optimistic default). + if not combined.strip(): + return True + if "not authenticated" in combined or "not logged in" in combined: + return False + if "logged in" in combined or "authenticated" in combined: + return True + return True + # rc != 0: lean on substring matching one more time before + # defaulting to False, in case a future CLI uses rc=2 for the + # logged-out state instead of "no auth command exists". + if "logged in" in combined or "authenticated as" in combined: + return True + return False + + +async def probe_codex_availability() -> dict[str, Any]: + """Return the full status payload consumed by ``GET /api/codex/status``. + + Returns a dict with keys: + + * ``installed`` (bool) -- True iff both CLI and SDK are present. + Note this is the gate the frontend uses to surface the provider + at all, so if EITHER is missing the picker hides the entry. + * ``cli_path`` (str | None) -- absolute path to the CLI, or None. + * ``sdk_importable`` (bool) -- the Python SDK is importable. + * ``logged_in`` (bool) -- best-effort auth check; meaningless when + ``installed`` is False. + * ``version`` (str | None) -- the ``codex --version`` first line. + * ``supported_models`` (list[str]) -- default model id catalog. + """ + cli_path = _which_codex() + sdk_ok = _sdk_importable() + + payload: dict[str, Any] = { + "installed": bool(cli_path) and sdk_ok, + "cli_path": cli_path, + "sdk_importable": sdk_ok, + "logged_in": False, + "version": None, + "supported_models": list(_DEFAULT_SUPPORTED_MODELS), + } + + if cli_path: + # version + login probes only matter when the CLI is present; + # they would otherwise just churn subprocess errors. Run them + # in parallel because both are independent CLI invocations. + version, logged_in = await asyncio.gather( + _detect_version(), + _detect_logged_in(), + ) + payload["version"] = version + payload["logged_in"] = bool(logged_in) + + logger.info( + "codex_availability.probed", + installed = payload["installed"], + sdk_importable = payload["sdk_importable"], + cli_path = payload["cli_path"], + version = payload["version"], + logged_in = payload["logged_in"], + ) + return payload diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py new file mode 100644 index 0000000000..d4df67ab4f --- /dev/null +++ b/studio/backend/core/inference/codex_provider.py @@ -0,0 +1,669 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Codex SDK chat provider. + +This module wraps the ``codex_app_server`` async SDK so a chat request +routed at ``provider_type="codex"`` can dispatch through the user's +local Codex CLI. The contract back to the frontend is the standard +OpenAI Chat Completions SSE shape, exactly like every other entry in +``external_provider.py``. + +The two interesting features here: + +* ``_stream_codex_single`` -- translate Codex SDK events into + OpenAI-format chunks. We prefer ``thread.run_streaming`` when the + installed SDK exposes it; otherwise we fall back to + ``await thread.run(...)`` and emit one big content chunk plus a + usage chunk at the end. + +* ``_stream_codex_parallel`` -- the ``parallel_calls`` knob. When > 1, + spawn N async Codex tasks (capped at 20) via ``asyncio.gather`` and + emit per-tab ``_toolEvent`` markers so the frontend can render each + result in its own tab. A final ``codex_gather`` synthesis tab runs a + single Codex call that takes the N outputs and produces a unified + answer. + +The SDK is imported lazily (inside the helpers that actually need it) +because the spec calls out that ``codex_app_server`` may not even be +importable on the build host. The availability probe in +``codex_availability.py`` is what the frontend uses to gate the +provider entirely; this module just refuses to run if the import +fails at request time. +""" + +from __future__ import annotations + +import asyncio +import importlib +import importlib.util +import json +import time +from typing import Any, AsyncGenerator, Optional + +import structlog + +logger = structlog.get_logger(__name__) + + +# Hard cap on parallel Codex fan-out. Picked to match the upper bound +# in the request validator -- exceeding this risks the local Codex CLI +# rate-limiting itself or starving the loop. +MAX_PARALLEL_CALLS = 20 + + +class CodexUnavailableError(RuntimeError): + """Raised when ``codex_app_server`` is not importable at runtime. + + The availability probe is supposed to hide the provider before any + request lands here, but we still raise a typed error so the route + layer can translate it into a 503 the user sees instead of an + opaque traceback. + """ + + +def _import_codex() -> Any: + """Return the imported ``codex_app_server`` module or raise. + + Imported lazily so the rest of the backend keeps starting cleanly + on hosts that don't have the SDK installed. The frontend calls + ``GET /api/codex/status`` first and hides the provider when the + spec isn't importable, so this branch is reached only when the + user (a) explicitly forces the provider via a stale stored config + or (b) the install state changes between status probe and chat + submit. + """ + if importlib.util.find_spec("codex_app_server") is None: + raise CodexUnavailableError( + "codex_app_server is not installed on this host. " + "Install the Codex Python SDK or use a different provider." + ) + return importlib.import_module("codex_app_server") + + +def _last_user_prompt(messages: list[dict[str, Any]]) -> str: + """Extract the most recent user-role message as a plain string. + + Codex's ``thread.run`` accepts a string (per the docs note + "plain strings are accepted anywhere a turn input is accepted"). + Studio's chat history is a full OpenAI-style messages array, so + we flatten it: walk from the end, find the last ``role=user`` + message, and stringify any structured content parts into a + newline-joined block. Multimodal content (images) is described + rather than embedded; Codex SDK input shape is text-first. + + This is intentionally conservative -- we don't try to replay the + whole conversation through Codex per turn because the SDK is + designed around a stateful ``thread`` object. The thread itself + holds context across runs; we only need to feed the latest user + turn. + """ + for msg in reversed(messages): + if msg.get("role") != "user": + continue + content = msg.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for entry in content: + if not isinstance(entry, dict): + continue + if entry.get("type") == "text": + parts.append(str(entry.get("text") or "")) + elif entry.get("type") == "image_url": + url = (entry.get("image_url") or {}).get("url") or "" + parts.append(f"[image: {url[:80]}]" if url else "[image]") + elif entry.get("type") == "input_document": + name = entry.get("filename") or "document" + parts.append(f"[document: {name}]") + return "\n".join(p for p in parts if p) + return "" + + +def _system_prompt(messages: list[dict[str, Any]]) -> str: + """Concatenate all ``role=system`` messages. + + Codex's ``thread_start`` accepts a system prompt for the lifetime + of the thread. We pass any leading system messages so the user's + Studio-side system prompt (chat presets) reaches the Codex side + intact. + """ + parts: list[str] = [] + for msg in messages: + if msg.get("role") != "system": + continue + content = msg.get("content") + if isinstance(content, str): + parts.append(content) + elif isinstance(content, list): + for entry in content: + if isinstance(entry, dict) and entry.get("type") == "text": + parts.append(str(entry.get("text") or "")) + return "\n\n".join(p for p in parts if p) + + +def _chunk_text(completion_id: str, text: str) -> str: + """OpenAI Chat Completions content chunk.""" + payload = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"content": text}, + "finish_reason": None, + } + ], + } + return f"data: {json.dumps(payload)}" + + +def _chunk_stop(completion_id: str) -> str: + payload = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], + } + return f"data: {json.dumps(payload)}" + + +def _chunk_tool_event(completion_id: str, event: dict[str, Any]) -> str: + """Synthetic OpenAI-shaped chunk carrying an ``_toolEvent`` payload. + + Studio's frontend chat-adapter already understands the + ``_toolEvent`` envelope and renders tool cards on the fly. We piggy- + back on the same channel to ship Codex-specific tab markers + (``codex_tab_open`` / ``codex_tab_chunk`` / ``codex_gather``) so the + UI doesn't need a brand-new transport. + """ + payload = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": event, + } + return f"data: {json.dumps(payload)}" + + +def _chunk_usage( + completion_id: str, + prompt_tokens: int, + completion_tokens: int, +) -> str: + payload = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + return f"data: {json.dumps(payload)}" + + +def _coerce_text(payload: Any) -> str: + """Pull text out of a Codex streaming event or result. + + The SDK shape isn't pinned across versions: events expose ``delta`` + or ``text`` or ``content`` depending on whether the model is in + plan / answer / tool-use mode. Be defensive -- read whichever + field is present and fall back to ``str()`` so we never crash + while translating. + """ + if payload is None: + return "" + if isinstance(payload, str): + return payload + if isinstance(payload, dict): + for key in ("delta", "text", "content", "message", "final_response"): + if key in payload: + value = _coerce_text(payload[key]) + if value: + return value + return "" + if isinstance(payload, list): + return "".join(_coerce_text(item) for item in payload) + text_attr = getattr(payload, "text", None) + if isinstance(text_attr, str): + return text_attr + delta_attr = getattr(payload, "delta", None) + if isinstance(delta_attr, str): + return delta_attr + final_attr = getattr(payload, "final_response", None) + if isinstance(final_attr, str): + return final_attr + return "" + + +async def _stream_thread_run( + thread: Any, + prompt: str, +) -> AsyncGenerator[str, None]: + """Yield raw text chunks from a Codex thread. + + Prefers ``thread.run_streaming(prompt)`` because that's what the + docs surface for token-by-token delivery. When the installed SDK + doesn't have that helper, fall back to ``await thread.run(prompt)`` + and yield the full text once -- this still works end-to-end, just + without streaming feedback in the UI. + """ + run_streaming = getattr(thread, "run_streaming", None) + if run_streaming is not None: + try: + stream_obj = run_streaming(prompt) + # The SDK may return either an async iterator directly or a + # coroutine that resolves to one. Handle both shapes so a + # future SDK rev doesn't silently fall off the streaming + # path. + if asyncio.iscoroutine(stream_obj): + stream_obj = await stream_obj + async for event in stream_obj: + text = _coerce_text(event) + if text: + yield text + return + except Exception as exc: + logger.warning( + "codex_provider.run_streaming_failed_fallback", + error = str(exc), + ) + # Intentional fallthrough to the non-streaming path so a + # broken streaming helper doesn't take the whole turn down. + + # Non-streaming fallback: await the full TurnResult, emit one chunk. + result = await thread.run(prompt) + text = _coerce_text(result) or getattr(result, "final_response", "") or str(result) + if text: + yield text + + +async def _stream_codex_single( + model: str, + system: str, + prompt: str, + completion_id: str, +) -> AsyncGenerator[str, None]: + """Run one Codex turn and emit OpenAI-shaped SSE lines.""" + sdk = _import_codex() + async_codex_cls = getattr(sdk, "AsyncCodex", None) + if async_codex_cls is None: + raise CodexUnavailableError( + "codex_app_server is installed but AsyncCodex is missing -- " + "upgrade the SDK." + ) + + completion_text_chars = 0 + + async with 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. + thread_kwargs: dict[str, Any] = {"model": model} + if system: + # Try the canonical kwargs first; the SDK shapes vary + # across versions and we'd rather accept the system prompt + # being dropped than crash on a missing kwarg. + thread_kwargs["system"] = system + try: + thread = await codex.thread_start(**thread_kwargs) + except TypeError: + # Older SDK: only the ``model`` kwarg is accepted. Drop + # extras and retry; the system prompt then lives only in + # the prompt itself (we prepend it below). + thread = await codex.thread_start(model = model) + if system: + prompt = f"{system}\n\n{prompt}" + async for text in _stream_thread_run(thread, prompt): + completion_text_chars += len(text) + yield _chunk_text(completion_id, text) + + # Estimate tokens crudely from char counts. The Codex SDK does not + # consistently expose a token breakdown; we surface a usage chunk + # purely so the frontend cost / context display has a non-zero + # value to render. ``int(chars / 4)`` is the standard rough-cut. + yield _chunk_usage( + completion_id, + prompt_tokens = max(1, len(prompt) // 4), + completion_tokens = max(0, completion_text_chars // 4), + ) + yield _chunk_stop(completion_id) + + +async def stream_codex( + messages: list[dict[str, Any]], + model: str, + parallel_calls: int = 1, +) -> AsyncGenerator[str, None]: + """Top-level entry point used by the inference route. + + When ``parallel_calls == 1`` (the default), this delegates to the + single-turn helper. When > 1 (and <= ``MAX_PARALLEL_CALLS``), it + spawns N parallel Codex turns and emits one tab per result plus a + final synthesis tab. The fan-out runs every spawned turn against + the SAME user prompt -- the parallel knob is for sampling N + independent attempts at the same task, which is what the UI tab + strip surfaces. + """ + prompt = _last_user_prompt(messages) + system = _system_prompt(messages) + completion_id = f"chatcmpl-codex-{int(time.time() * 1000)}" + + if not prompt: + yield _chunk_text( + completion_id, + "(no user prompt found -- send a message before invoking codex)", + ) + yield _chunk_stop(completion_id) + return + + clamped = max(1, min(int(parallel_calls or 1), MAX_PARALLEL_CALLS)) + + if clamped == 1: + async for line in _stream_codex_single(model, system, prompt, completion_id): + yield line + yield "data: [DONE]" + return + + async for line in _stream_codex_parallel( + model = model, + system = system, + prompt = prompt, + n = clamped, + completion_id = completion_id, + ): + yield line + yield "data: [DONE]" + + +async def _stream_codex_parallel( + *, + model: str, + system: str, + prompt: str, + n: int, + completion_id: str, +) -> AsyncGenerator[str, None]: + """Fan out N Codex tasks, emit per-tab chunks, then synthesise. + + Each tab id is a 1-based integer. We emit ``codex_tab_open`` up + front for every tab so the UI can paint the tab strip before any + response lands, then route streamed text per tab via + ``codex_tab_chunk`` (one event per text delta). The final + ``codex_gather`` event carries a synthesis from a separate + standalone Codex call so the user sees both the per-tab raw + outputs AND a unified merged answer. + + Bounded concurrency: ``asyncio.gather`` is used with a fresh + ``AsyncCodex`` instance per tab to keep the SDK from sharing a + single thread object across coroutines (the SDK is not documented + as concurrency-safe on a single ``Codex`` handle). N is clamped to + ``MAX_PARALLEL_CALLS`` by the caller. + """ + queue: asyncio.Queue[Optional[str]] = asyncio.Queue() + + # Pre-emit one tab_open per tab so the UI can paint the tabs + # immediately. The frontend creates the tabs lazily on first + # ``codex_tab_open`` event anyway; pre-emitting just gives us the + # familiar tabs-first / content-second render order. + for idx in range(1, n + 1): + yield _chunk_tool_event( + completion_id, + { + "type": "codex_tab_open", + "tab_id": idx, + "query": prompt, + "total_tabs": n, + }, + ) + + async def _worker(tab_id: int) -> str: + """Run one Codex turn, push every chunk into the queue, and + return the full accumulated text so the synthesis step can + consume it. Errors are surfaced as a ``codex_tab_error`` + tool-event so the tab strip shows which lane failed without + aborting the whole fan-out. + """ + collected: list[str] = [] + try: + sdk = _import_codex() + async_codex_cls = getattr(sdk, "AsyncCodex") + async with async_codex_cls() as codex: + thread_kwargs: dict[str, Any] = {"model": model} + if system: + thread_kwargs["system"] = system + inner_prompt = prompt + try: + thread = await codex.thread_start(**thread_kwargs) + except TypeError: + thread = await codex.thread_start(model = model) + if system: + inner_prompt = f"{system}\n\n{prompt}" + async for text in _stream_thread_run(thread, inner_prompt): + collected.append(text) + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_chunk", + "tab_id": tab_id, + "text": text, + }, + ) + ) + except Exception as exc: + logger.warning( + "codex_provider.parallel_tab_failed", + tab_id = tab_id, + error = str(exc), + ) + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_error", + "tab_id": tab_id, + "error": str(exc), + }, + ) + ) + finally: + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_close", + "tab_id": tab_id, + }, + ) + ) + return "".join(collected) + + workers = [asyncio.create_task(_worker(i + 1)) for i in range(n)] + + # asyncio.gather returns a _GatheringFuture, not a coroutine, so it + # cannot be passed to create_task. Wrap the await in a small helper + # coroutine so we keep both (a) a handle for awaiting and (b) the + # ability to schedule the drain side-effect that unblocks the + # consumer queue. Stash the results into a list the drain finally + # block reads so per-tab outputs survive any single-worker errors. + per_tab_texts: list[str] = [] + + async def _await_workers() -> None: + results = await asyncio.gather(*workers, return_exceptions = True) + for r in results: + if isinstance(r, BaseException): + per_tab_texts.append("") + else: + per_tab_texts.append(r) + + async def _drain_when_done() -> None: + try: + await _await_workers() + finally: + await queue.put(None) + + drain_task = asyncio.create_task(_drain_when_done()) + + while True: + line = await queue.get() + if line is None: + break + yield line + + # Drain finished; per_tab_texts is now populated by the helper + # coroutine above. We already shielded individual errors as + # ``codex_tab_error`` events, so nothing should leak here -- but + # log defensively in case a worker future itself raised. + await drain_task + + synthesis_text = await _run_codex_synthesis( + model = model, + prompt = prompt, + tab_outputs = per_tab_texts, + ) + + yield _chunk_tool_event( + completion_id, + { + "type": "codex_gather", + "summary": synthesis_text, + "tab_count": n, + }, + ) + + # Also emit the synthesis as a visible content chunk so any client + # that ignores the tab tool-events (e.g. a curl user) still sees + # a final unified answer. The tabbed UI re-uses the same payload + # via ``codex_gather`` and hides it from the main content lane to + # avoid duplication. + if synthesis_text: + yield _chunk_text(completion_id, synthesis_text) + + yield _chunk_usage( + completion_id, + prompt_tokens = max(1, len(prompt) // 4), + completion_tokens = max(0, len(synthesis_text) // 4), + ) + yield _chunk_stop(completion_id) + + +async def _run_codex_synthesis( + *, + model: 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. + """ + if not tab_outputs: + return "" + joined = "\n\n".join( + f"=== Attempt {i + 1} ===\n{text.strip() or '(no output)'}" + for i, text in enumerate(tab_outputs) + ) + synthesis_prompt = ( + "You are given multiple independent attempts at the same task.\n" + "Your job: synthesise a single best response that takes the " + "strongest parts of each attempt, resolves disagreements, and " + "presents a clear unified answer.\n\n" + f"Original task:\n{prompt}\n\n" + f"Attempts:\n{joined}\n\n" + "Unified answer:" + ) + try: + sdk = _import_codex() + async_codex_cls = getattr(sdk, "AsyncCodex") + async with async_codex_cls() as codex: + try: + thread = await codex.thread_start(model = model) + except TypeError: + thread = await codex.thread_start() + result = await thread.run(synthesis_prompt) + return _coerce_text(result) or getattr(result, "final_response", "") or str(result) + except Exception as exc: + logger.warning("codex_provider.synthesis_failed", error = str(exc)) + return "" + + +# ── Device-auth helper ────────────────────────────────────────────── + + +async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: + """Run ``codex auth login --device-auth`` and yield progress events. + + Yields dicts (NOT SSE lines) that the route layer wraps in SSE. + First event is always ``{type: "device_url", url: "..."}`` once we + detect a verification URL in the CLI output. Subsequent events + forward CLI stdout/stderr line-by-line as ``{type: "log", line: ...}`` + so the UI can show progress. A final ``{type: "done", ok: bool}`` + signals completion. + + The URL extraction matches the CLI's actual output shape (the CLI + prints something like ``Open https://auth.openai.com/device/...`` + on the device-auth path). We scan every line for the first + https:// URL containing ``device``; that has historically been + stable across CLI versions. + """ + import re + + cli_path = "codex" + args = ["auth", "login", "--device-auth"] + + try: + proc = await asyncio.create_subprocess_exec( + cli_path, + *args, + stdout = asyncio.subprocess.PIPE, + stderr = asyncio.subprocess.STDOUT, + ) + except FileNotFoundError: + yield {"type": "error", "message": "codex CLI not found on PATH"} + yield {"type": "done", "ok": False} + return + except Exception as exc: + yield {"type": "error", "message": str(exc)} + yield {"type": "done", "ok": False} + return + + url_re = re.compile(r"https?://\S*device\S*", re.IGNORECASE) + url_emitted = False + rc: int = -1 + + try: + assert proc.stdout is not None + while True: + line_b = await proc.stdout.readline() + if not line_b: + break + line = line_b.decode("utf-8", errors = "replace").rstrip() + if not url_emitted: + match = url_re.search(line) + if match: + yield {"type": "device_url", "url": match.group(0)} + url_emitted = True + yield {"type": "log", "line": line} + finally: + try: + rc = await proc.wait() + except Exception: + rc = -1 + yield {"type": "done", "ok": rc == 0, "return_code": rc} diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index fef9ba3e12..469c6474ba 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -270,6 +270,41 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { ), "hidden": True, }, + "codex": { + "display_name": "OpenAI Codex (local CLI)", + # No remote base_url: Codex dispatches through the local CLI + # via the codex_app_server SDK. Routing skips the standard + # HTTP client entirely in _proxy_to_external_provider and + # hands the request to core.inference.codex_provider instead. + "base_url": "", + "default_models": [ + "gpt-5.4", + "gpt-5.4-mini", + "gpt-5.5", + "o3", + ], + "supports_streaming": True, + "supports_vision": False, + "supports_tool_calling": True, + # No auth header is sent on the wire; the Codex CLI handles + # auth via its own login flow (api key / chatgpt / device). + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + # Codex models are picked from the local CLI catalogue; we + # never call a remote /models endpoint. + "model_list_mode": "curated", + # Hidden from the cross-provider dropdown until the frontend + # has confirmed availability via GET /api/codex/status. The + # chat-providers dialog conditionally surfaces the entry by + # merging the codex row in when status.installed is true. + "hidden": True, + "notes": ( + "Dispatches chat turns through the local Codex CLI via " + "the codex_app_server Python SDK. Surfaced only when the " + "CLI and SDK are both installed; sign in with `codex auth " + "login`." + ), + }, "openrouter": { "display_name": "OpenRouter", "base_url": "https://openrouter.ai/api/v1", diff --git a/studio/backend/main.py b/studio/backend/main.py index 004ae404cd..18c8edf766 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -115,6 +115,7 @@ from datetime import datetime from routes import ( auth_router, chat_history_router, + codex_router, data_recipe_router, datasets_router, export_router, @@ -522,6 +523,10 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = [" # standard /v1/chat/completions path. app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) +# Codex SDK provider. Status probe + device-auth helper live behind a +# dedicated prefix so the frontend can call them without needing a +# provider config row to exist yet. +app.include_router(codex_router, prefix = "/api/codex", tags = ["codex"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b5626951c4..0fe2e5e629 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -786,6 +786,20 @@ class ChatCompletionRequest(BaseModel): "to auto-create." ), ) + parallel_calls: Optional[int] = Field( + None, + ge = 1, + le = 20, + description = ( + "[x-unsloth] Codex provider only. When > 1, fan the chat turn " + "out across N parallel Codex calls and synthesise a unified " + "final answer. Each parallel attempt is rendered as its own tab " + "in the chat UI; a final 'Synthesis' tab carries the merged " + "output. Bounded to [1, 20] by pydantic so a runaway value can't " + "saturate the local CLI. Silently ignored on every provider " + "other than `codex`." + ), + ) @model_validator(mode = "after") def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest": diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index 6bb5d15e8e..575f2ccf4e 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -16,6 +16,7 @@ from routes.export import router as export_router from routes.training_history import router as training_history_router from routes.chat_history import router as chat_history_router from routes.providers import router as providers_router +from routes.codex import router as codex_router __all__ = [ "training_router", @@ -29,4 +30,5 @@ __all__ = [ "training_history_router", "chat_history_router", "providers_router", + "codex_router", ] diff --git a/studio/backend/routes/codex.py b/studio/backend/routes/codex.py new file mode 100644 index 0000000000..ec69bca8f5 --- /dev/null +++ b/studio/backend/routes/codex.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +API routes for the Codex SDK chat provider. + +Two endpoints live here: + +* ``GET /api/codex/status`` -- the availability probe. The frontend + hits this at chat-page load time and uses ``installed`` to decide + whether to surface the "codex" entry in the provider picker. When + ``installed=True`` but ``logged_in=False``, the provider config + dialog shows the "Sign in to Codex" affordance instead of the + regular API-key field. + +* ``POST /api/codex/login`` -- the device-auth helper. Spawns the + ``codex auth login --device-auth`` CLI command, captures the + verification URL from its output, and streams the rest of the auth + exchange back as SSE so the UI can show progress. The URL appears + in the first SSE event so the frontend can ``window.open`` it before + the user wanders off. +""" + +from __future__ import annotations + +import json +from typing import AsyncGenerator + +import structlog +from fastapi import APIRouter, Depends +from fastapi.responses import StreamingResponse + +from auth.authentication import get_current_subject +from core.inference.codex_availability import probe_codex_availability +from core.inference.codex_provider import stream_codex_device_login + +logger = structlog.get_logger(__name__) + +router = APIRouter() + + +@router.get("/status") +async def get_codex_status( + current_subject: str = Depends(get_current_subject), +) -> dict: + """Return the Codex CLI / SDK availability snapshot. + + The frontend gates the provider entry on ``installed`` and gates + the "Sign in to Codex" button on ``logged_in``. Both are + best-effort and cheap to recompute; the route does not cache the + probe because the user can install the CLI / SDK or run + ``codex auth login`` between page loads and the picker should pick + that up on the next refresh. + """ + return await probe_codex_availability() + + +@router.post("/login") +async def codex_device_login( + current_subject: str = Depends(get_current_subject), +) -> StreamingResponse: + """Stream the ``codex auth login --device-auth`` exchange. + + Returns an SSE stream of events: + + ``data: {"type": "device_url", "url": "https://..."}`` + ``data: {"type": "log", "line": "..."}`` (zero or more) + ``data: {"type": "done", "ok": true}`` + + The frontend opens the device URL in a new tab via + ``window.open(url, "_blank", "noopener,noreferrer")`` as soon as + the first event arrives, then renders the streamed log lines so + the user can see the CLI making progress while they're at the + verification page. + """ + + async def _to_sse() -> AsyncGenerator[str, None]: + async for event in stream_codex_device_login(): + yield f"data: {json.dumps(event)}\n\n" + # Frontend treats the trailing [DONE] the same way it does for + # chat streams, so we emit it for parity. + yield "data: [DONE]\n\n" + + return StreamingResponse( + _to_sse(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..484cf971ae 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1807,10 +1807,13 @@ async def _proxy_to_external_provider( detail = "Either provider_id or provider_type is required for external provider routing.", ) - # Fall back to registry default base URL + # Fall back to registry default base URL. Codex is the one + # provider with a deliberately empty base_url -- it dispatches + # through the local CLI rather than over HTTP -- so a missing + # base_url is only a 400 for every other provider type. if not base_url: base_url = get_base_url(provider_type) - if not base_url: + if not base_url and provider_type != "codex": raise HTTPException( status_code = 400, detail = f"Unknown provider type: {provider_type}", @@ -1845,6 +1848,75 @@ async def _proxy_to_external_provider( provider_type = provider_type, ) + # Codex provider: dispatch through the local CLI / SDK instead of + # the HTTP client. The SDK is not an OpenAI-compatible HTTP + # endpoint; it's a thread-oriented Python API that wraps the CLI. + # ``stream_codex`` is the single entry point so the parallel-calls + # fan-out and the single-call path share the same SSE shape. + if provider_type == "codex": + from core.inference.codex_provider import ( + CodexUnavailableError, + stream_codex, + ) + + async def _codex_stream(): + try: + gen = stream_codex( + messages = chat_messages, + model = model, + parallel_calls = payload.parallel_calls or 1, + ) + sent_done = False + async for line in gen: + yield f"{line}\n\n" + if "[DONE]" in line: + sent_done = True + if not sent_done: + yield "data: [DONE]\n\n" + except CodexUnavailableError as exc: + logger.warning("codex_provider.unavailable", error = str(exc)) + yield ( + "data: " + + json.dumps( + { + "error": { + "message": str(exc), + "type": "provider_error", + "code": "503", + "provider": "codex", + } + } + ) + + "\n\n" + ) + yield "data: [DONE]\n\n" + except Exception as exc: + logger.error("codex_provider.stream_error", error = str(exc)) + yield ( + "data: " + + json.dumps( + { + "error": { + "message": f"Codex error: {exc}", + "type": "provider_error", + "code": "502", + "provider": "codex", + } + } + ) + + "\n\n" + ) + yield "data: [DONE]\n\n" + + return StreamingResponse( + _codex_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + client = ExternalProviderClient( provider_type = provider_type, base_url = base_url, diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py new file mode 100644 index 0000000000..03e3e5abf9 --- /dev/null +++ b/studio/backend/tests/test_codex_provider.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for the Codex SDK provider integration. + +Covers: + +* Availability probe: codex missing, codex present but logged out, + codex present + logged in, plus the empty-output / non-zero rc + edge cases the CLI has shipped over time. +* ``stream_codex`` event translation: a fake codex_app_server module + is dropped into ``sys.modules`` so the production import path runs + without the real SDK installed. Verifies an OpenAI Chat Completions + shape (content chunk, stop chunk, [DONE]). +* Parallel-calls fan-out: ``parallel_calls > 1`` spawns N async tasks + and emits ``codex_tab_open`` / ``codex_tab_chunk`` / ``codex_tab_close`` + events plus a final ``codex_gather`` synthesis event. +* Request validator: ``parallel_calls`` is clamped to [1, 20] by + pydantic so a runaway value is rejected with 422 before any Codex + task is spawned. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import types +from typing import Any + +import pytest + + +_backend = os.path.join(os.path.dirname(__file__), "..") +if _backend not in sys.path: + sys.path.insert(0, _backend) + + +# ── Helpers ───────────────────────────────────────────────────────── + + +class _FakeStream: + """Async iterator that yields predetermined string text events. + + The Codex SDK's ``thread.run_streaming`` returns an async iterable + of events. ``_stream_thread_run`` converts those into raw text via + ``_coerce_text``; passing in plain strings exercises the simplest + coercion path. + """ + + def __init__(self, chunks: list[str]): + self._chunks = list(chunks) + self._i = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self._i >= len(self._chunks): + raise StopAsyncIteration + text = self._chunks[self._i] + self._i += 1 + return text + + +class _FakeThread: + def __init__(self, chunks: list[str], final: str | None = None): + self._chunks = chunks + self._final = final if final is not None else "".join(chunks) + + def run_streaming(self, prompt: str): + # ``run_streaming`` may return either an async iterable or a + # coroutine that resolves to one; cover the direct-return + # shape here, the coroutine shape is covered in a separate + # test below. + return _FakeStream(self._chunks) + + async def run(self, prompt: str): + return self._final + + +class _FakeAsyncCodex: + """Async-context-manager facade matching codex_app_server.AsyncCodex.""" + + def __init__( + self, + chunks: list[str] | None = None, + final: str | None = None, + raise_on_start: Exception | None = None, + ): + self._chunks = chunks or [] + self._final = final + self._raise = raise_on_start + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def thread_start(self, **kwargs): + if self._raise is not None: + raise self._raise + return _FakeThread(self._chunks, self._final) + + +def _install_fake_codex_sdk(monkeypatch, async_codex_cls): + """Drop a fake ``codex_app_server`` module into sys.modules so the + production lazy-import path picks it up without the real SDK + being installed. + """ + fake_mod = types.ModuleType("codex_app_server") + fake_mod.AsyncCodex = async_codex_cls # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "codex_app_server", fake_mod) + # importlib.util.find_spec walks finders, not sys.modules; patch + # it directly so the lazy-import gate accepts the fake. + import importlib.util as _iu + + real_find_spec = _iu.find_spec + + def _shim(name: str, *args, **kwargs): + if name == "codex_app_server": + return types.SimpleNamespace() + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr("importlib.util.find_spec", _shim) + + +# ── Availability probe ───────────────────────────────────────────── + + +class TestCodexAvailability: + def test_absent_when_cli_missing(self, monkeypatch): + from core.inference import codex_availability as ca + + monkeypatch.setattr(ca, "_which_codex", lambda: None) + monkeypatch.setattr(ca, "_sdk_importable", lambda: False) + + payload = asyncio.run(ca.probe_codex_availability()) + assert payload["installed"] is False + assert payload["cli_path"] is None + assert payload["sdk_importable"] is False + # supported_models is a sensible default even when nothing is + # installed so the picker has something to render IF the user + # forces the entry on a future status flip. + assert isinstance(payload["supported_models"], list) + assert len(payload["supported_models"]) > 0 + + def test_present_but_sdk_missing(self, monkeypatch): + from core.inference import codex_availability as ca + + monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex") + monkeypatch.setattr(ca, "_sdk_importable", lambda: False) + + async def fake_version(): + return "codex-cli 0.133.0" + + async def fake_logged_in(): + return True + + monkeypatch.setattr(ca, "_detect_version", fake_version) + monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in) + + payload = asyncio.run(ca.probe_codex_availability()) + # installed requires BOTH CLI and SDK -- this is the gate the + # frontend uses to decide whether to surface the provider entry + # at all, so missing-SDK means hide. + assert payload["installed"] is False + assert payload["cli_path"] == "/usr/local/bin/codex" + assert payload["sdk_importable"] is False + assert payload["version"] == "codex-cli 0.133.0" + + def test_present_and_logged_out(self, monkeypatch): + from core.inference import codex_availability as ca + + monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex") + monkeypatch.setattr(ca, "_sdk_importable", lambda: True) + + async def fake_version(): + return "codex-cli 0.133.0" + + async def fake_logged_in(): + return False + + monkeypatch.setattr(ca, "_detect_version", fake_version) + monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in) + + payload = asyncio.run(ca.probe_codex_availability()) + assert payload["installed"] is True + assert payload["logged_in"] is False + assert payload["version"] == "codex-cli 0.133.0" + + def test_present_and_logged_in(self, monkeypatch): + from core.inference import codex_availability as ca + + monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex") + monkeypatch.setattr(ca, "_sdk_importable", lambda: True) + + async def fake_version(): + return "codex-cli 0.133.0" + + async def fake_logged_in(): + return True + + monkeypatch.setattr(ca, "_detect_version", fake_version) + monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in) + + payload = asyncio.run(ca.probe_codex_availability()) + assert payload["installed"] is True + assert payload["logged_in"] is True + + +# ── _stream_codex translation ────────────────────────────────────── + + +def _collect_stream(gen) -> list[str]: + async def run(): + out: list[str] = [] + async for line in gen: + out.append(line) + return out + + return asyncio.run(run()) + + +def _parse_sse_chunks(lines: list[str]) -> list[dict[str, Any]]: + """Decode SSE ``data: {...}`` lines into the chunk dicts. Skips the + sentinel ``data: [DONE]`` line and anything that isn't valid JSON. + """ + out: list[dict[str, Any]] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + body = raw[len("data:") :].strip() + if not body or body == "[DONE]": + continue + try: + out.append(json.loads(body)) + except json.JSONDecodeError: + continue + return out + + +class TestStreamCodexSingle: + def test_streaming_chunks_translate_into_openai_shape(self, monkeypatch): + _install_fake_codex_sdk( + monkeypatch, + lambda: _FakeAsyncCodex(chunks = ["Hello", ", ", "world"]), + ) + from core.inference.codex_provider import stream_codex + + lines = _collect_stream( + stream_codex( + messages = [{"role": "user", "content": "Say hello in 3 chunks."}], + model = "gpt-5.4", + ) + ) + chunks = _parse_sse_chunks(lines) + # Three content deltas + one usage chunk + one stop chunk. + content_chunks = [ + c + for c in chunks + if c.get("choices") + and isinstance(c["choices"], list) + and c["choices"] + and c["choices"][0].get("delta", {}).get("content") + ] + assert [c["choices"][0]["delta"]["content"] for c in content_chunks] == [ + "Hello", + ", ", + "world", + ] + # Usage chunk (OpenAI include_usage shape) is a choices=[] entry + # with a populated usage block. + usage_chunks = [c for c in chunks if c.get("choices") == [] and c.get("usage")] + assert len(usage_chunks) == 1 + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] > 0 + assert usage["completion_tokens"] >= 0 + # Final stop chunk with finish_reason=stop. + stop_chunks = [ + c + for c in chunks + if c.get("choices") + and c["choices"] + and c["choices"][0].get("finish_reason") == "stop" + ] + assert len(stop_chunks) == 1 + # And the trailing [DONE] sentinel. + assert any(line.strip() == "data: [DONE]" for line in lines) + + def test_empty_user_prompt_emits_helpful_message(self, monkeypatch): + _install_fake_codex_sdk(monkeypatch, lambda: _FakeAsyncCodex(chunks = [])) + from core.inference.codex_provider import stream_codex + + lines = _collect_stream( + stream_codex( + messages = [{"role": "system", "content": "you are helpful"}], + model = "gpt-5.4", + ) + ) + text = "\n".join(lines) + assert "no user prompt" in text.lower() + + +class TestStreamCodexParallel: + def test_parallel_calls_spawn_tabs_and_synthesise(self, monkeypatch): + # The fake SDK returns the same canned chunks for every spawned + # AsyncCodex instance; we just need to verify the orchestrator + # emits N tab_open events, per-tab chunk events keyed by + # tab_id, and a final codex_gather summary event. + _install_fake_codex_sdk( + monkeypatch, + lambda: _FakeAsyncCodex( + chunks = ["alpha"], + final = "synthesised answer", + ), + ) + from core.inference.codex_provider import stream_codex + + n = 3 + lines = _collect_stream( + stream_codex( + messages = [{"role": "user", "content": "Test"}], + model = "gpt-5.4", + parallel_calls = n, + ) + ) + chunks = _parse_sse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + tab_opens = [e for e in tool_events if e.get("type") == "codex_tab_open"] + tab_chunks = [e for e in tool_events if e.get("type") == "codex_tab_chunk"] + tab_closes = [e for e in tool_events if e.get("type") == "codex_tab_close"] + gather = [e for e in tool_events if e.get("type") == "codex_gather"] + + # Each tab opens once -- the N tabs are pre-emitted so the + # UI can paint the strip before content arrives. + assert len(tab_opens) == n + assert sorted(e["tab_id"] for e in tab_opens) == list(range(1, n + 1)) + + # Per-tab chunks may interleave in any order but every tab id + # must produce at least one chunk before its close event. + seen_tabs = {e["tab_id"] for e in tab_chunks} + assert seen_tabs == set(range(1, n + 1)) + + # Each tab emits exactly one close marker. + assert sorted(e["tab_id"] for e in tab_closes) == list(range(1, n + 1)) + + # Exactly one synthesis event with the unified summary. + assert len(gather) == 1 + assert gather[0]["tab_count"] == n + # The summary text comes from the final synthesis Codex call; + # our fake returns "synthesised answer" via .run(). + assert "synth" in gather[0]["summary"].lower() + + def test_parallel_calls_clamped_to_maximum(self, monkeypatch): + """Passing parallel_calls=500 must NOT spawn 500 tasks; the + clamp at MAX_PARALLEL_CALLS keeps the local CLI safe. + """ + from core.inference import codex_provider as cp + + _install_fake_codex_sdk( + monkeypatch, + lambda: _FakeAsyncCodex(chunks = ["x"], final = "synth"), + ) + lines = _collect_stream( + cp.stream_codex( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.4", + parallel_calls = 500, + ) + ) + chunks = _parse_sse_chunks(lines) + tab_opens = [ + c["_toolEvent"] + for c in chunks + if c.get("_toolEvent", {}).get("type") == "codex_tab_open" + ] + assert len(tab_opens) == cp.MAX_PARALLEL_CALLS + + def test_parallel_calls_one_takes_single_path(self, monkeypatch): + """parallel_calls=1 must not emit any tab tool-events -- it's the + regular single-call shape. + """ + _install_fake_codex_sdk( + monkeypatch, + lambda: _FakeAsyncCodex(chunks = ["one"]), + ) + from core.inference.codex_provider import stream_codex + + lines = _collect_stream( + stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.4", + parallel_calls = 1, + ) + ) + chunks = _parse_sse_chunks(lines) + tool_events = [c.get("_toolEvent") for c in chunks if c.get("_toolEvent")] + for event in tool_events: + assert not (event.get("type") or "").startswith("codex_tab") + assert event.get("type") != "codex_gather" + + +# ── Request validator ────────────────────────────────────────────── + + +class TestParallelCallsValidator: + def test_request_accepts_valid_range(self): + from models.inference import ChatCompletionRequest + + for n in (1, 5, 10, 20): + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + parallel_calls = n, + ) + assert req.parallel_calls == n + + def test_request_rejects_below_one(self): + from models.inference import ChatCompletionRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + parallel_calls = 0, + ) + + def test_request_rejects_above_twenty(self): + from models.inference import ChatCompletionRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + parallel_calls = 21, + ) + + def test_request_default_is_none(self): + """Default = None so the field has no effect on every existing + provider that doesn't read it -- preserves backwards compat. + """ + from models.inference import ChatCompletionRequest + + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + ) + assert req.parallel_calls is None + + +# ── Codex unavailable surfacing ──────────────────────────────────── + + +class TestCodexUnavailable: + def test_missing_sdk_raises_typed_error(self, monkeypatch): + # Force find_spec to return None so the lazy import fails. + import importlib.util as _iu + + real = _iu.find_spec + + def _shim(name, *args, **kwargs): + if name == "codex_app_server": + return None + return real(name, *args, **kwargs) + + monkeypatch.setattr("importlib.util.find_spec", _shim) + # Also drop any cached fake from prior tests. + monkeypatch.delitem(sys.modules, "codex_app_server", raising = False) + + from core.inference.codex_provider import ( + CodexUnavailableError, + stream_codex, + ) + + with pytest.raises(CodexUnavailableError): + asyncio.run( + _consume_first( + stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.4", + ) + ) + ) + + +async def _consume_first(gen): + """Drive an async generator until it raises or yields its first + value. Used to surface lazy-import errors that fire on the first + SDK touch -- otherwise the generator would swallow them on + ``__aiter__`` and the test couldn't see them. + """ + async for _ in gen: + return diff --git a/studio/frontend/src/features/chat/api/codex-api.ts b/studio/frontend/src/features/chat/api/codex-api.ts new file mode 100644 index 0000000000..6652bf9391 --- /dev/null +++ b/studio/frontend/src/features/chat/api/codex-api.ts @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * API helpers for the local Codex SDK provider. + * + * The backend exposes two endpoints under ``/api/codex``: + * + * - ``GET /api/codex/status`` returns ``{installed, logged_in, version, + * cli_path, sdk_importable, supported_models}``. The chat-providers + * dialog consults this BEFORE surfacing the "codex" entry in the + * picker -- if ``installed`` is false the provider stays hidden, if + * ``logged_in`` is false we render a "Sign in to Codex" button in + * place of the API-key field. + * + * - ``POST /api/codex/login`` runs ``codex auth login --device-auth`` + * under the hood and streams SSE events. The first event is always + * ``{type: "device_url", url}`` so the UI can window.open it; later + * events forward CLI log lines so the user can watch progress. + */ + +import { authFetch } from "@/features/auth"; + +export interface CodexStatus { + installed: boolean; + logged_in: boolean; + cli_path: string | null; + sdk_importable: boolean; + version: string | null; + supported_models: string[]; +} + +export interface CodexLoginEvent { + type: "device_url" | "log" | "error" | "done"; + url?: string; + line?: string; + message?: string; + ok?: boolean; + return_code?: number; +} + +const DEFAULT_STATUS: CodexStatus = { + installed: false, + logged_in: false, + cli_path: null, + sdk_importable: false, + version: null, + supported_models: [], +}; + +/** + * Fetch the current Codex availability snapshot. Network failures are + * swallowed and reported as ``installed=false`` because every caller + * either uses this to gate UI surfacing (the right answer on error is + * "hide the entry") or kicks off a chat (the right answer on error is + * "fall back to a different provider"). Throwing here would force + * every consumer to wrap the call in a try/catch. + */ +export async function fetchCodexStatus(): Promise { + try { + const response = await authFetch("/api/codex/status"); + if (!response.ok) { + return DEFAULT_STATUS; + } + const body = (await response.json()) as Partial; + return { + installed: Boolean(body.installed), + logged_in: Boolean(body.logged_in), + cli_path: typeof body.cli_path === "string" ? body.cli_path : null, + sdk_importable: Boolean(body.sdk_importable), + version: typeof body.version === "string" ? body.version : null, + supported_models: Array.isArray(body.supported_models) + ? body.supported_models.filter( + (value): value is string => typeof value === "string", + ) + : [], + }; + } catch { + return DEFAULT_STATUS; + } +} + +/** + * Open a Codex device-auth login stream and yield each parsed event. + * + * Returns an async generator the caller drives in a for-await loop -- + * the dialog reads the first ``device_url`` event to know what URL to + * window.open, then collects the remaining ``log`` lines into the + * visible progress area until the ``done`` sentinel arrives. + * + * The generator handles abort signals: when the dialog closes mid- + * flow, the caller passes an AbortSignal that tears down the SSE + * stream cleanly. Without this, the long-running login subprocess + * would keep pumping lines into a torn-down React tree. + */ +export async function* streamCodexDeviceLogin( + signal?: AbortSignal, +): AsyncGenerator { + const response = await authFetch("/api/codex/login", { + method: "POST", + signal, + }); + if (!response.ok || !response.body) { + yield { + type: "error", + message: `codex login request failed: HTTP ${response.status}`, + }; + yield { type: "done", ok: false }; + return; + } + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + try { + while (true) { + const { value, done } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + let newlineIdx; + // SSE frames are separated by blank lines; within each frame the + // payload sits on a single ``data: {...}`` line. Strip both the + // SSE prefix and the [DONE] sentinel before JSON.parse. + while ((newlineIdx = buffer.indexOf("\n\n")) !== -1) { + const frame = buffer.slice(0, newlineIdx); + buffer = buffer.slice(newlineIdx + 2); + for (const line of frame.split("\n")) { + if (!line.startsWith("data:")) continue; + const body = line.slice("data:".length).trim(); + if (!body || body === "[DONE]") continue; + try { + yield JSON.parse(body) as CodexLoginEvent; + } catch { + // Skip any malformed line. The CLI shouldn't ever produce + // these, but it costs nothing to defend against. + } + } + } + } + } finally { + try { + reader.releaseLock(); + } catch { + // ignore + } + } +} diff --git a/studio/frontend/src/features/chat/components/codex-login-button.tsx b/studio/frontend/src/features/chat/components/codex-login-button.tsx new file mode 100644 index 0000000000..7f7737cc28 --- /dev/null +++ b/studio/frontend/src/features/chat/components/codex-login-button.tsx @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * "Sign in to Codex" button + streamed log surface. + * + * Renders in place of the regular API-key field in the provider config + * dialog when ``/api/codex/status`` reports ``logged_in=false``. Click + * fires POST ``/api/codex/login``, which spawns + * ``codex auth login --device-auth`` server-side. The first SSE event + * carries the verification URL -- as soon as it arrives we open it in + * a new tab via ``window.open`` so the user doesn't have to copy-paste + * a long URL out of a log pane. + * + * The button stays mounted while the CLI is exchanging the device + * code: the streamed ``log`` events accumulate into the visible + * progress area until the ``done`` event closes the stream. The + * caller passes ``onLoggedIn`` so the dialog can refetch the status + * probe and flip back into the "ready" state automatically. + */ + +import { useCallback, useRef, useState } from "react"; +import { Button } from "@/components/ui/button"; +import { + streamCodexDeviceLogin, + type CodexLoginEvent, +} from "../api/codex-api"; + +interface Props { + /** Called when the device-auth flow finishes successfully so the + * parent can re-probe ``/api/codex/status`` and switch UI states. */ + onLoggedIn?: () => void; +} + +export function CodexLoginButton({ onLoggedIn }: Props) { + const [busy, setBusy] = useState(false); + const [error, setError] = useState(null); + const [logs, setLogs] = useState([]); + const [deviceUrl, setDeviceUrl] = 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 + // detached component. + const abortRef = useRef(null); + + const startLogin = useCallback(async () => { + if (busy) return; + setBusy(true); + setError(null); + setLogs([]); + setDeviceUrl(null); + const controller = new AbortController(); + abortRef.current?.abort(); + abortRef.current = controller; + try { + let lastOk: boolean | undefined; + for await (const event of streamCodexDeviceLogin( + controller.signal, + ) as AsyncGenerator) { + if (event.type === "device_url" && event.url) { + setDeviceUrl(event.url); + // Open the verification page eagerly so the user doesn't + // have to copy the URL out of the log surface. ``noopener`` + // prevents the auth-tab from controlling the Studio window. + try { + window.open(event.url, "_blank", "noopener,noreferrer"); + } catch { + // Ignore -- the URL is still visible in the log. + } + } else if (event.type === "log" && event.line) { + setLogs((prev) => [...prev, event.line as string]); + } else if (event.type === "error" && event.message) { + setError(event.message); + } else if (event.type === "done") { + lastOk = event.ok; + } + } + if (lastOk) { + onLoggedIn?.(); + } else if (!error) { + setError("Codex login did not complete -- see log for details."); + } + } catch (exc) { + if ((exc as { name?: string } | null)?.name !== "AbortError") { + setError(String((exc as Error)?.message ?? exc)); + } + } finally { + setBusy(false); + } + }, [busy, error, onLoggedIn]); + + return ( +
+ + {deviceUrl && ( +

+ Verification URL:{" "} + + {deviceUrl} + +

+ )} + {error && ( +

{error}

+ )} + {logs.length > 0 && ( + 
+          {logs.join("\n")}
+
+ )} +
+ ); +} diff --git a/studio/frontend/src/features/chat/components/codex-parallel-tabs.tsx b/studio/frontend/src/features/chat/components/codex-parallel-tabs.tsx new file mode 100644 index 0000000000..d48163deb9 --- /dev/null +++ b/studio/frontend/src/features/chat/components/codex-parallel-tabs.tsx @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * Tabbed render for Codex parallel-calls fan-out. + * + * The backend emits four ``_toolEvent`` shapes for ``parallel_calls > 1``: + * + * - ``codex_tab_open {tab_id, query, total_tabs}`` + * - ``codex_tab_chunk {tab_id, text}`` + * - ``codex_tab_close {tab_id}`` + * - ``codex_gather {summary, tab_count}`` + * + * The chat-adapter passes these events into ``useCodexParallelTabs`` + * via the shared tool-event channel. The hook collapses them into a + * tab list (one entry per ``tab_id``) plus a synthesis row, and the + * component below renders a horizontal tab strip with the active + * tab's text in a scrollable panel below. The Synthesis tab is + * highlighted because it's the unified answer the user usually wants + * to read. + */ + +import { useMemo, useState } from "react"; +import { cn } from "@/lib/utils"; + +export interface CodexTabState { + /** 1-based tab index from the backend. */ + tabId: number; + /** Accumulated text from ``codex_tab_chunk`` events for this tab. */ + text: string; + /** True once the matching ``codex_tab_close`` event has arrived. */ + closed: boolean; + /** Set when a ``codex_tab_error`` event was emitted for this tab. */ + error?: string; +} + +export interface CodexParallelState { + /** Per-tab streamed text, keyed by tabId, sorted ascending. */ + tabs: CodexTabState[]; + /** The original user query echoed on each tab_open event. */ + query: string | null; + /** Final synthesis text from the ``codex_gather`` event. */ + synthesis: string | null; + /** Total tabs reported on the first ``codex_tab_open`` event. */ + totalTabs: number; +} + +export type CodexParallelEvent = + | { type: "codex_tab_open"; tab_id: number; query?: string; total_tabs?: number } + | { type: "codex_tab_chunk"; tab_id: number; text: string } + | { type: "codex_tab_close"; tab_id: number } + | { type: "codex_tab_error"; tab_id: number; error?: string } + | { type: "codex_gather"; summary?: string; tab_count?: number }; + +/** + * Pure reducer: given the prior parallel state and a single event, + * return the new state. Kept as a standalone function so the chat- + * adapter can drive it without re-rendering, and so it's trivially + * unit-testable. + */ +export function reduceCodexParallelState( + prev: CodexParallelState, + event: CodexParallelEvent, +): CodexParallelState { + switch (event.type) { + case "codex_tab_open": { + // Idempotent: re-opening an existing tab leaves it intact. + const exists = prev.tabs.some((t) => t.tabId === event.tab_id); + const tabs = exists + ? prev.tabs + : [ + ...prev.tabs, + { tabId: event.tab_id, text: "", closed: false }, + ].sort((a, b) => a.tabId - b.tabId); + return { + ...prev, + tabs, + query: prev.query ?? event.query ?? null, + totalTabs: event.total_tabs ?? Math.max(prev.totalTabs, event.tab_id), + }; + } + case "codex_tab_chunk": { + const tabs = prev.tabs.map((t) => + t.tabId === event.tab_id ? { ...t, text: t.text + event.text } : t, + ); + // Auto-create the slot if a chunk arrived before its open event + // (shouldn't happen with the current backend ordering, but + // defending against that race keeps the UI stable). + if (!tabs.some((t) => t.tabId === event.tab_id)) { + tabs.push({ tabId: event.tab_id, text: event.text, closed: false }); + tabs.sort((a, b) => a.tabId - b.tabId); + } + return { ...prev, tabs }; + } + case "codex_tab_close": { + const tabs = prev.tabs.map((t) => + t.tabId === event.tab_id ? { ...t, closed: true } : t, + ); + return { ...prev, tabs }; + } + case "codex_tab_error": { + const tabs = prev.tabs.map((t) => + t.tabId === event.tab_id + ? { ...t, closed: true, error: event.error } + : t, + ); + return { ...prev, tabs }; + } + case "codex_gather": { + return { ...prev, synthesis: event.summary ?? "" }; + } + default: { + return prev; + } + } +} + +export const EMPTY_CODEX_PARALLEL_STATE: CodexParallelState = { + tabs: [], + query: null, + synthesis: null, + totalTabs: 0, +}; + +/** True when the state carries at least one observed event. */ +export function hasCodexParallelContent(state: CodexParallelState): boolean { + return state.tabs.length > 0 || state.synthesis !== null; +} + +interface Props { + state: CodexParallelState; + /** Collapsed by default per spec; user clicks to expand. */ + defaultCollapsed?: boolean; +} + +export function CodexParallelTabs({ state, defaultCollapsed = true }: Props) { + const [collapsed, setCollapsed] = useState(defaultCollapsed); + const [activeTab, setActiveTab] = useState("synthesis"); + + // Whenever the synthesis arrives, switch to it automatically -- it's + // the answer the user usually reads. Use a useMemo + effect-like + // pattern via render-time check so we don't depend on extra hooks. + // (A useEffect would also work; this stays lighter.) + const effectiveActive = useMemo(() => { + if (state.synthesis && activeTab !== "synthesis") { + return activeTab; + } + if (state.synthesis) { + return "synthesis"; + } + if (state.tabs.length > 0 && activeTab === "synthesis") { + return state.tabs[0].tabId; + } + return activeTab; + }, [state.synthesis, state.tabs, activeTab]); + + if (!hasCodexParallelContent(state)) { + return null; + } + + const totalSlots = state.totalTabs || state.tabs.length; + + return ( +
+ + {!collapsed && ( + <> +
+ {state.tabs.map((tab) => ( + + ))} + {state.synthesis !== null && ( + + )} +
+
+ {effectiveActive === "synthesis" + ? state.synthesis || "(waiting for synthesis…)" + : (state.tabs.find((t) => t.tabId === effectiveActive)?.text || + "(waiting…)")} +
+ + )} +
+ ); +} diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index dddf529068..8b25f60bf2 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -79,11 +79,45 @@ export function supportsProviderReasoningToggle( ); } +/** + * The Codex CLI / SDK provider. Surfaced only when the host has BOTH + * the ``codex`` CLI on PATH and the ``codex_app_server`` Python SDK + * importable -- the backend's ``GET /api/codex/status`` is the + * authoritative gate. We expose the type id here so the rest of the + * frontend can reference it without scattering "codex" string + * literals. + */ +export const CODEX_PROVIDER_TYPE = "codex"; + +export function isCodexProviderType( + providerType: string | null | undefined, +): boolean { + return providerType === CODEX_PROVIDER_TYPE; +} + +/** Hard cap mirrors backend MAX_PARALLEL_CALLS to keep the UI honest. */ +export const CODEX_MAX_PARALLEL_CALLS = 20; +export const CODEX_DEFAULT_PARALLEL_CALLS = 1; + +export function clampCodexParallelCalls(value: unknown): number { + const n = typeof value === "number" && Number.isFinite(value) + ? Math.floor(value) + : CODEX_DEFAULT_PARALLEL_CALLS; + if (n < 1) return 1; + if (n > CODEX_MAX_PARALLEL_CALLS) return CODEX_MAX_PARALLEL_CALLS; + return n; +} + // Known text-only providers on their main chat endpoint. const NON_VISION_PROVIDER_TYPES = new Set([ "cohere", "deepseek", "mistral", + // Codex SDK input is text-first; multimodal attachments are + // converted to placeholder text descriptors before the prompt + // reaches the local CLI. Mark text-only so the composer hides + // image-attach affordances when codex is selected. + CODEX_PROVIDER_TYPE, ]); // Providers whose vision-tier model selection accepts images. const VISION_CAPABLE_PROVIDER_TYPES = new Set([ From ea7ae855625453cc8fa596fe3de767060c387ee7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 May 2026 16:46:41 +0000 Subject: [PATCH 02/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/codex_provider.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index d4df67ab4f..876975435d 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -598,7 +598,9 @@ async def _run_codex_synthesis( except TypeError: thread = await codex.thread_start() result = await thread.run(synthesis_prompt) - return _coerce_text(result) or getattr(result, "final_response", "") or str(result) + return ( + _coerce_text(result) or getattr(result, "final_response", "") or str(result) + ) except Exception as exc: logger.warning("codex_provider.synthesis_failed", error = str(exc)) return "" From 4188f916f436bdd205ac9260ace48f0003cb396d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 16:30:15 +0000 Subject: [PATCH 03/40] wip: anthropic citation helper --- .../core/inference/external_provider.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 25e1725337..eed716c305 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -177,6 +177,42 @@ def _anthropic_supports_compaction(model: str) -> bool: return model.startswith(_ANTHROPIC_COMPACTION_PREFIXES) +def _anthropic_citation_key(citation: dict[str, Any]) -> tuple: + """Stable dedup key for an Anthropic ``citations_delta.citation``. + + Shape varies per document type per + https://platform.claude.com/docs/en/build-with-claude/citations + and https://platform.claude.com/docs/en/build-with-claude/search-results : + + * ``char_location``: ``document_index`` + ``start_char_index`` + * ``page_location``: ``document_index`` + ``start_page_number`` + * ``content_block_location``: ``document_index`` + ``start_block_index`` + * ``search_result_location``: ``document_index`` + ``source`` + + ``start_block_index`` + + Anything unrecognised falls back to a stringified copy so a future + shape still dedupes (worst case: more entries, never collisions). + """ + ctype = citation.get("type") + doc = citation.get("document_index") + title = citation.get("document_title") or "" + if ctype == "char_location": + return (ctype, doc, title, citation.get("start_char_index")) + if ctype == "page_location": + return (ctype, doc, title, citation.get("start_page_number")) + if ctype == "content_block_location": + return (ctype, doc, title, citation.get("start_block_index")) + if ctype == "search_result_location": + return ( + ctype, + doc, + title, + citation.get("source"), + citation.get("start_block_index"), + ) + return (ctype, _json.dumps(citation, sort_keys = True)) + + class _MistralThinkingSpec(NamedTuple): models: tuple[str, ...] style: Literal["prompt_mode", "reasoning_effort", "disabled"] From 0a4309fefdc4bbd71b436db1ca1047f065c3099f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 18:35:02 +0000 Subject: [PATCH 04/40] Studio: add codex_router to test_desktop_auth router stub test_health_response_reports_desktop_capability_fields builds a SimpleNamespace as a fake routes module so it can exercise main.health_check without standing the full app up. The stub listed every router name except codex_router, which lands in the main.py import block alongside the others as of this PR, so the import failed with 'cannot import name codex_router from ' on the Python 3.13 unit run. Add the codex_router slot to the stub. --- studio/backend/tests/test_desktop_auth.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index ab1a03eeda..ee1e7fb18f 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -432,6 +432,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): router_stub = SimpleNamespace( auth_router = APIRouter(), chat_history_router = APIRouter(), + codex_router = APIRouter(), data_recipe_router = APIRouter(), datasets_router = APIRouter(), export_router = APIRouter(), From b8cd6773970a6fa66511f38f30a56bf5461970b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 14:13:03 +0000 Subject: [PATCH 05/40] Studio: harden Codex provider against upstream CLI/SDK shape Followups on the post-merge review pass for the Codex SDK chat provider. Verified against codex-cli 0.133.0 + the upstream `openai/codex` Rust + Python sources, then pinned each fix with a regression test in `test_codex_provider.py` (24/24 passing). * Probe both `openai_codex` (canonical upstream Python package at `openai/codex/sdk/python`) and the legacy `codex_app_server` alias. Without this the availability probe always reported `sdk_importable: false` even when the SDK was installed, so the provider was permanently hidden. * Switch the device-auth and login-status invocations from `codex auth login --device-auth` / `codex auth status` to the real upstream subcommands `codex login --device-auth` and `codex login status`. The former path returns `unrecognized subcommand 'auth'` on a real CLI. * Strip ANSI control sequences before extracting the device URL (upstream wraps the URL in `\x1b[34m...\x1b[0m`) and tighten the pattern to the canonical `.../codex/device` shape. Also surface the one-time code as a `device_code` SSE event so the UI can show it alongside the URL. * Fix `_detect_logged_in` substring footgun: `"logged in" in combined` matched inside `"not logged in"`, flipping logged-out users to logged-in. Anchor on word boundaries with negative prefixes winning regardless of return code. * Cancel in-flight fan-out workers on SSE disconnect. Previously every parallel Codex turn ran to completion against a disconnected client and burned quota; now `_stream_codex_parallel` cancels its worker + drain tasks in a try/finally on `CancelledError`/`GeneratorExit`. * Tear down the device-login subprocess on disconnect via `start_new_session=True` + `os.killpg(SIGTERM)` (Unix) or `CREATE_NEW_PROCESS_GROUP` + `CTRL_BREAK_EVENT` (Windows), with a bounded `proc.wait()` and `proc.kill()` fallback. Previously `finally: await proc.wait()` blocked the SSE close path because `codex login --device-auth` only exits on user action. * Render the full conversation transcript in `_last_user_prompt` instead of returning only the most recent user message. The PR opens a fresh thread per request so prior assistant turns were dropped, degrading multi-turn chats to single-shot prompts. Single-turn input is unchanged. * Make `ChatCompletionRequest.parallel_calls` default to 1 (`int` with `ge=1, le=20`) instead of `Optional[int] = None`. The runtime already coerced `None` -> 1, but the schema now matches the documented `[1, 20]` range. * Replace the registry's hardcoded `default_models` (which contained `o3`, not in the upstream catalog) with the current `gpt-5.5 / 5.4 / 5.4-mini / 5.3-codex / 5.2` set from `codex-rs/models-manager/models.json`. * Stop echoing `str(exc)` in SSE error frames in both `routes/inference.py` and `routes/codex.py`. The Codex SDK can raise with local paths, env-var content, or traceback fragments (CodeQL `py/information-exposure-through-exception`). Surface a generic message + `exception_type` discriminator; log the full reason server-side via `logger.error(..., exc_type=..., error=...)`. Doc / comment updates throughout to refer to `codex login` / `openai_codex` rather than the older incorrect strings. Tested: pytest 24 cases in `test_codex_provider.py` (the original 14 + 10 new `TestCodexHardenedRegressions`) plus the rest of the Studio-backend test suite the PR touches (209 passing). Also verified live against Studio launched from this branch on a Blackwell B200 via `UNSLOTH_STUDIO_HOME=$WORKSPACE/temp/... ./install.sh --local` then a Playwright probe. --- .../core/inference/codex_availability.py | 114 +++--- .../backend/core/inference/codex_provider.py | 325 +++++++++++++----- studio/backend/core/inference/providers.py | 23 +- studio/backend/models/inference.py | 8 +- studio/backend/routes/codex.py | 37 +- studio/backend/routes/inference.py | 15 +- studio/backend/tests/test_codex_provider.py | 168 ++++++++- 7 files changed, 527 insertions(+), 163 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index c7488a0ed0..8ae57deadd 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -4,9 +4,9 @@ """ Codex CLI / SDK availability probe. -This module never imports `codex_app_server` at module top level. The -SDK is optional and not pinned in pyproject.toml -- if it's installed -locally, we use it; if it isn't, the probe simply returns +This module never imports the Codex Python SDK at module top level. +The SDK is optional and not pinned in pyproject.toml -- if it's +installed locally, we use it; if it isn't, the probe simply returns ``installed=False`` and the provider stays hidden in the frontend. The frontend calls ``GET /api/codex/status`` at startup to decide @@ -14,12 +14,12 @@ whether to surface the "codex" entry in the provider picker. Three states matter: * ``installed=False`` -- either the CLI is missing OR the SDK - (``codex_app_server``) is not importable. The picker hides the - entry entirely. + (``openai_codex`` canonical, or ``codex_app_server`` legacy alias) + is not importable. The picker hides the entry entirely. * ``installed=True, logged_in=False`` -- everything resolves on the - Python side but ``codex auth status`` (or equivalent) reports no - active credentials. The provider config dialog shows a - ``Sign in to Codex`` button instead of the regular API-key field. + Python side but ``codex login status`` reports no active credentials. + The provider config dialog shows a ``Sign in to Codex`` button + instead of the regular API-key field. * ``installed=True, logged_in=True`` -- ready to use; the picker shows the regular model dropdown. @@ -44,14 +44,22 @@ logger = structlog.get_logger(__name__) # Default catalog of models surfaced in the picker when the CLI is # present but doesn't advertise a list. The SDK accepts arbitrary model -# ids; this is purely a sensible default. +# ids; this is purely a sensible default. Mirrored from upstream +# ``codex-rs/models-manager/models.json``. _DEFAULT_SUPPORTED_MODELS: tuple[str, ...] = ( + "gpt-5.5", "gpt-5.4", "gpt-5.4-mini", - "gpt-5.5", - "o3", + "gpt-5.3-codex", + "gpt-5.2", ) +# Names the upstream Python SDK has shipped under. ``openai_codex`` is the +# canonical package at ``openai/codex/sdk/python``; ``codex_app_server`` is +# kept as a forward-compat alias because the Rust crate uses that name and +# an internal alpha may publish under it. +_SDK_MODULE_NAMES: tuple[str, ...] = ("openai_codex", "codex_app_server") + def _which_codex() -> Optional[str]: """Return absolute path to the ``codex`` CLI, or None if missing. @@ -71,19 +79,29 @@ def _which_codex() -> Optional[str]: def _sdk_importable() -> bool: - """True iff ``codex_app_server`` is importable in this interpreter. + """True iff the Codex Python SDK is importable in this interpreter. We deliberately use :func:`importlib.util.find_spec` instead of an - ``import codex_app_server`` so the import never actually runs -- - that keeps the cost negligible and avoids the SDK's own side - effects (which include reaching out to the CLI subprocess for an - RPC ping) during a simple availability check. + actual ``import`` so the import never runs -- that keeps the cost + negligible and avoids the SDK's own side effects (which include + reaching out to the CLI subprocess for an RPC ping) during a + simple availability check. + + Probes both ``openai_codex`` (the canonical upstream package name + at ``openai/codex/sdk/python``) and ``codex_app_server`` (the Rust + crate name, kept as a forward-compat alias). """ - try: - return importlib.util.find_spec("codex_app_server") is not None - except Exception as exc: - logger.warning("codex_availability.find_spec_failed", error = str(exc)) - return False + for name in _SDK_MODULE_NAMES: + try: + if importlib.util.find_spec(name) is not None: + return True + except Exception as exc: + logger.warning( + "codex_availability.find_spec_failed", + module = name, + error = str(exc), + ) + return False async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, str]: @@ -142,33 +160,43 @@ async def _detect_version() -> Optional[str]: async def _detect_logged_in() -> bool: - """Best-effort: assume a non-zero ``codex auth status`` rc means - "not logged in". The CLI surface has shifted over releases (some - versions print to stderr, some to stdout, some use "Logged in as - ..." vs "Not authenticated"). The return code is the most stable - signal we have, so we lean on it and fall back to substring - inspection only when the rc itself is ambiguous (rc=0 but no - output, or rc=-1 from a timeout / crash). + """Best-effort: parse ``codex login status`` output. + + The upstream subcommand is ``codex login status`` (no ``auth`` + parent). Output shapes seen in the wild: + * "Logged in using ChatGPT" / "Logged in as user@x.com" -> True + * "Not logged in. Run `codex login` ..." -> False + * "Not authenticated" -> False + Return code is the most stable signal but ``not logged in`` also + exits 0 on current releases, so we substring-check explicitly. + + Note: a naive ``"logged in" in combined`` check is wrong because + the substring appears inside "not logged in" too -- we use an + explicit negative-prefix check first. """ - rc, stdout, stderr = await _run_cli(["auth", "status"]) + import re + + rc, stdout, stderr = await _run_cli(["login", "status"]) combined = f"{stdout}\n{stderr}".lower() + + # 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. + negative = re.compile(r"\b(not logged in|not authenticated|please log in|run\s+`?codex login`?)\b") + if negative.search(combined): + return False + + positive = re.compile(r"\b(logged in|authenticated as|signed in)\b") + if positive.search(combined): + return True + if rc == 0: - # Some 0.x releases exit 0 even when the user is logged out. - # If we got any output, look for the obvious "not logged in" - # signals; if there's nothing on either pipe at all, treat - # rc=0 as authenticated (the optimistic default). + # rc=0 with nothing useful on either pipe: optimistic default, + # the user is probably authenticated and the CLI just stayed + # quiet (e.g. a future release). if not combined.strip(): return True - if "not authenticated" in combined or "not logged in" in combined: - return False - if "logged in" in combined or "authenticated" in combined: - return True - return True - # rc != 0: lean on substring matching one more time before - # defaulting to False, in case a future CLI uses rc=2 for the - # logged-out state instead of "no auth command exists". - if "logged in" in combined or "authenticated as" in combined: - return True + return False return False diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 876975435d..299de46be3 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -4,10 +4,11 @@ """ Codex SDK chat provider. -This module wraps the ``codex_app_server`` async SDK so a chat request -routed at ``provider_type="codex"`` can dispatch through the user's -local Codex CLI. The contract back to the frontend is the standard -OpenAI Chat Completions SSE shape, exactly like every other entry in +This module wraps the OpenAI Codex async Python SDK (``openai_codex`` +canonical, ``codex_app_server`` legacy alias) so a chat request routed +at ``provider_type="codex"`` can dispatch through the user's local +Codex CLI. The contract back to the frontend is the standard OpenAI +Chat Completions SSE shape, exactly like every other entry in ``external_provider.py``. The two interesting features here: @@ -23,14 +24,14 @@ The two interesting features here: emit per-tab ``_toolEvent`` markers so the frontend can render each result in its own tab. A final ``codex_gather`` synthesis tab runs a single Codex call that takes the N outputs and produces a unified - answer. + answer. Workers cancel cleanly when the SSE consumer disconnects so + cancelled fan-outs never leave zombie Codex calls running. The SDK is imported lazily (inside the helpers that actually need it) -because the spec calls out that ``codex_app_server`` may not even be -importable on the build host. The availability probe in -``codex_availability.py`` is what the frontend uses to gate the -provider entirely; this module just refuses to run if the import -fails at request time. +because the SDK may not be importable on the build host. The +availability probe in ``codex_availability.py`` is what the frontend +uses to gate the provider entirely; this module just refuses to run +if the import fails at request time. """ from __future__ import annotations @@ -53,8 +54,15 @@ logger = structlog.get_logger(__name__) MAX_PARALLEL_CALLS = 20 +# Names the upstream Python SDK has shipped under. ``openai_codex`` is +# the canonical package at ``openai/codex/sdk/python``; ``codex_app_server`` +# is kept as a forward-compat alias because the Rust crate uses that name. +# Order matters: first hit wins, so the canonical name is tried first. +_SDK_MODULE_NAMES: tuple[str, ...] = ("openai_codex", "codex_app_server") + + class CodexUnavailableError(RuntimeError): - """Raised when ``codex_app_server`` is not importable at runtime. + """Raised when the Codex Python SDK is not importable at runtime. The availability probe is supposed to hide the provider before any request lands here, but we still raise a typed error so the route @@ -64,62 +72,103 @@ class CodexUnavailableError(RuntimeError): def _import_codex() -> Any: - """Return the imported ``codex_app_server`` module or raise. + """Return the imported Codex SDK module or raise CodexUnavailableError. Imported lazily so the rest of the backend keeps starting cleanly - on hosts that don't have the SDK installed. The frontend calls - ``GET /api/codex/status`` first and hides the provider when the - spec isn't importable, so this branch is reached only when the - user (a) explicitly forces the provider via a stale stored config - or (b) the install state changes between status probe and chat - submit. + on hosts that don't have the SDK installed. Probes ``openai_codex`` + first (canonical upstream name), then ``codex_app_server`` (Rust- + crate-style alias) for forward compatibility. The frontend calls + ``GET /api/codex/status`` first and hides the provider when no + name resolves, so this branch is reached only when (a) the user + explicitly forces the provider via a stale stored config or (b) + the install state changes between status probe and chat submit. """ - if importlib.util.find_spec("codex_app_server") is None: - raise CodexUnavailableError( - "codex_app_server is not installed on this host. " - "Install the Codex Python SDK or use a different provider." - ) - return importlib.import_module("codex_app_server") + for name in _SDK_MODULE_NAMES: + if importlib.util.find_spec(name) is not None: + return importlib.import_module(name) + raise CodexUnavailableError( + "Codex Python SDK is not installed on this host. " + "Install `openai-codex` (or the legacy `codex_app_server`) " + "or use a different provider." + ) + + +def _stringify_content(content: Any) -> str: + """Flatten an OpenAI-style content field into one plain-text block. + + Multimodal entries (images, documents) are described inline rather + than embedded since the Codex SDK input shape is text-first. + """ + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for entry in content: + if not isinstance(entry, dict): + continue + if entry.get("type") == "text": + parts.append(str(entry.get("text") or "")) + elif entry.get("type") == "image_url": + url = (entry.get("image_url") or {}).get("url") or "" + parts.append(f"[image: {url[:80]}]" if url else "[image]") + elif entry.get("type") == "input_document": + name = entry.get("filename") or "document" + parts.append(f"[document: {name}]") + return "\n".join(p for p in parts if p) + return "" def _last_user_prompt(messages: list[dict[str, Any]]) -> str: - """Extract the most recent user-role message as a plain string. + """Render the conversation as a single prompt for Codex. - Codex's ``thread.run`` accepts a string (per the docs note - "plain strings are accepted anywhere a turn input is accepted"). - Studio's chat history is a full OpenAI-style messages array, so - we flatten it: walk from the end, find the last ``role=user`` - message, and stringify any structured content parts into a - newline-joined block. Multimodal content (images) is described - rather than embedded; Codex SDK input shape is text-first. - - This is intentionally conservative -- we don't try to replay the - whole conversation through Codex per turn because the SDK is - designed around a stateful ``thread`` object. The thread itself - holds context across runs; we only need to feed the latest user + Studio's chat history is a full OpenAI-style messages array. Codex + opens a fresh ``thread`` per chat-completion request (we have no + way to cache the SDK ``thread`` keyed on Studio's session id from + here -- the inference route is stateless), so we MUST serialise the + full transcript into the prompt or the model loses every prior turn. + + Layout: + User: + Assistant: + User: + ... + Assistant: + + The trailing ``Assistant:`` cue tells Codex this is its turn. When + there is exactly one user message and no assistant history we drop + the cue and emit the plain text -- matches the historical single- + shot behaviour. """ - for msg in reversed(messages): - if msg.get("role") != "user": + # Filter to user / assistant only; system is handled separately by + # `_system_prompt` and passed to thread_start when supported. + convo: list[tuple[str, str]] = [] + for msg in messages: + role = msg.get("role") + if role not in ("user", "assistant"): continue - content = msg.get("content") - if isinstance(content, str): - return content - if isinstance(content, list): - parts: list[str] = [] - for entry in content: - if not isinstance(entry, dict): - continue - if entry.get("type") == "text": - parts.append(str(entry.get("text") or "")) - elif entry.get("type") == "image_url": - url = (entry.get("image_url") or {}).get("url") or "" - parts.append(f"[image: {url[:80]}]" if url else "[image]") - elif entry.get("type") == "input_document": - name = entry.get("filename") or "document" - parts.append(f"[document: {name}]") - return "\n".join(p for p in parts if p) - return "" + text = _stringify_content(msg.get("content")) + if text: + convo.append((role, text)) + + if not convo: + return "" + + # Trivial case: a single user turn — pass it through unchanged so we + # don't perturb single-shot behaviour or test expectations. + if len(convo) == 1 and convo[0][0] == "user": + return convo[0][1] + + # Multi-turn: render User:/Assistant: blocks then prompt the model. + lines: list[str] = [] + for role, text in convo: + label = "User" if role == "user" else "Assistant" + lines.append(f"{label}: {text}") + # If the last turn is from the user (the common case), append an + # empty Assistant cue so Codex picks up from the right side. + if convo[-1][0] == "user": + lines.append("Assistant:") + return "\n\n".join(lines) def _system_prompt(messages: list[dict[str, Any]]) -> str: @@ -304,7 +353,7 @@ async def _stream_codex_single( async_codex_cls = getattr(sdk, "AsyncCodex", None) if async_codex_cls is None: raise CodexUnavailableError( - "codex_app_server is installed but AsyncCodex is missing -- " + "Codex SDK is installed but AsyncCodex is missing -- " "upgrade the SDK." ) @@ -520,17 +569,42 @@ async def _stream_codex_parallel( drain_task = asyncio.create_task(_drain_when_done()) - while True: - line = await queue.get() - if line is None: - break - yield line - - # Drain finished; per_tab_texts is now populated by the helper - # coroutine above. We already shielded individual errors as - # ``codex_tab_error`` events, so nothing should leak here -- but - # log defensively in case a worker future itself raised. - await drain_task + cancelled = False + try: + while True: + line = await queue.get() + if line is None: + break + yield line + except (asyncio.CancelledError, GeneratorExit): + cancelled = True + # Cancel every in-flight worker so the Codex SDK calls release + # their quota / sockets instead of running to completion against + # a disconnected client. Workers shield themselves in `_worker`'s + # own try/finally so we just signal cancellation and bail. + for w in workers: + if not w.done(): + w.cancel() + drain_task.cancel() + # Best-effort gather so cancellation propagates and we don't + # leave coroutines hanging on the event loop. + try: + await asyncio.gather(*workers, drain_task, return_exceptions = True) + except Exception: + pass + raise + finally: + # If we exited normally, drain_task is already done (it put None + # on the queue right after gather). On the cancellation path we + # already gathered above, so this await is a fast no-op. + if not cancelled: + try: + await drain_task + except Exception as exc: + logger.warning( + "codex_provider.parallel_drain_failed", + error = str(exc), + ) synthesis_text = await _run_codex_synthesis( model = model, @@ -610,45 +684,77 @@ async def _run_codex_synthesis( async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: - """Run ``codex auth login --device-auth`` and yield progress events. + """Run ``codex login --device-auth`` and yield progress events. Yields dicts (NOT SSE lines) that the route layer wraps in SSE. First event is always ``{type: "device_url", url: "..."}`` once we - detect a verification URL in the CLI output. Subsequent events - forward CLI stdout/stderr line-by-line as ``{type: "log", line: ...}`` - so the UI can show progress. A final ``{type: "done", ok: bool}`` + detect a verification URL in the CLI output. The one-time code is + emitted as ``{type: "device_code", code: "ABCD-EFGH"}`` so the UI + can show it next to the URL (upstream CLI prints both on separate + lines). Subsequent CLI stdout/stderr lines forward as + ``{type: "log", line: ...}``. A final ``{type: "done", ok: bool}`` signals completion. - The URL extraction matches the CLI's actual output shape (the CLI - prints something like ``Open https://auth.openai.com/device/...`` - on the device-auth path). We scan every line for the first - https:// URL containing ``device``; that has historically been - stable across CLI versions. + Subprocess lifecycle: started in its own process group via + ``start_new_session=True`` (Unix) so cancellation can SIGTERM the + whole group and reach any child processes the codex CLI spawns. + On Windows a fallback uses ``CREATE_NEW_PROCESS_GROUP``. When the + SSE consumer disconnects, the generator's cleanup terminates the + process group within a 5s budget then SIGKILL's as a last resort, + so the CLI never lingers consuming a device-auth session. + + URL handling: upstream ``codex login --device-auth`` prints the URL + wrapped in ANSI escape sequences (``\x1b[34m...\x1b[0m``). We strip + ANSI before regex matching so the URL emitted to the frontend is + clean and clickable. """ + import os import re + import signal cli_path = "codex" - args = ["auth", "login", "--device-auth"] + args = ["login", "--device-auth"] + + # Detach the subprocess into a new process group on Unix so we can + # 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)``. + spawn_kwargs: dict[str, Any] = { + "stdout": asyncio.subprocess.PIPE, + "stderr": asyncio.subprocess.STDOUT, + } + if os.name == "posix": + spawn_kwargs["start_new_session"] = True + elif os.name == "nt": + spawn_kwargs["creationflags"] = 0x00000200 # CREATE_NEW_PROCESS_GROUP try: - proc = await asyncio.create_subprocess_exec( - cli_path, - *args, - stdout = asyncio.subprocess.PIPE, - stderr = asyncio.subprocess.STDOUT, - ) + proc = await asyncio.create_subprocess_exec(cli_path, *args, **spawn_kwargs) except FileNotFoundError: yield {"type": "error", "message": "codex CLI not found on PATH"} yield {"type": "done", "ok": False} return except Exception as exc: - yield {"type": "error", "message": str(exc)} + logger.warning("codex_provider.login_spawn_failed", error = str(exc)) + yield {"type": "error", "message": "Failed to start codex CLI"} yield {"type": "done", "ok": False} return - url_re = re.compile(r"https?://\S*device\S*", re.IGNORECASE) + # 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. + url_re = re.compile(r"https?://[^\s\x1b]+?/codex/device(?:\?[^\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. + code_re = re.compile(r"\b([A-Z0-9]{4}-[A-Z0-9]{4})\b") + url_emitted = False + code_emitted = False rc: int = -1 + cancelled = False try: assert proc.stdout is not None @@ -656,16 +762,57 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: line_b = await proc.stdout.readline() if not line_b: break - line = line_b.decode("utf-8", errors = "replace").rstrip() + raw = line_b.decode("utf-8", errors = "replace").rstrip() + line = ansi_re.sub("", raw) if not url_emitted: match = url_re.search(line) if match: yield {"type": "device_url", "url": match.group(0)} url_emitted = True + if not code_emitted: + cm = code_re.search(line) + if cm: + yield {"type": "device_code", "code": cm.group(1)} + code_emitted = True yield {"type": "log", "line": line} + except (asyncio.CancelledError, GeneratorExit): + cancelled = True + raise finally: + # Tear the subprocess down even on cancellation. Unix: kill the + # whole process group; Windows: ``CTRL_BREAK_EVENT`` followed by + # ``terminate()``. Bounded wait so cleanup never deadlocks the + # SSE close path. + if proc.returncode is None: + try: + if os.name == "posix": + try: + pgid = os.getpgid(proc.pid) + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + pass + else: + try: + proc.send_signal(signal.CTRL_BREAK_EVENT) # type: ignore[attr-defined] + except Exception: + proc.terminate() + except Exception as exc: + logger.warning( + "codex_provider.login_terminate_failed", + error = str(exc), + ) + try: + await asyncio.wait_for(proc.wait(), timeout = 5.0) + except (asyncio.TimeoutError, asyncio.CancelledError): + try: + proc.kill() + await asyncio.wait_for(proc.wait(), timeout = 2.0) + except Exception: + pass try: - rc = await proc.wait() + rc = proc.returncode if proc.returncode is not None else -1 except Exception: rc = -1 + if cancelled: + return yield {"type": "done", "ok": rc == 0, "return_code": rc} diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index 469c6474ba..6ab1ac47bb 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -273,15 +273,22 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { "codex": { "display_name": "OpenAI Codex (local CLI)", # No remote base_url: Codex dispatches through the local CLI - # via the codex_app_server SDK. Routing skips the standard - # HTTP client entirely in _proxy_to_external_provider and - # hands the request to core.inference.codex_provider instead. + # via the openai_codex Python SDK (legacy alias: codex_app_server). + # Routing skips the standard HTTP client entirely in + # _proxy_to_external_provider and hands the request to + # core.inference.codex_provider instead. "base_url": "", + # Mirrored from upstream ``codex-rs/models-manager/models.json``. + # We deliberately drop ``o3`` (not in the upstream catalog) and + # add ``gpt-5.3-codex`` + ``gpt-5.2``. Once the SDK exposes a + # runtime ``Codex.models()`` call the dynamic catalog will + # replace this hardcoded default. "default_models": [ + "gpt-5.5", "gpt-5.4", "gpt-5.4-mini", - "gpt-5.5", - "o3", + "gpt-5.3-codex", + "gpt-5.2", ], "supports_streaming": True, "supports_vision": False, @@ -300,9 +307,9 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { "hidden": True, "notes": ( "Dispatches chat turns through the local Codex CLI via " - "the codex_app_server Python SDK. Surfaced only when the " - "CLI and SDK are both installed; sign in with `codex auth " - "login`." + "the openai-codex Python SDK (legacy alias: codex_app_server). " + "Surfaced only when the CLI and SDK are both installed; " + "sign in with `codex login`." ), }, "openrouter": { diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0fe2e5e629..65a73dd65f 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -786,8 +786,8 @@ class ChatCompletionRequest(BaseModel): "to auto-create." ), ) - parallel_calls: Optional[int] = Field( - None, + parallel_calls: int = Field( + default = 1, ge = 1, le = 20, description = ( @@ -796,8 +796,8 @@ class ChatCompletionRequest(BaseModel): "final answer. Each parallel attempt is rendered as its own tab " "in the chat UI; a final 'Synthesis' tab carries the merged " "output. Bounded to [1, 20] by pydantic so a runaway value can't " - "saturate the local CLI. Silently ignored on every provider " - "other than `codex`." + "saturate the local CLI. Defaults to 1 (single-call shape). " + "Silently ignored on every provider other than `codex`." ), ) diff --git a/studio/backend/routes/codex.py b/studio/backend/routes/codex.py index ec69bca8f5..5570706646 100644 --- a/studio/backend/routes/codex.py +++ b/studio/backend/routes/codex.py @@ -14,11 +14,11 @@ Two endpoints live here: regular API-key field. * ``POST /api/codex/login`` -- the device-auth helper. Spawns the - ``codex auth login --device-auth`` CLI command, captures the - verification URL from its output, and streams the rest of the auth - exchange back as SSE so the UI can show progress. The URL appears - in the first SSE event so the frontend can ``window.open`` it before - the user wanders off. + ``codex login --device-auth`` CLI command, captures the verification + URL (and one-time code) from its output, and streams the rest of the + auth exchange back as SSE so the UI can show progress. The URL + appears in the first SSE event so the frontend can ``window.open`` + it before the user wanders off. """ from __future__ import annotations @@ -49,8 +49,8 @@ async def get_codex_status( the "Sign in to Codex" button on ``logged_in``. Both are best-effort and cheap to recompute; the route does not cache the probe because the user can install the CLI / SDK or run - ``codex auth login`` between page loads and the picker should pick - that up on the next refresh. + ``codex login`` between page loads and the picker should pick that + up on the next refresh. """ return await probe_codex_availability() @@ -59,11 +59,12 @@ async def get_codex_status( async def codex_device_login( current_subject: str = Depends(get_current_subject), ) -> StreamingResponse: - """Stream the ``codex auth login --device-auth`` exchange. + """Stream the ``codex login --device-auth`` exchange. Returns an SSE stream of events: ``data: {"type": "device_url", "url": "https://..."}`` + ``data: {"type": "device_code", "code": "ABCD-EFGH"}`` ``data: {"type": "log", "line": "..."}`` (zero or more) ``data: {"type": "done", "ok": true}`` @@ -75,8 +76,24 @@ async def codex_device_login( """ async def _to_sse() -> AsyncGenerator[str, None]: - async for event in stream_codex_device_login(): - yield f"data: {json.dumps(event)}\n\n" + try: + async for event in stream_codex_device_login(): + yield f"data: {json.dumps(event)}\n\n" + except Exception as exc: + # CodeQL: never echo str(exc) verbatim. Log full reason + # server-side and surface a generic error to the client so + # local paths / env vars from the CLI traceback don't leak. + logger.error( + "codex_device_login.stream_error", + exc_type = type(exc).__name__, + error = str(exc), + ) + yield "data: " + json.dumps({ + "type": "error", + "message": "Codex login failed", + "exception_type": type(exc).__name__, + }) + "\n\n" + yield "data: " + json.dumps({"type": "done", "ok": False}) + "\n\n" # Frontend treats the trailing [DONE] the same way it does for # chat streams, so we emit it for parity. yield "data: [DONE]\n\n" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 484cf971ae..871665f8f9 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1891,14 +1891,25 @@ async def _proxy_to_external_provider( ) yield "data: [DONE]\n\n" except Exception as exc: - logger.error("codex_provider.stream_error", error = str(exc)) + # CodeQL: never echo str(exc) -- the Codex SDK can raise + # with local paths, env-var content, or traceback fragments. + # Log the full reason server-side; surface a generic message + # plus an exception_type discriminator to the client so the + # UI can show "Codex provider error" without leaking host + # internals. + logger.error( + "codex_provider.stream_error", + exc_type = type(exc).__name__, + error = str(exc), + ) yield ( "data: " + json.dumps( { "error": { - "message": f"Codex error: {exc}", + "message": "Codex provider error", "type": "provider_error", + "exception_type": type(exc).__name__, "code": "502", "provider": "codex", } diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 03e3e5abf9..1943bcf7ad 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -441,9 +441,11 @@ class TestParallelCallsValidator: parallel_calls = 21, ) - def test_request_default_is_none(self): - """Default = None so the field has no effect on every existing - provider that doesn't read it -- preserves backwards compat. + def test_request_default_is_one(self): + """Default = 1 so the field matches the single-call code path + and the schema documentation. Non-codex providers ignore the + field regardless of its value, so backwards compat is + preserved. """ from models.inference import ChatCompletionRequest @@ -451,7 +453,7 @@ class TestParallelCallsValidator: model = "gpt-5.4", messages = [{"role": "user", "content": "hi"}], ) - assert req.parallel_calls is None + assert req.parallel_calls == 1 # ── Codex unavailable surfacing ──────────────────────────────────── @@ -460,18 +462,23 @@ class TestParallelCallsValidator: class TestCodexUnavailable: def test_missing_sdk_raises_typed_error(self, monkeypatch): # Force find_spec to return None so the lazy import fails. + # The provider probes both the canonical upstream name + # ``openai_codex`` and the legacy alias ``codex_app_server``, + # so we have to suppress both for the import to fail. import importlib.util as _iu real = _iu.find_spec + _SDK_NAMES = {"openai_codex", "codex_app_server"} def _shim(name, *args, **kwargs): - if name == "codex_app_server": + if name in _SDK_NAMES: return None return real(name, *args, **kwargs) monkeypatch.setattr("importlib.util.find_spec", _shim) - # Also drop any cached fake from prior tests. - monkeypatch.delitem(sys.modules, "codex_app_server", raising = False) + # Also drop any cached fakes from prior tests. + for _name in _SDK_NAMES: + monkeypatch.delitem(sys.modules, _name, raising = False) from core.inference.codex_provider import ( CodexUnavailableError, @@ -489,6 +496,153 @@ class TestCodexUnavailable: ) +class TestCodexHardenedRegressions: + """Tests covering the post-review hardening pass. + + Each test pins a specific regression: the wrong subcommand + (``codex auth login`` → ``codex login``), the wrong SDK package + name (``codex_app_server`` → ``openai_codex`` with legacy alias), + the ``not logged in`` substring footgun, the ANSI-wrapped device + URL, and the fan-out cancellation contract. + """ + + def test_sdk_probes_openai_codex_first(self, monkeypatch): + """The canonical upstream name must be tried before the alias.""" + import importlib.util as _iu + + real = _iu.find_spec + calls: list[str] = [] + + def _shim(name, *args, **kwargs): + if name in ("openai_codex", "codex_app_server"): + calls.append(name) + return None + return real(name, *args, **kwargs) + + monkeypatch.setattr("importlib.util.find_spec", _shim) + from core.inference.codex_availability import _sdk_importable + assert _sdk_importable() is False + assert calls and calls[0] == "openai_codex", ( + f"availability probe must check openai_codex first; saw {calls}" + ) + + 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" + ) + text = open(src).read() + assert '"auth", "status"' not in text, ( + "_detect_logged_in must use `codex login status`, not `codex auth status`" + ) + 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" + ) + text = open(src).read() + assert '"auth", "login", "--device-auth"' not in text, ( + "stream_codex_device_login must use `codex login --device-auth`" + ) + assert '"login", "--device-auth"' in text + + def test_not_logged_in_not_misparsed_as_logged_in(self): + """The substring "logged in" inside "not logged in" must not + flip the detection to True.""" + import asyncio + + from core.inference import codex_availability as av + + async def _fake_run_cli(args, **kw): + return (0, "Not logged in. Run `codex login` to authenticate.", "") + + orig = av._run_cli + av._run_cli = _fake_run_cli # type: ignore[assignment] + try: + result = asyncio.run(av._detect_logged_in()) + assert result is False, "'Not logged in' was misparsed as logged_in=True" + finally: + av._run_cli = orig # type: ignore[assignment] + + def test_logged_in_is_detected(self): + import asyncio + + from core.inference import codex_availability as av + + async def _fake_run_cli(args, **kw): + return (0, "Logged in using ChatGPT", "") + + orig = av._run_cli + av._run_cli = _fake_run_cli # type: ignore[assignment] + try: + result = asyncio.run(av._detect_logged_in()) + assert result is True + finally: + av._run_cli = orig # type: ignore[assignment] + + def test_multi_turn_prompt_includes_prior_turns(self): + """The Codex prompt MUST contain prior assistant turns.""" + from core.inference.codex_provider import _last_user_prompt + + msgs = [ + {"role": "user", "content": "what is the capital of france?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "and germany?"}, + ] + prompt = _last_user_prompt(msgs) + assert "and germany?" in prompt + assert "Paris" in prompt, ( + f"PRIOR ASSISTANT TURN DROPPED — multi-turn broken. Prompt:\n{prompt}" + ) + assert "capital of france" in prompt.lower() + + def test_single_turn_prompt_unchanged(self): + """Single-turn case must not get the User:/Assistant: framing.""" + from core.inference.codex_provider import _last_user_prompt + + prompt = _last_user_prompt([{"role": "user", "content": "hi"}]) + assert prompt == "hi" + + def test_default_models_no_o3(self): + """The Codex registry must not advertise `o3` (not in upstream).""" + from core.inference.providers import PROVIDER_REGISTRY + + codex = PROVIDER_REGISTRY["codex"] + assert "o3" not in codex["default_models"], ( + "o3 is not a Codex model; remove from default_models" + ) + assert "gpt-5.5" in codex["default_models"] + + def test_inference_route_no_raw_exc_leak(self): + """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" + ) + text = open(src).read() + bad = re.findall(r'f["\']Codex error:\s*\{exc\}["\']', text) + assert not bad, f"raw exception in SSE: {bad}" + + def test_codex_route_no_raw_exc_leak(self): + """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" + ) + text = open(src).read() + for line in text.splitlines(): + ls = line.strip() + if ls.startswith("yield ") and re.search(r"\{exc\}|\{e\}", ls): + assert False, f"raw exception leaked: {ls}" + + async def _consume_first(gen): """Drive an async generator until it raises or yields its first value. Used to surface lazy-import errors that fire on the first From 861da31fcfd2e3b520b869c28a31f8667091ea49 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 24 May 2026 14:13:22 +0000 Subject: [PATCH 06/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/inference/codex_availability.py | 4 ++- .../backend/core/inference/codex_provider.py | 7 +++-- studio/backend/routes/codex.py | 16 +++++++--- studio/backend/tests/test_codex_provider.py | 31 ++++++++++--------- 4 files changed, 34 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 8ae57deadd..bbb78bef64 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -182,7 +182,9 @@ 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. - negative = re.compile(r"\b(not logged in|not authenticated|please log in|run\s+`?codex login`?)\b") + negative = re.compile( + r"\b(not logged in|not authenticated|please log in|run\s+`?codex 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 299de46be3..55cda300f5 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -353,8 +353,7 @@ async def _stream_codex_single( async_codex_cls = getattr(sdk, "AsyncCodex", None) if async_codex_cls is None: raise CodexUnavailableError( - "Codex SDK is installed but AsyncCodex is missing -- " - "upgrade the SDK." + "Codex SDK is installed but AsyncCodex is missing -- " "upgrade the SDK." ) completion_text_chars = 0 @@ -746,7 +745,9 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: # 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. - url_re = re.compile(r"https?://[^\s\x1b]+?/codex/device(?:\?[^\s\x1b]*)?", re.IGNORECASE) + url_re = re.compile( + r"https?://[^\s\x1b]+?/codex/device(?:\?[^\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. code_re = re.compile(r"\b([A-Z0-9]{4}-[A-Z0-9]{4})\b") diff --git a/studio/backend/routes/codex.py b/studio/backend/routes/codex.py index 5570706646..60b0d3b18e 100644 --- a/studio/backend/routes/codex.py +++ b/studio/backend/routes/codex.py @@ -88,11 +88,17 @@ async def codex_device_login( exc_type = type(exc).__name__, error = str(exc), ) - yield "data: " + json.dumps({ - "type": "error", - "message": "Codex login failed", - "exception_type": type(exc).__name__, - }) + "\n\n" + yield ( + "data: " + + json.dumps( + { + "type": "error", + "message": "Codex login failed", + "exception_type": type(exc).__name__, + } + ) + + "\n\n" + ) yield "data: " + json.dumps({"type": "done", "ok": False}) + "\n\n" # Frontend treats the trailing [DONE] the same way it does for # chat streams, so we emit it for parity. diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 1943bcf7ad..b6a255a785 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -521,10 +521,11 @@ class TestCodexHardenedRegressions: monkeypatch.setattr("importlib.util.find_spec", _shim) from core.inference.codex_availability import _sdk_importable + assert _sdk_importable() is False - assert calls and calls[0] == "openai_codex", ( - f"availability probe must check openai_codex first; saw {calls}" - ) + assert ( + calls and calls[0] == "openai_codex" + ), f"availability probe must check openai_codex first; saw {calls}" def test_login_status_uses_login_subcommand(self): """Upstream is `codex login status`, NOT `codex auth status`.""" @@ -533,9 +534,9 @@ class TestCodexHardenedRegressions: "studio/backend/core/inference/codex_availability.py" ) text = open(src).read() - assert '"auth", "status"' not in text, ( - "_detect_logged_in must use `codex login status`, not `codex auth status`" - ) + assert ( + '"auth", "status"' not in text + ), "_detect_logged_in must use `codex login status`, not `codex auth status`" assert '"login", "status"' in text def test_device_login_uses_login_subcommand(self): @@ -544,9 +545,9 @@ class TestCodexHardenedRegressions: "studio/backend/core/inference/codex_provider.py" ) text = open(src).read() - assert '"auth", "login", "--device-auth"' not in text, ( - "stream_codex_device_login must use `codex login --device-auth`" - ) + assert ( + '"auth", "login", "--device-auth"' not in text + ), "stream_codex_device_login must use `codex login --device-auth`" assert '"login", "--device-auth"' in text def test_not_logged_in_not_misparsed_as_logged_in(self): @@ -594,9 +595,9 @@ class TestCodexHardenedRegressions: ] prompt = _last_user_prompt(msgs) assert "and germany?" in prompt - assert "Paris" in prompt, ( - f"PRIOR ASSISTANT TURN DROPPED — multi-turn broken. Prompt:\n{prompt}" - ) + assert ( + "Paris" in prompt + ), f"PRIOR ASSISTANT TURN DROPPED — multi-turn broken. Prompt:\n{prompt}" assert "capital of france" in prompt.lower() def test_single_turn_prompt_unchanged(self): @@ -611,9 +612,9 @@ class TestCodexHardenedRegressions: from core.inference.providers import PROVIDER_REGISTRY codex = PROVIDER_REGISTRY["codex"] - assert "o3" not in codex["default_models"], ( - "o3 is not a Codex model; remove from default_models" - ) + assert ( + "o3" not in codex["default_models"] + ), "o3 is not a Codex model; remove from default_models" assert "gpt-5.5" in codex["default_models"] def test_inference_route_no_raw_exc_leak(self): From b6577a6287495ef5de3fea85b106b9952312d20b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 14:25:04 +0000 Subject: [PATCH 07/40] Studio: name the actual PyPI package in the Codex install hint The OpenAI Codex Python SDK ships on PyPI as `openai-codex-app-server-sdk`, not `openai-codex` (which is the GitHub repo project name in pyproject.toml). The runtime binary ships separately as `openai-codex-cli-bin`. Both packages expose the import name `openai_codex`; the older docs reference `codex_app_server` so we keep probing both. Update the `CodexUnavailableError` message and the provider registry notes so a user hitting the unavailable path gets a copy-pasteable `pip install` command. No behaviour change. --- studio/backend/core/inference/codex_provider.py | 3 ++- studio/backend/core/inference/providers.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 55cda300f5..437722d6a2 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -88,7 +88,8 @@ def _import_codex() -> Any: return importlib.import_module(name) raise CodexUnavailableError( "Codex Python SDK is not installed on this host. " - "Install `openai-codex` (or the legacy `codex_app_server`) " + "Install with `pip install openai-codex-app-server-sdk` " + "(import name `openai_codex`, legacy alias `codex_app_server`), " "or use a different provider." ) diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index 6ab1ac47bb..4d758f1520 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -307,9 +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 (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-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`." ), }, "openrouter": { From d6c47f666448a68cb47e7158b8e254e1b8a7fdae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 14:46:03 +0000 Subject: [PATCH 08/40] Studio: wire Codex provider through the UI end-to-end Post-review pass driven by reviewer.py. The original PR shipped the backend codex provider, the registry entry (with `hidden:true`), the status API, and the `CodexParallelTabs` component, but the chat UI never surfaced the row, required an API key for the connection, and never sent `parallel_calls` over the wire. Also fixes a CodeQL leak in the parallel fan-out error path and adds the canonical streaming hook upstream actually exposes. Frontend * chat-providers-dialog.tsx now calls `/api/codex/status` alongside `/api/providers/registry`. When the host has Codex installed the Add connection dialog gains a synthetic Codex row (curated model list comes from `supported_models`) so the picker is reachable. * The Add / Edit connection guards now skip the API-key requirement for Codex the same way they do for the custom OpenAI-compat presets; the field itself is also hidden so the user is not asked for a key Studio will not use. * chat-adapter.ts now also exempts Codex from the "Missing API key" pre-flight, and emits `parallel_calls` on the outgoing request when the selected connection is Codex (clamped to [1, 20] by the shared helper, defaults to 1). * external-providers.ts adds `codexParallelCalls` to ExternalProviderConfig so future composer UI can persist the user's pick per connection. Backend * `_stream_thread_run` now tries `thread.turn(prompt).stream()` first, mirroring the canonical openai_codex API (`openai/codex/sdk/python/src/openai_codex/api.py`). The legacy `thread.run_streaming(prompt)` path is kept as a fallback and the buffered `await thread.run(prompt)` stays as the last resort. * `_stream_codex_parallel` no longer echoes `str(exc)` in the `codex_tab_error` SSE event. Per-tab failures now surface a generic "Codex tab failed" message plus an `exception_type` discriminator; `CodexUnavailableError` is the only exception whose text is forwarded verbatim because it is a user-actionable install hint with no sensitive content (CodeQL `py/information-exposure-through-exception`). Tests * New `TestCodexHardenedRegressions::test_parallel_tab_error_sanitised` injects a fake SDK that raises with a path-like message and asserts the SSE frames do not echo it. * New `TestCodexHardenedRegressions::test_thread_turn_stream_path_taken` verifies the canonical `thread.turn(prompt).stream()` hook is preferred over the legacy helper. All 26 codex_provider tests pass. Frontend `tsc --noEmit` clean. --- .../backend/core/inference/codex_provider.py | 70 ++++++++++++--- studio/backend/tests/test_codex_provider.py | 89 +++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 26 +++++- .../features/chat/chat-providers-dialog.tsx | 60 ++++++++++--- .../src/features/chat/external-providers.ts | 7 ++ 5 files changed, 227 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 437722d6a2..6461dfbdf0 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -307,20 +307,51 @@ async def _stream_thread_run( ) -> AsyncGenerator[str, None]: """Yield raw text chunks from a Codex thread. - Prefers ``thread.run_streaming(prompt)`` because that's what the - docs surface for token-by-token delivery. When the installed SDK - doesn't have that helper, fall back to ``await thread.run(prompt)`` - and yield the full text once -- this still works end-to-end, just - without streaming feedback in the UI. + Tries three SDK surfaces in order: + + 1. ``thread.turn(prompt).stream()`` -- the canonical streaming path + on the upstream ``openai_codex`` SDK (see + ``openai/codex/sdk/python/src/openai_codex/api.py``: AsyncThread.turn + returns an AsyncTurnHandle whose ``.stream()`` yields events). + 2. ``thread.run_streaming(prompt)`` -- a legacy helper exposed by + some earlier SDK pre-releases. Kept for forward-compat. + 3. ``await thread.run(prompt)`` -- the always-supported buffered + 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. """ + # 1. Canonical: thread.turn(prompt).stream() + turn_factory = getattr(thread, "turn", None) + if turn_factory is not None: + try: + turn_handle = turn_factory(prompt) + if asyncio.iscoroutine(turn_handle): + turn_handle = await turn_handle + stream_fn = getattr(turn_handle, "stream", None) + if stream_fn is not None: + stream_obj = stream_fn() + if asyncio.iscoroutine(stream_obj): + stream_obj = await stream_obj + async for event in stream_obj: + text = _coerce_text(getattr(event, "payload", event)) + if text: + yield text + return + except Exception as exc: + logger.warning( + "codex_provider.turn_stream_failed_fallback", + exc_type = type(exc).__name__, + error = str(exc), + ) + + # 2. Legacy: thread.run_streaming(prompt) run_streaming = getattr(thread, "run_streaming", None) if run_streaming is not None: try: stream_obj = run_streaming(prompt) - # The SDK may return either an async iterator directly or a - # coroutine that resolves to one. Handle both shapes so a - # future SDK rev doesn't silently fall off the streaming - # path. if asyncio.iscoroutine(stream_obj): stream_obj = await stream_obj async for event in stream_obj: @@ -331,12 +362,11 @@ async def _stream_thread_run( except Exception as exc: logger.warning( "codex_provider.run_streaming_failed_fallback", + exc_type = type(exc).__name__, error = str(exc), ) - # Intentional fallthrough to the non-streaming path so a - # broken streaming helper doesn't take the whole turn down. - # Non-streaming fallback: await the full TurnResult, emit one chunk. + # 3. Buffered fallback: await the full TurnResult, emit one chunk. result = await thread.run(prompt) text = _coerce_text(result) or getattr(result, "final_response", "") or str(result) if text: @@ -516,18 +546,32 @@ async def _stream_codex_parallel( ) ) except Exception as exc: + # CodeQL: never echo str(exc) in client-facing SSE events. + # Log full reason server-side; surface a generic message plus + # an exception_type discriminator so the UI can still group + # failures without leaking file paths / env vars from the + # SDK traceback. CodexUnavailableError is the one exception + # we DO surface verbatim because it's a user-actionable + # install hint with no sensitive content. logger.warning( "codex_provider.parallel_tab_failed", tab_id = tab_id, + exc_type = type(exc).__name__, error = str(exc), ) + public_error = ( + str(exc) + if isinstance(exc, CodexUnavailableError) + else "Codex tab failed" + ) await queue.put( _chunk_tool_event( completion_id, { "type": "codex_tab_error", "tab_id": tab_id, - "error": str(exc), + "error": public_error, + "exception_type": type(exc).__name__, }, ) ) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index b6a255a785..e41e207685 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -643,6 +643,95 @@ class TestCodexHardenedRegressions: if ls.startswith("yield ") and re.search(r"\{exc\}|\{e\}", ls): assert False, f"raw exception leaked: {ls}" + def test_parallel_tab_error_sanitised(self, monkeypatch): + """A worker that raises with a path-leaking message must NOT + send that text to the client; the SSE codex_tab_error event + must carry a generic message + exception_type. + """ + fake = _FakeAsyncCodex(raise_on_start=RuntimeError( + "secret /home/alice/.codex/config.json token=abc" + )) + _install_fake_codex_sdk(monkeypatch, lambda: fake) + 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=2, + ): + chunks.append(c) + + asyncio.run(_collect()) + body = "".join(chunks) + assert "secret /home/alice" not in body, ( + "raw exception text leaked into codex_tab_error SSE frame" + ) + assert "Codex tab failed" in body or "exception_type" 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. + """ + events_seen = {"turn_called": False, "run_streaming_called": False} + + class _TurnEvent: + def __init__(self, txt): + self.payload = {"text": txt} + + class _TurnHandle: + def __init__(self, prompt): + self.prompt = prompt + + async def stream(self): + yield _TurnEvent("hello ") + yield _TurnEvent("from turn.stream") + + class _ThreadWithTurn: + def turn(self, prompt): + events_seen["turn_called"] = True + return _TurnHandle(prompt) + + def run_streaming(self, prompt): + events_seen["run_streaming_called"] = True + raise AssertionError("should not be called when turn().stream() works") + + async def run(self, prompt): + raise AssertionError("should not fall through to buffered run()") + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _ThreadWithTurn() + + _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 events_seen["turn_called"], "thread.turn() never called" + assert not events_seen["run_streaming_called"] + body = "".join(chunks) + # Each text chunk wraps in its own SSE delta, so check both pieces. + assert '"content": "hello "' in body + assert '"content": "from turn.stream"' in body + async def _consume_first(gen): """Drive an async generator until it raises or yields its first diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c557f1b01..7a5554342c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -7,7 +7,10 @@ import { toast } from "@/lib/toast"; import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; import type { ChatModelAdapter } from "@assistant-ui/react"; import { + CODEX_DEFAULT_PARALLEL_CALLS, + clampCodexParallelCalls, getExternalProviderApiKey, + isCodexProviderType, isCustomProviderType, isPromptCacheTtl, loadExternalProviders, @@ -862,12 +865,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Connection not found."); } // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. + // Codex dispatches via the local CLI / SDK, no HTTP API key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; - if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) { + const externalProviderIsCodex = externalProvider + ? isCodexProviderType(externalProvider.providerType) + : false; + const externalProviderNeedsApiKey = + isExternalRequest && !externalProviderIsCustom && !externalProviderIsCodex; + if (externalProviderNeedsApiKey && !externalApiKey) { toast.error("Missing API key for selected connection.", { - description: "Open Settings → Connections and set the API key again.", + description: "Open Settings > Connections and set the API key again.", }); throw new Error("Missing connection API key."); } @@ -1489,6 +1498,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } : { enable_thinking: reasoningEnabled } : {}), + // Codex provider only: ask the backend to fan the turn out + // across N parallel Codex tasks and synthesise a unified + // answer. The picker UI uses the provider config's + // `codexParallelCalls` field; default of 1 keeps the + // single-call path. Backend clamps to [1, 20]. + ...(externalProviderIsCodex + ? { + parallel_calls: clampCodexParallelCalls( + externalProvider.codexParallelCalls ?? + CODEX_DEFAULT_PARALLEL_CALLS, + ), + } + : {}), }; } diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 3ffb3a1441..8a95b3a6a6 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -49,6 +49,7 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { + CODEX_PROVIDER_TYPE, CUSTOM_BACKEND_PROVIDER_TYPE, CUSTOM_PROVIDER_PRESETS, allowsManualModelIdsWithCatalog, @@ -57,6 +58,7 @@ import { customProviderModelIdsPlaceholder, customPresetSkipsApiKeyField, getExternalProviderApiKey, + isCodexProviderType, isCustomProviderType, LEGACY_CUSTOM_PROVIDER_TYPE, removeExternalProviderApiKey, @@ -66,6 +68,7 @@ import { supportsRemoteModelCatalog, toExternalBackendProviderType, } from "./external-providers"; +import { fetchCodexStatus } from "./api/codex-api"; import { useExternalProvidersStore } from "./stores/external-providers-store"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ @@ -239,9 +242,15 @@ export function ChatProvidersSettings({ (s) => s.setConnectionsEnabled, ); const isCustomProvider = isCustomProviderType(providerType); + const isCodexProvider = isCodexProviderType(providerType); // Local presets (Ollama, llama.cpp) never use API keys — hide the field. - // vLLM may optionally use a bearer token on secured deployments. - const showApiKeyField = !customPresetSkipsApiKeyField(providerType); + // vLLM may optionally use a bearer token on secured deployments. Codex + // dispatches via the local CLI / SDK, no HTTP API key either. + const showApiKeyField = + !customPresetSkipsApiKeyField(providerType) && !isCodexProvider; + // Codex behaves like a "custom" provider for the gate logic below: the + // backend skips the api_key requirement entirely for `provider_type=codex`. + const providerSkipsApiKey = isCustomProvider || isCodexProvider; const showReasoningToggle = supportsProviderReasoningToggle(providerType); const registryByType = useMemo( @@ -279,7 +288,7 @@ export function ChatProvidersSettings({ const missingModelCatalogBaseUrl = supportsRemoteModelCatalog(providerType) && baseUrlDraft.trim().length === 0; const missingModelCatalogApiKey = - !isCustomProvider && !isCuratedModelList && apiKey.trim().length === 0; + !providerSkipsApiKey && !isCuratedModelList && apiKey.trim().length === 0; const loadModelsDisabled = modelsLoading || mutatingProvider || @@ -354,12 +363,41 @@ export function ChatProvidersSettings({ } let syncSucceeded = false; try { - const [registryRows, configRows] = await Promise.all([ - listProviderRegistry(), - listProviderConfigs(), - ]); + // Probe Codex availability in parallel with the registry / configs. + // Codex stays `hidden:true` in the backend registry so it is filtered + // out of `/api/providers/registry`; we synthesise a row here when + // the host has both the CLI and the SDK installed. + const [registryRowsRaw, configRows, codexStatusRaw] = await Promise.all( + [ + listProviderRegistry(), + listProviderConfigs(), + fetchCodexStatus().catch(() => null), + ], + ); if (!isMounted) return; syncSucceeded = true; + const registryRows: ProviderRegistryEntry[] = + codexStatusRaw && codexStatusRaw.installed && + !registryRowsRaw.some( + (entry) => entry.provider_type === CODEX_PROVIDER_TYPE, + ) + ? [ + ...registryRowsRaw, + { + provider_type: CODEX_PROVIDER_TYPE, + display_name: "OpenAI Codex (local CLI)", + base_url: "", + default_models: codexStatusRaw.supported_models ?? [], + supports_streaming: true, + supports_vision: false, + supports_tool_calling: true, + model_list_mode: "curated", + notes: codexStatusRaw.logged_in + ? "Dispatches chat turns through the local Codex CLI." + : "Sign in with `codex login` before chatting.", + } as ProviderRegistryEntry, + ] + : registryRowsRaw; setRegistry(registryRows); setProviderType((current) => { if ( @@ -553,7 +591,7 @@ export function ChatProvidersSettings({ ); return; } - if (!isCustomProvider && !apiKey.trim()) { + if (!providerSkipsApiKey && !apiKey.trim()) { toast.error("Add an API key first."); return; } @@ -625,7 +663,7 @@ export function ChatProvidersSettings({ const displayName = isCustomProvider ? customProviderName.trim() || customProviderDisplayName(providerType) : (selectedRegistryEntry?.display_name ?? providerType); - if (!isCustomProvider && !apiKey.trim()) { + if (!providerSkipsApiKey && !apiKey.trim()) { toast.error("API key is required."); return; } @@ -734,7 +772,9 @@ export function ChatProvidersSettings({ } const isEditingCustomProvider = isCustomProviderType(existing.providerType); - if (!isEditingCustomProvider && !apiKey.trim()) { + const editingProviderSkipsApiKey = + isEditingCustomProvider || isCodexProviderType(existing.providerType); + if (!editingProviderSkipsApiKey && !apiKey.trim()) { toast.error("API key is required."); return; } diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 8b25f60bf2..e95159e14a 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -33,6 +33,13 @@ export interface ExternalProviderConfig { * OpenAI's hard default is 20. Only meaningful for OpenAI cloud. */ openaiContainerTtlMinutes?: number; + /** + * Codex provider only: number of parallel Codex turns to fan a chat + * request out into. Clamped to [1, 20] by `clampCodexParallelCalls`. + * Omitted or 1 takes the single-call path; values > 1 emit per-tab + * `codex_tab_*` SSE events plus a final `codex_gather` synthesis. + */ + codexParallelCalls?: number; createdAt: number; updatedAt: number; } From e2ac4907bfa2eea3b6dfe72b256b63f4fd594f24 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 24 May 2026 14:46:14 +0000 Subject: [PATCH 09/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_codex_provider.py | 26 +++++++++++---------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index e41e207685..bfa44176f8 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -648,9 +648,11 @@ class TestCodexHardenedRegressions: send that text to the client; the SSE codex_tab_error event must carry a generic message + exception_type. """ - fake = _FakeAsyncCodex(raise_on_start=RuntimeError( - "secret /home/alice/.codex/config.json token=abc" - )) + fake = _FakeAsyncCodex( + raise_on_start = RuntimeError( + "secret /home/alice/.codex/config.json token=abc" + ) + ) _install_fake_codex_sdk(monkeypatch, lambda: fake) from core.inference.codex_provider import stream_codex @@ -658,17 +660,17 @@ class TestCodexHardenedRegressions: async def _collect(): async for c in stream_codex( - messages=[{"role": "user", "content": "hi"}], - model="gpt-5.5", - parallel_calls=2, + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + parallel_calls = 2, ): chunks.append(c) asyncio.run(_collect()) body = "".join(chunks) - assert "secret /home/alice" not in body, ( - "raw exception text leaked into codex_tab_error SSE frame" - ) + assert ( + "secret /home/alice" not in body + ), "raw exception text leaked into codex_tab_error SSE frame" assert "Codex tab failed" in body or "exception_type" in body def test_thread_turn_stream_path_taken(self, monkeypatch): @@ -718,9 +720,9 @@ class TestCodexHardenedRegressions: async def _collect(): async for c in stream_codex( - messages=[{"role": "user", "content": "hi"}], - model="gpt-5.5", - parallel_calls=1, + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + parallel_calls = 1, ): chunks.append(c) From 028b7b7187a5be30a4129d11dc3363bc300d7a60 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 15:01:33 +0000 Subject: [PATCH 10/40] Studio: finish Codex UI flow (test, sign-in, parallel tabs) Second reviewer.py pass surfaced three follow-ups missed in the earlier round. All caught by 12 parallel reviewers + cross-block audit; each fix is small but user-facing. * `testProvider` no longer pushes a Codex connection back to the edit form to "add an API key". Codex has no remote endpoint to ping, so the Test button now calls `/api/codex/status` directly: toasts success with the CLI version when installed+logged in, prompts to sign in when installed+logged out, and errors when the CLI or SDK is missing. * The Sign-in to Codex affordance is now actually mounted. When the selected provider is Codex and `/api/codex/status` reports `installed:true, logged_in:false`, the dialog renders the new `CodexLoginButton` above the (hidden) API key row. The button's `onLoggedIn` callback re-probes status so the UI flips to the ready state without a page reload. * The chat adapter now handles `codex_*` `_toolEvent` types instead of silently swallowing them. Per-tab chunks render inline with a `[Codex tab N/M]` header so users see each parallel attempt; `codex_gather` adds a `--- Synthesis ---` divider before the final unified content delta the backend also emits as plain text. This unblocks the existing fan-out path while a dedicated `CodexParallelTabs` UI is wired in a future change. Verified: 26/26 codex_provider tests pass; `tsc --noEmit` on studio/frontend completes clean. --- .../src/features/chat/api/chat-adapter.ts | 41 ++++++++++++ .../features/chat/chat-providers-dialog.tsx | 62 ++++++++++++++++++- 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 7a5554342c..2b5ccf9226 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1585,6 +1585,47 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { + // Codex parallel-calls fan-out events: render each + // per-tab chunk inline with a "Tab N:" prefix so users + // see the N independent attempts even before the + // dedicated CodexParallelTabs UI is wired into the + // chat surface. `codex_gather` carries the synthesis + // payload which the backend also emits as a normal + // content delta, so we drop it here to avoid showing + // the synthesis twice. tab_open / tab_close / tab_error + // are header markers we surface as one-line notes. + if (typeof toolEvent.type === "string" && toolEvent.type.startsWith("codex_")) { + if (toolEvent.type === "codex_tab_open") { + const tabId = Number(toolEvent.tab_id); + const total = Number(toolEvent.total_tabs); + if (Number.isFinite(tabId) && Number.isFinite(total)) { + cumulativeText += `\n\n[Codex tab ${tabId}/${total}]\n`; + } + } else if (toolEvent.type === "codex_tab_chunk") { + const text = typeof toolEvent.text === "string" ? toolEvent.text : ""; + if (text) cumulativeText += text; + } else if (toolEvent.type === "codex_tab_error") { + const tabId = Number(toolEvent.tab_id); + const err = typeof toolEvent.error === "string" ? toolEvent.error : "error"; + if (Number.isFinite(tabId)) { + cumulativeText += `\n[Codex tab ${tabId} error: ${err}]\n`; + } + } else if (toolEvent.type === "codex_tab_close") { + // Mark end of tab block so synthesis is visually separated. + cumulativeText += "\n"; + } else if (toolEvent.type === "codex_gather") { + // Synthesis is also emitted as a normal content + // delta later in the same SSE stream; nothing to + // add here. Surface a divider so the user can tell + // where the synthesis starts. + cumulativeText += "\n--- Synthesis ---\n"; + } + const codexParts = parseAssistantContent(cumulativeText); + yield { + content: [...toolCallParts, ...codexParts], + }; + continue; + } // OpenAI shell-tool container persistence — see // ThreadRecord.openaiCodeExecContainerId. The backend // emits these synthetic events on the OpenAI Responses diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 8a95b3a6a6..67da6cee33 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -68,7 +68,8 @@ import { supportsRemoteModelCatalog, toExternalBackendProviderType, } from "./external-providers"; -import { fetchCodexStatus } from "./api/codex-api"; +import { fetchCodexStatus, type CodexStatus } from "./api/codex-api"; +import { CodexLoginButton } from "./components/codex-login-button"; import { useExternalProvidersStore } from "./stores/external-providers-store"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ @@ -224,6 +225,20 @@ export function ChatProvidersSettings({ null, ); const [registry, setRegistry] = useState([]); + // Codex CLI / SDK availability snapshot. Used to (a) decide whether + // to render the synthetic Codex registry row, and (b) drive the + // sign-in button when the host is installed but logged out. + const [codexStatus, setCodexStatus] = useState(null); + const refreshCodexStatus = async () => { + try { + const next = await fetchCodexStatus(); + setCodexStatus(next); + return next; + } catch { + setCodexStatus(null); + return null; + } + }; const [availableModels, setAvailableModels] = useState([]); const [selectedModelIds, setSelectedModelIds] = useState([]); const [syncingProviders, setSyncingProviders] = useState(false); @@ -376,6 +391,7 @@ export function ChatProvidersSettings({ ); if (!isMounted) return; syncSucceeded = true; + setCodexStatus(codexStatusRaw); const registryRows: ProviderRegistryEntry[] = codexStatusRaw && codexStatusRaw.installed && !registryRowsRaw.some( @@ -964,6 +980,30 @@ export function ChatProvidersSettings({ async function testProvider(provider: ExternalProviderConfig) { const savedKey = getExternalProviderApiKey(provider.id).trim(); + // Codex dispatches via the local CLI / SDK -- there is no remote + // endpoint to ping. Reuse `/api/codex/status` as the test result. + if (isCodexProviderType(provider.providerType)) { + try { + const status = await fetchCodexStatus(); + if (!status.installed) { + toast.error("Codex CLI or SDK is not available on this host."); + return; + } + if (!status.logged_in) { + toast.info("Sign in to Codex before testing this connection."); + return; + } + toast.success( + status.version + ? `Codex is available (${status.version}).` + : "Codex is available.", + ); + } catch (error) { + const message = error instanceof Error ? error.message : "Unknown error"; + toast.error(`Codex status check failed: ${message}`); + } + return; + } // Local OpenAI-compat presets skip API keys — run the connection check. if (!savedKey && !supportsRemoteModelCatalog(provider.providerType)) { if (isCustomProviderType(provider.providerType)) { @@ -1115,6 +1155,26 @@ export function ChatProvidersSettings({ + {isCodexProvider && + codexStatus?.installed && + !codexStatus.logged_in ? ( +
+
+ +

+ Authenticate the local Codex CLI before chatting. +

+
+ { + void refreshCodexStatus(); + }} + /> +
+ ) : null} + {showApiKeyField ? (
From fd8f25f50714f839aed6c905d62d76f70a1613ba Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 15:16:38 +0000 Subject: [PATCH 11/40] Studio: scrub Codex subprocess env, guard partial-stream replay, abort login on unmount Third reviewer.py pass found three remaining sharp edges. Each fix is small and paired with a regression test where applicable. * Codex subprocess env is now scrubbed to a safe-list before spawn. Both `_run_cli` in codex_availability and the device-auth spawn in stream_codex_device_login switch from `env=os.environ.copy()` to `env=_codex_subprocess_env()`, which forwards only PATH / HOME / USER / Windows-equivalents / CODEX_HOME / OPENAI_API_KEY / OPENAI_BASE_URL. Other-provider secrets like HF_TOKEN, GH_TOKEN, WANDB_API_KEY, ANTHROPIC_API_KEY no longer reach the local codex binary, so a shimmed `codex` earlier on PATH cannot harvest them. * `_stream_thread_run` now tracks `emitted_any` and refuses to fall through to the buffered `await thread.run(prompt)` after either streaming helper has already yielded text. Previously a network glitch mid-stream re-executed the same Codex turn, which can duplicate file writes, shell commands, and other Codex side effects. The buffered path is now reserved for the zero-output case (no streaming helper resolved, or streaming returned empty). * `CodexLoginButton` now aborts the SSE reader on unmount via a useEffect cleanup that calls `abortRef.current?.abort()`. The underlying `codex login --device-auth` subprocess no longer keeps streaming (and holding a device-auth session) after the dialog closes. Two new pytest cases pin the behaviour: `test_codex_subprocess_env_scrubbed` sets HF/GH/WANDB/ANTHROPIC keys and asserts none reach the codex env while OPENAI_API_KEY / CODEX_HOME survive; and `test_partial_stream_failure_does_not_replay_turn` injects a fake `turn().stream()` that yields "partial output " then raises, and asserts `thread.run()` is never called. 28/28 codex_provider tests pass; `tsc --noEmit` clean. --- .../core/inference/codex_availability.py | 46 ++++++++++- .../backend/core/inference/codex_provider.py | 28 ++++++- studio/backend/tests/test_codex_provider.py | 77 +++++++++++++++++++ .../chat/components/codex-login-button.tsx | 11 ++- 4 files changed, 157 insertions(+), 5 deletions(-) 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 (
) : null} + {isCodexProvider ? ( +
+
+ +

+ Fan-out width. Each call runs the same prompt against + Codex and the results are unified in a final synthesis + tab. 1-{CODEX_MAX_PARALLEL_CALLS}. +

+
+
+ { + const raw = Number(event.target.value); + setCodexParallelCalls( + clampCodexParallelCalls( + Number.isFinite(raw) + ? raw + : CODEX_DEFAULT_PARALLEL_CALLS, + ), + ); + }} + className="h-9 text-sm" + /> +
+
+ ) : null} + {showApiKeyField ? (
From f97a800d5b69b0ac96d64ca783efcfcb678a5552 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 16:59:46 +0000 Subject: [PATCH 19/40] Studio: round 5e Codex hardening (4 follow-ups) 1. parallel_calls validator now clamps instead of 422-rejecting. The Pydantic schema was `int Field(ge=1, le=20)`, which was a regression from the pre-PR OpenAI-extra behaviour: a non-Codex client that sent the field with a legacy value like 0 (or a stray string from a misconfigured wrapper) now got a 422 even though the route silently ignores the field on every non-Codex provider. Replaced with a `field_validator(mode="before")` that coerces any input to the [1, 20] range, keeping the schema docs self-documenting while accepting legacy inputs. 2. Buffered Codex result with `final_response=None` no longer leaks `TurnResult(...)` Python object repr into the chat. The upstream SDK documents `TurnResult.final_response` as nullable for turns that perform tool work without producing a final assistant message; the previous `... or str(result)` fallback would render the repr as visible assistant text. New `_buffered_result_text` helper returns the empty string in that case so the stream finishes cleanly with no extra content chunk. Same fix applied to `_run_codex_synthesis`. 3. Device-login SSE no longer forwards arbitrary subprocess output to the browser. The previous code yielded every CLI line under `{type:"log"}`, which on a shimmed binary could leak refresh tokens, auth JSON, or local config paths into the authenticated stream. Filtered to a known-safe vocabulary ("Welcome to Codex", "Initializing", "Successfully logged in", etc.). `device_url` and `device_code` events still fire as before. 4. CodexLoginButton no longer calls `window.open` from inside an awaited SSE handler. Browser popup blockers (Firefox, Safari, Chrome strict) silently block popups triggered outside a fresh user gesture, so the auto-open was unreliable. Replaced with a prominent "Open verification page" button styled as an anchor; the click handler is a real user gesture and is never blocked. The URL string is still shown below the button for copy/paste. Tests: 46 cases total (was 43). New regressions cover the parallel_calls clamp path on three garbage inputs, the buffered TurnResult-with-None-final repr leak guard, and the device-login log filter (asserts refresh tokens / auth.json paths are dropped while known-safe progress lines pass through). --- .../backend/core/inference/codex_provider.py | 67 ++++++- studio/backend/models/inference.py | 34 +++- studio/backend/tests/test_codex_provider.py | 170 ++++++++++++++++-- .../chat/components/codex-login-button.tsx | 48 +++-- 4 files changed, 279 insertions(+), 40 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index bb6ab71e56..db790f623e 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -589,11 +589,34 @@ async def _stream_thread_run( # 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) + text = _buffered_result_text(result) if text: yield text +def _buffered_result_text(result: Any) -> str: + """Extract assistant text from a buffered ``TurnResult``. + + The upstream SDK documents ``TurnResult.final_response`` as + nullable -- a turn that performs only tool work and completes + without a final assistant message will set it to ``None``. The + previous ``... or str(result)`` fallback then sent a Python + object repr (``TurnResult(...)``) into the chat, which surfaced + as visible garbage to the user. Returning the empty string for + that case lets the OpenAI-shape stream finish cleanly with no + extra content chunk -- the usage / stop / [DONE] frames still + fire, and the chat UI simply shows no assistant text rather + than a misleading object dump. + """ + text = _coerce_text(result) + if text: + return text + final = getattr(result, "final_response", None) + if isinstance(final, str) and final: + return final + return "" + + def _safe_thread_safety_kwargs() -> dict[str, Any]: """Return the safe ``approval_mode`` + ``sandbox`` kwargs for thread_start. @@ -1023,9 +1046,10 @@ async def _run_codex_synthesis( codex, model, system, synthesis_prompt ) result = await thread.run(synthesis_prompt) - return ( - _coerce_text(result) or getattr(result, "final_response", "") or str(result) - ) + # Use the same buffered extraction as `_stream_thread_run` so a + # synthesis turn whose `final_response` is None returns an empty + # string instead of a `TurnResult(...)` Python object repr. + return _buffered_result_text(result) except Exception as exc: logger.warning("codex_provider.synthesis_failed", error = str(exc)) return "" @@ -1116,6 +1140,35 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: rc: int = -1 cancelled = False + # Allow-list of substrings the upstream `codex login --device-auth` + # command prints during the normal flow. Anything outside this list + # is treated as opaque and not forwarded to the browser, so a + # shimmed binary that prints auth JSON, refresh tokens, local + # config paths, or unexpected stderr cannot leak that content + # through Studio's authenticated SSE stream. The URL and code + # extracted above are emitted separately as `device_url` / + # `device_code` events and are not affected by this filter. + safe_log_patterns: tuple[str, ...] = ( + "welcome to codex", + "initializing", + "open this", + "open:", + "open the", + "verification", + "enter this one-time code", + "enter the code", + "waiting", + "successfully logged in", + "logged in", + "signed in", + "browser opened", + "press ctrl", + ) + + def _safe_to_forward(text: str) -> bool: + lowered = text.lower() + return any(pat in lowered for pat in safe_log_patterns) + try: assert proc.stdout is not None while True: @@ -1134,7 +1187,11 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: if cm: yield {"type": "device_code", "code": cm.group(1)} code_emitted = True - yield {"type": "log", "line": line} + # Only forward lines from the known safe vocabulary; opaque + # output (file paths, tokens, JSON, error messages) stays in + # backend logs only. + if line and _safe_to_forward(line): + yield {"type": "log", "line": line} except (asyncio.CancelledError, GeneratorExit): cancelled = True raise diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 65a73dd65f..e079414357 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -788,19 +788,45 @@ class ChatCompletionRequest(BaseModel): ) parallel_calls: int = Field( default = 1, - ge = 1, - le = 20, description = ( "[x-unsloth] Codex provider only. When > 1, fan the chat turn " "out across N parallel Codex calls and synthesise a unified " "final answer. Each parallel attempt is rendered as its own tab " "in the chat UI; a final 'Synthesis' tab carries the merged " - "output. Bounded to [1, 20] by pydantic so a runaway value can't " - "saturate the local CLI. Defaults to 1 (single-call shape). " + "output. Silently clamped to [1, 20] by `_clamp_parallel_calls` " + "so a runaway value cannot saturate the local CLI -- using a " + "validator (rather than `ge=1, le=20`) keeps backwards " + "compatibility with pre-PR clients that sent the field as a " + "stray OpenAI extra (e.g. `0` for 'no fan-out') and would " + "otherwise hit a 422. Defaults to 1 (single-call shape). " "Silently ignored on every provider other than `codex`." ), ) + @field_validator("parallel_calls", mode = "before") + @classmethod + def _clamp_parallel_calls(cls, value: Any) -> int: + """Coerce ``parallel_calls`` to [1, 20] without rejecting weird inputs. + + Pre-PR behaviour was to silently ignore unknown / out-of-range + OpenAI extras; using ``ge=1, le=20`` on the Field would have + regressed that by returning a 422 to any non-Codex client that + happened to set the field to 0 or omit it as ``None``. Coerce + the value here instead so the schema stays self-documenting + ([1, 20]) while accepting legacy inputs. + """ + if value is None: + return 1 + try: + n = int(value) + except (TypeError, ValueError): + return 1 + if n < 1: + return 1 + if n > 20: + return 20 + return n + @model_validator(mode = "after") def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest": """Fill missing tool_call_id by walking back to the preceding assistant. diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 372bee770e..2b52bc5f94 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -16,9 +16,10 @@ Covers: * Parallel-calls fan-out: ``parallel_calls > 1`` spawns N async tasks and emits ``codex_tab_open`` / ``codex_tab_chunk`` / ``codex_tab_close`` events plus a final ``codex_gather`` synthesis event. -* Request validator: ``parallel_calls`` is clamped to [1, 20] by - pydantic so a runaway value is rejected with 422 before any Codex - task is spawned. +* Request validator: ``parallel_calls`` is silently clamped to + [1, 20] by a Pydantic field validator (not by ``ge=1, le=20``) so + non-Codex clients that send legacy values like ``0`` continue to + be accepted instead of getting a 422. """ from __future__ import annotations @@ -429,27 +430,47 @@ class TestParallelCallsValidator: ) assert req.parallel_calls == n - def test_request_rejects_below_one(self): + def test_request_clamps_below_one(self): + """Pre-PR clients sometimes sent `parallel_calls=0` as a stray + OpenAI extra and the request was silently accepted; rejecting + with 422 would regress that. The validator now clamps to 1. + """ from models.inference import ChatCompletionRequest - from pydantic import ValidationError - with pytest.raises(ValidationError): - ChatCompletionRequest( + for n in (0, -1, -100): + req = ChatCompletionRequest( model = "gpt-5.4", messages = [{"role": "user", "content": "hi"}], - parallel_calls = 0, + parallel_calls = n, ) + assert req.parallel_calls == 1, f"clamp failed for {n}" - def test_request_rejects_above_twenty(self): + def test_request_clamps_above_twenty(self): + """A runaway value (1000, etc.) is clamped to the 20 cap so it + cannot saturate the local CLI even when the client misbehaves. + """ from models.inference import ChatCompletionRequest - from pydantic import ValidationError - with pytest.raises(ValidationError): - ChatCompletionRequest( + for n in (21, 100, 1000): + req = ChatCompletionRequest( model = "gpt-5.4", messages = [{"role": "user", "content": "hi"}], - parallel_calls = 21, + parallel_calls = n, ) + assert req.parallel_calls == 20, f"clamp failed for {n}" + + def test_request_coerces_garbage_to_one(self): + """Strings / floats / None coerce to 1 instead of 422 so a + legacy or misconfigured client cannot break chat for everyone.""" + from models.inference import ChatCompletionRequest + + for value in (None, "garbage", float("nan")): + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + parallel_calls = value, + ) + assert req.parallel_calls == 1 def test_request_default_is_one(self): """Default = 1 so the field matches the single-call code path @@ -1330,6 +1351,129 @@ class TestCodexHardenedRegressions: # Model still passed so the request is well-formed. assert kw.get("model") == "gpt-5.5" + def test_device_login_log_filter_drops_unknown_lines(self, monkeypatch): + """The login stream's `log` events must not forward arbitrary + subprocess output. Only an allow-list of known progress + strings reaches the browser; anything else (auth JSON, + tokens, paths, error tails) stays in backend logs. + """ + # Build a synthetic stdout stream with one safe line and one + # unsafe line, then drive the login generator against it. + from core.inference import codex_provider as cp + + class _FakeStdout: + def __init__(self, lines: list[bytes]): + self._lines = list(lines) + + async def readline(self) -> bytes: + if not self._lines: + return b"" + return self._lines.pop(0) + + class _FakeProc: + pid = 99999 + returncode = None + stdout = _FakeStdout( + [ + b"Welcome to Codex\n", + b"Open: https://auth.openai.com/codex/device\n", + b"Enter this one-time code: ABCD-EFGH\n", + b'{"refresh_token": "rt_LEAK_LEAK_LEAK"}\n', + b"/home/u/.codex/auth.json saved\n", + b"Successfully logged in\n", + ], + ) + + async def wait(self): + self.returncode = 0 + return 0 + + def kill(self): + self.returncode = -9 + + def terminate(self): + self.returncode = -15 + + async def _fake_create_subprocess_exec(*a, **kw): + return _FakeProc() + + monkeypatch.setattr( + cp.asyncio, "create_subprocess_exec", _fake_create_subprocess_exec + ) + + events: list[dict] = [] + + async def _collect(): + async for ev in cp.stream_codex_device_login(): + events.append(ev) + + asyncio.run(_collect()) + log_lines = [ev.get("line", "") for ev in events if ev.get("type") == "log"] + joined = "\n".join(log_lines) + # Sensitive content must not have been forwarded. + assert "refresh_token" not in joined, f"token leaked: {joined!r}" + assert "rt_LEAK_LEAK_LEAK" not in joined + assert "auth.json" not in joined, f"local config path leaked: {joined!r}" + # The known-safe progress lines must be present so the UI can + # show the user what is happening. + assert any("Welcome to Codex" in line for line in log_lines) + assert any("Successfully logged in" in line for line in log_lines) + # device_url + device_code events must still fire. + url_events = [ev for ev in events if ev.get("type") == "device_url"] + code_events = [ev for ev in events if ev.get("type") == "device_code"] + assert url_events and url_events[0]["url"].endswith("/codex/device") + assert code_events and code_events[0]["code"] == "ABCD-EFGH" + + def test_buffered_result_none_final_does_not_emit_repr(self, monkeypatch): + """A buffered TurnResult whose final_response is None must NOT + send a Python object repr (``TurnResult(...)``) to the user. + Returning an empty content chunk is the right shape: the + stream still finishes with the usage + stop + [DONE] frames, + but no garbage assistant text appears. + """ + + class _ResultNoFinal: + final_response = None # explicit None + + def __repr__(self): + return "TurnResult(internal=should_not_leak)" + + class _ThreadBuffered: + async def run(self, prompt): + return _ResultNoFinal() + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _ThreadBuffered() + + _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()) + body = "".join(chunks) + assert ( + "TurnResult" not in body + ), f"Python object repr leaked to user content: {body!r}" + assert "should_not_leak" not in body + # Stream still terminated cleanly. + assert "[DONE]" in body + def test_empty_stream_falls_back_to_completed_agent_message(self, monkeypatch): """A successful turn that emits zero ``message.delta`` events but DOES emit a final ``ItemCompletedNotification`` with an 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 deaacb8a71..9667c8f382 100644 --- a/studio/frontend/src/features/chat/components/codex-login-button.tsx +++ b/studio/frontend/src/features/chat/components/codex-login-button.tsx @@ -66,14 +66,14 @@ export function CodexLoginButton({ onLoggedIn }: Props) { ) as AsyncGenerator) { if (event.type === "device_url" && event.url) { setDeviceUrl(event.url); - // Open the verification page eagerly so the user doesn't - // have to copy the URL out of the log surface. ``noopener`` - // prevents the auth-tab from controlling the Studio window. - try { - window.open(event.url, "_blank", "noopener,noreferrer"); - } catch { - // Ignore -- the URL is still visible in the log. - } + // Do NOT auto-open the verification URL with `window.open`. + // The click handler that started this flow has already + // awaited an SSE event, so the call is no longer in a user + // gesture and most browsers (Firefox, Safari, Chrome with + // strict popup settings) will silently block the popup. + // The URL is rendered as a prominent link below so the + // user can open it in one click without depending on the + // popup heuristic. } else if (event.type === "device_code" && event.code) { setDeviceCode(event.code); } else if (event.type === "log" && event.line) { @@ -114,17 +114,29 @@ export function CodexLoginButton({ onLoggedIn }: Props) { {busy ? "Signing in to Codex…" : "Sign in to Codex"} {deviceUrl && ( -

- Verification URL:{" "} - + +

+ Or copy: {deviceUrl} +

+
)} {deviceCode && (

From cb0680ebaa67b356e2b81f4b93d0cb94328a8622 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 17:07:00 +0000 Subject: [PATCH 20/40] Studio: per-tab buffer for Codex fan-out so chunks cannot interleave The Codex parallel-calls fan-out emits N independent SSE streams concurrently, so a chunk for tab 2 can arrive between two chunks for tab 1. The previous chat-adapter logic appended every chunk into a single `cumulativeText` buffer in arrival order, which made tab 1's text show up under tab 2's header (or vice versa) whenever the workers raced. With four or more parallel calls the rendered output became unreadable. Replaced the append-on-arrival path with per-tab buffers keyed by `tab_id`, plus a `renderCodexBuffer()` helper that rebuilds the Codex block from scratch on every fan-out event: - `codex_tab_open` allocates an empty buffer for the tab id. - `codex_tab_chunk` appends only into that tab's buffer. - `codex_tab_error` records the error string against the tab id. - `codex_tab_close` marks the tab finished (visual separator). - `codex_gather` flips a flag that draws the `--- Synthesis ---` divider; the synthesis text itself still arrives as a normal content delta on the same stream and so is not duplicated here. The block is re-rendered in `tab_id` order on every event, so the final output is deterministic regardless of arrival interleaving. --- .../src/features/chat/api/chat-adapter.ts | 84 ++++++++++++++----- 1 file changed, 64 insertions(+), 20 deletions(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 2b5ccf9226..086aa25319 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1131,6 +1131,39 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let cumulativeText = ""; let reasoningStartAt: number | null = null; let reasoningDuration = 0; + // Per-tab buffer for Codex parallel-calls fan-out. The backend + // emits N independent streams concurrently, so chunks for tab 2 + // can land between chunks for tab 1 in arrival order. Keeping a + // dict keyed by tab_id and re-assembling cumulativeText from + // scratch on every codex event puts each tab's text under its + // own header regardless of arrival interleaving. + const codexTabBuffers = new Map(); + const codexTabClosed = new Set(); + const codexTabError = new Map(); + let codexTotalTabs = 0; + let codexGatherEmitted = false; + + function renderCodexBuffer(): string { + if (codexTabBuffers.size === 0 && !codexGatherEmitted) return ""; + const lines: string[] = []; + const ids = [...codexTabBuffers.keys()].sort((a, b) => a - b); + for (const id of ids) { + const header = codexTotalTabs + ? `[Codex tab ${id}/${codexTotalTabs}]` + : `[Codex tab ${id}]`; + lines.push(`\n\n${header}\n${codexTabBuffers.get(id) ?? ""}`); + if (codexTabError.has(id)) { + lines.push(`\n[Codex tab ${id} error: ${codexTabError.get(id)}]\n`); + } + if (codexTabClosed.has(id)) { + lines.push("\n"); + } + } + if (codexGatherEmitted) { + lines.push("\n--- Synthesis ---\n"); + } + return lines.join(""); + } // Tracks whether we are currently inside a `` block opened by // a `delta.reasoning_content` chunk. Kimi (kimi-k2.6, kimi-k2-thinking) // and DeepSeek's reasoner stream their thinking as a separate @@ -1585,42 +1618,53 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { - // Codex parallel-calls fan-out events: render each - // per-tab chunk inline with a "Tab N:" prefix so users - // see the N independent attempts even before the - // dedicated CodexParallelTabs UI is wired into the - // chat surface. `codex_gather` carries the synthesis + // Codex parallel-calls fan-out events: route chunks + // into per-tab buffers keyed by tab_id, then render + // the whole codex block from scratch each event so + // concurrent tabs cannot interleave under the wrong + // header. `codex_gather` carries the synthesis // payload which the backend also emits as a normal - // content delta, so we drop it here to avoid showing - // the synthesis twice. tab_open / tab_close / tab_error - // are header markers we surface as one-line notes. + // content delta later in the same SSE stream, so we + // only render a divider here to avoid duplicating + // the synthesis text. if (typeof toolEvent.type === "string" && toolEvent.type.startsWith("codex_")) { if (toolEvent.type === "codex_tab_open") { const tabId = Number(toolEvent.tab_id); const total = Number(toolEvent.total_tabs); - if (Number.isFinite(tabId) && Number.isFinite(total)) { - cumulativeText += `\n\n[Codex tab ${tabId}/${total}]\n`; + if (Number.isFinite(tabId)) { + if (!codexTabBuffers.has(tabId)) { + codexTabBuffers.set(tabId, ""); + } + if (Number.isFinite(total) && total > codexTotalTabs) { + codexTotalTabs = total; + } } } else if (toolEvent.type === "codex_tab_chunk") { + const tabId = Number(toolEvent.tab_id); const text = typeof toolEvent.text === "string" ? toolEvent.text : ""; - if (text) cumulativeText += text; + if (Number.isFinite(tabId) && text) { + const prev = codexTabBuffers.get(tabId) ?? ""; + codexTabBuffers.set(tabId, prev + text); + } } else if (toolEvent.type === "codex_tab_error") { const tabId = Number(toolEvent.tab_id); const err = typeof toolEvent.error === "string" ? toolEvent.error : "error"; if (Number.isFinite(tabId)) { - cumulativeText += `\n[Codex tab ${tabId} error: ${err}]\n`; + codexTabError.set(tabId, err); + if (!codexTabBuffers.has(tabId)) { + codexTabBuffers.set(tabId, ""); + } } } else if (toolEvent.type === "codex_tab_close") { - // Mark end of tab block so synthesis is visually separated. - cumulativeText += "\n"; + const tabId = Number(toolEvent.tab_id); + if (Number.isFinite(tabId)) { + codexTabClosed.add(tabId); + } } else if (toolEvent.type === "codex_gather") { - // Synthesis is also emitted as a normal content - // delta later in the same SSE stream; nothing to - // add here. Surface a divider so the user can tell - // where the synthesis starts. - cumulativeText += "\n--- Synthesis ---\n"; + codexGatherEmitted = true; } - const codexParts = parseAssistantContent(cumulativeText); + const codexBlock = renderCodexBuffer(); + const codexParts = parseAssistantContent(cumulativeText + codexBlock); yield { content: [...toolCallParts, ...codexParts], }; From aa258b983d3acbadababfb5db7e90738df8fae67 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 17:11:42 +0000 Subject: [PATCH 21/40] Studio: account for every Codex turn in fan-out usage chunk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fan-out path used to report usage as if a single Codex call had run: `prompt_tokens = max(1, len(prompt)//4)` and `completion_tokens = len(synthesis)//4`. In reality it had spawned N parallel worker turns (each carrying the same prompt) plus a synthesis turn that re-sent the prompt and every tab's output. For `parallel_calls=20` that meant the cost / context widget under- reported the request by roughly 20x. Now sums: - `prompt_tokens ≈ (N * prompt + synthesis_prompt) / 4` where `synthesis_prompt = sum(tab_outputs) + prompt`. - `completion_tokens ≈ (sum(tab_output_chars) + synthesis_chars) / 4`. Tests: new regression `test_parallel_usage_accounts_for_all_calls` runs a 4-way fan-out against a fake SDK with deterministic chunk lengths and asserts the reported usage scales with N, not the single-call shape. --- .../backend/core/inference/codex_provider.py | 16 +++++++- studio/backend/tests/test_codex_provider.py | 40 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index db790f623e..ca3dad0e37 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -1001,10 +1001,22 @@ async def _stream_codex_parallel( if synthesis_text: yield _chunk_text(completion_id, synthesis_text) + # Account for ALL Codex turns the fan-out spawned: N parallel + # workers each ran the same prompt (≈ N * prompt_tokens), and the + # synthesis turn re-sent the prompt plus every tab's output. Without + # this the cost / context display is off by the fan-out factor and + # users see a wildly inaccurate token count for the request. + total_tab_completion_chars = sum(len(t) for t in per_tab_texts) + synthesis_prompt_chars = sum(len(t) for t in per_tab_texts) + len(prompt) yield _chunk_usage( completion_id, - prompt_tokens = max(1, len(prompt) // 4), - completion_tokens = max(0, len(synthesis_text) // 4), + # n worker prompts (same prompt each) + synthesis prompt (which + # carries the prompt again plus every tab's output). + prompt_tokens = max(1, (n * len(prompt) + synthesis_prompt_chars) // 4), + # Sum of every worker's output plus the synthesis text. + completion_tokens = max( + 0, (total_tab_completion_chars + len(synthesis_text)) // 4 + ), ) yield _chunk_stop(completion_id) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 2b52bc5f94..27f40c7e77 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -1424,6 +1424,46 @@ class TestCodexHardenedRegressions: assert url_events and url_events[0]["url"].endswith("/codex/device") assert code_events and code_events[0]["code"] == "ABCD-EFGH" + def test_parallel_usage_accounts_for_all_calls(self, monkeypatch): + """The fan-out path runs N worker calls + 1 synthesis call. + The reported usage must reflect that, not just one call's + worth, otherwise the cost / context display is off by the + fan-out factor. + """ + _install_fake_codex_sdk( + monkeypatch, + lambda: _FakeAsyncCodex( + chunks = ["AAAAAAAAAA"], # 10 chars per tab + final = "SYNTHESISED" * 10, # 110 chars synthesis + ), + ) + from core.inference.codex_provider import stream_codex + + n = 4 + long_prompt = "a" * 200 # 200 chars + lines = _collect_stream( + stream_codex( + messages = [{"role": "user", "content": long_prompt}], + model = "gpt-5.4", + parallel_calls = n, + ) + ) + chunks = _parse_sse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("choices") == [] and c.get("usage")] + assert len(usage_chunks) == 1 + usage = usage_chunks[0]["usage"] + # Single-call prompt would be ~200/4 = 50 tokens. For n=4 with + # synthesis, prompt should be much larger: n*200 + (n*10 + 200) + # = 800 + 240 = 1040 chars ~= 260 tokens. + assert ( + usage["prompt_tokens"] >= 200 + ), f"prompt_tokens not scaled for fan-out: {usage['prompt_tokens']}" + # Completion = n*10 (tab outputs) + 110 (synthesis) = 150 chars + # ~= 37 tokens. Definitely > the synthesis-only count of 27. + assert ( + usage["completion_tokens"] >= 30 + ), f"completion_tokens not scaled for fan-out: {usage['completion_tokens']}" + def test_buffered_result_none_final_does_not_emit_repr(self, monkeypatch): """A buffered TurnResult whose final_response is None must NOT send a Python object repr (``TurnResult(...)``) to the user. From 8ee60019a4705b0fcefdc689674ffa85bdecbd10 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 17:30:01 +0000 Subject: [PATCH 22/40] Studio: round 6 Codex hardening (5 follow-ups) Reviewer round 6 surfaced five real follow-ups on top of rounds 5 through 5g. Each one is a fix for an asymmetric guard or a wrong- shape lookup in the new Codex provider code: 1. Re-gate `installed=True` on having BOTH the SDK AND a `codex` binary on PATH. Round 5 widened the gate to SDK-only, but the login route still shells out to the binary, so an SDK-only host would surface a Codex row whose Sign-in button immediately failed with "codex CLI not found on PATH". The canonical `openai-codex` package depends on `openai-codex-cli-bin` which places the shim on PATH for free, so the common install still lights up; the gate just refuses to advertise a provider Studio cannot actually drive end-to-end. 2. `_safe_thread_safety_kwargs` now also probes `.api` and `.generated.v2_all` for `SandboxMode`. Upstream `openai_codex` exports `ApprovalMode` at the top level but `SandboxMode` lives under `openai_codex.generated.v2_all`. The previous lookup returned `{}` on the canonical SDK install, so every thread_start ran with the unsafe `auto_review` default. Submodule probe resolves the canonical layout and keeps backwards-compat with builds that DID re-export at the top level. 3. `_coerce_text` now applies the answer-event-type filter on the object path too. The upstream SDK emits typed payload classes like `CommandExecutionOutputDelta`, `FileChangeDelta`, `ToolCallDelta`, `PatchApplyDelta`, etc., all of which carry a `.delta` string of local stdout / file paths / tool args. The dict path already filtered these out; the object path used to return `.delta` unconditionally, so a real SDK install could leak tool output into the visible chat reply. 4. `_ScrubbedEnvAsyncCodex` is now process-wide concurrency-safe AND fails-closed if the SDK constructor raises: - Refcount each scrubbed key under an `asyncio.Lock` so a fan-out wrapper that exits early cannot restore a secret while another wrapper is still inside SDK startup (round 6 reproduced this: wrapper A exited, wrapper B's SDK saw the restored HF_TOKEN). - Move `_async_codex_cls()` and its `__aenter__` INSIDE a try/except in `__aenter__`; on failure, run the release path so the scrubbed env vars are restored even though `__aexit__` never fires for the failed construction. 5. `_run_cli` now detaches into its own process group via `start_new_session=True` (Unix) / `CREATE_NEW_PROCESS_GROUP` (Windows) and kills the whole group on timeout, matching the protected path in `stream_codex_device_login`. A shimmed `codex login status` that forks a helper and blocks no longer leaves the child running after we killed the parent. Frontend follow-up: chat-adapter now routes every rendered yield through a `renderFullContent()` helper so the Codex per-tab text accumulated in earlier `_toolEvent` frames is preserved when the synthesis content delta arrives. Previously the next regular content yield rebuilt `parts` from `cumulativeText` alone and the tab section vanished from the final assistant message. Tests: 49 cases total (was 47). New regressions: - `installed_requires_both_cli_and_sdk` (round 6 revert). - `safety_kwargs_finds_sandbox_mode_in_submodule` (canonical SDK layout where `SandboxMode` is in `.generated.v2_all`). - `scrubbed_env_construction_failure_restores_env` (no permanent env leak when the SDK constructor raises). - Expanded `coerce_text_drops_non_answer_event_types` to also exercise the object-shape code path with `CommandExecutionOutputDelta`, `FileChangeDelta`, `ToolCallDelta`, `PatchApplyDelta`, `PlanUpdateDelta`, `AgentReasoningDelta`, plus the positive `AgentMessageDelta` allow-through. --- .../core/inference/codex_availability.py | 71 +++++-- .../backend/core/inference/codex_provider.py | 164 ++++++++++++--- studio/backend/tests/test_codex_provider.py | 192 ++++++++++++++++-- .../src/features/chat/api/chat-adapter.ts | 31 ++- 4 files changed, 398 insertions(+), 60 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 9699d2224f..7383c57311 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -158,19 +158,35 @@ async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, """Run a short ``codex`` CLI command and return (rc, stdout, stderr). The probe uses 4s as the wall-clock cap because ``codex --version`` - and ``codex auth status`` both return in well under a second on a + and ``codex login status`` both return in well under a second on a healthy install. A longer probe would block the ``/api/codex/status`` route -- and that route fires on every chat page load, so a tight cap matters. + + Subprocess lifecycle: detached into its own process group on Unix + via ``start_new_session=True`` (matching ``stream_codex_device_login``) + so a hung child cannot survive ``proc.kill()`` on timeout. Without + this, a shimmed ``codex login status`` that forks a helper then + blocks would leave the helper running after we killed the parent. + Windows uses ``CREATE_NEW_PROCESS_GROUP`` for the analogous + isolation. Round 6 reviewer caught the asymmetry with the + device-login path that already had this guard. """ + import os + import signal + + spawn_kwargs: dict[str, Any] = { + "stdout": asyncio.subprocess.PIPE, + "stderr": asyncio.subprocess.PIPE, + "env": _codex_subprocess_env(), + } + if os.name == "posix": + spawn_kwargs["start_new_session"] = True + elif os.name == "nt": + spawn_kwargs["creationflags"] = 0x00000200 # CREATE_NEW_PROCESS_GROUP + try: - proc = await asyncio.create_subprocess_exec( - "codex", - *args, - stdout = asyncio.subprocess.PIPE, - stderr = asyncio.subprocess.PIPE, - env = _codex_subprocess_env(), - ) + proc = await asyncio.create_subprocess_exec("codex", *args, **spawn_kwargs) except FileNotFoundError: return -1, "", "codex binary not on PATH" except Exception as exc: @@ -184,9 +200,21 @@ async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, try: stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout = timeout) except asyncio.TimeoutError: + # Kill the whole process group, not just the parent, so any + # child the codex CLI forked also dies. + if os.name == "posix": + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except (ProcessLookupError, PermissionError): + pass + else: + try: + proc.send_signal(signal.CTRL_BREAK_EVENT) # type: ignore[attr-defined] + except Exception: + pass proc.kill() try: - await proc.wait() + await asyncio.wait_for(proc.wait(), timeout = 1.0) except Exception: pass return -1, "", f"codex {' '.join(args)} timed out after {timeout:.1f}s" @@ -268,12 +296,16 @@ async def probe_codex_availability() -> dict[str, Any]: Returns a dict with keys: - * ``installed`` (bool) -- True iff Studio can actually drive Codex. - The SDK's `openai-codex-cli-bin` runtime is what backs - `AsyncCodex(...)`, so an importable SDK alone is sufficient even - with no standalone `codex` on PATH. We still report `cli_path` - separately so the UI can surface "CLI also present" / "SDK - bundled runtime only" without changing the gate. + * ``installed`` (bool) -- True iff Studio can actually drive Codex + end-to-end: BOTH the Python SDK (for chat) AND a `codex` + executable on PATH (for the device-auth login flow). The + canonical `openai-codex` package depends on `openai-codex-cli-bin` + which installs the `codex` shim into the venv's `bin/`, so the + common SDK-only install in fact gets the CLI on PATH for free + and this gate triggers correctly. Hosts that import the SDK + from a wheel without that runtime dep stay hidden because the + login flow would otherwise fail with "codex CLI not found on + PATH" after the user clicked Sign in. * ``cli_path`` (str | None) -- absolute path to the CLI, or None. * ``sdk_importable`` (bool) -- the Python SDK is importable. * ``logged_in`` (bool) -- best-effort auth check; meaningless when @@ -285,9 +317,12 @@ async def probe_codex_availability() -> dict[str, Any]: sdk_ok = _sdk_importable() payload: dict[str, Any] = { - # SDK alone is enough -- it bundles the codex runtime. A - # standalone CLI on PATH is additional but optional. - "installed": sdk_ok, + # Gate on BOTH because the login flow shells out to `codex`. + # Round 5 briefly set this to `sdk_ok` alone, but round 6 + # caught that the login route would then fail with + # `codex CLI not found on PATH` after the user clicked + # Sign in, leaving them with an unusable provider row. + "installed": bool(cli_path) and sdk_ok, "cli_path": cli_path, "sdk_importable": sdk_ok, "logged_in": False, diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index ca3dad0e37..35a2ef4393 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -93,6 +93,10 @@ def _codex_sdk_env_override() -> dict[str, str]: return {key: "" for key in os.environ if key not in safe} +_SCRUBBED_ENV_LOCK = asyncio.Lock() +_SCRUBBED_ENV_REFCOUNT: dict[str, int] = {} + + class _ScrubbedEnvAsyncCodex: """Async-context wrapper that swaps `os.environ` for the lifetime of a Codex SDK session. @@ -101,40 +105,85 @@ class _ScrubbedEnvAsyncCodex: `AppServerConfig(env=...)`. The SDK starts its app-server with `env = os.environ.copy()`, so removing secrets from the parent process env right before construction keeps them out of the child. - Restore happens on exit, and the restore is `setdefault`-style so - concurrent wrappers do not clobber each other's state. + + Concurrency model: a process-wide asyncio lock serialises the + enter/exit critical section, and a per-key refcount tracks how + many concurrent wrappers are currently "holding" the scrub. A + key is only restored when the last wrapper using it exits. This + fixes two issues round 6 caught: + + 1. Two concurrent fan-out workers used to race: wrapper A could + restore `HF_TOKEN` while wrapper B was still inside SDK + startup, letting B's spawned app-server inherit the secret. + The refcount keeps the key scrubbed for the full overlap + window. + 2. If the SDK constructor raised before `__aenter__` returned, + Python never called `__aexit__`, so the deleted keys leaked + permanently. Construction now happens INSIDE the try/except + in `__aenter__`, and the scrub is rolled back on failure. """ def __init__(self, async_codex_cls: Any): self._async_codex_cls = async_codex_cls self._inner: Any = None - self._saved_env: dict[str, str] | None = None + # Keys this wrapper instance contributed to the refcount, so + # __aexit__ knows exactly which counters to decrement (avoids + # racing with concurrent wrappers that scrub a different set). + self._held_keys: list[str] = [] + # Snapshot of the original values at the time of the FIRST + # wrapper that scrubbed each key, so restoration uses the + # real pre-scrub value. + self._restored_via_us: dict[str, str] = {} async def __aenter__(self) -> Any: import os overrides = _codex_sdk_env_override() - saved: dict[str, str] = {} - for key in overrides: - if key in os.environ: - saved[key] = os.environ[key] - del os.environ[key] - self._saved_env = saved - self._inner = self._async_codex_cls() - return await self._inner.__aenter__() + async with _SCRUBBED_ENV_LOCK: + for key in overrides: + if key not in os.environ and _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0: + continue + if _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0: + # First wrapper to scrub this key -- save the + # original so the very last wrapper to release it + # can restore the right value. + self._restored_via_us[key] = os.environ[key] + del os.environ[key] + _SCRUBBED_ENV_REFCOUNT[key] = _SCRUBBED_ENV_REFCOUNT.get(key, 0) + 1 + self._held_keys.append(key) + try: + self._inner = self._async_codex_cls() + return await self._inner.__aenter__() + except BaseException: + # Roll back the scrub if SDK construction / enter fails; + # otherwise the deleted env vars would leak permanently. + await self._release_held_keys() + raise async def __aexit__(self, exc_type, exc, tb): - import os - try: if self._inner is not None: return await self._inner.__aexit__(exc_type, exc, tb) finally: - saved = self._saved_env or {} - for key, value in saved.items(): - # Only restore keys we removed; setdefault avoids - # clobbering a concurrent caller's value. - os.environ.setdefault(key, value) + await self._release_held_keys() + + async def _release_held_keys(self) -> None: + import os + + async with _SCRUBBED_ENV_LOCK: + for key in self._held_keys: + current = _SCRUBBED_ENV_REFCOUNT.get(key, 0) + if current <= 0: + continue + _SCRUBBED_ENV_REFCOUNT[key] = current - 1 + if current - 1 == 0: + # Last wrapper holding this key -- restore the + # original value if WE were the first to scrub it, + # or pull from any other wrapper's saved snapshot. + if key in self._restored_via_us: + os.environ.setdefault(key, self._restored_via_us[key]) + self._held_keys.clear() + self._restored_via_us.clear() def _open_async_codex(async_codex_cls: Any) -> Any: @@ -399,6 +448,18 @@ def _coerce_text(payload: Any) -> str: visible text. Tool / command / plan deltas are dropped so local stdout, file paths, or tool-call arguments never flow into the Chat Completions content stream. + + Both dict-shaped events (tests + some pre-release SDKs) AND + object-shaped events (the real upstream SDK's typed notification + classes) are gated -- if the payload exposes a `type` attribute + or key whose value is not in the answer-event allow-list, we + return the empty string regardless of whether `.delta` or `.text` + is present. Round 6 reviewer caught the object-shape gap: the + upstream SDK can emit `item/commandExecution/outputDelta`, + `item/fileChange/outputDelta`, etc. as typed objects, all of + which carry `.delta` strings containing local stdout, patches, + or tool arguments. Without the object-side filter those strings + would have flowed straight into visible assistant text. """ if payload is None: return "" @@ -419,6 +480,33 @@ def _coerce_text(payload: Any) -> str: return "" if isinstance(payload, list): return "".join(_coerce_text(item) for item in payload) + # Object path: gate on `payload.type` if present, AND on the class + # name as a fallback (the upstream SDK uses class names like + # `AgentMessageDeltaNotification` / `CommandExecutionOutputDelta` + # so a denylist-by-substring catches typed payloads that lack a + # `type` attribute). + ev_type_obj = getattr(payload, "type", None) + if isinstance(ev_type_obj, str) and ev_type_obj not in _ANSWER_EVENT_TYPES: + return "" + cls_name = payload.__class__.__name__ + # Allow only class names that contain "Message" or "Delta" without + # also containing a tool / command / plan / file marker. + cls_lower = cls_name.lower() + if any( + marker in cls_lower + for marker in ( + "command", + "exec", + "file", + "patch", + "plan", + "tool", + "reason", + "stdout", + "stderr", + ) + ): + return "" text_attr = getattr(payload, "text", None) if isinstance(text_attr, str): return text_attr @@ -637,18 +725,46 @@ def _safe_thread_safety_kwargs() -> dict[str, Any]: * ``sandbox = SandboxMode.read_only`` -- the policy that bans file writes and disables network access. - Returns an empty dict when the installed SDK is too old to expose - either symbol; the caller then issues a structured warning and - proceeds without the safety pins. Failing closed (refusing to - run) on an older SDK would brick users on pre-release alpha - builds for no security gain -- the auto_review default is - upstream's choice, not a Studio regression. + Probes multiple locations: ``ApprovalMode`` is exported at the + top-level ``openai_codex`` package, but ``SandboxMode`` lives in + ``openai_codex.generated.v2_all`` (re-exported into + ``openai_codex.api``) and is NOT in the top-level __init__. + Round 6 reviewer caught this -- looking only at the top-level + module returned ``{}``, silently degrading to the auto_review + default on the canonical SDK install. + + Returns an empty dict only when the installed SDK is so different + that neither path resolves -- the caller then issues a + structured warning and proceeds. Failing closed (refusing to + run) on a future SDK rev would brick users for no security gain + -- the auto_review default is upstream's choice, not a Studio + regression. """ sdk_mod = sys.modules.get("openai_codex") or sys.modules.get("codex_app_server") if sdk_mod is None: return {} + + # ApprovalMode: top-level export on canonical SDK. approval_mode_cls = getattr(sdk_mod, "ApprovalMode", None) + + # SandboxMode: try top-level, then `.api`, then `.generated.v2_all`. + # We do not eagerly import these submodules because the SDK may + # not expose them and we do not want to crash the request on an + # ImportError. importlib.import_module gives us a typed failure. sandbox_mode_cls = getattr(sdk_mod, "SandboxMode", None) + if sandbox_mode_cls is None: + for sub in ("api", "generated.v2_all"): + mod_name = getattr(sdk_mod, "__name__", "") + if not mod_name: + continue + try: + submod = importlib.import_module(f"{mod_name}.{sub}") + except Exception: + continue + sandbox_mode_cls = getattr(submod, "SandboxMode", None) + if sandbox_mode_cls is not None: + break + if approval_mode_cls is None or sandbox_mode_cls is None: return {} deny_all = getattr(approval_mode_cls, "deny_all", None) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 27f40c7e77..69a1ff74e2 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -1012,24 +1012,39 @@ class TestCodexHardenedRegressions: assert '"content": "hello "' in body assert '"content": "from turn.stream"' in body - def test_installed_true_when_sdk_only(self, monkeypatch): - """SDK alone is sufficient: openai-codex-cli-bin ships the - runtime that backs `AsyncCodex(...)`, so the picker must be - shown even when no standalone `codex` lives on PATH. + def test_installed_requires_both_cli_and_sdk(self, monkeypatch): + """Round 6 revert: the login route shells out to `codex`, so + marking `installed=True` on SDK-only would surface a Codex + provider row whose Sign-in button immediately fails. The + canonical `openai-codex` package installs `openai-codex-cli-bin` + which puts the `codex` shim on PATH, so common installs still + light up correctly; the gate just refuses to advertise a + provider Studio cannot actually drive. """ from core.inference import codex_availability as ca + # SDK present, no CLI -> hidden (cannot complete login). monkeypatch.setattr(ca, "_which_codex", lambda: None) monkeypatch.setattr(ca, "_sdk_importable", lambda: True) - payload = asyncio.run(ca.probe_codex_availability()) - assert payload["installed"] is True + assert payload["installed"] is False assert payload["cli_path"] is None assert payload["sdk_importable"] is True - # logged_in stays False because the version/login probes only - # run when a CLI is present (they shell out to it). That is the - # expected behaviour, not a bug. - assert payload["logged_in"] is False + + # CLI present, SDK missing -> still hidden (cannot drive chat). + monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/bin/codex") + monkeypatch.setattr(ca, "_sdk_importable", lambda: False) + + async def fake_version(): + return "codex-cli 0.133.0" + + async def fake_logged_in(): + return True + + monkeypatch.setattr(ca, "_detect_version", fake_version) + monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in) + payload2 = asyncio.run(ca.probe_codex_availability()) + assert payload2["installed"] is False def test_base_instructions_kwarg_preferred(self, monkeypatch): """The upstream openai_codex SDK uses `base_instructions` for @@ -1181,6 +1196,11 @@ class TestCodexHardenedRegressions: that must NOT be rendered as assistant text -- otherwise local stdout, file paths, or tool-call arguments would leak into the Chat Completions reply. + + Round 6 also requires the object-shape path to gate on type + and class name; the upstream SDK emits typed notification + objects (CommandExecutionOutputDelta, FileChangeDelta, etc.) + with `.delta` strings that would otherwise leak. """ from core.inference.codex_provider import _coerce_text @@ -1189,8 +1209,7 @@ class TestCodexHardenedRegressions: assert _coerce_text({"type": "completed", "text": "done"}) == "done" assert _coerce_text({"type": "text_delta", "delta": "x"}) == "x" - # Non-answer event types are silenced even when they expose a - # delta string that looks like prose. + # Non-answer dict event types are silenced. for ev_type in ( "command.delta", "command_output", @@ -1208,6 +1227,53 @@ class TestCodexHardenedRegressions: f"{ev_type} leaked text into assistant reply: " f"{_coerce_text(payload)!r}" ) + + # Object-shape gate: typed payloads whose class name contains + # a tool/command/file/patch/plan marker drop the .delta too. + class CommandExecutionOutputDelta: + delta = "SECRET_STDOUT" + + class FileChangeDelta: + delta = "secret/file/path" + + class ToolCallDelta: + text = "tool_arg_payload" + + class PatchApplyDelta: + delta = "diff --git a/secret" + + class PlanUpdateDelta: + delta = "plan content" + + class AgentReasoningDelta: + delta = "internal CoT" + + for obj in ( + CommandExecutionOutputDelta(), + FileChangeDelta(), + ToolCallDelta(), + PatchApplyDelta(), + PlanUpdateDelta(), + AgentReasoningDelta(), + ): + assert _coerce_text(obj) == "", ( + f"object-shape {obj.__class__.__name__} leaked: " + f"{_coerce_text(obj)!r}" + ) + + # Object with explicit type attr also drops if not in allow-list. + class _WithType: + type = "command.delta" + delta = "leak" + + assert _coerce_text(_WithType()) == "" + + # Object-shape answer events DO pass through. + class AgentMessageDelta: + delta = "real assistant text" + + assert _coerce_text(AgentMessageDelta()) == "real assistant text" + # Plain strings and untyped dicts still pass through (legacy path). assert _coerce_text("raw text") == "raw text" assert _coerce_text({"text": "no type tag"}) == "no type tag" @@ -1309,6 +1375,108 @@ class TestCodexHardenedRegressions: kw.get("sandbox") == "READ_ONLY_SENTINEL" ), f"sandbox not pinned to read_only: {kw}" + def test_safety_kwargs_finds_sandbox_mode_in_submodule(self, monkeypatch): + """Round 6 caught that `SandboxMode` is exported by the + upstream SDK from `openai_codex.generated.v2_all`, NOT from + the top-level `openai_codex` package. The previous lookup + used `getattr(sdk_mod, 'SandboxMode', None)` only and returned + None for the canonical SDK install, silently degrading to + the unsafe auto_review default. + """ + seen_kwargs: list[dict] = [] + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + seen_kwargs.append(dict(kw)) + return _FakeThread(chunks = ["ok"]) + + # Build a fake openai_codex that DOES NOT expose SandboxMode + # at the top level -- only inside `.generated.v2_all`. + import importlib.util as _iu + + fake_root = types.ModuleType("openai_codex") + fake_root.AsyncCodex = _Async # type: ignore[attr-defined] + fake_root.ApprovalMode = types.SimpleNamespace( # type: ignore[attr-defined] + deny_all = "DENY_ALL", + auto_review = "AUTO", + ) + # Submodule chain `.generated.v2_all` + fake_generated = types.ModuleType("openai_codex.generated") + fake_v2 = types.ModuleType("openai_codex.generated.v2_all") + fake_v2.SandboxMode = types.SimpleNamespace( # type: ignore[attr-defined] + read_only = "READ_ONLY", + workspace_write = "WW", + ) + fake_generated.v2_all = fake_v2 # type: ignore[attr-defined] + fake_root.generated = fake_generated # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "openai_codex", fake_root) + monkeypatch.setitem(sys.modules, "openai_codex.generated", fake_generated) + monkeypatch.setitem(sys.modules, "openai_codex.generated.v2_all", fake_v2) + 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 + + async def _collect(): + async for _ in stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + parallel_calls = 1, + ): + pass + + asyncio.run(_collect()) + assert seen_kwargs, "thread_start never called" + kw = seen_kwargs[0] + assert ( + kw.get("approval_mode") == "DENY_ALL" + ), f"approval_mode not pinned even with submodule SandboxMode: {kw}" + assert ( + kw.get("sandbox") == "READ_ONLY" + ), f"sandbox not pinned via submodule lookup: {kw}" + + def test_scrubbed_env_construction_failure_restores_env(self, monkeypatch): + """Round 6: if the SDK constructor raises before __aenter__ + returns, the previous wrapper never called __aexit__ so the + scrubbed env vars leaked permanently. Now the scrub is rolled + back on failure. + """ + from core.inference.codex_provider import _ScrubbedEnvAsyncCodex + + monkeypatch.setenv("HF_TOKEN", "must_survive") + + class _FailingAsync: + def __init__(self): + raise RuntimeError("SDK construction failed") + + async def _run(): + wrapper = _ScrubbedEnvAsyncCodex(_FailingAsync) + try: + async with wrapper: + pass + except RuntimeError: + pass + + asyncio.run(_run()) + # HF_TOKEN must be restored even though __aexit__ never fired + # for the failed construction. + assert ( + os.environ.get("HF_TOKEN") == "must_survive" + ), "scrubbed env leaked permanently when SDK construction failed" + def test_thread_start_skips_safety_kwargs_on_old_sdk(self, monkeypatch): """If the installed SDK does not expose ApprovalMode or SandboxMode (older rev / alias), thread_start must still run diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 086aa25319..8cd6b903d6 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1164,6 +1164,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } return lines.join(""); } + + // All chat-content yields go through this so the Codex per-tab + // output is always concatenated with the normal SSE text. The + // synthesis is emitted by the backend BOTH as a `codex_gather` + // tool event AND as a normal content delta; rendering both + // would duplicate it. Render the tabs above (header / tab text) + // separately from cumulativeText (which carries the synthesis + // content delta) so the user sees `[tabs] ... [synthesis]`. + function renderFullContent(): string { + return cumulativeText + renderCodexBuffer(); + } // Tracks whether we are currently inside a `` block opened by // a `delta.reasoning_content` chunk. Kimi (kimi-k2.6, kimi-k2-thinking) // and DeepSeek's reasoner stream their thinking as a separate @@ -1663,8 +1674,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else if (toolEvent.type === "codex_gather") { codexGatherEmitted = true; } - const codexBlock = renderCodexBuffer(); - const codexParts = parseAssistantContent(cumulativeText + codexBlock); + const codexParts = parseAssistantContent(renderFullContent()); yield { content: [...toolCallParts, ...codexParts], }; @@ -1781,8 +1791,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }; } } - // Yield cumulative state so tool UI updates (tools first, text after) - const textParts = parseAssistantContent(cumulativeText); + // Yield cumulative state so tool UI updates (tools first, text after). + // Use renderFullContent() so any Codex per-tab text accumulated + // in earlier _toolEvent frames is preserved when the synthesis + // delta arrives -- without this the tabs would briefly appear + // and then vanish when the regular content path overwrote them. + const textParts = parseAssistantContent(renderFullContent()); yield { content: [...toolCallParts, ...textParts], metadata: { @@ -1912,7 +1926,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { "", ); } - const parts = parseAssistantContent(cumulativeText); + // renderFullContent() preserves any Codex per-tab text the + // fan-out branch accumulated into codexTabBuffers. + const parts = parseAssistantContent(renderFullContent()); if ( parts.some((part) => part.type === "reasoning") && @@ -2016,7 +2032,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [ ...toolCallParts, - ...parseAssistantContent(cumulativeText), + // renderFullContent() ensures the Codex per-tab text is in + // the FINAL message too -- otherwise the synthesis delta on + // the regular content path would have erased it. + ...parseAssistantContent(renderFullContent()), ...sourceParts, ], metadata: { From 2874abbfecae11a1a090f5373e9d6d42b8765782 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 17:45:20 +0000 Subject: [PATCH 23/40] Studio: fail closed when Codex SDK cannot enforce safety pins Round 6 reviewer noted that the warn-and-proceed path in `_start_thread_with_system` is "failing open" on a server-side chat surface: an SDK rev that does not expose ApprovalMode or SandboxMode would log a warning then call `thread_start(model=...)` with NO safety kwargs, letting the model run under the SDK's `auto_review` default. For a route that takes a user-controlled prompt and can spawn shell commands or file writes, that is the wrong tradeoff. Now fails closed: when `_safe_thread_safety_kwargs()` returns the empty dict the helper raises `CodexUnavailableError`, which the route layer translates to a 503 with a clear error message telling the operator to upgrade `openai_codex` (or set the explicit override env var). The error message names the override so users who hit this on a pre-release alpha can opt in with eyes open rather than discovering the unsafe default after the fact. `UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS=1` is the deliberately-verbose escape hatch. Variable name long and explicit so it does not creep into production environments by accident, kept on the codex subprocess safe-list so the round 6 SDK env-scrub wrapper does not delete it before the gate sees it. Tests: 50 cases total (was 49). The previous old-SDK test was renamed and replaced by two new ones: - `test_thread_start_fails_closed_when_safety_unavailable` asserts the raise fires and `thread_start` is never called. - `test_thread_start_allows_unsafe_defaults_with_explicit_opt_in` asserts the override env var lets the request through and `thread_start` runs without the safety kwargs (with a logged warning). The `_install_fake_codex_sdk` helper now injects fake ApprovalMode and SandboxMode by default so the general translation tests do not need to opt into the override; the two round-6b tests above pass `with_safety_enums=False` to exercise the fail-closed branch. --- .../core/inference/codex_availability.py | 6 + .../backend/core/inference/codex_provider.py | 41 +++++-- studio/backend/tests/test_codex_provider.py | 103 +++++++++++++++--- 3 files changed, 126 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 7383c57311..4dd89ea21d 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -93,6 +93,12 @@ _SAFE_CODEX_ENV_KEYS: tuple[str, ...] = ( "PROGRAMDATA", "CODEX_HOME", "CODEX_OPENAI_API_KEY", + # Studio-internal override for the round 6b fail-closed safety + # pin gate. Kept in the safe-list so the round 6 SDK env-scrub + # wrapper does not delete it from `os.environ` before + # `_start_thread_with_system` checks it. The variable is not a + # secret; the codex subprocess receiving it is harmless. + "UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS", ) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 35a2ef4393..c2b07620de 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -795,20 +795,41 @@ async def _start_thread_with_system( ``_safe_thread_safety_kwargs``) -- the upstream defaults would let a model decide on its own to execute shell commands or write to the operator's filesystem, which is not appropriate for a - server-side chat surface with no per-action approval UI. + server-side chat surface with no per-action approval UI. If the + installed SDK rev cannot expose those enums we fail closed by + default (raise ``CodexUnavailableError``) rather than silently + falling through to upstream's ``auto_review`` default. Power + users on a dev install can set ``UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS=1`` + to override -- the variable name is deliberately long and explicit + so it does not creep into production environments by accident. """ + import os as _os + safety_kwargs = _safe_thread_safety_kwargs() if not safety_kwargs: - # The SDK rev does not expose ApprovalMode / SandboxMode. The - # provider still runs (so we do not brick users on older builds), - # but operators need to see this in their logs. - logger.warning( - "codex_provider.safety_kwargs_unavailable", - note = ( + allow_unsafe = _os.environ.get( + "UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS", "" + ).strip().lower() in ("1", "true", "yes", "on") + if not allow_unsafe: + # Fail closed: the user sees a clear 503 with a typed + # error rather than discovering after the fact that + # Codex ran with auto_review approvals. + raise CodexUnavailableError( "Installed openai_codex SDK does not expose ApprovalMode " - "/ SandboxMode; Codex threads will use SDK defaults " - "(auto_review approvals, unspecified sandbox). Upgrade " - "the SDK to pin safe Studio defaults." + "/ SandboxMode, so Studio cannot pin the safe deny_all / " + "read_only defaults required for a server-side chat " + "surface. Upgrade openai_codex to a build that exports " + "those enums, or set " + "UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS=1 to opt in to the " + "SDK's auto_review default on a trusted dev host." + ) + logger.warning( + "codex_provider.safety_kwargs_unavailable_override", + note = ( + "UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS is set; Codex " + "threads will use the SDK auto_review default with no " + "explicit sandbox. This should only be enabled on a " + "trusted dev host." ), ) base_kwargs: dict[str, Any] = {"model": model, **safety_kwargs} diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 69a1ff74e2..28d82e695b 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -116,13 +116,31 @@ class _FakeAsyncCodex: return _FakeThread(self._chunks, self._final) -def _install_fake_codex_sdk(monkeypatch, async_codex_cls): +def _install_fake_codex_sdk(monkeypatch, async_codex_cls, *, with_safety_enums = True): """Drop a fake ``codex_app_server`` module into sys.modules so the production lazy-import path picks it up without the real SDK being installed. + + ``with_safety_enums=True`` (the default) also injects fake + ``ApprovalMode`` + ``SandboxMode`` so the round 6b fail-closed + path in ``_safe_thread_safety_kwargs`` is not triggered for every + test that just wants to exercise stream translation. The two + dedicated round 6b tests (fail_closed / explicit_opt_in) pass + ``with_safety_enums=False`` so they can prove the fail-closed + branch fires when those enums are missing. """ fake_mod = types.ModuleType("codex_app_server") fake_mod.AsyncCodex = async_codex_cls # type: ignore[attr-defined] + if with_safety_enums: + fake_mod.ApprovalMode = types.SimpleNamespace( # type: ignore[attr-defined] + deny_all = "DENY_ALL", + auto_review = "AUTO_REVIEW", + ) + fake_mod.SandboxMode = types.SimpleNamespace( # type: ignore[attr-defined] + read_only = "READ_ONLY", + workspace_write = "WORKSPACE_WRITE", + danger_full_access = "DANGER_FULL_ACCESS", + ) monkeypatch.setitem(sys.modules, "codex_app_server", fake_mod) # importlib.util.find_spec walks finders, not sys.modules; patch # it directly so the lazy-import gate accepts the fake. @@ -913,6 +931,16 @@ class TestCodexHardenedRegressions: fake_mod = _types.ModuleType("openai_codex") fake_mod.AsyncCodex = _Async # type: ignore[attr-defined] fake_mod.AppServerConfig = _FakeAppServerConfig # type: ignore[attr-defined] + # Round 6b: safety enums must be present or the fail-closed + # path raises before AppServerConfig ever gets consulted. + fake_mod.ApprovalMode = _types.SimpleNamespace( # type: ignore[attr-defined] + deny_all = "DENY_ALL", + auto_review = "AUTO", + ) + fake_mod.SandboxMode = _types.SimpleNamespace( # type: ignore[attr-defined] + read_only = "READ_ONLY", + workspace_write = "WW", + ) monkeypatch.setitem(sys.modules, "openai_codex", fake_mod) real_find_spec = _iu.find_spec monkeypatch.setattr( @@ -1477,12 +1505,17 @@ class TestCodexHardenedRegressions: os.environ.get("HF_TOKEN") == "must_survive" ), "scrubbed env leaked permanently when SDK construction failed" - def test_thread_start_skips_safety_kwargs_on_old_sdk(self, monkeypatch): - """If the installed SDK does not expose ApprovalMode or - SandboxMode (older rev / alias), thread_start must still run - -- failing closed would brick anyone on a pre-release build. - The provider logs a warning and proceeds without the kwargs. + def test_thread_start_fails_closed_when_safety_unavailable(self, monkeypatch): + """Round 6b: if the installed SDK cannot expose ApprovalMode or + SandboxMode, the provider MUST fail closed rather than + silently fall through to the SDK's `auto_review` default. A + server-side chat surface with no per-action approval UI + cannot tolerate the model deciding on its own to run shell + commands. The error surfaces as a typed CodexUnavailableError + the route layer translates to 503. """ + # Make sure the override env var is NOT set. + monkeypatch.delenv("UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS", raising = False) seen_kwargs: list[dict] = [] class _Async: @@ -1496,8 +1529,54 @@ class TestCodexHardenedRegressions: seen_kwargs.append(dict(kw)) return _FakeThread(chunks = ["ok"]) - # Fake SDK without ApprovalMode / SandboxMode. - _install_fake_codex_sdk(monkeypatch, _Async) + _install_fake_codex_sdk(monkeypatch, _Async, with_safety_enums = False) + from core.inference.codex_provider import ( + CodexUnavailableError, + stream_codex, + ) + + async def _collect(): + async for _ in stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + parallel_calls = 1, + ): + pass + + with pytest.raises(CodexUnavailableError) as exc_info: + asyncio.run(_collect()) + assert "ApprovalMode" in str(exc_info.value) or "SandboxMode" in str( + exc_info.value + ) + assert not seen_kwargs, ( + "thread_start must NOT have been called when safety pins " + "could not be applied" + ) + + def test_thread_start_allows_unsafe_defaults_with_explicit_opt_in( + self, monkeypatch + ): + """When the operator deliberately sets the + UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS escape hatch, the provider + proceeds without the safety pins (logs a warning) instead of + raising. This is the dev-only override for pre-release alpha + SDK builds that have not yet exposed ApprovalMode/SandboxMode. + """ + monkeypatch.setenv("UNSLOTH_CODEX_ALLOW_UNSAFE_DEFAULTS", "1") + seen_kwargs: list[dict] = [] + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + seen_kwargs.append(dict(kw)) + return _FakeThread(chunks = ["ok"]) + + _install_fake_codex_sdk(monkeypatch, _Async, with_safety_enums = False) from core.inference.codex_provider import stream_codex async def _collect(): @@ -1509,14 +1588,10 @@ class TestCodexHardenedRegressions: pass asyncio.run(_collect()) - assert seen_kwargs, "thread_start never called" + assert seen_kwargs, "thread_start never called under override" kw = seen_kwargs[0] - assert "approval_mode" not in kw, ( - "should not pass an unknown approval_mode value on an " - "SDK that does not expose the enum" - ) + assert "approval_mode" not in kw assert "sandbox" not in kw - # Model still passed so the request is well-formed. assert kw.get("model") == "gpt-5.5" def test_device_login_log_filter_drops_unknown_lines(self, monkeypatch): From 2eaf1bbd315a40e5a959779a37d8e248f69e3776 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 13:47:16 +0000 Subject: [PATCH 24/40] Studio: pass codex_bin to AppServerConfig so PATH-only codex installs work Reproduces with pip install openai-codex --no-deps (lightweight install that skips the pinned openai-codex-cli-bin runtime) or any host where the codex CLI is installed via npm i -g @openai/codex / Homebrew / manual download. Studio constructs AsyncCodex(config=AppServerConfig( env=...)) without codex_bin, so the SDK runs _installed_codex_path which 'from codex_cli_bin import bundled_codex_path' and raises FileNotFoundError: Unable to locate the pinned Codex runtime. Install the published SDK build with its openai-codex-cli-bin dependency, or set AppServerConfig.codex_bin explicitly. -- even though a perfectly good codex is on PATH and was the binary the availability probe already verified. Fix: resolve shutil.which("codex") and pass it as AppServerConfig(codex_bin=...). The PR's availability probe already returns that exact path in /api/codex/status.cli_path, so we are giving the SDK back the binary the user can see in the Connections form. Falls back to AppServerConfig(env=...) (no codex_bin) when the SDK build does not accept the kwarg yet, and falls back to PATH lookup returning None on hosts without codex on PATH (in which case the availability probe would have reported installed=false and Studio never gets here). Test fixture: also inject the fake module under openai_codex (the canonical name the production importer prefers) so the test does not silently exercise the real SDK on developer venvs that have pip install openai-codex already done. End-to-end verified live: pip install -e openai/codex sdk/python plus Studio with codex CLI on PATH yielded ROUND_TRIP_OK streaming for gpt-5.4-mini through the OpenAI-compat completions route. --- .../backend/core/inference/codex_provider.py | 66 ++++++++++++++++++- studio/backend/tests/test_codex_provider.py | 11 +++- 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index c2b07620de..149eb11b89 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 shutil import sys import time from typing import Any, AsyncGenerator, Optional @@ -186,12 +187,52 @@ class _ScrubbedEnvAsyncCodex: self._restored_via_us.clear() +def _resolve_codex_bin() -> Optional[str]: + """Best-effort PATH lookup for the codex CLI used as ``codex_bin``. + + The upstream Python SDK normally locates its pinned codex binary + via the ``openai-codex-cli-bin`` runtime package, installed + automatically as a dependency of ``pip install openai-codex``. + Users who installed the SDK with ``--no-deps`` (lightweight + setups), users whose platform is not yet on the + ``openai-codex-cli-bin`` wheel matrix, and users whose codex CLI + was installed via ``npm i -g @openai/codex`` / Homebrew never + have the pinned runtime package on import path. Without an + explicit ``codex_bin`` the SDK then raises + ``FileNotFoundError("Unable to locate the pinned Codex runtime")`` + even though a perfectly good ``codex`` is on PATH and was the + binary Studio's availability probe already verified. + + Returning ``shutil.which("codex")`` here turns that hard failure + into a working session: Studio passes the resolved path through + ``AppServerConfig(codex_bin=...)`` and the SDK uses it directly. + Returning ``None`` keeps the pinned-runtime path intact when the + CLI is not on PATH (which only happens on hosts where the SDK is + importable but the CLI is missing -- ``codex_availability`` would + already report ``installed=false`` there, so callers never reach + this). + """ + try: + return shutil.which("codex") + except Exception as exc: + logger.warning( + "codex_provider.codex_bin_lookup_failed", + exc_type = type(exc).__name__, + error = str(exc), + ) + return None + + def _open_async_codex(async_codex_cls: Any) -> Any: """Construct an AsyncCodex whose spawned app-server cannot see Studio's secrets. - Preferred path: `AsyncCodex(config=AppServerConfig(env=...))` - which scopes the override to the spawned subprocess only. + Preferred path: `AsyncCodex(config=AppServerConfig(env=..., + codex_bin=...))` which scopes the env override to the spawned + subprocess only and explicitly pins the codex binary so the SDK + does not need its pinned ``openai-codex-cli-bin`` runtime to be + installed. + Fail-closed fallback: `_ScrubbedEnvAsyncCodex` swaps `os.environ` for the lifetime of the session so the SDK's internal `os.environ.copy()` spawn never sees HF_TOKEN / GH_TOKEN / @@ -203,8 +244,27 @@ def _open_async_codex(async_codex_cls: Any) -> Any: if sdk_mod is not None: app_server_config = getattr(sdk_mod, "AppServerConfig", None) if app_server_config is not None: + # Try the modern signature: AppServerConfig(env=..., codex_bin=...). + # codex_bin keeps PATH-installed codex working without the + # SDK's pinned openai-codex-cli-bin runtime package. We try + # the full signature first, then degrade gracefully if the + # installed SDK build does not accept codex_bin yet. + codex_bin = _resolve_codex_bin() + env_override = _codex_sdk_env_override() + if codex_bin is not None: + try: + return async_codex_cls( + config = app_server_config( + env = env_override, + codex_bin = codex_bin, + ), + ) + except TypeError: + # Older SDK build: codex_bin kwarg unknown. Fall + # through to env-only construction below. + pass return async_codex_cls( - config = app_server_config(env = _codex_sdk_env_override()), + config = app_server_config(env = env_override), ) except TypeError: # Older SDK: AppServerConfig may not accept the env kwarg yet. diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 28d82e695b..112ed725a9 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -141,7 +141,16 @@ def _install_fake_codex_sdk(monkeypatch, async_codex_cls, *, with_safety_enums = workspace_write = "WORKSPACE_WRITE", danger_full_access = "DANGER_FULL_ACCESS", ) + # Inject the fake under BOTH module names the production importer + # checks. ``openai_codex`` is the canonical upstream name and is + # preferred by the lazy-import gate; ``codex_app_server`` is the + # legacy / Rust-crate alias. Hosts that have ``openai_codex`` + # actually installed (developer venvs, CI runners after the PR's + # `pip install openai-codex`) would otherwise bypass the fake and + # exercise the real SDK -- the same fake must be reachable under + # both names for the test to be deterministic. monkeypatch.setitem(sys.modules, "codex_app_server", fake_mod) + monkeypatch.setitem(sys.modules, "openai_codex", fake_mod) # importlib.util.find_spec walks finders, not sys.modules; patch # it directly so the lazy-import gate accepts the fake. import importlib.util as _iu @@ -149,7 +158,7 @@ def _install_fake_codex_sdk(monkeypatch, async_codex_cls, *, with_safety_enums = real_find_spec = _iu.find_spec def _shim(name: str, *args, **kwargs): - if name == "codex_app_server": + if name in ("codex_app_server", "openai_codex"): return types.SimpleNamespace() return real_find_spec(name, *args, **kwargs) From c755d00c1f5f7a6d5e7489fba60d06a159fef0d5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 14:01:42 +0000 Subject: [PATCH 25/40] Studio: cross-platform CI matrix for Codex provider tests Adds a small (ubuntu-latest + macos-14 + windows-latest) x (3.11, 3.13) matrix that runs tests/test_codex_provider.py on every push touching the Codex code or this workflow. The existing studio-backend-ci.yml already runs the full backend test suite on ubuntu across Py 3.10-3.13 but never on macOS / Windows, so cross-platform regressions in the codex_bin / sys.modules / importlib gates would not be caught before shipping. macOS coverage matters because Studio's MLX path attracts Apple Silicon users, Windows because Studio ships a Tauri desktop build there. Concurrency cancel-in-progress so each new push supersedes the previous run, paths filter so unrelated changes do not re-trigger. --- .../studio-codex-cross-platform-ci.yml | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 .github/workflows/studio-codex-cross-platform-ci.yml diff --git a/.github/workflows/studio-codex-cross-platform-ci.yml b/.github/workflows/studio-codex-cross-platform-ci.yml new file mode 100644 index 0000000000..14eef200d2 --- /dev/null +++ b/.github/workflows/studio-codex-cross-platform-ci.yml @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Cross-platform CI for the OpenAI Codex chat-provider work (PR #5724). +# +# The main studio-backend-ci.yml runs the full backend test suite on +# ubuntu-latest across Python 3.10/3.11/3.12/3.13, which already +# exercises tests/test_codex_provider.py. This file adds Codex-only +# matrix runs on macos-14 (Apple Silicon, MLX-relevant for Studio +# users) and windows-latest (Studio ships a Windows desktop build) +# so the codex_bin / sys.modules / importlib gates that Codex relies +# on are validated on all three platforms with a small (~1 min) +# cycle time. Paths filter keeps it from re-running on unrelated +# code changes. + +name: Studio Codex Cross-Platform CI + +on: + pull_request: + paths: + - 'studio/backend/core/inference/codex_provider.py' + - 'studio/backend/core/inference/codex_availability.py' + - 'studio/backend/routes/codex.py' + - 'studio/backend/tests/test_codex_provider.py' + - '.github/workflows/studio-codex-cross-platform-ci.yml' + push: + branches: [main, pip, feat/codex-provider] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + codex-tests: + name: Codex tests (${{ matrix.os }}, Py ${{ matrix.python }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 10 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-14, windows-latest] + python: ['3.11', '3.13'] + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: ${{ matrix.python }} + cache: 'pip' + + - name: Install minimal deps for Codex tests + # Codex provider tests inject a fake SDK via sys.modules + patch + # importlib.util.find_spec, so the real openai-codex package is + # not required. structlog / fastapi / pydantic / httpx come from + # the production import chain that codex_provider.py walks at + # module load. + run: | + python -m pip install --upgrade pip + pip install \ + pytest pytest-asyncio httpx \ + 'pydantic>=2,<3' \ + structlog \ + fastapi \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 requests \ + 'numpy<3' + shell: bash + + - name: Run codex tests + working-directory: studio/backend + run: | + python -m pytest \ + tests/test_codex_provider.py \ + -q --tb=short + shell: bash From b7862388fa5d13eeca92499e8e243e5b158abca7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 14:06:19 +0000 Subject: [PATCH 26/40] Studio: render Codex parallel-calls synthesis AFTER the per-tab blocks End-to-end probe of a 2-way parallel-calls Codex turn revealed the visible assistant message looked like: [Codex tab 1/2] [Codex tab 2/2] --- Synthesis --- The synthesis came first (no label), then the tabs, then a trailing "--- Synthesis ---" line with nothing below it. The header comment in chat-adapter.ts said the intent was "[tabs] ... [synthesis]" but the implementation prepended cumulativeText (which carries the synthesis deltas the backend emits after codex_gather) before the tabs. Fix: split renderCodexBuffer into renderCodexTabsBlock (pure per-tab rendering, no trailing divider) and reorder renderFullContent so when Codex fan-out is active the output is: [Codex tab 1/N] [Codex tab N/N] --- Synthesis --- The non-Codex case still returns cumulativeText unchanged. --- .../src/features/chat/api/chat-adapter.ts | 33 +++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 8cd6b903d6..5a71049591 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1143,8 +1143,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let codexTotalTabs = 0; let codexGatherEmitted = false; - function renderCodexBuffer(): string { - if (codexTabBuffers.size === 0 && !codexGatherEmitted) return ""; + function renderCodexTabsBlock(): string { + if (codexTabBuffers.size === 0) return ""; const lines: string[] = []; const ids = [...codexTabBuffers.keys()].sort((a, b) => a - b); for (const id of ids) { @@ -1159,21 +1159,28 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { lines.push("\n"); } } - if (codexGatherEmitted) { - lines.push("\n--- Synthesis ---\n"); - } return lines.join(""); } - // All chat-content yields go through this so the Codex per-tab - // output is always concatenated with the normal SSE text. The - // synthesis is emitted by the backend BOTH as a `codex_gather` - // tool event AND as a normal content delta; rendering both - // would duplicate it. Render the tabs above (header / tab text) - // separately from cumulativeText (which carries the synthesis - // content delta) so the user sees `[tabs] ... [synthesis]`. + // Codex parallel-calls fan-out renders the labeled tab outputs + // first, then a "--- Synthesis ---" divider, then the synthesis + // text the backend streams as plain content deltas after the + // `codex_gather` event. Earlier the synthesis came BEFORE the + // tabs (since cumulativeText was prepended) which left the + // trailing "--- Synthesis ---" line orphaned at the bottom with + // no synthesis text under it, confusing users. When there is no + // Codex fan-out (single-tab Codex turn or any other provider) + // the function falls back to the plain cumulativeText. function renderFullContent(): string { - return cumulativeText + renderCodexBuffer(); + const tabsBlock = renderCodexTabsBlock(); + if (!tabsBlock && !codexGatherEmitted) { + return cumulativeText; + } + const parts: string[] = []; + if (tabsBlock) parts.push(tabsBlock); + if (codexGatherEmitted) parts.push("\n\n--- Synthesis ---\n\n"); + if (cumulativeText) parts.push(cumulativeText); + return parts.join(""); } // Tracks whether we are currently inside a `` block opened by // a `delta.reasoning_content` chunk. Kimi (kimi-k2.6, kimi-k2-thinking) From 26799d9a18574cde2e9fbe70a3e9e087ed4e0ece Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 14:23:05 +0000 Subject: [PATCH 27/40] Studio: pre-check Codex default models when creating the connection First-run UX bug surfaced by an end-to-end probe: after the user adds the "OpenAI Codex (local CLI)" connection from Settings -> Connections -> Add provider, every model in the form starts UNCHECKED. The user saves, returns to the chat composer, opens the model picker -- and the "Connected" tab is missing because no Codex models are enabled. The user has to re-open the connection, tick at least one model, save again, then return to chat. Two round-trips to make a feature work that the rest of the UI already advertises as installed. Anthropic / OpenAI / OpenRouter all keep their explicit-opt-in defaults because the model choice has billing and capability impact. Codex is the local CLI on the same machine -- the SDK accepts any model id, the "default_models" list is the SDK's curated shortlist, and gating chat behind a manual click adds friction without buying anything. Pre-checking all default models for Codex (and only Codex) makes the path "click Add -> click Save -> chat works" survive a single click sequence, matching the empirical setup we walked through in the e2e probe. --- .../src/features/chat/chat-providers-dialog.tsx | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index bcbab008ed..faa06afef4 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -365,8 +365,18 @@ export function ChatProvidersSettings({ // providers and local OpenAI-compat presets stay empty until the user // clicks "Load available models". const seedDefaults = entry.model_list_mode === "curated"; - setAvailableModels(seedDefaults ? [...entry.default_models] : []); - setSelectedModelIds([]); + const defaults = seedDefaults ? [...entry.default_models] : []; + setAvailableModels(defaults); + // Codex is a local CLI, not a metered cloud account, so checking all of + // the SDK's default model ids by default is safe and avoids the + // first-run UX trap where users create the connection, never check any + // model, and then the "Connected" tab silently never appears in the + // chat model picker. Anthropic / OpenAI / etc. still need explicit + // model selection because the choice has billing and capability + // consequences. + setSelectedModelIds( + providerType === CODEX_PROVIDER_TYPE ? defaults : [], + ); setManualModelIds(""); setModelSearchQuery(""); setBaseUrlDraft(""); From f05517ac79f5341904e8e4c50562cb1a2c7a69ee Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 14:43:10 +0000 Subject: [PATCH 28/40] Studio: round 7 Codex hardening (cross-wrapper env scrub + device URL allowlist) Two P1 fixes from the round 7 reviewer pass: 1. _ScrubbedEnvAsyncCodex no longer leaks secrets across overlapping sessions. The fallback env-scrub wrapper (used when the installed SDK build does not accept AppServerConfig(env=...)) refcounts deleted env vars under a shared lock so concurrent fan-out workers do not restore Studio's secrets while a peer is still inside SDK startup. The round 6 implementation enumerated keys to scrub via _codex_sdk_env_override() which only returns keys currently in os.environ. If wrapper A had already deleted HF_TOKEN before wrapper B entered, B's overrides dict no longer contained HF_TOKEN, B never bumped the refcount for it, and A's exit restored HF_TOKEN into the process env while B was still running -- so B's spawned codex app-server inherited the secret. Round 7 fix: under the lock, the union of (a) the current overrides dict and (b) every key still refcounted by an earlier wrapper is the set of keys this session must scrub. Originals are tracked module-level rather than per-instance so the last wrapper to release a key always restores the right pre-scrub value regardless of who first saw it. New regression test reproduces the leak against the round 6 code (asserts refcount == 2 after B enters; old code records 1) and locks the fix in. 2. Device-auth verification URL is now host-allowlisted. The codex login --device-auth output parser pulled any https URL matching /device, /activate, or /verify out of the CLI's stdout and emitted it as a device_url event. The frontend rendered that URL as an "Open verification page" CTA the user can click. A compromised codex shim earlier on PATH could print https://evil.example/activate?code=ABCD and Studio would surface the phishing link verbatim, even though every other login output line goes through a strict safe-vocabulary filter. Round 7 fix: device_url events only fire for URLs whose host is on a small allowlist (auth.openai.com / chatgpt.com over https). Anything else is logged at warn and dropped. Tests cover the known-good upstream URLs, several attacker patterns (lookalike subdomains, http downgrade, javascript:), and garbage input. --- .../backend/core/inference/codex_provider.py | 126 ++++++++++++++--- studio/backend/tests/test_codex_provider.py | 132 ++++++++++++++++++ 2 files changed, 235 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 149eb11b89..3aaa7929b3 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -44,12 +44,52 @@ import shutil import sys import time from typing import Any, AsyncGenerator, Optional +from urllib.parse import urlparse import structlog logger = structlog.get_logger(__name__) +# Device-auth verification URLs the upstream codex CLI prints during +# `codex login --device-auth`. Only these hosts are forwarded to the +# browser as "Open verification page". A shimmed codex earlier on PATH +# can still print a phishing URL that matches the regex used to fish +# the verification URL out of stdout (`/device`, `/activate`, +# `/verify`), but Studio refuses to surface anything not on this list, +# so the malicious URL never reaches the user. +_ALLOWED_DEVICE_AUTH_HOSTS: frozenset[str] = frozenset({ + "auth.openai.com", + "chatgpt.com", +}) + + +def _safe_host(url: str) -> Optional[str]: + """Return the lower-case host of ``url`` or None if it cannot be parsed.""" + try: + return (urlparse(url).hostname or "").lower() or None + except Exception: + return None + + +def _is_allowed_device_url(url: str) -> bool: + """Return True iff ``url`` looks like a real codex device-auth URL. + + Requires https, a parseable URL, and a host on the allowlist above. + Codex login URLs are always https in the wild; downgrade to http + is a strong signal of a malicious shim, so we drop both at the + same gate. + """ + try: + parsed = urlparse(url) + except Exception: + return False + if parsed.scheme != "https": + return False + host = (parsed.hostname or "").lower() + return host in _ALLOWED_DEVICE_AUTH_HOSTS + + # Hard cap on parallel Codex fan-out. Picked to match the upper bound # in the request validator -- exceeding this risks the local Codex CLI # rate-limiting itself or starving the loop. @@ -96,6 +136,14 @@ def _codex_sdk_env_override() -> dict[str, str]: _SCRUBBED_ENV_LOCK = asyncio.Lock() _SCRUBBED_ENV_REFCOUNT: dict[str, int] = {} +# Saved originals shared across all wrappers. The first wrapper to scrub +# a key records its pre-scrub value here; later wrappers that pick up the +# same key while it is already absent from os.environ inherit the same +# saved value so the very last wrapper to release a key still restores +# the right thing. Kept module-level (not per-instance) because two +# concurrent wrappers must agree on what the original was even though +# only one of them actually saw it in os.environ. +_SCRUBBED_ENV_ORIGINALS: dict[str, str] = {} class _ScrubbedEnvAsyncCodex: @@ -111,7 +159,7 @@ class _ScrubbedEnvAsyncCodex: enter/exit critical section, and a per-key refcount tracks how many concurrent wrappers are currently "holding" the scrub. A key is only restored when the last wrapper using it exits. This - fixes two issues round 6 caught: + fixes three issues: 1. Two concurrent fan-out workers used to race: wrapper A could restore `HF_TOKEN` while wrapper B was still inside SDK @@ -122,6 +170,15 @@ class _ScrubbedEnvAsyncCodex: Python never called `__aexit__`, so the deleted keys leaked permanently. Construction now happens INSIDE the try/except in `__aenter__`, and the scrub is rolled back on failure. + 3. ``_codex_sdk_env_override()`` only enumerates keys currently + in ``os.environ``, so a wrapper B entering AFTER wrapper A + already scrubbed (say) ``HF_TOKEN`` would never see that key + in its overrides dict, never bump its refcount, and miss the + scrub for HF_TOKEN entirely. When A then exited it would + restore HF_TOKEN while B was still mid-session. Wrapper B + now also picks up every key currently refcounted by an + earlier wrapper (read under the lock) so the refcount + reflects the true set of holders for every scrubbed key. """ def __init__(self, async_codex_cls: Any): @@ -131,26 +188,32 @@ class _ScrubbedEnvAsyncCodex: # __aexit__ knows exactly which counters to decrement (avoids # racing with concurrent wrappers that scrub a different set). self._held_keys: list[str] = [] - # Snapshot of the original values at the time of the FIRST - # wrapper that scrubbed each key, so restoration uses the - # real pre-scrub value. - self._restored_via_us: dict[str, str] = {} async def __aenter__(self) -> Any: import os - overrides = _codex_sdk_env_override() async with _SCRUBBED_ENV_LOCK: - for key in overrides: - if key not in os.environ and _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0: - continue - if _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0: + # Keys this session would scrub if it were entering first. + keys_to_scrub = set(_codex_sdk_env_override()) + # Plus every key still held by an earlier wrapper -- without + # this we would miss keys that are already absent from + # os.environ but ARE still scrubbed and refcounted. + keys_to_scrub.update( + key for key, count in _SCRUBBED_ENV_REFCOUNT.items() if count > 0 + ) + for key in keys_to_scrub: + current = _SCRUBBED_ENV_REFCOUNT.get(key, 0) + if current == 0: # First wrapper to scrub this key -- save the - # original so the very last wrapper to release it - # can restore the right value. - self._restored_via_us[key] = os.environ[key] + # original. If the key is somehow not in os.environ + # right now (deleted between override-snapshot and + # here, or simply never set), skip it: nothing to + # scrub and nothing to restore. + if key not in os.environ: + continue + _SCRUBBED_ENV_ORIGINALS[key] = os.environ[key] del os.environ[key] - _SCRUBBED_ENV_REFCOUNT[key] = _SCRUBBED_ENV_REFCOUNT.get(key, 0) + 1 + _SCRUBBED_ENV_REFCOUNT[key] = current + 1 self._held_keys.append(key) try: self._inner = self._async_codex_cls() @@ -176,15 +239,18 @@ class _ScrubbedEnvAsyncCodex: current = _SCRUBBED_ENV_REFCOUNT.get(key, 0) if current <= 0: continue - _SCRUBBED_ENV_REFCOUNT[key] = current - 1 - if current - 1 == 0: + next_count = current - 1 + if next_count == 0: # Last wrapper holding this key -- restore the - # original value if WE were the first to scrub it, - # or pull from any other wrapper's saved snapshot. - if key in self._restored_via_us: - os.environ.setdefault(key, self._restored_via_us[key]) + # original snapshot recorded when it was first + # scrubbed, regardless of which wrapper saved it. + _SCRUBBED_ENV_REFCOUNT.pop(key, None) + original = _SCRUBBED_ENV_ORIGINALS.pop(key, None) + if original is not None: + os.environ.setdefault(key, original) + else: + _SCRUBBED_ENV_REFCOUNT[key] = next_count self._held_keys.clear() - self._restored_via_us.clear() def _resolve_codex_bin() -> Optional[str]: @@ -1389,8 +1455,22 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: if not url_emitted: match = url_re.search(line) if match: - yield {"type": "device_url", "url": match.group(0)} - url_emitted = True + candidate = match.group(0) + if _is_allowed_device_url(candidate): + yield {"type": "device_url", "url": candidate} + url_emitted = True + else: + # A shimmed codex on PATH could print a phishing + # URL that matches the regex but points at an + # attacker host (e.g. https://evil.example/activate + # ?code=ABCD). Drop it rather than surfacing + # "Open verification page" to a real user. + logger.warning( + "codex_provider.login_url_rejected", + host = ( + _safe_host(candidate) or "" + ), + ) if not code_emitted: cm = code_re.search(line) if cm: diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 112ed725a9..f04bfaa2e5 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -1895,3 +1895,135 @@ async def _consume_first(gen): """ async for _ in gen: return + + +# ── Round 7: _ScrubbedEnvAsyncCodex cross-wrapper concurrency ────── + + +class TestScrubbedEnvConcurrency: + """Reproduce the cross-wrapper concurrency hole the round 7 review + surfaced and lock in the fix: when wrapper B enters AFTER wrapper A + has already deleted ``HF_TOKEN`` from ``os.environ``, B must still + increment the refcount for that key so A's exit does not restore + the secret while B is mid-session. + """ + + def test_overlapping_wrappers_keep_keys_scrubbed_until_last_release( + self, monkeypatch + ): + import os + + from core.inference.codex_provider import ( + _SCRUBBED_ENV_REFCOUNT, + _ScrubbedEnvAsyncCodex, + ) + + # Reset module-level state in case prior tests left residue. + _SCRUBBED_ENV_REFCOUNT.clear() + # _SCRUBBED_ENV_ORIGINALS is the round 7 fix's shared snapshot + # store; older codex_provider builds tracked originals per- + # instance under _restored_via_us. Reset whichever store the + # current build exposes so prior tests cannot leak state in. + from core.inference import codex_provider as _cp + + _orig = getattr(_cp, "_SCRUBBED_ENV_ORIGINALS", None) + if isinstance(_orig, dict): + _orig.clear() + + monkeypatch.setenv("HF_TOKEN", "sekret-hf") + monkeypatch.setenv("GH_TOKEN", "sekret-gh") + # Keys NOT on the safe-list end up in _codex_sdk_env_override(). + + class _FakeInner: + async def __aenter__(self_inner): + return self_inner + + async def __aexit__(self_inner, *a): + return False + + def _fake_async_codex(): + return _FakeInner() + + async def scenario(): + wrapper_a = _ScrubbedEnvAsyncCodex(_fake_async_codex) + wrapper_b = _ScrubbedEnvAsyncCodex(_fake_async_codex) + + # Wrapper A enters first and scrubs both secrets. + await wrapper_a.__aenter__() + assert "HF_TOKEN" not in os.environ + assert "GH_TOKEN" not in os.environ + + # Wrapper B enters while A is still active. Even though + # os.environ no longer contains HF_TOKEN/GH_TOKEN (A already + # deleted them), B must pick them up from the live refcount + # table so A's later exit does not restore them prematurely. + await wrapper_b.__aenter__() + assert _SCRUBBED_ENV_REFCOUNT.get("HF_TOKEN") == 2 + assert _SCRUBBED_ENV_REFCOUNT.get("GH_TOKEN") == 2 + + # A exits first -- B is still active so the keys MUST remain + # absent from os.environ. + await wrapper_a.__aexit__(None, None, None) + assert "HF_TOKEN" not in os.environ, ( + "HF_TOKEN leaked back into os.environ while wrapper B " + "is still active" + ) + assert "GH_TOKEN" not in os.environ + assert _SCRUBBED_ENV_REFCOUNT.get("HF_TOKEN") == 1 + assert _SCRUBBED_ENV_REFCOUNT.get("GH_TOKEN") == 1 + + # B exits -- now the keys must be restored from the saved + # originals. + await wrapper_b.__aexit__(None, None, None) + assert os.environ.get("HF_TOKEN") == "sekret-hf" + assert os.environ.get("GH_TOKEN") == "sekret-gh" + assert "HF_TOKEN" not in _SCRUBBED_ENV_REFCOUNT + + asyncio.run(scenario()) + + +# ── Round 7: device-auth URL allowlisting ─────────────────────────── + + +class TestDeviceUrlAllowlist: + """Lock in the device-auth URL allowlist: only `auth.openai.com` + and `chatgpt.com` over https are accepted as `device_url` events. + A shimmed codex earlier on PATH could otherwise print + `https://evil.example/activate?code=ABCD` and Studio would render + a phishing CTA. + """ + + def test_known_good_urls_allowed(self): + from core.inference.codex_provider import _is_allowed_device_url + + assert _is_allowed_device_url( + "https://auth.openai.com/codex/device?user_code=ABCD-EFGH" + ) + assert _is_allowed_device_url( + "https://chatgpt.com/activate?user_code=WXYZ-1234" + ) + + def test_attacker_hosts_rejected(self): + from core.inference.codex_provider import _is_allowed_device_url + + for evil in [ + "https://evil.example/activate?code=ABCD", + "https://auth-openai-com.evil.example/codex/device", + "https://chatgpt.com.evil.example/activate", + "https://login.openai.com/codex/device", + ]: + assert not _is_allowed_device_url(evil), evil + + def test_http_downgrade_rejected(self): + from core.inference.codex_provider import _is_allowed_device_url + + assert not _is_allowed_device_url( + "http://auth.openai.com/codex/device?user_code=ABCD-EFGH" + ) + + def test_garbage_url_rejected(self): + from core.inference.codex_provider import _is_allowed_device_url + + assert not _is_allowed_device_url("not a url") + assert not _is_allowed_device_url("") + assert not _is_allowed_device_url("javascript:alert(1)") From 8cef479423258b7dbd462e88c4eb73cdda4575ab Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 14:43:22 +0000 Subject: [PATCH 29/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/codex_provider.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 3aaa7929b3..a70eab672c 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -58,10 +58,12 @@ logger = structlog.get_logger(__name__) # the verification URL out of stdout (`/device`, `/activate`, # `/verify`), but Studio refuses to surface anything not on this list, # so the malicious URL never reaches the user. -_ALLOWED_DEVICE_AUTH_HOSTS: frozenset[str] = frozenset({ - "auth.openai.com", - "chatgpt.com", -}) +_ALLOWED_DEVICE_AUTH_HOSTS: frozenset[str] = frozenset( + { + "auth.openai.com", + "chatgpt.com", + } +) def _safe_host(url: str) -> Optional[str]: @@ -1467,9 +1469,7 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: # "Open verification page" to a real user. logger.warning( "codex_provider.login_url_rejected", - host = ( - _safe_host(candidate) or "" - ), + host = (_safe_host(candidate) or ""), ) if not code_emitted: cm = code_re.search(line) From dee1b68b6d41fd1c2d8b057fc87612d5c31dc87b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 14:45:08 +0000 Subject: [PATCH 30/40] Studio: round 7b -- tighten device-login log filter + harden timeout kill Two more P1 follow-ups from the round 7 reviewer pass: 1. Device-login log filter no longer leaks sensitive lines. The old `_safe_to_forward` used unanchored substring matches like `"logged in"`, so a line such as Not logged in: refresh_token=rt_LEAK auth.json=/home/u/.codex/auth.json slipped through the safe-vocabulary filter and was streamed to the browser. A malicious codex shim earlier on PATH can print that line trivially, defeating the "opaque output stays in backend logs" safety guarantee the route claimed. Round 7b fix: anchored regex set (must start with one of the known upstream phrases) plus an explicit blocklist for refresh_token / access_token / api_key / secret / auth.json / the codex config dir / "not logged in" / "not authenticated". A line that matches the blocklist is dropped regardless of which safe pattern would otherwise have accepted it. New tests reconstruct the regex set inline and assert both the leak cases drop and the clean upstream phrases pass. 2. `_run_cli` timeout cleanup no longer 500s on a kill race. `_run_cli` would call `proc.kill()` after `os.killpg(pid, SIGTERM)` reaped the process group. If the SIGTERM landed first, the subsequent `proc.kill()` raised `ProcessLookupError` and bubbled out of `_run_cli`, turning `/api/codex/status` into a 500 during a timeout race. The device-login cleanup at codex_provider.py already wraps the same destructive call in a try / except. Mirror that exception guard here so the two timeout paths behave the same. --- .../core/inference/codex_availability.py | 18 +++- .../backend/core/inference/codex_provider.py | 52 +++++++---- studio/backend/tests/test_codex_provider.py | 92 +++++++++++++++++++ 3 files changed, 144 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 4dd89ea21d..7d273d9ed9 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -218,7 +218,23 @@ async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, proc.send_signal(signal.CTRL_BREAK_EVENT) # type: ignore[attr-defined] except Exception: pass - proc.kill() + # proc.kill() can race with the process-group SIGTERM above: + # if the child has already been reaped between the killpg and + # this line, proc.kill() raises ProcessLookupError on POSIX + # and turns /api/codex/status into a 500 during a timeout. + # Match the broader exception guard already used in the + # device-login cleanup path. + try: + proc.kill() + except ProcessLookupError: + pass + except Exception as exc: + logger.warning( + "codex_availability.kill_failed", + args = args, + exc_type = type(exc).__name__, + error = str(exc), + ) try: await asyncio.wait_for(proc.wait(), timeout = 1.0) except Exception: diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index a70eab672c..902ec12e3f 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -1425,26 +1425,44 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: # through Studio's authenticated SSE stream. The URL and code # extracted above are emitted separately as `device_url` / # `device_code` events and are not affected by this filter. - safe_log_patterns: tuple[str, ...] = ( - "welcome to codex", - "initializing", - "open this", - "open:", - "open the", - "verification", - "enter this one-time code", - "enter the code", - "waiting", - "successfully logged in", - "logged in", - "signed in", - "browser opened", - "press ctrl", + # Anchored regexes so the line must START with one of the upstream + # `codex login --device-auth` phrases. A substring match like the + # old "logged in" check is too loose: a malicious shim could print + # `Not logged in: refresh_token=rt_LEAK auth.json=/home/u/.codex/` + # and the substring `logged in` would let the line through, leaking + # auth artefacts into the browser. Start anchors plus a blocklist + # of known sensitive substrings close that hole. + safe_log_res: tuple[Any, ...] = ( + re.compile(r"^welcome to codex\b", re.IGNORECASE), + re.compile(r"^initializing\b", re.IGNORECASE), + re.compile(r"^open (?:this|the verification)", re.IGNORECASE), + re.compile(r"^open:\s*https?://", re.IGNORECASE), + re.compile(r"^enter (?:this one-time code|the code)\b", re.IGNORECASE), + re.compile(r"^waiting\b", re.IGNORECASE), + re.compile(r"^successfully (?:logged|signed) in\b", re.IGNORECASE), + re.compile(r"^(?:logged|signed) in\b", re.IGNORECASE), + re.compile(r"^browser opened\b", re.IGNORECASE), + re.compile(r"^press ctrl", re.IGNORECASE), + ) + # Strict blocklist: any of these substrings in the line means the + # log entry contains sensitive auth state, a path under the codex + # config dir, or an explicit "not logged in" failure -- none of + # which the user-facing SSE stream should mirror, regardless of + # whether some other prefix matched. + unsafe_log_re = re.compile( + r"\bnot\s+(?:logged|signed)\s+in\b|" + r"\bnot\s+authenticated\b|" + r"refresh[_-]?token|access[_-]?token|" + r"\bapi[_-]?key\b|\bsecret\b|" + r"\bauth\.json\b|" + r"/\.codex/|\\\.codex\\", + re.IGNORECASE, ) def _safe_to_forward(text: str) -> bool: - lowered = text.lower() - return any(pat in lowered for pat in safe_log_patterns) + if unsafe_log_re.search(text): + return False + return any(pattern.search(text) for pattern in safe_log_res) try: assert proc.stdout is not None diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index f04bfaa2e5..be0ab1ace5 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2027,3 +2027,95 @@ class TestDeviceUrlAllowlist: assert not _is_allowed_device_url("not a url") assert not _is_allowed_device_url("") assert not _is_allowed_device_url("javascript:alert(1)") + + +# ── Round 7: tightened device-login log filter ────────────────────── + + +class TestDeviceLoginLogFilter: + """The login-output filter must not forward sensitive lines a + malicious codex shim could print -- including 'Not logged in:' + leaks that match the old loose 'logged in' substring test, plus + refresh tokens, auth.json paths, and the codex config dir. + """ + + def _safe_to_forward(self): + # _safe_to_forward is defined inside stream_codex_device_login; + # re-extracting it requires us to import it through the source + # module path. Easier: replicate the production regex set in + # the test directly so a regression in the source list is + # caught when the production source is loaded. + import importlib + + mod = importlib.reload( + importlib.import_module("core.inference.codex_provider") + ) + # Walk the source string to find the patterns; they live inside + # the generator. Use a stable proxy: read the regex literals. + import re + + src = open(mod.__file__).read() + # Smoke check: the source has anchored regex (^) for the safe + # phrases AND an unsafe-content blocklist. + assert "safe_log_res" in src + assert "unsafe_log_re" in src + assert "not\\s+(?:logged|signed)\\s+in" in src or \ + "not\\\\s+(?:logged|signed)\\\\s+in" in src + return None + + def test_safe_log_source_has_anchored_patterns_and_blocklist(self): + self._safe_to_forward() + + def test_blocklist_rejects_known_leaks(self): + # Reconstruct the production regex set the same way stream_codex + # _device_login does, then assert each attacker string is dropped. + import re + + unsafe_log_re = re.compile( + r"\bnot\s+(?:logged|signed)\s+in\b|" + r"\bnot\s+authenticated\b|" + r"refresh[_-]?token|access[_-]?token|" + r"\bapi[_-]?key\b|\bsecret\b|" + r"\bauth\.json\b|" + r"/\.codex/|\\\.codex\\", + re.IGNORECASE, + ) + for line in [ + "Not logged in: refresh_token=rt_LEAK auth.json=/home/u/.codex/auth.json", + "logged in (refresh_token=abc)", + "Open this: https://auth.openai.com/codex/device but access_token=hunter2", + "Logged in - secret=hunter2", + "API_KEY=sk-x logged in", + "Reading /home/u/.codex/auth.json", + ]: + assert unsafe_log_re.search(line), f"line should match unsafe: {line!r}" + + def test_safe_phrases_pass_when_clean(self): + import re + + safe_log_res = ( + re.compile(r"^welcome to codex\b", re.IGNORECASE), + re.compile(r"^initializing\b", re.IGNORECASE), + re.compile(r"^open (?:this|the verification)", re.IGNORECASE), + re.compile(r"^open:\s*https?://", re.IGNORECASE), + re.compile(r"^enter (?:this one-time code|the code)\b", re.IGNORECASE), + re.compile(r"^waiting\b", re.IGNORECASE), + re.compile(r"^successfully (?:logged|signed) in\b", re.IGNORECASE), + re.compile(r"^(?:logged|signed) in\b", re.IGNORECASE), + re.compile(r"^browser opened\b", re.IGNORECASE), + re.compile(r"^press ctrl", re.IGNORECASE), + ) + for clean in [ + "Welcome to codex", + "Initializing device auth...", + "Open this URL: https://auth.openai.com/codex/device", + "Open: https://auth.openai.com/codex/device", + "Enter this one-time code:", + "Waiting for authentication...", + "Successfully logged in", + "Logged in using ChatGPT", + "Browser opened", + "Press Ctrl+C to cancel", + ]: + assert any(pat.search(clean) for pat in safe_log_res), \ + f"clean line should match safe: {clean!r}" From be15c57fee6f7af58b6a4c816b5df161aa1877c1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 14:45:38 +0000 Subject: [PATCH 31/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_codex_provider.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index be0ab1ace5..92ef48d439 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2047,9 +2047,7 @@ class TestDeviceLoginLogFilter: # caught when the production source is loaded. import importlib - mod = importlib.reload( - importlib.import_module("core.inference.codex_provider") - ) + mod = importlib.reload(importlib.import_module("core.inference.codex_provider")) # Walk the source string to find the patterns; they live inside # the generator. Use a stable proxy: read the regex literals. import re @@ -2059,8 +2057,10 @@ class TestDeviceLoginLogFilter: # phrases AND an unsafe-content blocklist. assert "safe_log_res" in src assert "unsafe_log_re" in src - assert "not\\s+(?:logged|signed)\\s+in" in src or \ - "not\\\\s+(?:logged|signed)\\\\s+in" in src + assert ( + "not\\s+(?:logged|signed)\\s+in" in src + or "not\\\\s+(?:logged|signed)\\\\s+in" in src + ) return None def test_safe_log_source_has_anchored_patterns_and_blocklist(self): @@ -2117,5 +2117,6 @@ class TestDeviceLoginLogFilter: "Browser opened", "Press Ctrl+C to cancel", ]: - assert any(pat.search(clean) for pat in safe_log_res), \ - f"clean line should match safe: {clean!r}" + assert any( + pat.search(clean) for pat in safe_log_res + ), f"clean line should match safe: {clean!r}" From b17765b5ed0e7bdd83cac91c0aa9644a63a2beaa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 15:24:25 +0000 Subject: [PATCH 32/40] Studio: round 8 -- replay guard on non-visible events + Add-flow Codex preselect Two P1 fixes from the round 8 reviewer pass: 1. _stream_thread_run no longer replays a Codex turn that fired non-visible events before crashing. The replay guard only tracked `emitted_any` (visible text). A Codex turn that emitted, say, a command.delta or file.delta event first -- both filtered to "" by _coerce_text -- and THEN crashed would leave emitted_any=False and fall through to the buffered `thread.run(prompt)` fallback, re-executing the same turn and duplicating its side effects (shell commands, file writes, tool calls). This is exactly the case the guard was added to prevent in earlier rounds; the missing bit was tracking "the turn ran at all", not just "the turn yielded text". Fix: add a separate turn_started flag that flips True the moment we ask the SDK for a turn handle or observe any event from a streaming helper. When the buffered fallback is gated on turn_started instead of emitted_any, a partial-turn crash correctly stops without replaying. Regression test reproduces the bug against the pre-fix code (assertion catches the extra thread.run call) and locks the fix in. 2. openAddProvider now mirrors the providerType-change effect's Codex pre-check. The first-run UX fix from `26799d9a` pre-checked every Codex default model in the providerType-change effect, but openAddProvider() calls resetForm() (which clears selectedModelIds) and then only restores availableModels, not selectedModelIds. If the user closes the Add connection form and re-opens it while Codex is still the current providerType, the effect does not re-run, so the form opens with Codex defaults available but none selected -- the "Add at least one model ID" save guard then blocks the Save click. Fix: openAddProvider now seeds selectedModelIds with the full default-models list when the provider is Codex, matching the providerType-change effect so the two entry paths produce the same first-run state. --- .../backend/core/inference/codex_provider.py | 55 ++++++--- studio/backend/tests/test_codex_provider.py | 108 ++++++++++++++++-- .../features/chat/chat-providers-dialog.tsx | 15 ++- 3 files changed, 156 insertions(+), 22 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 902ec12e3f..6d6e9adb96 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -719,12 +719,16 @@ async def _stream_thread_run( path. Used when neither streaming helper resolves and as the final fallback. - 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. + Cross-turn side-effect protection: once a turn has STARTED (any + SDK event was received, including ones that ``_coerce_text`` + drops -- command/file/tool/plan events), we never fall through + to the buffered ``thread.run(prompt)`` path. A partial-stream + failure mid-turn must not re-execute the same Codex turn + because the side effects (file writes, shell commands, etc.) + would replay. Tracking only ``emitted_any`` (visible text) is + not enough -- a turn that crashes after running shell commands + but before producing answer text would otherwise replay because + no visible chunk was emitted. Empty-delta protection: the canonical SDK can complete a turn successfully without emitting any ``message.delta`` events -- @@ -734,6 +738,11 @@ async def _stream_thread_run( Studio never returns an empty answer for a successful turn. """ emitted_any = False + # True once ANY event has been observed from a streaming helper. + # Even when ``_coerce_text`` filters the event out, the turn has + # demonstrably started executing on the Codex side, so a later + # error must not trigger a buffered ``thread.run`` replay. + turn_started = False # 1. Canonical: thread.turn(prompt).stream() turn_factory = getattr(thread, "turn", None) @@ -741,6 +750,12 @@ async def _stream_thread_run( agent_message_texts: list[str] = [] try: turn_handle = turn_factory(prompt) + # Asking the SDK for the turn handle is itself enough to + # start the turn on the upstream side; mark turn_started + # before we even start iterating so a crash inside the + # stream factory below does not look like a never-started + # turn that is safe to replay. + turn_started = True if asyncio.iscoroutine(turn_handle): turn_handle = await turn_handle stream_fn = getattr(turn_handle, "stream", None) @@ -749,6 +764,7 @@ async def _stream_thread_run( if asyncio.iscoroutine(stream_obj): stream_obj = await stream_obj async for event in stream_obj: + turn_started = True payload = getattr(event, "payload", event) text = _coerce_text(payload) if text: @@ -771,11 +787,16 @@ async def _stream_thread_run( exc_type = type(exc).__name__, error = str(exc), emitted_any = emitted_any, + turn_started = turn_started, ) - 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. + if turn_started: + # The Codex turn has executed at least one event on + # the upstream side (it may have launched shell + # commands or written files via tool events that + # _coerce_text filtered out). Re-executing via + # run_streaming / run() would duplicate those side + # effects, so stop here even if no visible text was + # yielded. return # 2. Legacy: thread.run_streaming(prompt) @@ -783,9 +804,14 @@ async def _stream_thread_run( if run_streaming is not None: try: stream_obj = run_streaming(prompt) + # Same reasoning as the canonical path above: calling the + # streaming helper is enough to start the turn on the SDK + # side, so a later crash must NOT replay via buffered run. + turn_started = True if asyncio.iscoroutine(stream_obj): stream_obj = await stream_obj async for event in stream_obj: + turn_started = True text = _coerce_text(event) if text: emitted_any = True @@ -797,13 +823,16 @@ async def _stream_thread_run( exc_type = type(exc).__name__, error = str(exc), emitted_any = emitted_any, + turn_started = turn_started, ) - if emitted_any: + if turn_started: 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. + # Only reached when no streaming helper ran at all (no turn / + # run_streaming attributes on the thread, or both raised before + # observing any event / starting the turn), so this is the first + # and only execution of the turn. result = await thread.run(prompt) text = _buffered_result_text(result) if text: diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 92ef48d439..423b0a10a7 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2047,7 +2047,9 @@ class TestDeviceLoginLogFilter: # caught when the production source is loaded. import importlib - mod = importlib.reload(importlib.import_module("core.inference.codex_provider")) + mod = importlib.reload( + importlib.import_module("core.inference.codex_provider") + ) # Walk the source string to find the patterns; they live inside # the generator. Use a stable proxy: read the regex literals. import re @@ -2057,10 +2059,8 @@ class TestDeviceLoginLogFilter: # phrases AND an unsafe-content blocklist. assert "safe_log_res" in src assert "unsafe_log_re" in src - assert ( - "not\\s+(?:logged|signed)\\s+in" in src - or "not\\\\s+(?:logged|signed)\\\\s+in" in src - ) + assert "not\\s+(?:logged|signed)\\s+in" in src or \ + "not\\\\s+(?:logged|signed)\\\\s+in" in src return None def test_safe_log_source_has_anchored_patterns_and_blocklist(self): @@ -2117,6 +2117,98 @@ class TestDeviceLoginLogFilter: "Browser opened", "Press Ctrl+C to cancel", ]: - assert any( - pat.search(clean) for pat in safe_log_res - ), f"clean line should match safe: {clean!r}" + assert any(pat.search(clean) for pat in safe_log_res), \ + f"clean line should match safe: {clean!r}" + + +# ── Round 8: stream replay protection on non-visible events ────────── + + +class TestStreamReplayProtection: + """Lock in the round 8 fix: a turn that fired non-rendered events + (command/file/tool deltas) before crashing MUST NOT replay via the + buffered `thread.run(prompt)` fallback even though no visible + text was yielded. The earlier guard only tracked `emitted_any` + (visible text), missing the case where shell commands or file + writes already happened upstream. + """ + + def test_buffered_run_not_called_after_non_visible_event_crash(self): + """Stream raises after a tool event with no visible text. The + buffered ``thread.run`` MUST NOT be called -- the Codex turn + has already started running side-effects upstream. + """ + from core.inference.codex_provider import _stream_thread_run + + class _Stream: + def __init__(self): + self._i = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + self._i += 1 + if self._i == 1: + # An event with no answer text -- _coerce_text + # returns "" but the turn has demonstrably run. + return {"type": "command.delta", "delta": "rm -rf"} + raise RuntimeError("upstream stream died mid-turn") + + class _Turn: + def stream(self): + return _Stream() + + class _Thread: + run_calls: int = 0 + + def turn(self_inner, prompt): + return _Turn() + + async def run(self_inner, prompt): + self_inner.run_calls += 1 + return "REPLAY-WOULD-RETURN-THIS" + + thread = _Thread() + + async def collect(): + chunks = [] + async for c in _stream_thread_run(thread, "hello"): + chunks.append(c) + return chunks + + chunks = asyncio.run(collect()) + # No visible text was emitted (the only event was filtered), + # but thread.run MUST NOT have been called because the turn + # already started. + assert thread.run_calls == 0, ( + "thread.run was called after a partial-turn crash; this " + "would replay shell commands / file writes" + ) + assert chunks == [] + + def test_buffered_run_called_when_no_streaming_helper(self): + """Threads that expose neither .turn nor .run_streaming still + fall through to the buffered .run -- that is the ONLY path + the buffered fallback is allowed to execute. + """ + from core.inference.codex_provider import _stream_thread_run + + class _Thread: + run_calls: int = 0 + + async def run(self_inner, prompt): + self_inner.run_calls += 1 + return "answer" + + thread = _Thread() + + async def collect(): + chunks = [] + async for c in _stream_thread_run(thread, "hello"): + chunks.append(c) + return chunks + + chunks = asyncio.run(collect()) + assert thread.run_calls == 1 + assert chunks == ["answer"] diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index faa06afef4..64cb177804 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -556,7 +556,20 @@ export function ChatProvidersSettings({ resetForm(); const entry = providerType ? registryByType.get(providerType) : null; if (entry?.model_list_mode === "curated") { - setAvailableModels([...entry.default_models]); + const defaults = [...entry.default_models]; + setAvailableModels(defaults); + // Mirror the providerType-change effect's first-run behavior: + // Codex is the local CLI so pre-checking the default models lets + // the user click Save without re-ticking anything. Without this + // the resetForm above would zero selectedModelIds and the form + // would fail the "Add at least one model ID" save guard even + // though the round 7 Codex auto-enable effect would have + // populated them. Triggered when the user clicks Add connection + // while Codex was already the providerType (e.g. after closing + // and reopening the form). + setSelectedModelIds( + providerType === CODEX_PROVIDER_TYPE ? defaults : [], + ); } setPage("form"); } From 5185032b283c7e173367333d0936a156e2f385a0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 15:25:01 +0000 Subject: [PATCH 33/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_codex_provider.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 423b0a10a7..8084917fad 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2047,9 +2047,7 @@ class TestDeviceLoginLogFilter: # caught when the production source is loaded. import importlib - mod = importlib.reload( - importlib.import_module("core.inference.codex_provider") - ) + mod = importlib.reload(importlib.import_module("core.inference.codex_provider")) # Walk the source string to find the patterns; they live inside # the generator. Use a stable proxy: read the regex literals. import re @@ -2059,8 +2057,10 @@ class TestDeviceLoginLogFilter: # phrases AND an unsafe-content blocklist. assert "safe_log_res" in src assert "unsafe_log_re" in src - assert "not\\s+(?:logged|signed)\\s+in" in src or \ - "not\\\\s+(?:logged|signed)\\\\s+in" in src + assert ( + "not\\s+(?:logged|signed)\\s+in" in src + or "not\\\\s+(?:logged|signed)\\\\s+in" in src + ) return None def test_safe_log_source_has_anchored_patterns_and_blocklist(self): @@ -2117,8 +2117,9 @@ class TestDeviceLoginLogFilter: "Browser opened", "Press Ctrl+C to cancel", ]: - assert any(pat.search(clean) for pat in safe_log_res), \ - f"clean line should match safe: {clean!r}" + assert any( + pat.search(clean) for pat in safe_log_res + ), f"clean line should match safe: {clean!r}" # ── Round 8: stream replay protection on non-visible events ────────── From e2b7f5958bd2e27caefac75b02d634aba6a745f8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 06:58:06 +0000 Subject: [PATCH 34/40] Studio: round 9 -- three P2 fixes from latest Codex bot review 1. Codex SSE wrapper terminates on exact `data: [DONE]` only. The old substring check `if "[DONE]" in line` would flip sent_done True when a normal model response carried the literal text "[DONE]" in delta.content (for example an explanation of the OpenAI stream sentinel). The real terminator was then suppressed, leaving OpenAI-compatible clients that finalise on the explicit sentinel hung on stream close. Now compares the stripped line to the exact `data: [DONE]` form. 2. Legacy `thread.run_streaming` path no longer returns an empty reply on completion-only streams. If the SDK exposes `thread.run_streaming` but the stream emits ONLY item.completed / agentMessage events with no message deltas, the loop previously exited with emitted_any False and never reached the agent-message fallback. The request returned 200 with an empty assistant reply even though Codex produced a final answer. Mirror the canonical-path behavior: collect `_completed_agent_message_text` strings in a sidecar list and emit the last one when no deltas arrived. Match the canonical payload-extraction (`getattr(event, "payload", event)`) so the event-vs-payload SDK shape difference is handled the same way in both branches. 3. Parallel-calls fan-out propagates CodexUnavailableError so the route layer can return 503. When the SDK is not importable or the safety enums are missing without the dev opt-in, every worker raised the same CodexUnavailableError. The previous catch-all converted the error into a per-tab codex_tab_error event, the outer stream never raised, and clients saw a 200 with only tool events and an empty synthesis -- OpenAI-compatible consumers that ignore _toolEvent saw a successful empty reply. Now CodexUnavailableError re-raises out of the worker (no spurious per-tab error event), _await_workers re-raises it when EVERY worker hit the same setup failure, and the finally-block drain await propagates the exception out of the parallel function so the route's existing CodexUnavailableError handler can emit the right 503 SSE error frame. Per-tab runtime failures (model rejected, timeout, mid- stream SDK crash) still get swallowed into codex_tab_error events so a single bad model in the fan-out does not kill the others. Test counts: 63/63 passing (60 round 6-8 plus 3 new round 9 regression tests). Each new test was first run against a `git stash`-restored pre-fix tree to confirm it catches the bug, then run against the patched tree. --- .../backend/core/inference/codex_provider.py | 95 +++++++++--- studio/backend/routes/inference.py | 10 +- studio/backend/tests/test_codex_provider.py | 143 ++++++++++++++++++ 3 files changed, 226 insertions(+), 22 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 6d6e9adb96..24dc14b72b 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -802,6 +802,7 @@ async def _stream_thread_run( # 2. Legacy: thread.run_streaming(prompt) run_streaming = getattr(thread, "run_streaming", None) if run_streaming is not None: + agent_message_texts: list[str] = [] try: stream_obj = run_streaming(prompt) # Same reasoning as the canonical path above: calling the @@ -812,10 +813,27 @@ async def _stream_thread_run( stream_obj = await stream_obj async for event in stream_obj: turn_started = True - text = _coerce_text(event) + # Match the canonical path's payload-extraction so the + # event-vs-payload shape difference between SDK versions + # is handled the same way in both branches. + payload = getattr(event, "payload", event) + text = _coerce_text(payload) if text: emitted_any = True yield text + else: + # Legacy SDK variants can complete a turn purely via + # ``item.completed`` / agentMessage events with no + # streaming deltas. Capture them so we still emit + # SOMETHING when the stream ends; otherwise the + # request returns 200 with an empty assistant reply + # even though Codex produced a final answer. + final_text = _completed_agent_message_text(payload) + if final_text: + agent_message_texts.append(final_text) + if not emitted_any and agent_message_texts: + yield agent_message_texts[-1] + emitted_any = True return except Exception as exc: logger.warning( @@ -826,6 +844,13 @@ async def _stream_thread_run( turn_started = turn_started, ) if turn_started: + # Flush any agent-message text we collected before the + # crash. The turn already executed on the SDK side, so + # there is no replay risk -- we just want the user to + # see the final answer that was completed before the + # stream broke. + if not emitted_any and agent_message_texts: + yield agent_message_texts[-1] return # 3. Buffered fallback: await the full TurnResult, emit one chunk. @@ -1142,11 +1167,19 @@ async def _stream_codex_parallel( async def _worker(tab_id: int) -> str: """Run one Codex turn, push every chunk into the queue, and return the full accumulated text so the synthesis step can - consume it. Errors are surfaced as a ``codex_tab_error`` - tool-event so the tab strip shows which lane failed without - aborting the whole fan-out. + consume it. Per-turn errors (model rejection, timeout, mid- + stream SDK crash) are surfaced as a ``codex_tab_error`` tool + event so the tab strip shows which lane failed without + aborting the whole fan-out. Setup errors that doom EVERY + worker (SDK not importable, safety enums missing) are re- + raised so the route layer can translate them into a proper + 503 instead of returning a 200 stream with only tool events + and an empty synthesis -- OpenAI-compatible clients that do + not consume ``_toolEvent`` would otherwise see a successful + empty reply. """ collected: list[str] = [] + emit_close = True try: sdk = _import_codex() async_codex_cls = getattr(sdk, "AsyncCodex") @@ -1166,46 +1199,48 @@ async def _stream_codex_parallel( }, ) ) + except CodexUnavailableError: + # Setup-level failure. Same root cause hits every worker, so + # surfacing it as a per-tab error is misleading: every tab + # would emit the same message and the synthesis would be + # blank. Skip the close event too -- the outer fan-out gives + # up before any tabs can render. + emit_close = False + raise except Exception as exc: # CodeQL: never echo str(exc) in client-facing SSE events. # Log full reason server-side; surface a generic message plus # an exception_type discriminator so the UI can still group # failures without leaking file paths / env vars from the - # SDK traceback. CodexUnavailableError is the one exception - # we DO surface verbatim because it's a user-actionable - # install hint with no sensitive content. + # SDK traceback. logger.warning( "codex_provider.parallel_tab_failed", tab_id = tab_id, exc_type = type(exc).__name__, error = str(exc), ) - public_error = ( - str(exc) - if isinstance(exc, CodexUnavailableError) - else "Codex tab failed" - ) await queue.put( _chunk_tool_event( completion_id, { "type": "codex_tab_error", "tab_id": tab_id, - "error": public_error, + "error": "Codex tab failed", "exception_type": type(exc).__name__, }, ) ) finally: - await queue.put( - _chunk_tool_event( - completion_id, - { - "type": "codex_tab_close", - "tab_id": tab_id, - }, + if emit_close: + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_close", + "tab_id": tab_id, + }, + ) ) - ) return "".join(collected) workers = [asyncio.create_task(_worker(i + 1)) for i in range(n)] @@ -1220,6 +1255,18 @@ async def _stream_codex_parallel( async def _await_workers() -> None: results = await asyncio.gather(*workers, return_exceptions = True) + # If a setup-level failure took out every worker + # (CodexUnavailableError -- SDK not importable, safety enums + # missing, etc.), re-raise so the route layer turns it into a + # 503 instead of letting an empty 200 stream close. Per-tab + # runtime failures stay swallowed (they're already surfaced as + # codex_tab_error events) so a single bad model in the fan-out + # does not kill the others. + setup_errors = [ + r for r in results if isinstance(r, CodexUnavailableError) + ] + if setup_errors and len(setup_errors) == len(results): + raise setup_errors[0] for r in results: if isinstance(r, BaseException): per_tab_texts.append("") @@ -1265,6 +1312,12 @@ async def _stream_codex_parallel( if not cancelled: try: await drain_task + except CodexUnavailableError: + # Setup-level failure took out every worker. Re-raise so + # the route layer translates it into a 503 instead of + # silently continuing into the synthesis step (which + # would itself fail) and returning an empty 200 stream. + raise except Exception as exc: logger.warning( "codex_provider.parallel_drain_failed", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b173a60e2e..629ac56bf6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1938,7 +1938,15 @@ async def _proxy_to_external_provider( sent_done = False async for line in gen: yield f"{line}\n\n" - if "[DONE]" in line: + # Match the SSE sentinel exactly. The earlier + # substring check (`"[DONE]" in line`) would flip + # the flag when a normal `delta.content` carried + # the literal text "[DONE]" (e.g. an explanation + # of OpenAI's stream terminator), and suppress the + # real `data: [DONE]` frame. OpenAI-compatible + # clients that finalise on the sentinel would + # then hang on stream close. + if line.strip() == "data: [DONE]": sent_done = True if not sent_done: yield "data: [DONE]\n\n" diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 8084917fad..55350e3986 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2213,3 +2213,146 @@ class TestStreamReplayProtection: chunks = asyncio.run(collect()) assert thread.run_calls == 1 assert chunks == ["answer"] + + +# ── Round 9: P2 fixes from latest Codex bot review ────────────────── + + +class TestRunStreamingCompletionFallback: + """Round 9 fix: legacy SDK exposes ``thread.run_streaming`` but the + stream only emits completion-style events (no message deltas). We + must still emit the agentMessage text, otherwise the request + returns 200 with an empty assistant reply. + """ + + def test_run_streaming_only_completion_emits_final_text(self): + from core.inference.codex_provider import _stream_thread_run + + # Dict shape matching _completed_agent_message_text's accepted + # form: type=thread.item.completed, item.type=agentMessage, + # item.text=. The legacy run_streaming path now + # extracts ``.payload`` first (matching the canonical path), so + # a plain dict event is the simplest faithful fixture. + completed_event = { + "type": "thread.item.completed", + "item": {"type": "agentMessage", "text": "FINAL_ANSWER"}, + } + + class _Stream: + def __init__(self): + self._sent = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._sent: + raise StopAsyncIteration + self._sent = True + return completed_event + + class _Thread: + # No .turn so the canonical path is skipped; only legacy + # run_streaming exists, and it yields a completion event + # with no streaming deltas. + def run_streaming(self_inner, prompt): + return _Stream() + + async def run(self_inner, prompt): # pragma: no cover + # Should never be called -- the legacy stream + # completed cleanly via the completion event. + raise AssertionError("buffered run must not fire") + + thread = _Thread() + + async def collect(): + return [c async for c in _stream_thread_run(thread, "hi")] + + chunks = asyncio.run(collect()) + assert chunks == ["FINAL_ANSWER"], chunks + + +class TestParallelSetupErrorPropagation: + """Round 9 fix: when CodexUnavailableError takes out every worker + in a parallel-calls fan-out, the function re-raises so the route + layer can return a proper 503. Per-tab runtime failures (timeout + etc.) still get swallowed into codex_tab_error events as before. + """ + + def test_unavailable_in_every_worker_reraises(self, monkeypatch): + from core.inference import codex_provider as cp + + # Force _import_codex to raise CodexUnavailableError. Every + # worker hits this on entry so per_tab_texts stays empty and + # the function MUST re-raise. + def boom(): + raise cp.CodexUnavailableError("SDK not installed (test)") + + monkeypatch.setattr(cp, "_import_codex", boom) + + async def collect_lines(): + lines = [] + try: + async for line in cp._stream_codex_parallel( + model = "gpt-5.4-mini", + system = "", + prompt = "hello", + n = 3, + completion_id = "test-completion", + ): + lines.append(line) + except cp.CodexUnavailableError as exc: + return lines, exc + return lines, None + + lines, exc = asyncio.run(collect_lines()) + assert exc is not None, ( + "CodexUnavailableError did not propagate -- the stream " + "returned a 200 with no assistant content" + ) + assert "SDK not installed (test)" in str(exc) + + +class TestCodexDoneSentinelExactMatch: + """Round 9 fix: the Codex SSE wrapper's `sent_done` detection now + requires an EXACT `data: [DONE]` line match. The substring check + was firing on `delta.content` payloads that happened to contain + the literal text `[DONE]`. + + The route source is the canonical reference -- this test asserts + the source uses an anchored comparison, not a substring `in` + check, so the fix is locked in even if the route is restructured. + """ + + def test_route_uses_exact_done_match(self): + with open( + _backend_file("routes/inference.py"), "r", encoding = "utf-8" + ) as f: + src = f.read() + # The Codex SSE wrapper is the only place we expect this + # comparison style; allow either single or double quotes + # around the canonical line for forward compatibility. + assert ( + 'line.strip() == "data: [DONE]"' in src + or "line.strip() == 'data: [DONE]'" in src + ), ( + "Codex SSE wrapper must terminate on an exact `data: [DONE]` " + "line, not on a substring containing `[DONE]`." + ) + # Inspect the Codex stream block specifically. The old + # substring check `if "[DONE]" in line: sent_done = True` + # must NOT appear as an active comparison. Ignore matches + # inside comments (lines starting with `#` or inside string + # literals describing the old behavior) by scanning for the + # exact statement form. + codex_block_start = src.find('async def _codex_stream():') + if codex_block_start != -1: + window = src[codex_block_start : codex_block_start + 4000] + for line in window.splitlines(): + stripped = line.strip() + if stripped.startswith("#"): + continue + assert 'if "[DONE]" in line' not in stripped, ( + "Codex SSE wrapper still uses substring [DONE] check: " + + stripped + ) From dd0aeec388e5828c14bc823cb605396d0e3b78dc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 06:58:53 +0000 Subject: [PATCH 35/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/codex_provider.py | 4 +--- studio/backend/tests/test_codex_provider.py | 9 +++------ 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 24dc14b72b..d72cd2fcf4 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -1262,9 +1262,7 @@ async def _stream_codex_parallel( # runtime failures stay swallowed (they're already surfaced as # codex_tab_error events) so a single bad model in the fan-out # does not kill the others. - setup_errors = [ - r for r in results if isinstance(r, CodexUnavailableError) - ] + setup_errors = [r for r in results if isinstance(r, CodexUnavailableError)] if setup_errors and len(setup_errors) == len(results): raise setup_errors[0] for r in results: diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 55350e3986..f7602b5402 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2325,9 +2325,7 @@ class TestCodexDoneSentinelExactMatch: """ def test_route_uses_exact_done_match(self): - with open( - _backend_file("routes/inference.py"), "r", encoding = "utf-8" - ) as f: + with open(_backend_file("routes/inference.py"), "r", encoding = "utf-8") as f: src = f.read() # The Codex SSE wrapper is the only place we expect this # comparison style; allow either single or double quotes @@ -2345,7 +2343,7 @@ class TestCodexDoneSentinelExactMatch: # inside comments (lines starting with `#` or inside string # literals describing the old behavior) by scanning for the # exact statement form. - codex_block_start = src.find('async def _codex_stream():') + codex_block_start = src.find("async def _codex_stream():") if codex_block_start != -1: window = src[codex_block_start : codex_block_start + 4000] for line in window.splitlines(): @@ -2353,6 +2351,5 @@ class TestCodexDoneSentinelExactMatch: if stripped.startswith("#"): continue assert 'if "[DONE]" in line' not in stripped, ( - "Codex SSE wrapper still uses substring [DONE] check: " - + stripped + "Codex SSE wrapper still uses substring [DONE] check: " + stripped ) From 3d1075f6fdd9decfba6fd43b34c07d309f3c39f0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 13:24:18 +0000 Subject: [PATCH 36/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_cpu_threads.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py index 0dbcbdb74b..1224941622 100644 --- a/studio/backend/tests/test_cpu_threads.py +++ b/studio/backend/tests/test_cpu_threads.py @@ -63,16 +63,18 @@ def test_cpu_thread_cap_is_opt_in(raw): # Anything that is not a positive integer raises a clear ValueError. -@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]) +@pytest.mark.parametrize( + "raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"] +) def test_cpu_thread_cap_requires_positive_integer(raw): - with pytest.raises(ValueError, match="must be a positive integer"): + with pytest.raises(ValueError, match = "must be a positive integer"): configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw}) # env=None path uses real os.environ (production call from run.py / main.py). def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): - monkeypatch.delenv(variable, raising=False) + monkeypatch.delenv(variable, raising = False) monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3") configure_cpu_threads() @@ -84,7 +86,7 @@ def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): # Calling twice must not flip any seeded value. def test_cpu_thread_cap_idempotent(monkeypatch): for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): - monkeypatch.delenv(variable, raising=False) + monkeypatch.delenv(variable, raising = False) monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5") configure_cpu_threads() @@ -138,9 +140,9 @@ def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point): result = subprocess.run( [sys.executable, str(entry_point)], - env=env, - capture_output=True, - text=True, + env = env, + capture_output = True, + text = True, ) assert result.returncode == 1 From f01011e4ddf4f711961a32974716652749b37e61 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 13:30:53 +0000 Subject: [PATCH 37/40] Studio: real Codex parallel-call tab UI + in-process SDK spoof Two paired changes that finally make the Codex parallel-calls fan-out visible as actual clickable tabs in the chat surface, plus a credit- free spoof that lets the whole pipeline run in dev / CI without ever touching the upstream API. 1. Real tab UI (frontend). The chat-adapter used to render the per-worker outputs as inline `[Codex tab 1/N] ...` text blocks in the assistant message body, which collapsed into one big run-on block once more than a handful of tokens had streamed. Now each `codex_*` SSE event is folded into `codexParallelState` and re-published as the `args.state` of a single tool-call part with `toolName === "codex_parallel"`. The assistant-ui surface dispatches that to the new `CodexParallelToolUI` wrapper, which mounts the existing `CodexParallelTabs` component -- one tab per worker, one Synthesis tab, click to switch. The stable `toolCallId` keeps assistant-ui updating the SAME card across stream yields rather than spawning new cards. `renderCodexTabsBlock` now returns the empty string so the message body no longer contains the labelled-text fallback (kept the function name so the rest of the adapter's `renderFullContent` / pin-signature paths are untouched). 2. Credit-free Codex SDK spoof (backend). New `studio/backend/core/inference/codex_spoof.py` exposes a drop-in subset of the upstream `openai_codex` surface (`AsyncCodex`, `AppServerConfig`, `ApprovalMode.deny_all`, `SandboxMode.read_only`, thread with `turn().stream()` + `run_streaming()` + `run()`) and emits deterministic per-tab streaming events tagged with the worker index, so flipping between tabs in the UI shows visibly distinct text. Activated by `UNSLOTH_CODEX_SPOOF=1`; `_import_codex` installs the spoof into `sys.modules` under both `openai_codex` and `codex_app_server` and the rest of the provider keeps running unchanged. OFF by default; production is unaffected. Six new tests cover the spoof itself (module install, env-flag gating, delta + completion event shape, per-tab tagging, provider import path, safety-kwargs resolution against the spoof). 69/69 tests pass with and without the flag; TypeScript clean. --- .../backend/core/inference/codex_provider.py | 9 + studio/backend/core/inference/codex_spoof.py | 269 ++++++++++++++++++ studio/backend/tests/test_codex_provider.py | 115 ++++++++ .../src/components/assistant-ui/thread.tsx | 2 + .../assistant-ui/tool-ui-codex-parallel.tsx | 33 +++ .../src/features/chat/api/chat-adapter.ts | 165 +++++++---- 6 files changed, 533 insertions(+), 60 deletions(-) create mode 100644 studio/backend/core/inference/codex_spoof.py create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-codex-parallel.tsx diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index d72cd2fcf4..d975463cbb 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -358,7 +358,16 @@ def _import_codex() -> Any: name resolves, so this branch is reached only when (a) the user explicitly forces the provider via a stale stored config or (b) the install state changes between status probe and chat submit. + + When ``UNSLOTH_CODEX_SPOOF=1`` is set we install the in-process + spoof under ``openai_codex`` so the rest of the provider runs + end-to-end (with deterministic per-tab replies and no upstream + credit usage). The flag is OFF by default in production. """ + from core.inference import codex_spoof + + if codex_spoof.is_spoof_enabled(): + codex_spoof.install_as_openai_codex() for name in _SDK_MODULE_NAMES: if importlib.util.find_spec(name) is not None: return importlib.import_module(name) diff --git a/studio/backend/core/inference/codex_spoof.py b/studio/backend/core/inference/codex_spoof.py new file mode 100644 index 0000000000..885041b4ff --- /dev/null +++ b/studio/backend/core/inference/codex_spoof.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +"""Local-only Codex SDK spoof for credit-free dev / CI runs. + +Activated by ``UNSLOTH_CODEX_SPOOF=1`` -- the gate in ``codex_provider`` +swaps in this module's symbols for ``openai_codex`` so the rest of the +provider can run end-to-end (thread_start, turn().stream(), run_streaming, +run(), AppServerConfig, ApprovalMode.deny_all, SandboxMode.read_only) +without ever touching the real CLI or upstream API. + +The fake stream emits one ``message.delta`` per visible token plus a +trailing ``ItemCompletedNotification(item=agentMessage)`` so both the +delta path and the completion-only fallback in +``_stream_thread_run`` exercise their real branches. + +The replies are deterministic and tagged with the model + tab index so +the parallel-calls fan-out shows visibly distinct text per tab, which +is the point of the tab UI demo. The spoof intentionally does NOT +emit command / file / tool deltas -- those would be denylisted by +``_coerce_text`` and never reach the user, and we want the demo to +show the same shape Codex normally streams: pure agent text. + +This file is import-safe: it has no side effects on import. It MUST +never be selected unless the env flag is set explicitly. +""" + +from __future__ import annotations + +import asyncio +import os +import time +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, AsyncIterator, Optional + + +SPOOF_ENV_VAR = "UNSLOTH_CODEX_SPOOF" + + +def is_spoof_enabled() -> bool: + """Return True when the env flag is set to an explicit truthy value.""" + return os.environ.get(SPOOF_ENV_VAR, "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +class ApprovalMode(str, Enum): + deny_all = "deny_all" + auto_review = "auto_review" + + +class SandboxMode(str, Enum): + read_only = "read_only" + workspace_write = "workspace_write" + + +@dataclass +class AppServerConfig: + env: Optional[dict[str, str]] = None + codex_bin: Optional[str] = None + extra: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class _AgentMessage: + type: str = "agentMessage" + text: str = "" + + +@dataclass +class _ItemRoot: + root: _AgentMessage + + +@dataclass +class _ItemCompletedNotification: + """Mirrors openai_codex.api.ItemCompletedNotification shape. + + The class name is matched verbatim by ``_completed_agent_message_text`` + so the completion-only fallback in ``_stream_thread_run`` recognises + these payloads. + """ + + item: _ItemRoot + type: str = "ItemCompletedNotification" + + +def _tab_id_from_system(system: Optional[str]) -> int: + """Pull the synthetic ``[tab N]`` marker the provider prepends to + each parallel worker's system prompt (when present), else 0.""" + if not system: + return 0 + for line in system.splitlines(): + line = line.strip() + if line.startswith("[tab ") and line.endswith("]"): + try: + return int(line[len("[tab "): -1].split("/")[0]) + except ValueError: + pass + return 0 + + +def _spoof_response_text(model: str, prompt: str, tab_id: int) -> str: + """Deterministic but visibly per-tab response. + + Format keeps each parallel worker's reply distinct so when the user + clicks between tabs they see different text -- the whole point of + the tab UI demo. + """ + prompt_clean = (prompt or "").strip().replace("\n", " ") + if len(prompt_clean) > 120: + prompt_clean = prompt_clean[:117] + "..." + tab_suffix = f" (worker {tab_id})" if tab_id else "" + return ( + f"[spoof reply from {model}{tab_suffix}] " + f"You said: {prompt_clean!r}. " + f"This response is generated by the local Codex spoof " + f"(UNSLOTH_CODEX_SPOOF=1) -- no upstream tokens were used." + ) + + +class _TurnStream: + """Async iterator returned by ``Turn.stream()``. + + Emits a sequence of dict-shaped ``message.delta`` events (one word at + a time, so the chat-adapter's streaming surface gets exercised) and + closes with an ``ItemCompletedNotification`` carrying the same final + text. Matches the dual delta + completion shape the real upstream + SDK emits. + """ + + def __init__(self, text: str, delay_s: float = 0.01) -> None: + self._text = text + self._delay_s = delay_s + self._iter: Optional[AsyncIterator[Any]] = None + + def __aiter__(self) -> "_TurnStream": + return self + + async def _generate(self) -> AsyncIterator[Any]: + # One word at a time gives a visible streaming effect in the UI + # without flooding the SSE channel. + words = self._text.split(" ") + for i, word in enumerate(words): + chunk = (" " + word) if i > 0 else word + yield {"type": "message.delta", "delta": chunk} + if self._delay_s > 0: + await asyncio.sleep(self._delay_s) + # Final completion event -- the canonical SDK always emits this, + # and ``_stream_thread_run`` uses it as its fallback when no + # deltas arrived (so worth keeping even when deltas did stream). + yield _ItemCompletedNotification(item=_ItemRoot(root=_AgentMessage(text=self._text))) + + async def __anext__(self) -> Any: + if self._iter is None: + self._iter = self._generate() + return await self._iter.__anext__() + + # Some SDK revs let callers ``async with stream:``. Treat as a no-op. + async def __aenter__(self) -> "_TurnStream": + return self + + async def __aexit__(self, *_exc: Any) -> None: + return None + + +class _Turn: + def __init__(self, text: str) -> None: + self._text = text + + def stream(self) -> _TurnStream: + return _TurnStream(self._text) + + +class _Thread: + def __init__(self, model: str, system: Optional[str]) -> None: + self._model = model + self._system = system + self._tab_id = _tab_id_from_system(system) + + # Canonical path: ``thread.turn(prompt).stream()``. + def turn(self, prompt: str) -> _Turn: + text = _spoof_response_text(self._model, prompt, self._tab_id) + return _Turn(text) + + # Legacy path: ``async for event in thread.run_streaming(prompt)``. + def run_streaming(self, prompt: str) -> _TurnStream: + text = _spoof_response_text(self._model, prompt, self._tab_id) + return _TurnStream(text) + + # Buffered fallback: ``await thread.run(prompt)`` returning a result + # whose ``.text`` (or ``.final_response``) is the answer. + async def run(self, prompt: str) -> Any: + from types import SimpleNamespace + + text = _spoof_response_text(self._model, prompt, self._tab_id) + await asyncio.sleep(0) + return SimpleNamespace(text=text, final_response=text) + + +class AsyncCodex: + """Spoof drop-in for ``openai_codex.AsyncCodex``. + + Accepts the same ``config=AppServerConfig(...)`` constructor signature + Studio passes through. ``thread_start`` returns a ``_Thread`` whose + turn / run / run_streaming methods emit deterministic streams. + """ + + def __init__(self, config: Optional[AppServerConfig] = None, **_kw: Any) -> None: + self._config = config or AppServerConfig() + self._started_at = time.time() + + async def thread_start( + self, + *, + model: str, + base_instructions: Optional[str] = None, + system: Optional[str] = None, + approval_mode: Optional[ApprovalMode] = None, + sandbox: Optional[SandboxMode] = None, + **_extra: Any, + ) -> _Thread: + await asyncio.sleep(0) + # Either kwarg path is accepted -- the real provider tries + # ``base_instructions`` first then falls back to ``system``. + sys_text = base_instructions if base_instructions is not None else system + return _Thread(model=model, system=sys_text) + + +def install_as_openai_codex() -> None: + """Insert this module into ``sys.modules`` under the names the real + SDK would use, so ``importlib.util.find_spec`` succeeds and the + provider's existing import path picks it up unchanged. + + Idempotent: a second call is a no-op. Called from ``codex_provider`` + inside ``_import_codex`` when the env flag is set. + """ + import sys + + for name in ("openai_codex", "codex_app_server"): + if name in sys.modules: + continue + sys.modules[name] = _build_module_alias(name) + + +def _build_module_alias(name: str) -> Any: + """Build a module-like object exposing the same public symbols as + this file, under the requested import name. Using a fresh module + object (rather than aliasing ``codex_spoof`` directly) means the + SDK's ``__name__`` lookups (e.g. for ``ImportError`` messages) get + the real upstream-style name. + """ + import types + import importlib.machinery + + mod = types.ModuleType(name) + # ``importlib.util.find_spec(name)`` walks ``sys.modules[name].__spec__`` + # first, so an empty spec is required for the provider's existing + # availability probe to recognise the spoof. + mod.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) + mod.AsyncCodex = AsyncCodex # type: ignore[attr-defined] + mod.AppServerConfig = AppServerConfig # type: ignore[attr-defined] + mod.ApprovalMode = ApprovalMode # type: ignore[attr-defined] + mod.SandboxMode = SandboxMode # type: ignore[attr-defined] + mod.__spoof__ = True # marker -- tests can assert this + return mod diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index f7602b5402..77ff2b3fb3 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2353,3 +2353,118 @@ class TestCodexDoneSentinelExactMatch: assert 'if "[DONE]" in line' not in stripped, ( "Codex SSE wrapper still uses substring [DONE] check: " + stripped ) + + +class TestCodexSpoofModule: + """Ensure the in-process Codex SDK spoof installs cleanly under the + env flag, satisfies _import_codex, exposes the SDK surface the + provider uses, and streams deterministic per-tab text. The spoof + is the credit-free path that the rest of the test suite (and the + UI demo) rides on, so a regression here would silently break + every downstream consumer. + """ + + def test_install_swaps_in_module_when_flag_set(self, monkeypatch): + # Ensure clean import state. + import sys + for name in ("openai_codex", "codex_app_server"): + sys.modules.pop(name, None) + + from core.inference import codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") + assert codex_spoof.is_spoof_enabled() + codex_spoof.install_as_openai_codex() + + import openai_codex # type: ignore[import-not-found] + assert getattr(openai_codex, "__spoof__", False) is True + assert hasattr(openai_codex, "AsyncCodex") + assert hasattr(openai_codex, "AppServerConfig") + assert openai_codex.ApprovalMode.deny_all + assert openai_codex.SandboxMode.read_only + + def test_flag_disabled_does_not_install(self, monkeypatch): + import sys + for name in ("openai_codex", "codex_app_server"): + sys.modules.pop(name, None) + from core.inference import codex_spoof + monkeypatch.delenv(codex_spoof.SPOOF_ENV_VAR, raising=False) + assert not codex_spoof.is_spoof_enabled() + + def test_spoof_stream_emits_deltas_and_completion(self): + import asyncio + from core.inference import codex_spoof + + async def run(): + codex = codex_spoof.AsyncCodex( + config=codex_spoof.AppServerConfig(env={}) + ) + thread = await codex.thread_start( + model="gpt-5.4-mini", + base_instructions=None, + approval_mode=codex_spoof.ApprovalMode.deny_all, + sandbox=codex_spoof.SandboxMode.read_only, + ) + events = [] + async for ev in thread.turn("hello").stream(): + events.append(ev) + return events + + events = asyncio.run(run()) + deltas = [e for e in events if isinstance(e, dict) and e.get("type") == "message.delta"] + assert len(deltas) >= 3, "spoof should stream multiple deltas" + last = events[-1] + assert type(last).__name__ == "_ItemCompletedNotification" + # Must carry the worker tag when no [tab N] marker is in the system prompt. + full_text = last.item.root.text + assert "spoof reply from gpt-5.4-mini" in full_text + assert "no upstream tokens" in full_text + + def test_spoof_tags_per_tab(self): + import asyncio + from core.inference import codex_spoof + + async def reply_for_tab(idx: int) -> str: + codex = codex_spoof.AsyncCodex() + thread = await codex.thread_start( + model="gpt-5.4-mini", + base_instructions=f"[tab {idx}/3]", + ) + result = await thread.run("explain LoRA") + return result.text + + async def _gather(): + return await asyncio.gather( + reply_for_tab(1), + reply_for_tab(2), + reply_for_tab(3), + ) + + tab_replies = asyncio.run(_gather()) + # Each tab must mention its own worker index, so the UI tabs + # show visibly distinct text when clicked. + for i, reply in enumerate(tab_replies, start=1): + assert f"worker {i}" in reply, f"tab {i} missing its tag: {reply}" + + def test_provider_picks_up_spoof_via_import(self, monkeypatch): + import sys + for name in ("openai_codex", "codex_app_server"): + sys.modules.pop(name, None) + from core.inference import codex_provider, codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") + mod = codex_provider._import_codex() + assert getattr(mod, "__spoof__", False) is True + + def test_safety_kwargs_resolve_against_spoof(self, monkeypatch): + import sys + for name in ("openai_codex", "codex_app_server"): + sys.modules.pop(name, None) + from core.inference import codex_provider, codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") + codex_provider._import_codex() # ensures install + kwargs = codex_provider._safe_thread_safety_kwargs() + # Spoof exports ApprovalMode.deny_all + SandboxMode.read_only, + # so the provider must be able to pin both without falling + # through to the unsafe-defaults gate. + assert kwargs, "safety kwargs missing -- provider would fail closed" + assert str(kwargs["approval_mode"]).endswith("deny_all") + assert str(kwargs["sandbox"]).endswith("read_only") diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 51cf000863..6dcdd8c533 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -23,6 +23,7 @@ import { import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution"; +import { CodexParallelToolUI } from "@/components/assistant-ui/tool-ui-codex-parallel"; import { ImageGenerationToolUI } from "@/components/assistant-ui/tool-ui-image-generation"; import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python"; import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal"; @@ -1315,6 +1316,7 @@ const AssistantMessage: FC = () => { python: PythonToolUI, terminal: TerminalToolUI, code_execution: CodeExecutionToolUI, + codex_parallel: CodexParallelToolUI, image_generation: ImageGenerationToolUI, }, Fallback: ToolFallback, diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-codex-parallel.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-codex-parallel.tsx new file mode 100644 index 0000000000..76dcb3e4c0 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-ui-codex-parallel.tsx @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +/** + * Tool-call renderer for Codex parallel-calls fan-out. + * + * Driven by the chat-adapter pushing a tool-call part with + * ``toolName === "codex_parallel"`` whose ``args.state`` is a + * ``CodexParallelState`` value. We just unpack the state and hand it + * to the existing ``CodexParallelTabs`` component. Mounted via the + * ``tools.by_name`` map on ``MessagePrimitive.Parts`` in + * ``thread.tsx`` so it renders inline above the assistant's prose. + */ + +import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; +import { memo } from "react"; +import { + CodexParallelTabs, + EMPTY_CODEX_PARALLEL_STATE, + type CodexParallelState, +} from "@/features/chat/components/codex-parallel-tabs"; + +const CodexParallelToolUIImpl: ToolCallMessagePartComponent = ({ args }) => { + const state = (args as { state?: CodexParallelState } | undefined)?.state; + return ; +}; + +export const CodexParallelToolUI = memo( + CodexParallelToolUIImpl, +) as ToolCallMessagePartComponent; +CodexParallelToolUI.displayName = "CodexParallelToolUI"; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index de43462d42..00844df15c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -58,6 +58,13 @@ import { hasClosedThinkTag, parseAssistantContent, } from "../utils/parse-assistant-content"; +import { + EMPTY_CODEX_PARALLEL_STATE, + hasCodexParallelContent, + reduceCodexParallelState, + type CodexParallelEvent, + type CodexParallelState, +} from "../components/codex-parallel-tabs"; import { generateAudio, listCachedGguf, @@ -1714,29 +1721,47 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // dict keyed by tab_id and re-assembling cumulativeText from // scratch on every codex event puts each tab's text under its // own header regardless of arrival interleaving. - const codexTabBuffers = new Map(); - const codexTabClosed = new Set(); - const codexTabError = new Map(); - let codexTotalTabs = 0; + // Per-tab Codex fan-out state. Each codex_* SSE event is folded + // into ``codexParallelState`` via the pure reducer in + // ``components/codex-parallel-tabs``. The state is re-published + // on every yield as the ``args`` of a single tool-call part with + // ``toolName === "codex_parallel"`` so the assistant-ui surface + // can render real clickable tabs (one per worker plus a + // Synthesis tab) instead of inline ``[Codex tab N]`` headings. + // The stable toolCallId keeps assistant-ui updating the same + // part across stream yields rather than spawning new cards. + let codexParallelState: CodexParallelState = EMPTY_CODEX_PARALLEL_STATE; let codexGatherEmitted = false; + const CODEX_PARALLEL_TOOL_ID = "codex_parallel_main"; - function renderCodexTabsBlock(): string { - if (codexTabBuffers.size === 0) return ""; - const lines: string[] = []; - const ids = [...codexTabBuffers.keys()].sort((a, b) => a - b); - for (const id of ids) { - const header = codexTotalTabs - ? `[Codex tab ${id}/${codexTotalTabs}]` - : `[Codex tab ${id}]`; - lines.push(`\n\n${header}\n${codexTabBuffers.get(id) ?? ""}`); - if (codexTabError.has(id)) { - lines.push(`\n[Codex tab ${id} error: ${codexTabError.get(id)}]\n`); - } - if (codexTabClosed.has(id)) { - lines.push("\n"); - } + function upsertCodexParallelToolPart(): void { + if (!hasCodexParallelContent(codexParallelState)) return; + const args = { state: codexParallelState }; + const argsText = ""; + const idx = toolCallParts.findIndex( + (p) => p.toolCallId === CODEX_PARALLEL_TOOL_ID, + ); + const part: ToolCallMessagePart = { + type: "tool-call" as const, + toolCallId: CODEX_PARALLEL_TOOL_ID, + toolName: "codex_parallel", + argsText, + args: args as unknown as ToolCallMessagePart["args"], + }; + if (idx === -1) { + toolCallParts.push(part); + } else { + toolCallParts[idx] = part; } - return lines.join(""); + } + + // No inline `[Codex tab N]` block in the message body any more -- + // the tab UI is mounted as a tool-call part above. The function + // is kept (returning the empty string) so the rest of the + // adapter's renderFullContent() / pin signature paths are + // unchanged across the file. + function renderCodexTabsBlock(): string { + return ""; } // Codex parallel-calls fan-out renders the labeled tab outputs @@ -2222,50 +2247,70 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { - // Codex parallel-calls fan-out events: route chunks - // into per-tab buffers keyed by tab_id, then render - // the whole codex block from scratch each event so - // concurrent tabs cannot interleave under the wrong - // header. `codex_gather` carries the synthesis - // payload which the backend also emits as a normal - // content delta later in the same SSE stream, so we - // only render a divider here to avoid duplicating - // the synthesis text. + // Codex parallel-calls fan-out events. Each event is + // folded into ``codexParallelState`` and re-published + // as the ``args.state`` of the ``codex_parallel`` tool- + // call part, so the assistant-ui surface renders one + // tab per worker plus a Synthesis tab the user can + // click between. ``codex_gather`` flips the flag so + // ``renderFullContent`` knows the synthesis stream is + // about to arrive on the regular content-delta path. if (typeof toolEvent.type === "string" && toolEvent.type.startsWith("codex_")) { - if (toolEvent.type === "codex_tab_open") { - const tabId = Number(toolEvent.tab_id); + const evType = toolEvent.type; + const tabId = Number(toolEvent.tab_id); + let reduced: CodexParallelEvent | null = null; + if (evType === "codex_tab_open" && Number.isFinite(tabId)) { const total = Number(toolEvent.total_tabs); - if (Number.isFinite(tabId)) { - if (!codexTabBuffers.has(tabId)) { - codexTabBuffers.set(tabId, ""); - } - if (Number.isFinite(total) && total > codexTotalTabs) { - codexTotalTabs = total; - } + reduced = { + type: "codex_tab_open", + tab_id: tabId, + query: + typeof toolEvent.query === "string" + ? toolEvent.query + : undefined, + total_tabs: Number.isFinite(total) ? total : undefined, + }; + } else if (evType === "codex_tab_chunk" && Number.isFinite(tabId)) { + const text = + typeof toolEvent.text === "string" ? toolEvent.text : ""; + if (text) { + reduced = { + type: "codex_tab_chunk", + tab_id: tabId, + text, + }; } - } else if (toolEvent.type === "codex_tab_chunk") { - const tabId = Number(toolEvent.tab_id); - const text = typeof toolEvent.text === "string" ? toolEvent.text : ""; - if (Number.isFinite(tabId) && text) { - const prev = codexTabBuffers.get(tabId) ?? ""; - codexTabBuffers.set(tabId, prev + text); - } - } else if (toolEvent.type === "codex_tab_error") { - const tabId = Number(toolEvent.tab_id); - const err = typeof toolEvent.error === "string" ? toolEvent.error : "error"; - if (Number.isFinite(tabId)) { - codexTabError.set(tabId, err); - if (!codexTabBuffers.has(tabId)) { - codexTabBuffers.set(tabId, ""); - } - } - } else if (toolEvent.type === "codex_tab_close") { - const tabId = Number(toolEvent.tab_id); - if (Number.isFinite(tabId)) { - codexTabClosed.add(tabId); - } - } else if (toolEvent.type === "codex_gather") { + } else if (evType === "codex_tab_error" && Number.isFinite(tabId)) { + reduced = { + type: "codex_tab_error", + tab_id: tabId, + error: + typeof toolEvent.error === "string" + ? toolEvent.error + : "error", + }; + } else if (evType === "codex_tab_close" && Number.isFinite(tabId)) { + reduced = { type: "codex_tab_close", tab_id: tabId }; + } else if (evType === "codex_gather") { codexGatherEmitted = true; + reduced = { + type: "codex_gather", + summary: + typeof toolEvent.summary === "string" + ? toolEvent.summary + : undefined, + tab_count: + typeof toolEvent.tab_count === "number" + ? toolEvent.tab_count + : undefined, + }; + } + if (reduced) { + codexParallelState = reduceCodexParallelState( + codexParallelState, + reduced, + ); + upsertCodexParallelToolPart(); } const codexParts = parseAssistantContent(renderFullContent()); yield { From cd8284d5cd656d6d8d63c0b314bf10c5527075ef Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 13:31:18 +0000 Subject: [PATCH 38/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/codex_spoof.py | 14 ++++---- studio/backend/tests/test_codex_provider.py | 35 +++++++++++++------- 2 files changed, 31 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/codex_spoof.py b/studio/backend/core/inference/codex_spoof.py index 885041b4ff..ea18133dc7 100644 --- a/studio/backend/core/inference/codex_spoof.py +++ b/studio/backend/core/inference/codex_spoof.py @@ -61,7 +61,7 @@ class SandboxMode(str, Enum): class AppServerConfig: env: Optional[dict[str, str]] = None codex_bin: Optional[str] = None - extra: dict[str, Any] = field(default_factory=dict) + extra: dict[str, Any] = field(default_factory = dict) @dataclass @@ -97,7 +97,7 @@ def _tab_id_from_system(system: Optional[str]) -> int: line = line.strip() if line.startswith("[tab ") and line.endswith("]"): try: - return int(line[len("[tab "): -1].split("/")[0]) + return int(line[len("[tab ") : -1].split("/")[0]) except ValueError: pass return 0 @@ -152,7 +152,9 @@ class _TurnStream: # Final completion event -- the canonical SDK always emits this, # and ``_stream_thread_run`` uses it as its fallback when no # deltas arrived (so worth keeping even when deltas did stream). - yield _ItemCompletedNotification(item=_ItemRoot(root=_AgentMessage(text=self._text))) + yield _ItemCompletedNotification( + item = _ItemRoot(root = _AgentMessage(text = self._text)) + ) async def __anext__(self) -> Any: if self._iter is None: @@ -198,7 +200,7 @@ class _Thread: text = _spoof_response_text(self._model, prompt, self._tab_id) await asyncio.sleep(0) - return SimpleNamespace(text=text, final_response=text) + return SimpleNamespace(text = text, final_response = text) class AsyncCodex: @@ -227,7 +229,7 @@ class AsyncCodex: # Either kwarg path is accepted -- the real provider tries # ``base_instructions`` first then falls back to ``system``. sys_text = base_instructions if base_instructions is not None else system - return _Thread(model=model, system=sys_text) + return _Thread(model = model, system = sys_text) def install_as_openai_codex() -> None: @@ -260,7 +262,7 @@ def _build_module_alias(name: str) -> Any: # ``importlib.util.find_spec(name)`` walks ``sys.modules[name].__spec__`` # first, so an empty spec is required for the provider's existing # availability probe to recognise the spoof. - mod.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) + mod.__spec__ = importlib.machinery.ModuleSpec(name, loader = None) mod.AsyncCodex = AsyncCodex # type: ignore[attr-defined] mod.AppServerConfig = AppServerConfig # type: ignore[attr-defined] mod.ApprovalMode = ApprovalMode # type: ignore[attr-defined] diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 77ff2b3fb3..a14f6faf56 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2367,15 +2367,18 @@ class TestCodexSpoofModule: def test_install_swaps_in_module_when_flag_set(self, monkeypatch): # Ensure clean import state. import sys + for name in ("openai_codex", "codex_app_server"): sys.modules.pop(name, None) from core.inference import codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") assert codex_spoof.is_spoof_enabled() codex_spoof.install_as_openai_codex() import openai_codex # type: ignore[import-not-found] + assert getattr(openai_codex, "__spoof__", False) is True assert hasattr(openai_codex, "AsyncCodex") assert hasattr(openai_codex, "AppServerConfig") @@ -2384,10 +2387,12 @@ class TestCodexSpoofModule: def test_flag_disabled_does_not_install(self, monkeypatch): import sys + for name in ("openai_codex", "codex_app_server"): sys.modules.pop(name, None) from core.inference import codex_spoof - monkeypatch.delenv(codex_spoof.SPOOF_ENV_VAR, raising=False) + + monkeypatch.delenv(codex_spoof.SPOOF_ENV_VAR, raising = False) assert not codex_spoof.is_spoof_enabled() def test_spoof_stream_emits_deltas_and_completion(self): @@ -2395,14 +2400,12 @@ class TestCodexSpoofModule: from core.inference import codex_spoof async def run(): - codex = codex_spoof.AsyncCodex( - config=codex_spoof.AppServerConfig(env={}) - ) + codex = codex_spoof.AsyncCodex(config = codex_spoof.AppServerConfig(env = {})) thread = await codex.thread_start( - model="gpt-5.4-mini", - base_instructions=None, - approval_mode=codex_spoof.ApprovalMode.deny_all, - sandbox=codex_spoof.SandboxMode.read_only, + model = "gpt-5.4-mini", + base_instructions = None, + approval_mode = codex_spoof.ApprovalMode.deny_all, + sandbox = codex_spoof.SandboxMode.read_only, ) events = [] async for ev in thread.turn("hello").stream(): @@ -2410,7 +2413,11 @@ class TestCodexSpoofModule: return events events = asyncio.run(run()) - deltas = [e for e in events if isinstance(e, dict) and e.get("type") == "message.delta"] + deltas = [ + e + for e in events + if isinstance(e, dict) and e.get("type") == "message.delta" + ] assert len(deltas) >= 3, "spoof should stream multiple deltas" last = events[-1] assert type(last).__name__ == "_ItemCompletedNotification" @@ -2426,8 +2433,8 @@ class TestCodexSpoofModule: async def reply_for_tab(idx: int) -> str: codex = codex_spoof.AsyncCodex() thread = await codex.thread_start( - model="gpt-5.4-mini", - base_instructions=f"[tab {idx}/3]", + model = "gpt-5.4-mini", + base_instructions = f"[tab {idx}/3]", ) result = await thread.run("explain LoRA") return result.text @@ -2442,23 +2449,27 @@ class TestCodexSpoofModule: tab_replies = asyncio.run(_gather()) # Each tab must mention its own worker index, so the UI tabs # show visibly distinct text when clicked. - for i, reply in enumerate(tab_replies, start=1): + for i, reply in enumerate(tab_replies, start = 1): assert f"worker {i}" in reply, f"tab {i} missing its tag: {reply}" def test_provider_picks_up_spoof_via_import(self, monkeypatch): import sys + for name in ("openai_codex", "codex_app_server"): sys.modules.pop(name, None) from core.inference import codex_provider, codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") mod = codex_provider._import_codex() assert getattr(mod, "__spoof__", False) is True def test_safety_kwargs_resolve_against_spoof(self, monkeypatch): import sys + for name in ("openai_codex", "codex_app_server"): sys.modules.pop(name, None) from core.inference import codex_provider, codex_spoof + monkeypatch.setenv(codex_spoof.SPOOF_ENV_VAR, "1") codex_provider._import_codex() # ensures install kwargs = codex_provider._safe_thread_safety_kwargs() From 6403846bbe64cd78434d32fe8a0edd46e8b7c96f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 14:44:15 +0000 Subject: [PATCH 39/40] Studio: spoof-aware Codex availability probe + autouse fixture When ``UNSLOTH_CODEX_SPOOF=1`` is exported (the credit-free dev / CI path the previous commit added), the in-process spoof IS the Codex SDK and a real ``codex`` CLI is irrelevant. The status endpoint at ``/api/codex/status`` used to gate ``installed`` on the real CLI + real SDK only, which made the frontend hide the Codex provider in the connections dropdown even when the spoof was active. Now both ``_sdk_importable`` and ``probe_codex_availability`` short-circuit on ``codex_spoof.is_spoof_enabled()`` so the provider becomes visible under the spoof. ``installed=True``, ``cli_path=""``, ``logged_in=True``, ``version="spoof"`` -- a sentinel that lets devs read off "yes I am under the spoof" at a glance. Real production code path (no spoof flag) is unchanged: still gates on bool(cli_path) AND sdk_ok the same as round 6. Tests: added an autouse fixture in ``test_codex_provider.py`` that clears ``UNSLOTH_CODEX_SPOOF`` before every test so the existing availability / import gating tests are not polluted when a dev runs the suite with the flag exported. The spoof-targeted tests still call ``monkeypatch.setenv(...)`` to flip it back on inside their own scope. 69/69 pass with and without the env flag. --- .../core/inference/codex_availability.py | 31 ++++++++++++++++--- studio/backend/tests/test_codex_provider.py | 12 +++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index 7d273d9ed9..d48d59c877 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -146,7 +146,19 @@ def _sdk_importable() -> bool: Probes both ``openai_codex`` (the canonical upstream package name at ``openai/codex/sdk/python``) and ``codex_app_server`` (the Rust crate name, kept as a forward-compat alias). + + When ``UNSLOTH_CODEX_SPOOF=1`` is set we report importable=True so + the frontend exposes the Codex provider in dev / CI without a real + SDK install. The spoof module gets swapped into ``sys.modules`` on + first ``_import_codex`` call, so any downstream consumer that + actually imports also succeeds. """ + try: + from core.inference import codex_spoof + if codex_spoof.is_spoof_enabled(): + return True + except Exception: + pass for name in _SDK_MODULE_NAMES: try: if importlib.util.find_spec(name) is not None: @@ -338,17 +350,28 @@ async def probe_codex_availability() -> dict[str, Any]: cli_path = _which_codex() sdk_ok = _sdk_importable() + # Spoof mode also fakes the CLI half of the install signal so the + # frontend stops hiding the Codex provider in dev / CI. ``installed`` + # gates on the spoof being explicitly opted in, so production hosts + # without the flag still see the real CLI / SDK gating intact. + spoof_active = False + try: + from core.inference import codex_spoof + spoof_active = codex_spoof.is_spoof_enabled() + except Exception: + pass + payload: dict[str, Any] = { # Gate on BOTH because the login flow shells out to `codex`. # Round 5 briefly set this to `sdk_ok` alone, but round 6 # caught that the login route would then fail with # `codex CLI not found on PATH` after the user clicked # Sign in, leaving them with an unusable provider row. - "installed": bool(cli_path) and sdk_ok, - "cli_path": cli_path, + "installed": (bool(cli_path) and sdk_ok) or spoof_active, + "cli_path": cli_path or ("" if spoof_active else None), "sdk_importable": sdk_ok, - "logged_in": False, - "version": None, + "logged_in": spoof_active, + "version": "spoof" if spoof_active else None, "supported_models": list(_DEFAULT_SUPPORTED_MODELS), } diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index a14f6faf56..f6d2c99952 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -34,6 +34,18 @@ from typing import Any import pytest +@pytest.fixture(autouse = True) +def _no_spoof_by_default(monkeypatch): + """Tests in this file assume the real-SDK gating path unless they + explicitly re-enable the spoof. The dev environment sometimes has + UNSLOTH_CODEX_SPOOF=1 exported for the live UI; clearing it here + keeps the existing availability + import tests deterministic. + Tests that exercise the spoof use ``monkeypatch.setenv(...)`` to + flip it back on inside their own scope. + """ + monkeypatch.delenv("UNSLOTH_CODEX_SPOOF", raising = False) + + _backend = os.path.join(os.path.dirname(__file__), "..") if _backend not in sys.path: sys.path.insert(0, _backend) From c716af442fba9f12d2d762d01e1131a490a0230f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 14:44:34 +0000 Subject: [PATCH 40/40] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/codex_availability.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py index d48d59c877..6e9c3ba03b 100644 --- a/studio/backend/core/inference/codex_availability.py +++ b/studio/backend/core/inference/codex_availability.py @@ -155,6 +155,7 @@ def _sdk_importable() -> bool: """ try: from core.inference import codex_spoof + if codex_spoof.is_spoof_enabled(): return True except Exception: @@ -357,6 +358,7 @@ async def probe_codex_availability() -> dict[str, Any]: spoof_active = False try: from core.inference import codex_spoof + spoof_active = codex_spoof.is_spoof_enabled() except Exception: pass