ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: add Codex SDK as a chat provider with parallel-calls fan-out

Browse source 
Wires the OpenAI Codex CLI / Python SDK (codex_app_server) into Studio
as a new chat provider type. Hosts that don't have the CLI or the SDK
installed never see the entry; on logged-out hosts the provider config
dialog renders a device-auth Sign-in button that surfaces the
verification URL and streams CLI progress back over SSE.

Backend
- new core/inference/codex_availability.py probes the CLI + SDK and
  reports {installed, logged_in, version, supported_models}; it never
  imports codex_app_server at module top level so the rest of the
  backend keeps starting cleanly on hosts that don't have the SDK.
- new core/inference/codex_provider.py wraps AsyncCodex and translates
  Codex events into OpenAI chat-completion chunks. Supports the
  thread.run_streaming path with a non-streaming fallback for older
  SDK revs.
- parallel_calls > 1 fans the turn out across N tasks (capped at 20)
  via asyncio.gather and emits codex_tab_open / codex_tab_chunk /
  codex_tab_close tool-events per attempt plus a final codex_gather
  synthesis event. A separate standalone Codex call produces the
  unified answer.
- new routes/codex.py exposes GET /api/codex/status and POST
  /api/codex/login. The login route shells out to
  codex auth login --device-auth and streams events; the first event
  carries the verification URL so the frontend can window.open it.
- ChatCompletionRequest gains a parallel_calls field bounded [1, 20]
  by pydantic. The codex registry entry stays hidden by default; the
  /api/codex/status probe is the authoritative gate.
- routes/inference.py dispatches provider_type=codex through the
  local CLI/SDK pipeline instead of the standard HTTP client, with
  graceful error surfacing for CodexUnavailableError.

Frontend
- new api/codex-api.ts exposes fetchCodexStatus() and an async
  generator streamCodexDeviceLogin() that drives the SSE stream and
  yields parsed events.
- new components/codex-parallel-tabs.tsx renders the tabbed parallel-
  calls UI with a Synthesis tab highlighted once the codex_gather
  event arrives. Pure reducer keeps the state transitions unit-
  testable.
- new components/codex-login-button.tsx posts to /api/codex/login,
  opens the verification URL in a new tab via window.open, and shows
  the streamed CLI log as it lands.
- external-providers.ts exports CODEX_PROVIDER_TYPE,
  CODEX_MAX_PARALLEL_CALLS, isCodexProviderType, and
  clampCodexParallelCalls. Codex is marked text-only so the composer
  hides image-attach affordances when selected.

Tests
- tests/test_codex_provider.py (14 cases) covers the availability
  probe across the four install / login states, the streaming +
  parallel-calls translation against a fake codex_app_server module
  injected into sys.modules, the [1, 20] pydantic clamp, the
  CodexUnavailableError surfacing path, and the parallel_calls=1
  single-call shape (no tab tool-events).
This commit is contained in:
Daniel Han 2026-05-22 16:46:00 +00:00 • committed by danielhanchen
commit cbc3c43655
13 changed files with 2138 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

221
studio/backend/core/inference/codex_availability.py Normal file
View file

