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 diff --git a/studio/backend/core/inference/codex_availability.py b/studio/backend/core/inference/codex_availability.py new file mode 100644 index 0000000000..6e9c3ba03b --- /dev/null +++ b/studio/backend/core/inference/codex_availability.py @@ -0,0 +1,399 @@ +# 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 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 +whether to surface the "codex" entry in the provider picker. Three +states matter: + +* ``installed=False`` -- either the CLI is missing OR the SDK + (``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 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. + +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. 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.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") + +# 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 needs to spawn its own helpers +# (PATH), resolve its auth/config dir (HOME / USER / Windows equivalents +# plus CODEX_HOME), and emit log output in the user's locale. +# +# OPENAI_API_KEY is DELIBERATELY excluded. The codex CLI authenticates +# via its own `codex login --device-auth` ChatGPT flow or via stdin +# (`--with-api-key`); Studio's stored OpenAI key belongs to the OpenAI +# provider, not Codex. Forwarding it would let a shimmed `codex` binary +# on PATH exfiltrate the user's OpenAI credential. Users who want to +# wire the same key into Codex should set CODEX_OPENAI_API_KEY or feed +# the key via `codex login --with-api-key` themselves. +_SAFE_CODEX_ENV_KEYS: tuple[str, ...] = ( + "PATH", + "HOME", + "USER", + "USERNAME", + "SHELL", + "LANG", + "LC_ALL", + "TMPDIR", + "TEMP", + "TMP", + "SYSTEMROOT", + "WINDIR", + "APPDATA", + "LOCALAPPDATA", + "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", +) + + +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. + + 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 the Codex Python SDK is importable in this interpreter. + + We deliberately use :func:`importlib.util.find_spec` instead of an + 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). + + 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: + 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]: + """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 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, **spawn_kwargs) + 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: + # 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() 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: + 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: 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. + """ + 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. + # Covers the variants seen across CLI releases and locales. + negative = re.compile( + r"\b(not\s+(?:logged|signed)\s+in|" + r"not\s+authenticated|" + r"please\s+(?:log|sign)\s+in|" + r"run\s+`?codex\s+login`?)\b" + ) + if negative.search(combined): + return False + + positive = re.compile( + r"\b(" + r"logged in|" + r"authenticated as|" + r"authenticated:\s*yes|" + r"signed in" + r")\b" + ) + if positive.search(combined): + return True + + if rc == 0: + # 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 + return False + 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 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 + ``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() + + # 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) or spoof_active, + "cli_path": cli_path or ("" if spoof_active else None), + "sdk_importable": sdk_ok, + "logged_in": spoof_active, + "version": "spoof" if spoof_active else 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..d975463cbb --- /dev/null +++ b/studio/backend/core/inference/codex_provider.py @@ -0,0 +1,1631 @@ +# 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 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: + +* ``_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. 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 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 + +import asyncio +import importlib +import importlib.util +import json +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. +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 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 + layer can translate it into a 503 the user sees instead of an + opaque traceback. + """ + + +def _codex_sdk_env_override() -> dict[str, str]: + """Return an env update dict that scrubs sensitive vars from the + codex app-server subprocess env. + + The upstream openai_codex SDK's `AppServerConfig.env` is merged on + top of `os.environ.copy()` (see openai/codex/sdk/python/src/openai_codex/client.py), + so providing an empty-string mapping for every non-safe key + effectively overrides them in the spawn env. Combined with + `_codex_subprocess_env()` (used for direct CLI calls) this gives + parity between the CLI and SDK code paths: neither sees HF_TOKEN, + GH_TOKEN, WANDB_API_KEY, ANTHROPIC_API_KEY, or any other secret + that lives in the Studio parent environment. + """ + import os + + from core.inference.codex_availability import _SAFE_CODEX_ENV_KEYS + + safe = set(_SAFE_CODEX_ENV_KEYS) + return {key: "" for key in os.environ if key not in safe} + + +_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: + """Async-context wrapper that swaps `os.environ` for the lifetime + of a Codex SDK session. + + Used as the fail-closed fallback when the SDK does not expose + `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. + + 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 three issues: + + 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. + 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): + self._async_codex_cls = async_codex_cls + self._inner: Any = 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] = [] + + async def __aenter__(self) -> Any: + import os + + async with _SCRUBBED_ENV_LOCK: + # 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. 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] = current + 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): + try: + if self._inner is not None: + return await self._inner.__aexit__(exc_type, exc, tb) + finally: + 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 + next_count = current - 1 + if next_count == 0: + # Last wrapper holding this key -- restore the + # 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() + + +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=..., + 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 / + WANDB_API_KEY / etc. There is no code path that lets the SDK + inherit those secrets. + """ + try: + sdk_mod = sys.modules.get("openai_codex") or sys.modules.get("codex_app_server") + if sdk_mod is not None: + app_server_config = getattr(sdk_mod, "AppServerConfig", None) + if app_server_config is not None: + # 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 = env_override), + ) + except TypeError: + # Older SDK: AppServerConfig may not accept the env kwarg yet. + # Fall through to the os.environ-swap wrapper. + pass + except Exception as exc: + logger.warning( + "codex_provider.env_scrub_config_failed", + exc_type = type(exc).__name__, + error = str(exc), + ) + return _ScrubbedEnvAsyncCodex(async_codex_cls) + + +def _import_codex() -> Any: + """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. 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. + + 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) + raise CodexUnavailableError( + "Codex Python SDK is not installed on this host. " + "Install with `pip install openai-codex` (canonical upstream " + "name, imports as `openai_codex`; legacy alias `codex_app_server` " + "is also accepted), or use a different provider." + ) + + +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: + """Render the conversation as a single prompt for Codex. + + 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. + """ + # 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 + 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: + """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)}" + + +# Event types Codex emits that carry the assistant's natural-language +# answer (or its stream-time deltas). Other event types -- command +# execution, file edits, tool calls, plan steps -- have their own +# `delta` fields that the OpenAI Chat Completions surface must NOT +# render as visible assistant text or local stdout / paths would leak +# into the chat reply. +_ANSWER_EVENT_TYPES: frozenset[str] = frozenset( + { + "message.delta", + "message.completed", + "assistant.message.delta", + "assistant.message.completed", + "thread.message.delta", + "thread.message.completed", + "text_delta", + "completed", + # Some SDK revs use a bare "message" / "delta" wrapper without + # qualifying the role; we accept those too because the legacy + # tests rely on the shape. + "message", + "delta", + } +) + + +def _coerce_text(payload: Any) -> str: + """Pull text out of a Codex streaming event or result. + + Only events whose ``type`` is in `_ANSWER_EVENT_TYPES` (or have no + ``type`` field at all, i.e. raw text containers) are translated to + 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 "" + if isinstance(payload, str): + return payload + if isinstance(payload, dict): + # If the dict carries a typed event tag, gate on it: only + # answer-bearing types contribute visible text. Untyped dicts + # (legacy / raw text wrappers) fall through to the field walk. + ev_type = payload.get("type") + if isinstance(ev_type, str) and ev_type not in _ANSWER_EVENT_TYPES: + return "" + 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) + # 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 + 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 "" + + +def _completed_agent_message_text(payload: Any) -> str: + """Return the assistant text from an ``ItemCompletedNotification``. + + The canonical openai_codex SDK sometimes finishes a turn without + emitting any ``message.delta`` events: the final answer arrives + only as ``ItemCompletedNotification(item=AgentMessage(text=...))`` + at the end of the stream. Without recognising that shape, + ``_stream_thread_run`` would loop through the stream, see no + ``delta`` text, and return an empty Chat Completions response. + + Returns the empty string for any other event shape so the caller + can ignore it. Matches by class name + structural shape so the + function works on both real upstream events and the dict / fake + shapes the tests use. + """ + if payload is None: + return "" + + # Dict shape: tests + some pre-release SDK revs. + if isinstance(payload, dict): + if payload.get("type") not in ( + "ItemCompletedNotification", + "item.completed", + "thread.item.completed", + ): + return "" + item = payload.get("item") + # The upstream model wraps the item in a discriminated-union + # `root` field; some pre-release shapes drop the wrapper. Look + # both ways. + if isinstance(item, dict): + inner = item.get("root", item) + if not isinstance(inner, dict): + return "" + if inner.get("type") not in ("agentMessage", "agent_message"): + return "" + text = inner.get("text") + return text if isinstance(text, str) else "" + return "" + + # Object shape: upstream events with `.item.root.text`. + if payload.__class__.__name__ not in ( + "ItemCompletedNotification", + "ThreadItemCompletedNotification", + ): + return "" + item = getattr(payload, "item", None) + item = getattr(item, "root", item) + if getattr(item, "type", None) not in ("agentMessage", "agent_message"): + return "" + text = getattr(item, "text", None) + return text if isinstance(text, str) else "" + + +async def _stream_thread_run( + thread: Any, + prompt: str, +) -> AsyncGenerator[str, None]: + """Yield raw text chunks from a Codex thread. + + 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. + + 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 -- + the final text arrives only as an ``ItemCompletedNotification`` + whose ``item`` is an ``agentMessage``. We collect those during the + stream loop and emit the last one if no deltas came through, so + 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) + if turn_factory is not None: + 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) + 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: + turn_started = True + payload = getattr(event, "payload", event) + text = _coerce_text(payload) + if text: + emitted_any = True + yield text + else: + final_text = _completed_agent_message_text(payload) + if final_text: + agent_message_texts.append(final_text) + if not emitted_any and agent_message_texts: + # The stream completed cleanly but only via a final + # ItemCompletedNotification -- emit the last agent + # message text so the chat reply is not blank. + yield agent_message_texts[-1] + emitted_any = True + return + except Exception as exc: + logger.warning( + "codex_provider.turn_stream_failed_fallback", + exc_type = type(exc).__name__, + error = str(exc), + emitted_any = emitted_any, + turn_started = turn_started, + ) + 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) + 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 + # 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 + # 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( + "codex_provider.run_streaming_failed_fallback", + exc_type = type(exc).__name__, + error = str(exc), + emitted_any = emitted_any, + 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. + # 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: + 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. + + The upstream ``openai_codex.AsyncCodex.thread_start`` defaults + ``approval_mode`` to ``ApprovalMode.auto_review`` -- which the SDK + docs describe as "automatically execute tools when permission + escalations occur, without user intervention" -- and leaves + ``sandbox`` as ``None``. Studio drives Codex from a server-side + chat request with no per-action UI, so leaving those at the + defaults would let a model decide on its own to run shell + commands, write files, or hit the network on the operator's + machine. + + We pin both to the strictest values the SDK exposes: + + * ``approval_mode = ApprovalMode.deny_all`` -- reject any tool / + command request rather than auto-approving it. + * ``sandbox = SandboxMode.read_only`` -- the policy that bans + file writes and disables network access. + + 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) + read_only = getattr(sandbox_mode_cls, "read_only", None) + if deny_all is None or read_only is None: + return {} + return {"approval_mode": deny_all, "sandbox": read_only} + + +async def _start_thread_with_system( + codex: Any, + model: str, + system: str, + prompt: str, +) -> tuple[Any, str]: + """Start a Codex thread carrying the system prompt and safe defaults. + + Upstream ``openai_codex.AsyncCodex.thread_start`` accepts the system + prompt under the kwarg ``base_instructions``. Some pre-release / alias + SDK revisions historically used ``system`` instead. We try the + canonical kwarg first, then the legacy one, then drop both and + prepend the system text to the user prompt so the model still sees + it. The returned (thread, prompt) tuple lets the caller use the + possibly-rewritten prompt. + + We always pin ``approval_mode`` to ``deny_all`` and ``sandbox`` to + ``read_only`` when the SDK exposes them (see + ``_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. 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: + 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, 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} + + if not system: + thread = await codex.thread_start(**base_kwargs) + return thread, prompt + + for kw_name in ("base_instructions", "system"): + try: + thread = await codex.thread_start(**base_kwargs, **{kw_name: system}) + return thread, prompt + except TypeError: + continue + except Exception: + raise + # Last-resort fallback: inline the system text in the user prompt so + # the role intent reaches Codex even on an SDK with no kwarg for it. + thread = await codex.thread_start(**base_kwargs) + return thread, f"{system}\n\n{prompt}" + + +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 SDK is installed but AsyncCodex is missing -- " "upgrade the SDK." + ) + + completion_text_chars = 0 + + async with _open_async_codex(async_codex_cls) as codex: + # ``thread_start`` accepts a model id; system prompts are + # passed when supported by the SDK rev (older revs ignore the + # Upstream `openai_codex.AsyncCodex.thread_start` uses + # `base_instructions` for the system prompt (see + # openai/codex/sdk/python/src/openai_codex/api.py). Older / alias + # SDKs may use `system` instead. We try `base_instructions` + # first, then `system`, and finally fall through to inlining + # the system text in the user prompt if neither kwarg is + # accepted. + thread, prompt = await _start_thread_with_system(codex, model, system, 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. 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") + async with _open_async_codex(async_codex_cls) as codex: + thread, inner_prompt = await _start_thread_with_system( + codex, model, system, 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 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. + logger.warning( + "codex_provider.parallel_tab_failed", + tab_id = tab_id, + exc_type = type(exc).__name__, + error = str(exc), + ) + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_error", + "tab_id": tab_id, + "error": "Codex tab failed", + "exception_type": type(exc).__name__, + }, + ) + ) + finally: + 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)] + + # 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) + # 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("") + 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()) + + 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 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", + error = str(exc), + ) + + synthesis_text = await _run_codex_synthesis( + model = model, + system = system, + 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) + + # 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, + # 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) + + +async def _run_codex_synthesis( + *, + model: str, + system: str, + prompt: str, + tab_outputs: list[str], +) -> str: + """Run one extra Codex call that consumes the N per-tab outputs and + returns a unified synthesis. Returns the empty string on failure -- + the caller already surfaced the per-tab outputs so an empty + synthesis is recoverable. The Studio system prompt is forwarded to + the synthesis thread so style/role instructions like "Always answer + in Spanish" survive the fan-out. + """ + if not tab_outputs: + return "" + 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 _open_async_codex(async_codex_cls) as codex: + thread, synthesis_prompt = await _start_thread_with_system( + codex, model, system, synthesis_prompt + ) + result = await thread.run(synthesis_prompt) + # 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 "" + + +# ── Device-auth helper ────────────────────────────────────────────── + + +async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]: + """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. 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. + + 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 = ["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)``. + # 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 + elif os.name == "nt": + spawn_kwargs["creationflags"] = 0x00000200 # CREATE_NEW_PROCESS_GROUP + + try: + 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: + 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 + + # 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]") + # Accept any plausible device-auth URL the CLI prints. Upstream has + # used `.../codex/device`, `chatgpt.com/activate`, and + # `auth.openai.com/device`; rather than guess we look for any + # https URL whose path mentions `device`, `activate`, or `verify`. + url_re = re.compile( + r"https?://[^\s\x1b]+/(?:codex/)?(?:device|activate|verify)\b[^\s\x1b]*", + re.IGNORECASE, + ) + # One-time-code format from upstream device_code_auth.rs: 4 chars, + # dash, 4 chars. Pattern is tolerant of any uppercase alphanum. + 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 + + # 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. + # 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: + 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 + while True: + line_b = await proc.stdout.readline() + if not line_b: + break + 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: + 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: + yield {"type": "device_code", "code": cm.group(1)} + code_emitted = True + # 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 + 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 = 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/codex_spoof.py b/studio/backend/core/inference/codex_spoof.py new file mode 100644 index 0000000000..ea18133dc7 --- /dev/null +++ b/studio/backend/core/inference/codex_spoof.py @@ -0,0 +1,271 @@ +# 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/core/inference/providers.py b/studio/backend/core/inference/providers.py index 785f1dec3b..2688f656ee 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -319,6 +319,49 @@ 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 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.3-codex", + "gpt-5.2", + ], + "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 OpenAI Codex Python SDK (pip install `openai-codex`, " + "imports as `openai_codex`; legacy alias `codex_app_server` " + "is accepted). Surfaced only when the CLI and SDK are both " + "installed; sign in with `codex login`." + ), + }, "openrouter": { "display_name": "OpenRouter", "base_url": "https://openrouter.ai/api/v1", diff --git a/studio/backend/main.py b/studio/backend/main.py index 8457da0b2d..2eeeb9541b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -128,6 +128,7 @@ from datetime import datetime from routes import ( auth_router, chat_history_router, + codex_router, data_recipe_router, datasets_router, export_router, @@ -536,6 +537,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(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index bb1bd394d1..81a6756912 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -863,6 +863,22 @@ class ChatCompletionRequest(BaseModel): "to auto-create." ), ) + parallel_calls: int = Field( + default = 1, + 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. 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`." + ), + ) fast_mode: Optional[bool] = Field( None, description = ( @@ -874,6 +890,30 @@ class ChatCompletionRequest(BaseModel): ), ) + @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/routes/__init__.py b/studio/backend/routes/__init__.py index ee3ab61b6e..4f9019f73d 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 from routes.mcp_servers import router as mcp_servers_router __all__ = [ @@ -30,5 +31,6 @@ __all__ = [ "training_history_router", "chat_history_router", "providers_router", + "codex_router", "mcp_servers_router", ] diff --git a/studio/backend/routes/codex.py b/studio/backend/routes/codex.py new file mode 100644 index 0000000000..60b0d3b18e --- /dev/null +++ b/studio/backend/routes/codex.py @@ -0,0 +1,115 @@ +# 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 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 + +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 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 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}`` + + 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]: + 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" + + 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 259337616c..d94fb402f0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2077,10 +2077,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}", @@ -2116,6 +2119,94 @@ async def _proxy_to_external_provider( base_url = base_url, ) + # 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" + # 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" + 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: + # 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": "Codex provider error", + "type": "provider_error", + "exception_type": type(exc).__name__, + "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..f6d2c99952 --- /dev/null +++ b/studio/backend/tests/test_codex_provider.py @@ -0,0 +1,2493 @@ +# 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 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 + +import asyncio +import json +import os +import sys +import types +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) + +# Resolved relative to this file so the source-inspection tests work in any +# checkout location (CI, dev machines, the review worker, etc.). +_BACKEND_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "..")) + + +def _backend_file(rel: str) -> str: + """Return an absolute path inside the backend tree, regardless of cwd.""" + return os.path.join(_BACKEND_DIR, rel) + + +# ── Helpers ───────────────────────────────────────────────────────── + + +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, *, 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", + ) + # 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 + + real_find_spec = _iu.find_spec + + def _shim(name: str, *args, **kwargs): + if name in ("codex_app_server", "openai_codex"): + 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()) + # The SDK is what backs `AsyncCodex(...)`, so installed=False + # when the SDK is missing -- even if a standalone CLI is on + # PATH there is no way for Studio to drive it without the + # Python bindings. + 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_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 + + for n in (0, -1, -100): + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + parallel_calls = n, + ) + assert req.parallel_calls == 1, f"clamp failed for {n}" + + 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 + + for n in (21, 100, 1000): + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + 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 + and the schema documentation. Non-codex providers ignore the + field regardless of its value, so backwards compat is + preserved. + """ + from models.inference import ChatCompletionRequest + + req = ChatCompletionRequest( + model = "gpt-5.4", + messages = [{"role": "user", "content": "hi"}], + ) + assert req.parallel_calls == 1 + + +# ── 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. + # 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 in _SDK_NAMES: + return None + return real(name, *args, **kwargs) + + monkeypatch.setattr("importlib.util.find_spec", _shim) + # 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, + stream_codex, + ) + + with pytest.raises(CodexUnavailableError): + asyncio.run( + _consume_first( + stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.4", + ) + ) + ) + + +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 = _backend_file("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 = _backend_file("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 = _backend_file("routes/inference.py") + text = open(src).read() + bad = re.findall(r'f["\']Codex error:\s*\{exc\}["\']', text) + assert not bad, f"raw exception in SSE: {bad}" + + def test_codex_route_no_raw_exc_leak(self): + """codex.py SSE stream wrapping must also not leak str(exc).""" + import re + + src = _backend_file("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}" + + 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_codex_subprocess_env_scrubbed(self, monkeypatch): + """The codex subprocess env must not include other-provider secrets. + + OPENAI_API_KEY is intentionally excluded too: a shimmed `codex` + binary on PATH must not receive Studio's stored OpenAI provider + key. Users wire Codex auth via `codex login` or the + codex-specific CODEX_OPENAI_API_KEY override instead. + """ + 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_provider_key_not_for_codex") + monkeypatch.setenv("CODEX_OPENAI_API_KEY", "codex_specific_key") + 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", + # OPENAI_API_KEY belongs to the OpenAI provider, not Codex. + "OPENAI_API_KEY", + ): + assert secret not in env, f"{secret} leaked into codex env" + # Codex-relevant keys must be preserved. + assert env.get("CODEX_OPENAI_API_KEY") == "codex_specific_key" + 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_not_signed_in_wording_also_handled(self): + """`Not signed in` (alternative localisation) must also be + treated as logged-out, not as positive match. + """ + import asyncio + + from core.inference import codex_availability as av + + async def _fake_run_cli(args, **kw): + return (0, "Not signed in.", "") + + orig = av._run_cli + av._run_cli = _fake_run_cli # type: ignore[assignment] + try: + assert asyncio.run(av._detect_logged_in()) is False + finally: + av._run_cli = orig # type: ignore[assignment] + + def test_device_url_accepts_generic_verification_url(self): + """The login parser must accept upstream's chatgpt.com/activate + URL as well as the canonical /codex/device shape. + """ + import re + + src = _backend_file("core/inference/codex_provider.py") + text = open(src).read() + # Find the url_re pattern literal and compile it. + m = re.search(r"url_re\s*=\s*re\.compile\(\s*\n?\s*r\"([^\"]+)\"", text) + assert m, "url_re definition not found" + pattern = re.compile(m.group(1), re.IGNORECASE) + # Upstream device URLs we expect to match. + for u in ( + "https://auth.openai.com/codex/device", + "https://chatgpt.com/activate", + "https://auth.openai.com/device/verify?code=ABCD", + ): + assert pattern.search(u), f"device URL regex missed: {u}" + + def test_synthesis_call_forwards_system_prompt(self, monkeypatch): + """`_run_codex_synthesis` must pass the system prompt so a + fan-out style instruction ("Always answer in Spanish") survives + the unification step. + """ + seen_kwargs: list[dict] = [] + seen_prompts: list[str] = [] + + class _SynThread: + async def run(self, prompt): + seen_prompts.append(prompt) + return "synth ok" + + def turn(self, prompt): + # Force buffered path via no `stream` attr. + class _T: + pass + + return _T() + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + seen_kwargs.append(kw) + return _SynThread() + + _install_fake_codex_sdk(monkeypatch, _Async) + from core.inference.codex_provider import _run_codex_synthesis + + out = asyncio.run( + _run_codex_synthesis( + model = "gpt-5.5", + system = "Always answer in Spanish.", + prompt = "What is the capital of France?", + tab_outputs = ["Paris", "Paris."], + ) + ) + # The upstream openai_codex SDK uses `base_instructions` for the + # system prompt; the legacy alias accepts `system`; the last-resort + # fallback inlines the system text into the user prompt. Accept + # any of those paths. + system_seen = ( + any("Spanish" in (kw.get("base_instructions") or "") for kw in seen_kwargs) + or any("Spanish" in (kw.get("system") or "") for kw in seen_kwargs) + or any("Always answer in Spanish" in p for p in seen_prompts) + ) + assert system_seen, ( + f"system prompt dropped in synthesis. kwargs={seen_kwargs} " + f"prompts={seen_prompts}" + ) + # And the synthesis still returned the model's text. + assert "synth" in out.lower() + + def test_sdk_env_scrubbed_via_appserverconfig(self, monkeypatch): + """The SDK construction path must wire AppServerConfig(env=...) + when the SDK exposes it, so HF_TOKEN / GH_TOKEN are not leaked + to the codex app-server subprocess. + """ + monkeypatch.setenv("HF_TOKEN", "should_be_scrubbed") + monkeypatch.setenv("GH_TOKEN", "should_be_scrubbed") + # OPENAI_API_KEY is now ALSO scrubbed -- it belongs to the + # OpenAI provider, not Codex. CODEX_OPENAI_API_KEY is the + # codex-specific override that survives. + monkeypatch.setenv("OPENAI_API_KEY", "openai_provider_key_not_for_codex") + monkeypatch.setenv("CODEX_OPENAI_API_KEY", "codex_specific_key") + + seen_configs: list[Any] = [] + + class _FakeAppServerConfig: + def __init__(self, env = None, **kw): + self.env = env or {} + + class _Async: + def __init__(self, config = None): + seen_configs.append(config) + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _FakeThread(chunks = ["ok"]) + + # Inject a fake openai_codex module exposing AppServerConfig. + import importlib.util as _iu + import types as _types + + fake_mod = _types.ModuleType("openai_codex") + fake_mod.AsyncCodex = _Async # type: ignore[attr-defined] + fake_mod.AppServerConfig = _FakeAppServerConfig # type: ignore[attr-defined] + # 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( + "importlib.util.find_spec", + lambda n, *a, **kw: ( + _types.SimpleNamespace() + if n in ("openai_codex", "codex_app_server") + else real_find_spec(n, *a, **kw) + ), + ) + + from core.inference.codex_provider import stream_codex + + asyncio.run( + _consume_first( + stream_codex( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + parallel_calls = 1, + ) + ) + ) + assert seen_configs, "AsyncCodex was never instantiated" + cfg = seen_configs[0] + assert cfg is not None, "AppServerConfig was not passed to AsyncCodex" + assert ( + "HF_TOKEN" in cfg.env and cfg.env["HF_TOKEN"] == "" + ), "HF_TOKEN not overridden to empty in SDK env" + assert "GH_TOKEN" in cfg.env and cfg.env["GH_TOKEN"] == "" + # OPENAI_API_KEY is intentionally overridden to empty in the + # SDK env so the app-server cannot use it as a Codex credential + # by accident. The OpenAI provider still reads its own key from + # Studio's storage; nothing in this path needs the env var. + assert cfg.env.get("OPENAI_API_KEY") == "" + # CODEX_OPENAI_API_KEY is the Codex-specific override and must + # survive untouched so users can wire that key into Codex. + assert "CODEX_OPENAI_API_KEY" not in cfg.env + + def test_thread_turn_stream_path_taken(self, monkeypatch): + """The canonical openai_codex API uses thread.turn(prompt).stream(); + the provider must prefer that over the legacy run_streaming hook. + """ + 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 + + 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 False + assert payload["cli_path"] is None + assert payload["sdk_importable"] is True + + # 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 + the system prompt. The provider must try that name first; only + if the SDK rejects it should it fall back to `system`. + """ + 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) + from core.inference.codex_provider import stream_codex + + async def _collect(): + async for _ in stream_codex( + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"}, + ], + model = "gpt-5.5", + parallel_calls = 1, + ): + pass + + asyncio.run(_collect()) + # The first (and only, since this fake accepts any kwargs) + # call must use base_instructions, not the legacy `system`. + assert seen_kwargs, "thread_start was never called" + assert ( + "base_instructions" in seen_kwargs[0] + ), f"upstream-canonical kwarg not used: {seen_kwargs[0]}" + assert seen_kwargs[0]["base_instructions"] == "You are helpful." + assert ( + "system" not in seen_kwargs[0] + ), "legacy `system` kwarg was sent even though base_instructions worked" + + def test_base_instructions_falls_back_to_system(self, monkeypatch): + """When the SDK rejects `base_instructions` with TypeError the + helper must retry with the legacy `system` kwarg before giving + up and inlining the system text in the prompt. + """ + call_log: list[dict] = [] + + class _StrictSDK: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + call_log.append(dict(kw)) + if "base_instructions" in kw: + raise TypeError( + "thread_start() got an unexpected keyword 'base_instructions'" + ) + return _FakeThread(chunks = ["ok"]) + + _install_fake_codex_sdk(monkeypatch, _StrictSDK) + from core.inference.codex_provider import stream_codex + + async def _collect(): + async for _ in stream_codex( + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"}, + ], + model = "gpt-5.5", + parallel_calls = 1, + ): + pass + + asyncio.run(_collect()) + assert len(call_log) >= 2, "fallback to `system` kwarg never tried" + assert "base_instructions" in call_log[0] + assert "system" in call_log[1] and call_log[1]["system"] == "You are helpful." + + def test_scrubbed_env_wrapper_strips_secrets_before_construction(self, monkeypatch): + """When AppServerConfig is missing the fail-closed wrapper must + remove secret env vars BEFORE the SDK constructor runs (the + SDK starts its app-server with `env = os.environ.copy()`). + """ + observed_env_during_init: dict[str, str | None] = {} + + class _NoConfigAsync: + def __init__(self): + # Capture the environment exactly as the SDK would see + # it at construction time. + observed_env_during_init["HF_TOKEN"] = os.environ.get("HF_TOKEN") + observed_env_during_init["GH_TOKEN"] = os.environ.get("GH_TOKEN") + observed_env_during_init["WANDB_API_KEY"] = os.environ.get( + "WANDB_API_KEY" + ) + observed_env_during_init["PATH"] = os.environ.get("PATH") + observed_env_during_init["CODEX_HOME"] = os.environ.get("CODEX_HOME") + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _FakeThread(chunks = ["ok"]) + + monkeypatch.setenv("HF_TOKEN", "should_be_gone") + monkeypatch.setenv("GH_TOKEN", "should_be_gone") + monkeypatch.setenv("WANDB_API_KEY", "should_be_gone") + monkeypatch.setenv("PATH", "/usr/bin") + monkeypatch.setenv("CODEX_HOME", "/home/u/.codex") + # No AppServerConfig in the fake module -- forces the wrapper path. + _install_fake_codex_sdk(monkeypatch, _NoConfigAsync) + 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()) + # Secrets must have been removed from os.environ BEFORE the + # SDK constructor captured the env. + assert ( + observed_env_during_init["HF_TOKEN"] is None + ), "HF_TOKEN visible to SDK constructor -- env scrub failed" + assert observed_env_during_init["GH_TOKEN"] is None + assert observed_env_during_init["WANDB_API_KEY"] is None + # Safe-listed keys must survive. + assert observed_env_during_init["PATH"] == "/usr/bin" + assert observed_env_during_init["CODEX_HOME"] == "/home/u/.codex" + # And the wrapper must restore them after exit. + assert os.environ.get("HF_TOKEN") == "should_be_gone" + assert os.environ.get("GH_TOKEN") == "should_be_gone" + + def test_coerce_text_drops_non_answer_event_types(self): + """Tool / command / plan deltas have their own `delta` fields + 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 + + # Allowed answer-bearing event types contribute text. + assert _coerce_text({"type": "message.delta", "delta": "hello"}) == "hello" + assert _coerce_text({"type": "completed", "text": "done"}) == "done" + assert _coerce_text({"type": "text_delta", "delta": "x"}) == "x" + + # Non-answer dict event types are silenced. + for ev_type in ( + "command.delta", + "command_output", + "file_write.delta", + "tool_call.delta", + "plan.update", + "exec.stdout", + "exec.stderr", + "patch.apply", + "thread.tool_call", + "agent_reasoning", + ): + payload = {"type": ev_type, "delta": "this should NOT leak"} + assert _coerce_text(payload) == "", ( + 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" + + def test_authenticated_yes_wording_is_detected(self): + """An `Authenticated: Yes` line (a wording the CLI ships in + some locales / versions) must be parsed as logged-in. + """ + from core.inference import codex_availability as av + + async def _fake_run_cli(args, **kw): + return (0, "Authenticated: Yes\nuser@example.com", "") + + orig = av._run_cli + av._run_cli = _fake_run_cli # type: ignore[assignment] + try: + assert asyncio.run(av._detect_logged_in()) is True + finally: + av._run_cli = orig # type: ignore[assignment] + + def test_codex_openai_api_key_overrides_openai_provider_key(self, monkeypatch): + """Studio's `OPENAI_API_KEY` must NOT reach codex -- but the + codex-specific `CODEX_OPENAI_API_KEY` MUST be forwarded so + users can deliberately wire a key into Codex. + """ + from core.inference.codex_availability import _codex_subprocess_env + + monkeypatch.setenv("OPENAI_API_KEY", "belongs_to_openai_provider") + monkeypatch.setenv("CODEX_OPENAI_API_KEY", "explicit_codex_key") + + env = _codex_subprocess_env() + assert ( + "OPENAI_API_KEY" not in env + ), "OpenAI provider key leaked into codex subprocess env" + assert env.get("CODEX_OPENAI_API_KEY") == "explicit_codex_key" + + def test_thread_start_uses_safe_approval_and_sandbox(self, monkeypatch): + """When the SDK exposes ApprovalMode + SandboxMode, the + provider MUST pin approval to `deny_all` and sandbox to + `read_only`. The upstream SDK default + (`auto_review` approvals, unspecified sandbox) would let the + model auto-execute commands and write files on the server. + """ + 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"]) + + # Drop a fake openai_codex with ApprovalMode + SandboxMode enums. + import importlib.util as _iu + + fake_mod = types.ModuleType("openai_codex") + fake_mod.AsyncCodex = _Async # type: ignore[attr-defined] + fake_mod.ApprovalMode = types.SimpleNamespace( # type: ignore[attr-defined] + deny_all = "DENY_ALL_SENTINEL", + auto_review = "AUTO_REVIEW_SENTINEL", + ) + fake_mod.SandboxMode = types.SimpleNamespace( # type: ignore[attr-defined] + read_only = "READ_ONLY_SENTINEL", + workspace_write = "WS_WRITE_SENTINEL", + danger_full_access = "DANGER_SENTINEL", + ) + monkeypatch.setitem(sys.modules, "openai_codex", fake_mod) + real_find_spec = _iu.find_spec + monkeypatch.setattr( + "importlib.util.find_spec", + lambda n, *a, **kw: ( + types.SimpleNamespace() + if n in ("openai_codex", "codex_app_server") + else real_find_spec(n, *a, **kw) + ), + ) + + from core.inference.codex_provider import stream_codex + + 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_SENTINEL" + ), f"approval_mode not pinned to deny_all: {kw}" + assert ( + 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_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: + 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 ( + 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(): + 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 under override" + kw = seen_kwargs[0] + assert "approval_mode" not in kw + assert "sandbox" not in kw + 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_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. + 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 + agent message must surface that text. Without the fallback the + Chat Completions reply would be empty even though Codex + produced a complete answer. + """ + + class _CompletedEvent: + payload = { + "type": "item.completed", + "item": { + "root": { + "type": "agentMessage", + "text": "final answer from completion", + }, + }, + } + + class _Turn: + async def stream(self): + yield _CompletedEvent() + + class _ThreadEmptyDeltas: + def turn(self, prompt): + return _Turn() + + async def run(self, prompt): + raise AssertionError( + "must not fall through to buffered run() when " + "the stream completes successfully" + ) + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _ThreadEmptyDeltas() + + _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 ( + "final answer from completion" in body + ), f"agent message text from completion event was dropped; body={body!r}" + + def test_synthesis_also_pins_safety_kwargs(self, monkeypatch): + """The synthesis turn that unifies parallel fan-out outputs + must use the same safety pins -- a fan-out tab could otherwise + sneak an unsafe approval into the final synthesis prompt. + """ + seen_kwargs: list[dict] = [] + + class _SynThread: + async def run(self, prompt): + return "synth ok" + + 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 _SynThread() + + import importlib.util as _iu + + fake_mod = types.ModuleType("openai_codex") + fake_mod.AsyncCodex = _Async # type: ignore[attr-defined] + fake_mod.ApprovalMode = types.SimpleNamespace( # type: ignore[attr-defined] + deny_all = "DENY_ALL_SENTINEL", + ) + fake_mod.SandboxMode = types.SimpleNamespace( # type: ignore[attr-defined] + read_only = "READ_ONLY_SENTINEL", + ) + monkeypatch.setitem(sys.modules, "openai_codex", fake_mod) + real_find_spec = _iu.find_spec + monkeypatch.setattr( + "importlib.util.find_spec", + lambda n, *a, **kw: ( + types.SimpleNamespace() + if n in ("openai_codex", "codex_app_server") + else real_find_spec(n, *a, **kw) + ), + ) + + from core.inference.codex_provider import _run_codex_synthesis + + asyncio.run( + _run_codex_synthesis( + model = "gpt-5.5", + system = "Always answer in Spanish.", + prompt = "What is the capital of France?", + tab_outputs = ["Paris", "Paris."], + ) + ) + assert seen_kwargs, "synthesis thread_start never called" + kw = seen_kwargs[0] + assert kw.get("approval_mode") == "DENY_ALL_SENTINEL" + assert kw.get("sandbox") == "READ_ONLY_SENTINEL" + + +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 + + +# ── 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)") + + +# ── 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}" + + +# ── 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"] + + +# ── 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 + ) + + +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/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 diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index b1522dd382..435e8be896 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(), 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 a74a4a2dea..28d3b014ba 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, @@ -55,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, @@ -1337,10 +1347,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { clearSelectedImageEditReference(); throw new Error("Connection not found."); } - // Local providers and custom Gemini bases allow an empty key. + // Local providers, custom Gemini bases, and Codex (local CLI / SDK) all allow an empty key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; + const externalProviderIsCodex = externalProvider + ? isCodexProviderType(externalProvider.providerType) + : false; const externalProviderIsGeminiCustomBase = Boolean( externalProvider && externalProvider.providerType === "gemini" && @@ -1350,10 +1363,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { isExternalRequest && !externalApiKey && !externalProviderIsCustom && + !externalProviderIsCodex && !externalProviderIsGeminiCustomBase ) { 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.", }); clearSelectedImageEditReference(); throw new Error("Missing connection API key."); @@ -1702,9 +1716,85 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let cumulativeText = ""; let reasoningStartAt: number | null = null; let reasoningDuration = 0; - // True while wrapping a `delta.reasoning_content` stream in - // ... for parseAssistantContent. Lives outside - // the SSE loop because the close tag fires when content arrives. + // 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. + // 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 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; + } + } + + // 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 + // 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 { + 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) + // and DeepSeek's reasoner stream their thinking as a separate + // `reasoning_content` field on the chat-completion delta — not as + // `content`, not as a structured part. We wrap those chunks with + // inline `...` so the existing parseAssistantContent + // lifts them into the reasoning panel the same way it does for + // local Harmony models. State has to live outside the SSE loop + // because the close tag fires when the next chunk carries content + // (or when the stream ends). let reasoningContentOpen = false; // Tool call parts, cumulative; result lands on tool_end. const toolCallParts: ToolCallMessagePart[] = []; @@ -2066,6 +2156,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, + ), + } + : {}), }; } @@ -2146,7 +2249,82 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { - // Persist container_id onto the thread (OpenAI / Anthropic). + // 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_")) { + 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); + 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 (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 { + content: [...toolCallParts, ...codexParts], + }; + continue; + } + // OpenAI shell-tool container persistence — see + // ThreadRecord.openaiCodeExecContainerId. The backend + // emits these synthetic events on the OpenAI Responses + // SSE stream after capturing the container_id from a + // response, or detecting an expired-container error. if (toolEvent.type === "container_ready") { const newContainerId = toolEvent.container_id as | string @@ -2364,10 +2542,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }; } } - // Cumulative yield. orderAssistantContent puts search/ - // code before text and generated images after. + // Cumulative yield so tool UI updates. orderAssistantContent + // puts search / code before text and generated images after. + // renderFullContent() preserves any Codex per-tab text from + // earlier _toolEvent frames; pinTextThoughtSignature attaches + // Gemini thoughtSignature onto the final text part. const textParts = pinTextThoughtSignature( - parseAssistantContent(cumulativeText), + parseAssistantContent(renderFullContent()), ); yield { content: orderAssistantContent(textParts), @@ -2628,8 +2809,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { "", ); } + // renderFullContent() preserves any Codex per-tab text the + // fan-out branch accumulated into codexTabBuffers; + // pinTextThoughtSignature attaches Gemini thoughtSignature. const parts = pinTextThoughtSignature( - parseAssistantContent(cumulativeText), + parseAssistantContent(renderFullContent()), ); if ( @@ -2746,8 +2930,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [ + // 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. + // pinTextThoughtSignature attaches Gemini thoughtSignature. ...orderAssistantContent( - pinTextThoughtSignature(parseAssistantContent(cumulativeText)), + pinTextThoughtSignature(parseAssistantContent(renderFullContent())), ), ...sourceParts, ...documentCitationParts, 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..50d71cc009 --- /dev/null +++ b/studio/frontend/src/features/chat/api/codex-api.ts @@ -0,0 +1,150 @@ +// 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 { + // `device_code` is the one-time code the verification page asks for + // (separate from the URL); the backend extracts it from the CLI + // stdout via a dedicated regex and emits it as a structured event. + type: "device_url" | "device_code" | "log" | "error" | "done"; + url?: string; + code?: string; + line?: string; + message?: string; + ok?: boolean; + 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/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 3ffb3a1441..64cb177804 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -49,14 +49,19 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { + CODEX_PROVIDER_TYPE, CUSTOM_BACKEND_PROVIDER_TYPE, CUSTOM_PROVIDER_PRESETS, allowsManualModelIdsWithCatalog, + CODEX_DEFAULT_PARALLEL_CALLS, + CODEX_MAX_PARALLEL_CALLS, + clampCodexParallelCalls, customProviderBaseUrlPlaceholder, customProviderDisplayName, customProviderModelIdsPlaceholder, customPresetSkipsApiKeyField, getExternalProviderApiKey, + isCodexProviderType, isCustomProviderType, LEGACY_CUSTOM_PROVIDER_TYPE, removeExternalProviderApiKey, @@ -66,6 +71,8 @@ import { supportsRemoteModelCatalog, toExternalBackendProviderType, } from "./external-providers"; +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) */ @@ -221,6 +228,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); @@ -231,6 +252,14 @@ export function ChatProvidersSettings({ const [modelSearchQuery, setModelSearchQuery] = useState(""); const [customProviderName, setCustomProviderName] = useState("Custom"); const [isReasoningModel, setIsReasoningModel] = useState(false); + // Per-Codex-connection fan-out width. Stored on the provider so + // restoring it after a refresh / page reload does not collapse back + // to single-call. Clamped to [1, MAX] at every write because the + // input is a plain `` and a hand-edited + // localStorage entry could otherwise overflow. + const [codexParallelCalls, setCodexParallelCalls] = useState( + CODEX_DEFAULT_PARALLEL_CALLS, + ); const reduceMotion = useReducedMotion(); const connectionsEnabled = useExternalProvidersStore( (s) => s.connectionsEnabled, @@ -239,9 +268,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 +314,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 || @@ -330,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(""); @@ -354,12 +399,42 @@ 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; + setCodexStatus(codexStatusRaw); + 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 ( @@ -408,6 +483,15 @@ export function ChatProvidersSettings({ isReasoningModel: supportsProviderReasoningToggle(uiProviderType) ? existing?.isReasoningModel === true : undefined, + // Preserve the per-Codex fan-out width on sync. The + // backend provider row does not carry it (it lives in + // localStorage only), so we read it from `existing` and + // skip the field entirely for non-Codex providers. + codexParallelCalls: isCodexProviderType(uiProviderType) + ? clampCodexParallelCalls( + existing?.codexParallelCalls ?? CODEX_DEFAULT_PARALLEL_CALLS, + ) + : undefined, createdAt: existing?.createdAt ?? createdAt, updatedAt, }; @@ -465,13 +549,27 @@ export function ChatProvidersSettings({ setModelSearchQuery(""); setCustomProviderName(customProviderDisplayName(providerType)); setIsReasoningModel(false); + setCodexParallelCalls(CODEX_DEFAULT_PARALLEL_CALLS); } function openAddProvider() { 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"); } @@ -553,7 +651,7 @@ export function ChatProvidersSettings({ ); return; } - if (!isCustomProvider && !apiKey.trim()) { + if (!providerSkipsApiKey && !apiKey.trim()) { toast.error("Add an API key first."); return; } @@ -625,7 +723,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; } @@ -702,6 +800,11 @@ export function ChatProvidersSettings({ isReasoningModel: supportsProviderReasoningToggle(uiProviderType) ? isReasoningModel : undefined, + // Persist the fan-out width on the Codex provider only; other + // providers must not carry the field through normalization. + codexParallelCalls: isCodexProviderType(uiProviderType) + ? clampCodexParallelCalls(codexParallelCalls) + : undefined, createdAt, updatedAt, }; @@ -734,7 +837,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; } @@ -820,6 +925,12 @@ export function ChatProvidersSettings({ ) ? isReasoningModel : undefined, + // Carry through the fan-out width for Codex; clear it on + // every other provider type so a left-over value cannot + // hitchhike on the persisted record. + codexParallelCalls: isCodexProviderType(existing.providerType) + ? clampCodexParallelCalls(codexParallelCalls) + : undefined, updatedAt, } : provider, @@ -852,6 +963,13 @@ export function ChatProvidersSettings({ ? provider.isReasoningModel === true : false, ); + setCodexParallelCalls( + isCodexProviderType(provider.providerType) + ? clampCodexParallelCalls( + provider.codexParallelCalls ?? CODEX_DEFAULT_PARALLEL_CALLS, + ) + : CODEX_DEFAULT_PARALLEL_CALLS, + ); if ( isCustomProviderType(provider.providerType) && !supportsRemoteModelCatalog(provider.providerType) @@ -924,6 +1042,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)) { @@ -1075,6 +1217,66 @@ export function ChatProvidersSettings({ + {isCodexProvider && + codexStatus?.installed && + !codexStatus.logged_in ? ( +
+
+ +

