unsloth/studio/backend/core/inference/codex_spoof.py

# 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