@ -0,0 +1,221 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Codex CLI / SDK availability probe.
This module never imports `codex_app_server` at module top level. The
SDK is optional and not pinned in pyproject.toml -- if it's installed
locally, we use it; if it isn't, the probe simply returns
``installed=False`` and the provider stays hidden in the frontend.
The frontend calls ``GET /api/codex/status`` at startup to decide
whether to surface the "codex" entry in the provider picker. Three
states matter:
* ``installed=False`` -- either the CLI is missing OR the SDK
(``codex_app_server``) is not importable. The picker hides the
entry entirely.
* ``installed=True, logged_in=False`` -- everything resolves on the
Python side but ``codex auth status`` (or equivalent) reports no
active credentials. The provider config dialog shows a
``Sign in to Codex`` button instead of the regular API-key field.
* ``installed=True, logged_in=True`` -- ready to use; the picker
shows the regular model dropdown.
Detection is best-effort and cheap: we shell out to ``which codex``
plus ``codex --version`` for the CLI and use ``importlib.util.find_spec``
for the SDK. No long-running CLI commands are invoked here so the
status endpoint is safe to poll on every page load.
"""
from __future__ import annotations
import asyncio
import importlib.util
import os
import shutil
from typing import Any, Optional
import structlog
logger = structlog.get_logger(__name__)
# Default catalog of models surfaced in the picker when the CLI is
# present but doesn't advertise a list. The SDK accepts arbitrary model
# ids; this is purely a sensible default.
_DEFAULT_SUPPORTED_MODELS: tuple[str, ...] = (
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.5",
"o3",
)
def _which_codex() -> Optional[str]:
"""Return absolute path to the ``codex`` CLI, or None if missing.
Uses :func:`shutil.which` so the lookup honours ``PATH`` exactly
the way the user's shell would. Returns ``None`` on any failure
so callers can treat "missing" and "broken probe" the same way.
"""
try:
return shutil.which("codex")
except Exception as exc:
# shutil.which itself is documented as raising only on
# genuinely unusual conditions, but a hardened wrapper costs
# nothing and keeps the status endpoint from 500'ing.
logger.warning("codex_availability.which_failed", error = str(exc))
return None
def _sdk_importable() -> bool:
"""True iff ``codex_app_server`` is importable in this interpreter.
We deliberately use :func:`importlib.util.find_spec` instead of an
``import codex_app_server`` so the import never actually runs --
that keeps the cost negligible and avoids the SDK's own side
effects (which include reaching out to the CLI subprocess for an
RPC ping) during a simple availability check.
"""
try:
return importlib.util.find_spec("codex_app_server") is not None
except Exception as exc:
logger.warning("codex_availability.find_spec_failed", error = str(exc))
return False
async def _run_cli(args: list[str], *, timeout: float = 4.0) -> tuple[int, str, str]:
"""Run a short ``codex`` CLI command and return (rc, stdout, stderr).
The probe uses 4s as the wall-clock cap because ``codex --version``
and ``codex auth status`` both return in well under a second on a
healthy install. A longer probe would block the
``/api/codex/status`` route -- and that route fires on every chat
page load, so a tight cap matters.
"""
try:
proc = await asyncio.create_subprocess_exec(
"codex",
*args,
stdout = asyncio.subprocess.PIPE,
stderr = asyncio.subprocess.PIPE,
env = os.environ.copy(),
)
except FileNotFoundError:
return -1, "", "codex binary not on PATH"
except Exception as exc:
logger.warning(
"codex_availability.spawn_failed",
args = args,
error = str(exc),
)
return -1, "", str(exc)
try:
stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout = timeout)
except asyncio.TimeoutError:
proc.kill()
try:
await proc.wait()
except Exception:
pass
return -1, "", f"codex {' '.join(args)} timed out after {timeout:.1f}s"
return (
proc.returncode if proc.returncode is not None else -1,
stdout_b.decode("utf-8", errors = "replace").strip(),
stderr_b.decode("utf-8", errors = "replace").strip(),
)
async def _detect_version() -> Optional[str]:
rc, stdout, stderr = await _run_cli(["--version"])
if rc != 0:
return None
# ``codex --version`` prints something like "codex-cli 0.133.0".
# Surface the whole line so the UI can show the exact build the
# user has installed -- it's useful when troubleshooting.
text = stdout or stderr
return text.split("\n", 1)[0].strip() if text else None
async def _detect_logged_in() -> bool:
"""Best-effort: assume a non-zero ``codex auth status`` rc means
"not logged in". The CLI surface has shifted over releases (some
versions print to stderr, some to stdout, some use "Logged in as
..." vs "Not authenticated"). The return code is the most stable
signal we have, so we lean on it and fall back to substring
inspection only when the rc itself is ambiguous (rc=0 but no
output, or rc=-1 from a timeout / crash).
"""
rc, stdout, stderr = await _run_cli(["auth", "status"])
combined = f"{stdout}\n{stderr}".lower()
if rc == 0:
# Some 0.x releases exit 0 even when the user is logged out.
# If we got any output, look for the obvious "not logged in"
# signals; if there's nothing on either pipe at all, treat
# rc=0 as authenticated (the optimistic default).
if not combined.strip():
return True
if "not authenticated" in combined or "not logged in" in combined:
return False
if "logged in" in combined or "authenticated" in combined:
return True
return True
# rc != 0: lean on substring matching one more time before
# defaulting to False, in case a future CLI uses rc=2 for the
# logged-out state instead of "no auth command exists".
if "logged in" in combined or "authenticated as" in combined:
return True
return False
async def probe_codex_availability() -> dict[str, Any]:
"""Return the full status payload consumed by ``GET /api/codex/status``.
Returns a dict with keys:
* ``installed`` (bool) -- True iff both CLI and SDK are present.
Note this is the gate the frontend uses to surface the provider
at all, so if EITHER is missing the picker hides the entry.
* ``cli_path`` (str | None) -- absolute path to the CLI, or None.
* ``sdk_importable`` (bool) -- the Python SDK is importable.
* ``logged_in`` (bool) -- best-effort auth check; meaningless when
``installed`` is False.
* ``version`` (str | None) -- the ``codex --version`` first line.
* ``supported_models`` (list[str]) -- default model id catalog.
"""
cli_path = _which_codex()
sdk_ok = _sdk_importable()
payload: dict[str, Any] = {
"installed": bool(cli_path) and sdk_ok,
"cli_path": cli_path,
"sdk_importable": sdk_ok,
"logged_in": False,
"version": None,
"supported_models": list(_DEFAULT_SUPPORTED_MODELS),
}
if cli_path:
# version + login probes only matter when the CLI is present;
# they would otherwise just churn subprocess errors. Run them
# in parallel because both are independent CLI invocations.
version, logged_in = await asyncio.gather(
_detect_version(),
_detect_logged_in(),
)
payload["version"] = version
payload["logged_in"] = bool(logged_in)
logger.info(
"codex_availability.probed",
installed = payload["installed"],
sdk_importable = payload["sdk_importable"],
cli_path = payload["cli_path"],
version = payload["version"],
logged_in = payload["logged_in"],
)
return payload

669
studio/backend/core/inference/codex_provider.py Normal file
View file

@ -0,0 +1,669 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Codex SDK chat provider.
This module wraps the ``codex_app_server`` async SDK so a chat request
routed at ``provider_type="codex"`` can dispatch through the user's
local Codex CLI. The contract back to the frontend is the standard
OpenAI Chat Completions SSE shape, exactly like every other entry in
``external_provider.py``.
The two interesting features here:
* ``_stream_codex_single`` -- translate Codex SDK events into
OpenAI-format chunks. We prefer ``thread.run_streaming`` when the
installed SDK exposes it; otherwise we fall back to
``await thread.run(...)`` and emit one big content chunk plus a
usage chunk at the end.
* ``_stream_codex_parallel`` -- the ``parallel_calls`` knob. When > 1,
spawn N async Codex tasks (capped at 20) via ``asyncio.gather`` and
emit per-tab ``_toolEvent`` markers so the frontend can render each
result in its own tab. A final ``codex_gather`` synthesis tab runs a
single Codex call that takes the N outputs and produces a unified
answer.
The SDK is imported lazily (inside the helpers that actually need it)
because the spec calls out that ``codex_app_server`` may not even be
importable on the build host. The availability probe in
``codex_availability.py`` is what the frontend uses to gate the
provider entirely; this module just refuses to run if the import
fails at request time.
"""
from __future__ import annotations
import asyncio
import importlib
import importlib.util
import json
import time
from typing import Any, AsyncGenerator, Optional
import structlog
logger = structlog.get_logger(__name__)
# Hard cap on parallel Codex fan-out. Picked to match the upper bound
# in the request validator -- exceeding this risks the local Codex CLI
# rate-limiting itself or starving the loop.
MAX_PARALLEL_CALLS = 20
class CodexUnavailableError(RuntimeError):
"""Raised when ``codex_app_server`` is not importable at runtime.
The availability probe is supposed to hide the provider before any
request lands here, but we still raise a typed error so the route
layer can translate it into a 503 the user sees instead of an
opaque traceback.
"""
def _import_codex() -> Any:
"""Return the imported ``codex_app_server`` module or raise.
Imported lazily so the rest of the backend keeps starting cleanly
on hosts that don't have the SDK installed. The frontend calls
``GET /api/codex/status`` first and hides the provider when the
spec isn't importable, so this branch is reached only when the
user (a) explicitly forces the provider via a stale stored config
or (b) the install state changes between status probe and chat
submit.
"""
if importlib.util.find_spec("codex_app_server") is None:
raise CodexUnavailableError(
"codex_app_server is not installed on this host. "
"Install the Codex Python SDK or use a different provider."
)
return importlib.import_module("codex_app_server")
def _last_user_prompt(messages: list[dict[str, Any]]) -> str:
"""Extract the most recent user-role message as a plain string.
Codex's ``thread.run`` accepts a string (per the docs note
"plain strings are accepted anywhere a turn input is accepted").
Studio's chat history is a full OpenAI-style messages array, so
we flatten it: walk from the end, find the last ``role=user``
message, and stringify any structured content parts into a
newline-joined block. Multimodal content (images) is described
rather than embedded; Codex SDK input shape is text-first.
This is intentionally conservative -- we don't try to replay the
whole conversation through Codex per turn because the SDK is
designed around a stateful ``thread`` object. The thread itself
holds context across runs; we only need to feed the latest user
turn.
"""
for msg in reversed(messages):
if msg.get("role") != "user":
continue
content = msg.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for entry in content:
if not isinstance(entry, dict):
continue
if entry.get("type") == "text":
parts.append(str(entry.get("text") or ""))
elif entry.get("type") == "image_url":
url = (entry.get("image_url") or {}).get("url") or ""
parts.append(f"[image: {url[:80]}]" if url else "[image]")
elif entry.get("type") == "input_document":
name = entry.get("filename") or "document"
parts.append(f"[document: {name}]")
return "\n".join(p for p in parts if p)
return ""
def _system_prompt(messages: list[dict[str, Any]]) -> str:
"""Concatenate all ``role=system`` messages.
Codex's ``thread_start`` accepts a system prompt for the lifetime
of the thread. We pass any leading system messages so the user's
Studio-side system prompt (chat presets) reaches the Codex side
intact.
"""
parts: list[str] = []
for msg in messages:
if msg.get("role") != "system":
continue
content = msg.get("content")
if isinstance(content, str):
parts.append(content)
elif isinstance(content, list):
for entry in content:
if isinstance(entry, dict) and entry.get("type") == "text":
parts.append(str(entry.get("text") or ""))
return "\n\n".join(p for p in parts if p)
def _chunk_text(completion_id: str, text: str) -> str:
"""OpenAI Chat Completions content chunk."""
payload = {
"id": completion_id,
"object": "chat.completion.chunk",
"choices": [
{
"index": 0,
"delta": {"content": text},
"finish_reason": None,
}
],
}
return f"data: {json.dumps(payload)}"
def _chunk_stop(completion_id: str) -> str:
payload = {
"id": completion_id,
"object": "chat.completion.chunk",
"choices": [
{
"index": 0,
"delta": {},
"finish_reason": "stop",
}
],
}
return f"data: {json.dumps(payload)}"
def _chunk_tool_event(completion_id: str, event: dict[str, Any]) -> str:
"""Synthetic OpenAI-shaped chunk carrying an ``_toolEvent`` payload.
Studio's frontend chat-adapter already understands the
``_toolEvent`` envelope and renders tool cards on the fly. We piggy-
back on the same channel to ship Codex-specific tab markers
(``codex_tab_open`` / ``codex_tab_chunk`` / ``codex_gather``) so the
UI doesn't need a brand-new transport.
"""
payload = {
"id": completion_id,
"object": "chat.completion.chunk",
"choices": [
{
"index": 0,
"delta": {},
"finish_reason": None,
}
],
"_toolEvent": event,
}
return f"data: {json.dumps(payload)}"
def _chunk_usage(
completion_id: str,
prompt_tokens: int,
completion_tokens: int,
) -> str:
payload = {
"id": completion_id,
"object": "chat.completion.chunk",
"choices": [],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
return f"data: {json.dumps(payload)}"
def _coerce_text(payload: Any) -> str:
"""Pull text out of a Codex streaming event or result.
The SDK shape isn't pinned across versions: events expose ``delta``
or ``text`` or ``content`` depending on whether the model is in
plan / answer / tool-use mode. Be defensive -- read whichever
field is present and fall back to ``str()`` so we never crash
while translating.
"""
if payload is None:
return ""
if isinstance(payload, str):
return payload
if isinstance(payload, dict):
for key in ("delta", "text", "content", "message", "final_response"):
if key in payload:
value = _coerce_text(payload[key])
if value:
return value
return ""
if isinstance(payload, list):
return "".join(_coerce_text(item) for item in payload)
text_attr = getattr(payload, "text", None)
if isinstance(text_attr, str):
return text_attr
delta_attr = getattr(payload, "delta", None)
if isinstance(delta_attr, str):
return delta_attr
final_attr = getattr(payload, "final_response", None)
if isinstance(final_attr, str):
return final_attr
return ""
async def _stream_thread_run(
thread: Any,
prompt: str,
) -> AsyncGenerator[str, None]:
"""Yield raw text chunks from a Codex thread.
Prefers ``thread.run_streaming(prompt)`` because that's what the
docs surface for token-by-token delivery. When the installed SDK
doesn't have that helper, fall back to ``await thread.run(prompt)``
and yield the full text once -- this still works end-to-end, just
without streaming feedback in the UI.
"""
run_streaming = getattr(thread, "run_streaming", None)
if run_streaming is not None:
try:
stream_obj = run_streaming(prompt)
# The SDK may return either an async iterator directly or a
# coroutine that resolves to one. Handle both shapes so a
# future SDK rev doesn't silently fall off the streaming
# path.
if asyncio.iscoroutine(stream_obj):
stream_obj = await stream_obj
async for event in stream_obj:
text = _coerce_text(event)
if text:
yield text
return
except Exception as exc:
logger.warning(
"codex_provider.run_streaming_failed_fallback",
error = str(exc),
)
# Intentional fallthrough to the non-streaming path so a
# broken streaming helper doesn't take the whole turn down.
# Non-streaming fallback: await the full TurnResult, emit one chunk.
result = await thread.run(prompt)
text = _coerce_text(result) or getattr(result, "final_response", "") or str(result)
if text:
yield text
async def _stream_codex_single(
model: str,
system: str,
prompt: str,
completion_id: str,
) -> AsyncGenerator[str, None]:
"""Run one Codex turn and emit OpenAI-shaped SSE lines."""
sdk = _import_codex()
async_codex_cls = getattr(sdk, "AsyncCodex", None)
if async_codex_cls is None:
raise CodexUnavailableError(
"codex_app_server is installed but AsyncCodex is missing -- "
"upgrade the SDK."
)
completion_text_chars = 0
async with async_codex_cls() as codex:
# ``thread_start`` accepts a model id; system prompts are
# passed when supported by the SDK rev (older revs ignore the
# extra kwarg). Be tolerant about kwargs that may not exist.
thread_kwargs: dict[str, Any] = {"model": model}
if system:
# Try the canonical kwargs first; the SDK shapes vary
# across versions and we'd rather accept the system prompt
# being dropped than crash on a missing kwarg.
thread_kwargs["system"] = system
try:
thread = await codex.thread_start(**thread_kwargs)
except TypeError:
# Older SDK: only the ``model`` kwarg is accepted. Drop
# extras and retry; the system prompt then lives only in
# the prompt itself (we prepend it below).
thread = await codex.thread_start(model = model)
if system:
prompt = f"{system}\n\n{prompt}"
async for text in _stream_thread_run(thread, prompt):
completion_text_chars += len(text)
yield _chunk_text(completion_id, text)
# Estimate tokens crudely from char counts. The Codex SDK does not
# consistently expose a token breakdown; we surface a usage chunk
# purely so the frontend cost / context display has a non-zero
# value to render. ``int(chars / 4)`` is the standard rough-cut.
yield _chunk_usage(
completion_id,
prompt_tokens = max(1, len(prompt) // 4),
completion_tokens = max(0, completion_text_chars // 4),
)
yield _chunk_stop(completion_id)
async def stream_codex(
messages: list[dict[str, Any]],
model: str,
parallel_calls: int = 1,
) -> AsyncGenerator[str, None]:
"""Top-level entry point used by the inference route.
When ``parallel_calls == 1`` (the default), this delegates to the
single-turn helper. When > 1 (and <= ``MAX_PARALLEL_CALLS``), it
spawns N parallel Codex turns and emits one tab per result plus a
final synthesis tab. The fan-out runs every spawned turn against
the SAME user prompt -- the parallel knob is for sampling N
independent attempts at the same task, which is what the UI tab
strip surfaces.
"""
prompt = _last_user_prompt(messages)
system = _system_prompt(messages)
completion_id = f"chatcmpl-codex-{int(time.time() * 1000)}"
if not prompt:
yield _chunk_text(
completion_id,
"(no user prompt found -- send a message before invoking codex)",
)
yield _chunk_stop(completion_id)
return
clamped = max(1, min(int(parallel_calls or 1), MAX_PARALLEL_CALLS))
if clamped == 1:
async for line in _stream_codex_single(model, system, prompt, completion_id):
yield line
yield "data: [DONE]"
return
async for line in _stream_codex_parallel(
model = model,
system = system,
prompt = prompt,
n = clamped,
completion_id = completion_id,
):
yield line
yield "data: [DONE]"
async def _stream_codex_parallel(
*,
model: str,
system: str,
prompt: str,
n: int,
completion_id: str,
) -> AsyncGenerator[str, None]:
"""Fan out N Codex tasks, emit per-tab chunks, then synthesise.
Each tab id is a 1-based integer. We emit ``codex_tab_open`` up
front for every tab so the UI can paint the tab strip before any
response lands, then route streamed text per tab via
``codex_tab_chunk`` (one event per text delta). The final
``codex_gather`` event carries a synthesis from a separate
standalone Codex call so the user sees both the per-tab raw
outputs AND a unified merged answer.
Bounded concurrency: ``asyncio.gather`` is used with a fresh
``AsyncCodex`` instance per tab to keep the SDK from sharing a
single thread object across coroutines (the SDK is not documented
as concurrency-safe on a single ``Codex`` handle). N is clamped to
``MAX_PARALLEL_CALLS`` by the caller.
"""
queue: asyncio.Queue[Optional[str]] = asyncio.Queue()
# Pre-emit one tab_open per tab so the UI can paint the tabs
# immediately. The frontend creates the tabs lazily on first
# ``codex_tab_open`` event anyway; pre-emitting just gives us the
# familiar tabs-first / content-second render order.
for idx in range(1, n + 1):
yield _chunk_tool_event(
completion_id,
{
"type": "codex_tab_open",
"tab_id": idx,
"query": prompt,
"total_tabs": n,
},
)
async def _worker(tab_id: int) -> str:
"""Run one Codex turn, push every chunk into the queue, and
return the full accumulated text so the synthesis step can
consume it. Errors are surfaced as a ``codex_tab_error``
tool-event so the tab strip shows which lane failed without
aborting the whole fan-out.
"""
collected: list[str] = []
try:
sdk = _import_codex()
async_codex_cls = getattr(sdk, "AsyncCodex")
async with async_codex_cls() as codex:
thread_kwargs: dict[str, Any] = {"model": model}
if system:
thread_kwargs["system"] = system
inner_prompt = prompt
try:
thread = await codex.thread_start(**thread_kwargs)
except TypeError:
thread = await codex.thread_start(model = model)
if system:
inner_prompt = f"{system}\n\n{prompt}"
async for text in _stream_thread_run(thread, inner_prompt):
collected.append(text)
await queue.put(
_chunk_tool_event(
completion_id,
{
"type": "codex_tab_chunk",
"tab_id": tab_id,
"text": text,
},
)
)
except Exception as exc:
logger.warning(
"codex_provider.parallel_tab_failed",
tab_id = tab_id,
error = str(exc),
)
await queue.put(
_chunk_tool_event(
completion_id,
{
"type": "codex_tab_error",
"tab_id": tab_id,
"error": str(exc),
},
)
)
finally:
await queue.put(
_chunk_tool_event(
completion_id,
{
"type": "codex_tab_close",
"tab_id": tab_id,
},
)
)
return "".join(collected)
workers = [asyncio.create_task(_worker(i + 1)) for i in range(n)]
# asyncio.gather returns a _GatheringFuture, not a coroutine, so it
# cannot be passed to create_task. Wrap the await in a small helper
# coroutine so we keep both (a) a handle for awaiting and (b) the
# ability to schedule the drain side-effect that unblocks the
# consumer queue. Stash the results into a list the drain finally
# block reads so per-tab outputs survive any single-worker errors.
per_tab_texts: list[str] = []
async def _await_workers() -> None:
results = await asyncio.gather(*workers, return_exceptions = True)
for r in results:
if isinstance(r, BaseException):
per_tab_texts.append("")
else:
per_tab_texts.append(r)
async def _drain_when_done() -> None:
try:
await _await_workers()
finally:
await queue.put(None)
drain_task = asyncio.create_task(_drain_when_done())
while True:
line = await queue.get()
if line is None:
break
yield line
# Drain finished; per_tab_texts is now populated by the helper
# coroutine above. We already shielded individual errors as
# ``codex_tab_error`` events, so nothing should leak here -- but
# log defensively in case a worker future itself raised.
await drain_task
synthesis_text = await _run_codex_synthesis(
model = model,
prompt = prompt,
tab_outputs = per_tab_texts,
)
yield _chunk_tool_event(
completion_id,
{
"type": "codex_gather",
"summary": synthesis_text,
"tab_count": n,
},
)
# Also emit the synthesis as a visible content chunk so any client
# that ignores the tab tool-events (e.g. a curl user) still sees
# a final unified answer. The tabbed UI re-uses the same payload
# via ``codex_gather`` and hides it from the main content lane to
# avoid duplication.
if synthesis_text:
yield _chunk_text(completion_id, synthesis_text)
yield _chunk_usage(
completion_id,
prompt_tokens = max(1, len(prompt) // 4),
completion_tokens = max(0, len(synthesis_text) // 4),
)
yield _chunk_stop(completion_id)
async def _run_codex_synthesis(
*,
model: str,
prompt: str,
tab_outputs: list[str],
) -> str:
"""Run one extra Codex call that consumes the N per-tab outputs and
returns a unified synthesis. Returns the empty string on failure --
the caller already surfaced the per-tab outputs so an empty
synthesis is recoverable.
"""
if not tab_outputs:
return ""
joined = "\n\n".join(
f"=== Attempt {i + 1} ===\n{text.strip() or '(no output)'}"
for i, text in enumerate(tab_outputs)
)
synthesis_prompt = (
"You are given multiple independent attempts at the same task.\n"
"Your job: synthesise a single best response that takes the "
"strongest parts of each attempt, resolves disagreements, and "
"presents a clear unified answer.\n\n"
f"Original task:\n{prompt}\n\n"
f"Attempts:\n{joined}\n\n"
"Unified answer:"
)
try:
sdk = _import_codex()
async_codex_cls = getattr(sdk, "AsyncCodex")
async with async_codex_cls() as codex:
try:
thread = await codex.thread_start(model = model)
except TypeError:
thread = await codex.thread_start()
result = await thread.run(synthesis_prompt)
return _coerce_text(result) or getattr(result, "final_response", "") or str(result)
except Exception as exc:
logger.warning("codex_provider.synthesis_failed", error = str(exc))
return ""
# ── Device-auth helper ──────────────────────────────────────────────
async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]:
"""Run ``codex auth login --device-auth`` and yield progress events.
Yields dicts (NOT SSE lines) that the route layer wraps in SSE.
First event is always ``{type: "device_url", url: "..."}`` once we
detect a verification URL in the CLI output. Subsequent events
forward CLI stdout/stderr line-by-line as ``{type: "log", line: ...}``
so the UI can show progress. A final ``{type: "done", ok: bool}``
signals completion.
The URL extraction matches the CLI's actual output shape (the CLI
prints something like ``Open https://auth.openai.com/device/...``
on the device-auth path). We scan every line for the first
https:// URL containing ``device``; that has historically been
stable across CLI versions.
"""
import re
cli_path = "codex"
args = ["auth", "login", "--device-auth"]
try:
proc = await asyncio.create_subprocess_exec(
cli_path,
*args,
stdout = asyncio.subprocess.PIPE,
stderr = asyncio.subprocess.STDOUT,
)
except FileNotFoundError:
yield {"type": "error", "message": "codex CLI not found on PATH"}
yield {"type": "done", "ok": False}
return
except Exception as exc:
yield {"type": "error", "message": str(exc)}
yield {"type": "done", "ok": False}
return
url_re = re.compile(r"https?://\S*device\S*", re.IGNORECASE)
url_emitted = False
rc: int = -1
try:
assert proc.stdout is not None
while True:
line_b = await proc.stdout.readline()
if not line_b:
break
line = line_b.decode("utf-8", errors = "replace").rstrip()
if not url_emitted:
match = url_re.search(line)
if match:
yield {"type": "device_url", "url": match.group(0)}
url_emitted = True
yield {"type": "log", "line": line}
finally:
try:
rc = await proc.wait()
except Exception:
rc = -1
yield {"type": "done", "ok": rc == 0, "return_code": rc}

35
studio/backend/core/inference/providers.py
View file

@ -270,6 +270,41 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
),
"hidden": True,
},
"codex": {
"display_name": "OpenAI Codex (local CLI)",
# No remote base_url: Codex dispatches through the local CLI
# via the codex_app_server SDK. Routing skips the standard
# HTTP client entirely in _proxy_to_external_provider and
# hands the request to core.inference.codex_provider instead.
"base_url": "",
"default_models": [
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.5",
"o3",
],
"supports_streaming": True,
"supports_vision": False,
"supports_tool_calling": True,
# No auth header is sent on the wire; the Codex CLI handles
# auth via its own login flow (api key / chatgpt / device).
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
# Codex models are picked from the local CLI catalogue; we
# never call a remote /models endpoint.
"model_list_mode": "curated",
# Hidden from the cross-provider dropdown until the frontend
# has confirmed availability via GET /api/codex/status. The
# chat-providers dialog conditionally surfaces the entry by
# merging the codex row in when status.installed is true.
"hidden": True,
"notes": (
"Dispatches chat turns through the local Codex CLI via "
"the codex_app_server Python SDK. Surfaced only when the "
"CLI and SDK are both installed; sign in with `codex auth "
"login`."
),
},
"openrouter": {
"display_name": "OpenRouter",
"base_url": "https://openrouter.ai/api/v1",

5
studio/backend/main.py
View file

@ -115,6 +115,7 @@ from datetime import datetime
from routes import (
auth_router,
chat_history_router,
codex_router,
data_recipe_router,
datasets_router,
export_router,
@ -522,6 +523,10 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
# standard /v1/chat/completions path.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
# Codex SDK provider. Status probe + device-auth helper live behind a
# dedicated prefix so the frontend can call them without needing a
# provider config row to exist yet.
app.include_router(codex_router, prefix = "/api/codex", tags = ["codex"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])

14
studio/backend/models/inference.py
View file

@ -786,6 +786,20 @@ class ChatCompletionRequest(BaseModel):
"to auto-create."
),
)
parallel_calls: Optional[int] = Field(
None,
ge = 1,
le = 20,
description = (
"[x-unsloth] Codex provider only. When > 1, fan the chat turn "
"out across N parallel Codex calls and synthesise a unified "
"final answer. Each parallel attempt is rendered as its own tab "
"in the chat UI; a final 'Synthesis' tab carries the merged "
"output. Bounded to [1, 20] by pydantic so a runaway value can't "
"saturate the local CLI. Silently ignored on every provider "
"other than `codex`."
),
)
@model_validator(mode = "after")
def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest":

2
studio/backend/routes/__init__.py
View file

@ -16,6 +16,7 @@ from routes.export import router as export_router
from routes.training_history import router as training_history_router
from routes.chat_history import router as chat_history_router
from routes.providers import router as providers_router
from routes.codex import router as codex_router
__all__ = [
"training_router",
@ -29,4 +30,5 @@ __all__ = [
"training_history_router",
"chat_history_router",
"providers_router",
"codex_router",
]

92
studio/backend/routes/codex.py Normal file
View file

@ -0,0 +1,92 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
API routes for the Codex SDK chat provider.
Two endpoints live here:
* ``GET /api/codex/status`` -- the availability probe. The frontend
hits this at chat-page load time and uses ``installed`` to decide
whether to surface the "codex" entry in the provider picker. When
``installed=True`` but ``logged_in=False``, the provider config
dialog shows the "Sign in to Codex" affordance instead of the
regular API-key field.
* ``POST /api/codex/login`` -- the device-auth helper. Spawns the
``codex auth login --device-auth`` CLI command, captures the
verification URL from its output, and streams the rest of the auth
exchange back as SSE so the UI can show progress. The URL appears
in the first SSE event so the frontend can ``window.open`` it before
the user wanders off.
"""
from __future__ import annotations
import json
from typing import AsyncGenerator
import structlog
from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse
from auth.authentication import get_current_subject
from core.inference.codex_availability import probe_codex_availability
from core.inference.codex_provider import stream_codex_device_login
logger = structlog.get_logger(__name__)
router = APIRouter()
@router.get("/status")
async def get_codex_status(
current_subject: str = Depends(get_current_subject),
) -> dict:
"""Return the Codex CLI / SDK availability snapshot.
The frontend gates the provider entry on ``installed`` and gates
the "Sign in to Codex" button on ``logged_in``. Both are
best-effort and cheap to recompute; the route does not cache the
probe because the user can install the CLI / SDK or run
``codex auth login`` between page loads and the picker should pick
that up on the next refresh.
"""
return await probe_codex_availability()
@router.post("/login")
async def codex_device_login(
current_subject: str = Depends(get_current_subject),
) -> StreamingResponse:
"""Stream the ``codex auth login --device-auth`` exchange.
Returns an SSE stream of events:
``data: {"type": "device_url", "url": "https://..."}``
``data: {"type": "log", "line": "..."}`` (zero or more)
``data: {"type": "done", "ok": true}``
The frontend opens the device URL in a new tab via
``window.open(url, "_blank", "noopener,noreferrer")`` as soon as
the first event arrives, then renders the streamed log lines so
the user can see the CLI making progress while they're at the
verification page.
"""
async def _to_sse() -> AsyncGenerator[str, None]:
async for event in stream_codex_device_login():
yield f"data: {json.dumps(event)}\n\n"
# Frontend treats the trailing [DONE] the same way it does for
# chat streams, so we emit it for parity.
yield "data: [DONE]\n\n"
return StreamingResponse(
_to_sse(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)

76
studio/backend/routes/inference.py
View file

@ -1807,10 +1807,13 @@ async def _proxy_to_external_provider(
detail = "Either provider_id or provider_type is required for external provider routing.",
)
# Fall back to registry default base URL
# Fall back to registry default base URL. Codex is the one
# provider with a deliberately empty base_url -- it dispatches
# through the local CLI rather than over HTTP -- so a missing
# base_url is only a 400 for every other provider type.
if not base_url:
base_url = get_base_url(provider_type)
if not base_url:
if not base_url and provider_type != "codex":
raise HTTPException(
status_code = 400,
detail = f"Unknown provider type: {provider_type}",
@ -1845,6 +1848,75 @@ async def _proxy_to_external_provider(
provider_type = provider_type,
)
# Codex provider: dispatch through the local CLI / SDK instead of
# the HTTP client. The SDK is not an OpenAI-compatible HTTP
# endpoint; it's a thread-oriented Python API that wraps the CLI.
# ``stream_codex`` is the single entry point so the parallel-calls
# fan-out and the single-call path share the same SSE shape.
if provider_type == "codex":
from core.inference.codex_provider import (
CodexUnavailableError,
stream_codex,
)
async def _codex_stream():
try:
gen = stream_codex(
messages = chat_messages,
model = model,
parallel_calls = payload.parallel_calls or 1,
)
sent_done = False
async for line in gen:
yield f"{line}\n\n"
if "[DONE]" in line:
sent_done = True
if not sent_done:
yield "data: [DONE]\n\n"
except CodexUnavailableError as exc:
logger.warning("codex_provider.unavailable", error = str(exc))
yield (
"data: "
+ json.dumps(
{
"error": {
"message": str(exc),
"type": "provider_error",
"code": "503",
"provider": "codex",
}
}
)
+ "\n\n"
)
yield "data: [DONE]\n\n"
except Exception as exc:
logger.error("codex_provider.stream_error", error = str(exc))
yield (
"data: "
+ json.dumps(
{
"error": {
"message": f"Codex error: {exc}",
"type": "provider_error",
"code": "502",
"provider": "codex",
}
}
)
+ "\n\n"
)
yield "data: [DONE]\n\n"
return StreamingResponse(
_codex_stream(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
},
)
client = ExternalProviderClient(
provider_type = provider_type,
base_url = base_url,

499
studio/backend/tests/test_codex_provider.py Normal file
View file

@ -0,0 +1,499 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Unit tests for the Codex SDK provider integration.
Covers:
* Availability probe: codex missing, codex present but logged out,
codex present + logged in, plus the empty-output / non-zero rc
edge cases the CLI has shipped over time.
* ``stream_codex`` event translation: a fake codex_app_server module
is dropped into ``sys.modules`` so the production import path runs
without the real SDK installed. Verifies an OpenAI Chat Completions
shape (content chunk, stop chunk, [DONE]).
* Parallel-calls fan-out: ``parallel_calls > 1`` spawns N async tasks
and emits ``codex_tab_open`` / ``codex_tab_chunk`` / ``codex_tab_close``
events plus a final ``codex_gather`` synthesis event.
* Request validator: ``parallel_calls`` is clamped to [1, 20] by
pydantic so a runaway value is rejected with 422 before any Codex
task is spawned.
"""
from __future__ import annotations
import asyncio
import json
import os
import sys
import types
from typing import Any
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
if _backend not in sys.path:
sys.path.insert(0, _backend)
# ── Helpers ─────────────────────────────────────────────────────────
class _FakeStream:
"""Async iterator that yields predetermined string text events.
The Codex SDK's ``thread.run_streaming`` returns an async iterable
of events. ``_stream_thread_run`` converts those into raw text via
``_coerce_text``; passing in plain strings exercises the simplest
coercion path.
"""
def __init__(self, chunks: list[str]):
self._chunks = list(chunks)
self._i = 0
def __aiter__(self):
return self
async def __anext__(self):
if self._i >= len(self._chunks):
raise StopAsyncIteration
text = self._chunks[self._i]
self._i += 1
return text
class _FakeThread:
def __init__(self, chunks: list[str], final: str | None = None):
self._chunks = chunks
self._final = final if final is not None else "".join(chunks)
def run_streaming(self, prompt: str):
# ``run_streaming`` may return either an async iterable or a
# coroutine that resolves to one; cover the direct-return
# shape here, the coroutine shape is covered in a separate
# test below.
return _FakeStream(self._chunks)
async def run(self, prompt: str):
return self._final
class _FakeAsyncCodex:
"""Async-context-manager facade matching codex_app_server.AsyncCodex."""
def __init__(
self,
chunks: list[str] | None = None,
final: str | None = None,
raise_on_start: Exception | None = None,
):
self._chunks = chunks or []
self._final = final
self._raise = raise_on_start
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def thread_start(self, **kwargs):
if self._raise is not None:
raise self._raise
return _FakeThread(self._chunks, self._final)
def _install_fake_codex_sdk(monkeypatch, async_codex_cls):
"""Drop a fake ``codex_app_server`` module into sys.modules so the
production lazy-import path picks it up without the real SDK
being installed.
"""
fake_mod = types.ModuleType("codex_app_server")
fake_mod.AsyncCodex = async_codex_cls # type: ignore[attr-defined]
monkeypatch.setitem(sys.modules, "codex_app_server", fake_mod)
# importlib.util.find_spec walks finders, not sys.modules; patch
# it directly so the lazy-import gate accepts the fake.
import importlib.util as _iu
real_find_spec = _iu.find_spec
def _shim(name: str, *args, **kwargs):
if name == "codex_app_server":
return types.SimpleNamespace()
return real_find_spec(name, *args, **kwargs)
monkeypatch.setattr("importlib.util.find_spec", _shim)
# ── Availability probe ─────────────────────────────────────────────
class TestCodexAvailability:
def test_absent_when_cli_missing(self, monkeypatch):
from core.inference import codex_availability as ca
monkeypatch.setattr(ca, "_which_codex", lambda: None)
monkeypatch.setattr(ca, "_sdk_importable", lambda: False)
payload = asyncio.run(ca.probe_codex_availability())
assert payload["installed"] is False
assert payload["cli_path"] is None
assert payload["sdk_importable"] is False
# supported_models is a sensible default even when nothing is
# installed so the picker has something to render IF the user
# forces the entry on a future status flip.
assert isinstance(payload["supported_models"], list)
assert len(payload["supported_models"]) > 0
def test_present_but_sdk_missing(self, monkeypatch):
from core.inference import codex_availability as ca
monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex")
monkeypatch.setattr(ca, "_sdk_importable", lambda: False)
async def fake_version():
return "codex-cli 0.133.0"
async def fake_logged_in():
return True
monkeypatch.setattr(ca, "_detect_version", fake_version)
monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in)
payload = asyncio.run(ca.probe_codex_availability())
# installed requires BOTH CLI and SDK -- this is the gate the
# frontend uses to decide whether to surface the provider entry
# at all, so missing-SDK means hide.
assert payload["installed"] is False
assert payload["cli_path"] == "/usr/local/bin/codex"
assert payload["sdk_importable"] is False
assert payload["version"] == "codex-cli 0.133.0"
def test_present_and_logged_out(self, monkeypatch):
from core.inference import codex_availability as ca
monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex")
monkeypatch.setattr(ca, "_sdk_importable", lambda: True)
async def fake_version():
return "codex-cli 0.133.0"
async def fake_logged_in():
return False
monkeypatch.setattr(ca, "_detect_version", fake_version)
monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in)
payload = asyncio.run(ca.probe_codex_availability())
assert payload["installed"] is True
assert payload["logged_in"] is False
assert payload["version"] == "codex-cli 0.133.0"
def test_present_and_logged_in(self, monkeypatch):
from core.inference import codex_availability as ca
monkeypatch.setattr(ca, "_which_codex", lambda: "/usr/local/bin/codex")
monkeypatch.setattr(ca, "_sdk_importable", lambda: True)
async def fake_version():
return "codex-cli 0.133.0"
async def fake_logged_in():
return True
monkeypatch.setattr(ca, "_detect_version", fake_version)
monkeypatch.setattr(ca, "_detect_logged_in", fake_logged_in)
payload = asyncio.run(ca.probe_codex_availability())
assert payload["installed"] is True
assert payload["logged_in"] is True
# ── _stream_codex translation ──────────────────────────────────────
def _collect_stream(gen) -> list[str]:
async def run():
out: list[str] = []
async for line in gen:
out.append(line)
return out
return asyncio.run(run())
def _parse_sse_chunks(lines: list[str]) -> list[dict[str, Any]]:
"""Decode SSE ``data: {...}`` lines into the chunk dicts. Skips the
sentinel ``data: [DONE]`` line and anything that isn't valid JSON.
"""
out: list[dict[str, Any]] = []
for raw in lines:
if not raw.startswith("data:"):
continue
body = raw[len("data:") :].strip()
if not body or body == "[DONE]":
continue
try:
out.append(json.loads(body))
except json.JSONDecodeError:
continue
return out
class TestStreamCodexSingle:
def test_streaming_chunks_translate_into_openai_shape(self, monkeypatch):
_install_fake_codex_sdk(
monkeypatch,
lambda: _FakeAsyncCodex(chunks = ["Hello", ", ", "world"]),
)
from core.inference.codex_provider import stream_codex
lines = _collect_stream(
stream_codex(
messages = [{"role": "user", "content": "Say hello in 3 chunks."}],
model = "gpt-5.4",
)
)
chunks = _parse_sse_chunks(lines)
# Three content deltas + one usage chunk + one stop chunk.
content_chunks = [
c
for c in chunks
if c.get("choices")
and isinstance(c["choices"], list)
and c["choices"]
and c["choices"][0].get("delta", {}).get("content")
]
assert [c["choices"][0]["delta"]["content"] for c in content_chunks] == [
"Hello",
", ",
"world",
]
# Usage chunk (OpenAI include_usage shape) is a choices=[] entry
# with a populated usage block.
usage_chunks = [c for c in chunks if c.get("choices") == [] and c.get("usage")]
assert len(usage_chunks) == 1
usage = usage_chunks[0]["usage"]
assert usage["prompt_tokens"] > 0
assert usage["completion_tokens"] >= 0
# Final stop chunk with finish_reason=stop.
stop_chunks = [
c
for c in chunks
if c.get("choices")
and c["choices"]
and c["choices"][0].get("finish_reason") == "stop"
]
assert len(stop_chunks) == 1
# And the trailing [DONE] sentinel.
assert any(line.strip() == "data: [DONE]" for line in lines)
def test_empty_user_prompt_emits_helpful_message(self, monkeypatch):
_install_fake_codex_sdk(monkeypatch, lambda: _FakeAsyncCodex(chunks = []))
from core.inference.codex_provider import stream_codex
lines = _collect_stream(
stream_codex(
messages = [{"role": "system", "content": "you are helpful"}],
model = "gpt-5.4",
)
)
text = "\n".join(lines)
assert "no user prompt" in text.lower()
class TestStreamCodexParallel:
def test_parallel_calls_spawn_tabs_and_synthesise(self, monkeypatch):
# The fake SDK returns the same canned chunks for every spawned
# AsyncCodex instance; we just need to verify the orchestrator
# emits N tab_open events, per-tab chunk events keyed by
# tab_id, and a final codex_gather summary event.
_install_fake_codex_sdk(
monkeypatch,
lambda: _FakeAsyncCodex(
chunks = ["alpha"],
final = "synthesised answer",
),
)
from core.inference.codex_provider import stream_codex
n = 3
lines = _collect_stream(
stream_codex(
messages = [{"role": "user", "content": "Test"}],
model = "gpt-5.4",
parallel_calls = n,
)
)
chunks = _parse_sse_chunks(lines)
tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
tab_opens = [e for e in tool_events if e.get("type") == "codex_tab_open"]
tab_chunks = [e for e in tool_events if e.get("type") == "codex_tab_chunk"]
tab_closes = [e for e in tool_events if e.get("type") == "codex_tab_close"]
gather = [e for e in tool_events if e.get("type") == "codex_gather"]
# Each tab opens once -- the N tabs are pre-emitted so the
# UI can paint the strip before content arrives.
assert len(tab_opens) == n
assert sorted(e["tab_id"] for e in tab_opens) == list(range(1, n + 1))
# Per-tab chunks may interleave in any order but every tab id
# must produce at least one chunk before its close event.
seen_tabs = {e["tab_id"] for e in tab_chunks}
assert seen_tabs == set(range(1, n + 1))
# Each tab emits exactly one close marker.
assert sorted(e["tab_id"] for e in tab_closes) == list(range(1, n + 1))
# Exactly one synthesis event with the unified summary.
assert len(gather) == 1
assert gather[0]["tab_count"] == n
# The summary text comes from the final synthesis Codex call;
# our fake returns "synthesised answer" via .run().
assert "synth" in gather[0]["summary"].lower()
def test_parallel_calls_clamped_to_maximum(self, monkeypatch):
"""Passing parallel_calls=500 must NOT spawn 500 tasks; the
clamp at MAX_PARALLEL_CALLS keeps the local CLI safe.
"""
from core.inference import codex_provider as cp
_install_fake_codex_sdk(
monkeypatch,
lambda: _FakeAsyncCodex(chunks = ["x"], final = "synth"),
)
lines = _collect_stream(
cp.stream_codex(
messages = [{"role": "user", "content": "x"}],
model = "gpt-5.4",
parallel_calls = 500,
)
)
chunks = _parse_sse_chunks(lines)
tab_opens = [
c["_toolEvent"]
for c in chunks
if c.get("_toolEvent", {}).get("type") == "codex_tab_open"
]
assert len(tab_opens) == cp.MAX_PARALLEL_CALLS
def test_parallel_calls_one_takes_single_path(self, monkeypatch):
"""parallel_calls=1 must not emit any tab tool-events -- it's the
regular single-call shape.
"""
_install_fake_codex_sdk(
monkeypatch,
lambda: _FakeAsyncCodex(chunks = ["one"]),
)
from core.inference.codex_provider import stream_codex
lines = _collect_stream(
stream_codex(
messages = [{"role": "user", "content": "hi"}],
model = "gpt-5.4",
parallel_calls = 1,
)
)
chunks = _parse_sse_chunks(lines)
tool_events = [c.get("_toolEvent") for c in chunks if c.get("_toolEvent")]
for event in tool_events:
assert not (event.get("type") or "").startswith("codex_tab")
assert event.get("type") != "codex_gather"
# ── Request validator ──────────────────────────────────────────────
class TestParallelCallsValidator:
def test_request_accepts_valid_range(self):
from models.inference import ChatCompletionRequest
for n in (1, 5, 10, 20):
req = ChatCompletionRequest(
model = "gpt-5.4",
messages = [{"role": "user", "content": "hi"}],
parallel_calls = n,
)
assert req.parallel_calls == n
def test_request_rejects_below_one(self):
from models.inference import ChatCompletionRequest
from pydantic import ValidationError
with pytest.raises(ValidationError):
ChatCompletionRequest(
model = "gpt-5.4",
messages = [{"role": "user", "content": "hi"}],
parallel_calls = 0,
)
def test_request_rejects_above_twenty(self):
from models.inference import ChatCompletionRequest
from pydantic import ValidationError
with pytest.raises(ValidationError):
ChatCompletionRequest(
model = "gpt-5.4",
messages = [{"role": "user", "content": "hi"}],
parallel_calls = 21,
)
def test_request_default_is_none(self):
"""Default = None so the field has no effect on every existing
provider that doesn't read it -- preserves backwards compat.
"""
from models.inference import ChatCompletionRequest
req = ChatCompletionRequest(
model = "gpt-5.4",
messages = [{"role": "user", "content": "hi"}],
)
assert req.parallel_calls is None
# ── Codex unavailable surfacing ────────────────────────────────────
class TestCodexUnavailable:
def test_missing_sdk_raises_typed_error(self, monkeypatch):
# Force find_spec to return None so the lazy import fails.
import importlib.util as _iu
real = _iu.find_spec
def _shim(name, *args, **kwargs):
if name == "codex_app_server":
return None
return real(name, *args, **kwargs)
monkeypatch.setattr("importlib.util.find_spec", _shim)
# Also drop any cached fake from prior tests.
monkeypatch.delitem(sys.modules, "codex_app_server", raising = False)
from core.inference.codex_provider import (
CodexUnavailableError,
stream_codex,
)
with pytest.raises(CodexUnavailableError):
asyncio.run(
_consume_first(
stream_codex(
messages = [{"role": "user", "content": "hi"}],
model = "gpt-5.4",
)
)
)
async def _consume_first(gen):
"""Drive an async generator until it raises or yields its first
value. Used to surface lazy-import errors that fire on the first
SDK touch -- otherwise the generator would swallow them on
``__aiter__`` and the test couldn't see them.
"""
async for _ in gen:
return

146
studio/frontend/src/features/chat/api/codex-api.ts Normal file
View file

@ -0,0 +1,146 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* API helpers for the local Codex SDK provider.
*
* The backend exposes two endpoints under ``/api/codex``:
*
* - ``GET /api/codex/status`` returns ``{installed, logged_in, version,
* cli_path, sdk_importable, supported_models}``. The chat-providers
* dialog consults this BEFORE surfacing the "codex" entry in the
* picker -- if ``installed`` is false the provider stays hidden, if
* ``logged_in`` is false we render a "Sign in to Codex" button in
* place of the API-key field.
*
* - ``POST /api/codex/login`` runs ``codex auth login --device-auth``
* under the hood and streams SSE events. The first event is always
* ``{type: "device_url", url}`` so the UI can window.open it; later
* events forward CLI log lines so the user can watch progress.
*/
import { authFetch } from "@/features/auth";
export interface CodexStatus {
installed: boolean;
logged_in: boolean;
cli_path: string | null;
sdk_importable: boolean;
version: string | null;
supported_models: string[];
}
export interface CodexLoginEvent {
type: "device_url" | "log" | "error" | "done";
url?: string;
line?: string;
message?: string;
ok?: boolean;
return_code?: number;
}
const DEFAULT_STATUS: CodexStatus = {
installed: false,
logged_in: false,
cli_path: null,
sdk_importable: false,
version: null,
supported_models: [],
};
/**
* Fetch the current Codex availability snapshot. Network failures are
* swallowed and reported as ``installed=false`` because every caller
* either uses this to gate UI surfacing (the right answer on error is
* "hide the entry") or kicks off a chat (the right answer on error is
* "fall back to a different provider"). Throwing here would force
* every consumer to wrap the call in a try/catch.
*/
export async function fetchCodexStatus(): Promise<CodexStatus> {
try {
const response = await authFetch("/api/codex/status");
if (!response.ok) {
return DEFAULT_STATUS;
}
const body = (await response.json()) as Partial<CodexStatus>;
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<CodexLoginEvent, void, void> {
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
}
}
}

120
studio/frontend/src/features/chat/components/codex-login-button.tsx Normal file
View file

@ -0,0 +1,120 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* "Sign in to Codex" button + streamed log surface.
*
* Renders in place of the regular API-key field in the provider config
* dialog when ``/api/codex/status`` reports ``logged_in=false``. Click
* fires POST ``/api/codex/login``, which spawns
* ``codex auth login --device-auth`` server-side. The first SSE event
* carries the verification URL -- as soon as it arrives we open it in
* a new tab via ``window.open`` so the user doesn't have to copy-paste
* a long URL out of a log pane.
*
* The button stays mounted while the CLI is exchanging the device
* code: the streamed ``log`` events accumulate into the visible
* progress area until the ``done`` event closes the stream. The
* caller passes ``onLoggedIn`` so the dialog can refetch the status
* probe and flip back into the "ready" state automatically.
*/
import { useCallback, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import {
streamCodexDeviceLogin,
type CodexLoginEvent,
} from "../api/codex-api";
interface Props {
/** Called when the device-auth flow finishes successfully so the
* parent can re-probe ``/api/codex/status`` and switch UI states. */
onLoggedIn?: () => void;
}
export function CodexLoginButton({ onLoggedIn }: Props) {
const [busy, setBusy] = useState(false);
const [error, setError] = useState<string | null>(null);
const [logs, setLogs] = useState<string[]>([]);
const [deviceUrl, setDeviceUrl] = useState<string | null>(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<AbortController | null>(null);
const startLogin = useCallback(async () => {
if (busy) return;
setBusy(true);
setError(null);
setLogs([]);
setDeviceUrl(null);
const controller = new AbortController();
abortRef.current?.abort();
abortRef.current = controller;
try {
let lastOk: boolean | undefined;
for await (const event of streamCodexDeviceLogin(
controller.signal,
) as AsyncGenerator<CodexLoginEvent>) {
if (event.type === "device_url" && event.url) {
setDeviceUrl(event.url);
// Open the verification page eagerly so the user doesn't
// have to copy the URL out of the log surface. ``noopener``
// prevents the auth-tab from controlling the Studio window.
try {
window.open(event.url, "_blank", "noopener,noreferrer");
} catch {
// Ignore -- the URL is still visible in the log.
}
} else if (event.type === "log" && event.line) {
setLogs((prev) => [...prev, event.line as string]);
} else if (event.type === "error" && event.message) {
setError(event.message);
} else if (event.type === "done") {
lastOk = event.ok;
}
}
if (lastOk) {
onLoggedIn?.();
} else if (!error) {
setError("Codex login did not complete -- see log for details.");
}
} catch (exc) {
if ((exc as { name?: string } | null)?.name !== "AbortError") {
setError(String((exc as Error)?.message ?? exc));
}
} finally {
setBusy(false);
}
}, [busy, error, onLoggedIn]);
return (
<div className="space-y-2">
<Button type="button" disabled={busy} onClick={startLogin}>
{busy ? "Signing in to Codex…" : "Sign in to Codex"}
</Button>
{deviceUrl && (
<p className="text-xs text-muted-foreground">
Verification URL:{" "}
<a
href={deviceUrl}
target="_blank"
rel="noopener noreferrer"
className="underline"
>
{deviceUrl}
</a>
</p>
)}
{error && (
<p className="text-xs text-destructive">{error}</p>
)}
{logs.length > 0 && (
<pre className="max-h-40 overflow-auto rounded bg-muted/50 p-2 text-[11px]">
{logs.join("\n")}
</pre>
)}
</div>
);
}

227
studio/frontend/src/features/chat/components/codex-parallel-tabs.tsx Normal file
View file

@ -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<number | "synthesis">("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<number | "synthesis">(() => {
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 (
<div
className={cn(
"codex-parallel-card my-2 rounded-md border bg-muted/30 p-2 text-sm",
)}
>
<button
type="button"
className="flex w-full items-center justify-between gap-2 rounded px-1 py-1 text-left text-xs font-medium text-muted-foreground hover:bg-accent/50"
onClick={() => setCollapsed((v) => !v)}
aria-expanded={!collapsed}
>
<span>
Codex parallel calls
{totalSlots > 0 ? ` (${state.tabs.length}/${totalSlots})` : null}
{state.synthesis ? " — synthesis ready" : ""}
</span>
<span aria-hidden>{collapsed ? "+" : "−"}</span>
</button>
{!collapsed && (
<>
<div className="mt-2 flex flex-wrap gap-1 border-b pb-2">
{state.tabs.map((tab) => (
<button
key={tab.tabId}
type="button"
className={cn(
"rounded-t px-2 py-1 text-xs font-medium",
effectiveActive === tab.tabId
? "bg-background text-foreground"
: "text-muted-foreground hover:bg-accent/50",
tab.error && "text-destructive",
)}
onClick={() => setActiveTab(tab.tabId)}
>
Tab {tab.tabId}
{tab.error ? " (error)" : tab.closed ? "" : " …"}
</button>
))}
{state.synthesis !== null && (
<button
type="button"
className={cn(
"rounded-t px-2 py-1 text-xs font-semibold",
effectiveActive === "synthesis"
? "bg-primary/15 text-primary"
: "text-primary/70 hover:bg-primary/10",
)}
onClick={() => setActiveTab("synthesis")}
>
Synthesis
</button>
)}
</div>
<div className="mt-2 max-h-72 overflow-auto whitespace-pre-wrap rounded bg-background/50 p-2 text-xs">
{effectiveActive === "synthesis"
? state.synthesis || "(waiting for synthesis…)"
: (state.tabs.find((t) => t.tabId === effectiveActive)?.text ||
"(waiting…)")}
</div>
</>
)}
</div>
);
}

34
studio/frontend/src/features/chat/external-providers.ts
View file

@ -79,11 +79,45 @@ export function supportsProviderReasoningToggle(
);
}
/**
* The Codex CLI / SDK provider. Surfaced only when the host has BOTH
* the ``codex`` CLI on PATH and the ``codex_app_server`` Python SDK
* importable -- the backend's ``GET /api/codex/status`` is the
* authoritative gate. We expose the type id here so the rest of the
* frontend can reference it without scattering "codex" string
* literals.
*/
export const CODEX_PROVIDER_TYPE = "codex";
export function isCodexProviderType(
providerType: string | null | undefined,
): boolean {
return providerType === CODEX_PROVIDER_TYPE;
}
/** Hard cap mirrors backend MAX_PARALLEL_CALLS to keep the UI honest. */
export const CODEX_MAX_PARALLEL_CALLS = 20;
export const CODEX_DEFAULT_PARALLEL_CALLS = 1;
export function clampCodexParallelCalls(value: unknown): number {
const n = typeof value === "number" && Number.isFinite(value)
? Math.floor(value)
: CODEX_DEFAULT_PARALLEL_CALLS;
if (n < 1) return 1;
if (n > CODEX_MAX_PARALLEL_CALLS) return CODEX_MAX_PARALLEL_CALLS;
return n;
}
// Known text-only providers on their main chat endpoint.
const NON_VISION_PROVIDER_TYPES = new Set<string>([
"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<string>([