+ Authenticate the local Codex CLI before chatting. +

+
+ { + void refreshCodexStatus(); + }} + /> +
+ ) : 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 ? (
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..9667c8f382 --- /dev/null +++ b/studio/frontend/src/features/chat/components/codex-login-button.tsx @@ -0,0 +1,157 @@ +// 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, useEffect, 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); + const [deviceCode, setDeviceCode] = useState(null); + // Track the active stream's abort controller so a second click + // (or an unmount) tears the SSE reader down cleanly. Without this + // the long-running login subprocess would keep streaming into a + // detached component. + const abortRef = useRef(null); + + const startLogin = useCallback(async () => { + if (busy) return; + setBusy(true); + setError(null); + setLogs([]); + setDeviceUrl(null); + setDeviceCode(null); + const controller = new AbortController(); + abortRef.current?.abort(); + abortRef.current = controller; + // Track the specific backend error inside the closure so the + // generic fallback message does not overwrite it: setError is + // async and reading `error` after `setError(event.message)` would + // still see the stale pre-stream value. + let lastStreamError: string | null = null; + try { + let lastOk: boolean | undefined; + for await (const event of streamCodexDeviceLogin( + controller.signal, + ) as AsyncGenerator) { + if (event.type === "device_url" && event.url) { + setDeviceUrl(event.url); + // 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) { + setLogs((prev) => [...prev, event.line as string]); + } else if (event.type === "error" && event.message) { + lastStreamError = event.message; + setError(event.message); + } else if (event.type === "done") { + lastOk = event.ok; + } + } + if (lastOk) { + onLoggedIn?.(); + } else if (!lastStreamError) { + 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, 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 ( +
+ + {deviceUrl && ( +
+ +

+ Or copy: {deviceUrl} +

+
+ )} + {deviceCode && ( +

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

+ )} + {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 b10253c280..c3c7d34be5 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; } @@ -86,11 +93,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([