From 8389b0d3a71e2dc6114374e1bea7bec3cf801db4 Mon Sep 17 00:00:00 2001 From: tibetyalman <222756182+tibetyalman@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:06:59 +0300 Subject: [PATCH 1/5] feat: add Codex ChatGPT authentication runtime --- README.md | 32 ++ agent/codex_cli.py | 213 ++++++++++ agent/core/codex_models.py | 27 ++ agent/core/codex_runtime.py | 596 ++++++++++++++++++++++++++++ agent/core/llm_params.py | 8 + agent/core/model_switcher.py | 21 +- agent/main.py | 50 ++- tests/unit/test_cli_local_models.py | 20 + tests/unit/test_cli_rendering.py | 72 ++++ tests/unit/test_codex_runtime.py | 145 +++++++ 10 files changed, 1181 insertions(+), 3 deletions(-) create mode 100644 agent/codex_cli.py create mode 100644 agent/core/codex_models.py create mode 100644 agent/core/codex_runtime.py create mode 100644 tests/unit/test_codex_runtime.py diff --git a/README.md b/README.md index e0bf7384..f552b6f3 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,38 @@ model prefixes). Hosted inference is billed to the active Hugging Face user. See below on how to run `ml-intern` with local models. +#### OpenAI authentication through Codex + +The local CLI can use an existing Codex login instead of sending model calls +through Hugging Face Router: + +```bash +codex login +ml-intern --model codex/default +``` + +`codex/default` lets Codex select the current default model for the signed-in +account. To request a specific model available to that account: + +```bash +ml-intern --model codex/ +``` + +This mode starts the local `codex app-server`; ML Intern never reads or copies +Codex's cached credentials. If `codex login` used **Sign in with ChatGPT**, +Codex usage follows that ChatGPT plan's Codex allowance. If Codex was logged in +with an OpenAI API key, standard API billing applies instead. + +Codex supplies the main agent loop and local repository tools. ML Intern's +Hugging Face documentation, papers, datasets, Hub, Jobs, web research, and +optional sandbox tools are exposed to Codex under the `ml_intern` tool +namespace. `HF_TOKEN` is optional for the Codex model itself, but individual +Hub/Jobs tools and `--sandbox-tools` still require Hugging Face authentication. + +OpenAI-authenticated Codex mode currently targets the local CLI. The hosted web +app continues to use Hugging Face OAuth/Router because a local Codex login is a +device credential and must not be forwarded to the server. + #### Local models Local model support uses OpenAI-compatible HTTP endpoints through LiteLLM. The diff --git a/agent/codex_cli.py b/agent/codex_cli.py new file mode 100644 index 00000000..bb7a4e5d --- /dev/null +++ b/agent/codex_cli.py @@ -0,0 +1,213 @@ +"""CLI presentation for the OpenAI-authenticated Codex runtime.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from typing import Any + +from prompt_toolkit import PromptSession + +from agent.config import Config +from agent.core.approval_policy import is_scheduled_operation +from agent.core.codex_runtime import CodexAppServerRuntime, CodexRuntimeError +from agent.core.session import Event +from agent.core.tools import ToolRouter +from agent.utils.terminal_display import ( + get_console, + print_banner, + print_error, + print_markdown, + print_tool_call, + print_tool_log, + print_tool_output, +) + + +def _is_scheduled_job(tool_name: str, arguments: dict[str, Any]) -> bool: + return tool_name == "hf_jobs" and is_scheduled_operation(arguments.get("operation")) + + +def _available_mcp_servers(config: Config, hf_token: str | None) -> dict: + """Skip the default HF OAuth MCP when no local HF identity exists.""" + if hf_token: + return config.mcpServers + + available = {} + for name, server in config.mcpServers.items(): + data = server.model_dump() + url = str(data.get("url") or "") + if "huggingface.co/mcp" in url: + continue + available[name] = server + return available + + +async def run_codex_interactive( + *, + config: Config, + prompt_session: PromptSession, + hf_token: str | None, + hf_user: str | None, + local_mode: bool, +) -> None: + """Run the interactive ML Intern CLI on the authenticated Codex runtime.""" + console = get_console() + print_banner( + model=config.model_name, + hf_user=hf_user, + tool_runtime="local filesystem" if local_mode else "HF sandbox", + ) + + streamed = [False] + + async def on_delta(delta: str) -> None: + streamed[0] = True + console.file.write(delta) + console.file.flush() + + async def on_tool( + name: str, + arguments: dict[str, Any], + output: str | None, + success: bool | None, + ) -> None: + if output is None: + print_tool_call(name, json.dumps(arguments)[:120]) + else: + print_tool_output(output, bool(success), truncate=True) + + async def on_event(event: Event) -> None: + if event.event_type == "tool_log" and event.data: + print_tool_log( + str(event.data.get("tool") or ""), + str(event.data.get("log") or ""), + ) + + async def approve_tool(name: str, arguments: dict[str, Any]) -> bool: + console.print(f"\n[bold yellow]Approval required:[/bold yellow] {name}") + console.print_json(data=arguments) + answer = await prompt_session.prompt_async("Approve this tool call? [y/N] ") + return answer.strip().lower() in {"y", "yes"} + + tool_router = ToolRouter( + _available_mcp_servers(config, hf_token), + hf_token=hf_token, + local_mode=local_mode, + ) + + try: + async with CodexAppServerRuntime( + config=config, + tool_router=tool_router, + hf_token=hf_token, + local_mode=local_mode, + cwd=Path.cwd(), + autonomous_mode=False, + approve_tool=approve_tool, + on_delta=on_delta, + on_tool=on_tool, + on_event=on_event, + ) as runtime: + console.print(f"[dim]{runtime.auth_status}[/dim]") + while True: + try: + user_input = await prompt_session.prompt_async("\nYou: ") + except (EOFError, KeyboardInterrupt): + break + + stripped = user_input.strip() + if not stripped: + continue + if stripped.lower() in {"exit", "quit", "/quit", "/exit"}: + break + if stripped == "/status": + console.print(f"[bold]Model:[/bold] {config.model_name}") + console.print(f"[bold]Auth:[/bold] {runtime.auth_status}") + continue + if stripped == "/new": + await runtime.new_thread() + console.print("[green]Started a new Codex conversation.[/green]") + continue + if stripped.startswith("/model"): + console.print( + "[dim]Codex runtime models are selected at startup. " + "Restart with `ml-intern --model codex/default` or " + "`ml-intern --model codex/`.[/dim]" + ) + continue + + streamed[0] = False + try: + final_text = await runtime.run_turn(stripped, stream=True) + except KeyboardInterrupt: + await runtime.interrupt() + console.print("\n[yellow]Interrupted.[/yellow]") + continue + if streamed[0]: + console.file.write("\n") + console.file.flush() + elif final_text: + await print_markdown(final_text, instant=True) + except CodexRuntimeError as exc: + print_error(str(exc)) + finally: + console.print("\n[dim]Bye.[/dim]\n") + + +async def run_codex_headless( + prompt: str, + *, + config: Config, + hf_token: str | None, + local_mode: bool, + stream: bool, +) -> None: + """Run one prompt through Codex and exit.""" + streamed = [False] + + async def on_delta(delta: str) -> None: + streamed[0] = True + sys.stdout.write(delta) + sys.stdout.flush() + + async def on_tool( + name: str, + arguments: dict[str, Any], + output: str | None, + success: bool | None, + ) -> None: + if output is None: + print_tool_call(name, json.dumps(arguments)[:120]) + else: + print_tool_output(output, bool(success), truncate=True) + + async def approve_tool(name: str, arguments: dict[str, Any]) -> bool: + # Match the existing headless policy: scheduled Jobs never receive + # automatic approval because they can create recurring spend. + return not _is_scheduled_job(name, arguments) + + tool_router = ToolRouter( + _available_mcp_servers(config, hf_token), + hf_token=hf_token, + local_mode=local_mode, + ) + async with CodexAppServerRuntime( + config=config, + tool_router=tool_router, + hf_token=hf_token, + local_mode=local_mode, + cwd=Path.cwd(), + autonomous_mode=True, + approve_tool=approve_tool, + on_delta=on_delta, + on_tool=on_tool, + ) as runtime: + print(f"Codex auth: {runtime.auth_status}", file=sys.stderr) + final_text = await runtime.run_turn(prompt, stream=stream) + if streamed[0]: + sys.stdout.write("\n") + sys.stdout.flush() + elif final_text: + await print_markdown(final_text, instant=True) diff --git a/agent/core/codex_models.py b/agent/core/codex_models.py new file mode 100644 index 00000000..eb140e49 --- /dev/null +++ b/agent/core/codex_models.py @@ -0,0 +1,27 @@ +"""Model-id helpers for the Codex app-server runtime. + +``codex/`` is intentionally separate from the OpenAI-compatible +LiteLLM providers. Codex owns its authentication session and may use either a +ChatGPT subscription or an OpenAI API key, depending on how ``codex login`` was +completed. ML Intern never reads or forwards Codex's cached credentials. +""" + +CODEX_MODEL_PREFIX = "codex/" +CODEX_DEFAULT_MODEL_ID = f"{CODEX_MODEL_PREFIX}default" + + +def is_codex_model_id(model_id: str | None) -> bool: + """Return ``True`` for a well-formed Codex runtime model id.""" + if not model_id or any(char.isspace() for char in model_id): + return False + return model_id.startswith(CODEX_MODEL_PREFIX) and bool( + model_id.removeprefix(CODEX_MODEL_PREFIX) + ) + + +def codex_model_name(model_id: str) -> str | None: + """Return the model passed to Codex, or ``None`` for its current default.""" + if not is_codex_model_id(model_id): + raise ValueError(f"Unsupported Codex model id: {model_id}") + name = model_id.removeprefix(CODEX_MODEL_PREFIX) + return None if name == "default" else name diff --git a/agent/core/codex_runtime.py b/agent/core/codex_runtime.py new file mode 100644 index 00000000..bedf8f93 --- /dev/null +++ b/agent/core/codex_runtime.py @@ -0,0 +1,596 @@ +"""OpenAI-authenticated Codex runtime for ML Intern. + +This module talks to the locally installed ``codex app-server`` over its stdio +JSON-RPC transport. Authentication remains fully owned by Codex. In +particular, ML Intern never reads ``~/.codex/auth.json`` or treats a ChatGPT +session token as an OpenAI API key. + +Codex remains the agent loop in this mode. ML Intern's Hugging Face research, +Hub, Jobs, and sandbox tools are exposed to it as an experimental dynamic-tool +namespace supported by the Codex app-server. +""" + +from __future__ import annotations + +import asyncio +import inspect +import json +import logging +import re +import shutil +from collections import deque +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Any + +from agent.config import Config +from agent.core.agent_loop import _base_needs_approval +from agent.core.codex_models import codex_model_name +from agent.core.session import Event, Session +from agent.core.tools import ToolRouter + +logger = logging.getLogger(__name__) + +CODEX_TOOL_NAMESPACE = "ml_intern" +_DYNAMIC_TOOL_NAME = re.compile(r"^[A-Za-z0-9_-]+$") +_CODEX_BUILTIN_LOCAL_TOOLS = {"bash", "read", "write", "edit"} +_UNSUPPORTED_CODEX_TOOLS = { + # This tool creates a nested LiteLLM research loop using the active model. + # Codex already has its own independent context and can call the underlying + # research tools directly. + "research", +} + +ToolApprovalCallback = Callable[[str, dict[str, Any]], Awaitable[bool] | bool] +DeltaCallback = Callable[[str], Awaitable[None] | None] +ToolCallback = Callable[ + [str, dict[str, Any], str | None, bool | None], + Awaitable[None] | None, +] +EventCallback = Callable[[Event], Awaitable[None] | None] + + +class CodexRuntimeError(RuntimeError): + """Raised when Codex authentication or app-server execution fails.""" + + +async def _call_maybe_async(callback: Callable[..., Any] | None, *args: Any) -> Any: + if callback is None: + return None + result = callback(*args) + if inspect.isawaitable(result): + return await result + return result + + +async def codex_login_status(codex_bin: str = "codex") -> str: + """Return Codex's public login status without reading cached credentials.""" + resolved = shutil.which(codex_bin) + if resolved is None: + raise CodexRuntimeError( + "Codex CLI is not installed. Install it, then run `codex login`." + ) + + process = await asyncio.create_subprocess_exec( + resolved, + "login", + "status", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + text = (stdout or stderr).decode(errors="replace").strip() + if process.returncode != 0: + detail = text or "no active Codex login" + raise CodexRuntimeError( + f"Codex authentication is unavailable: {detail}. Run `codex login`." + ) + return text or "Codex authentication active" + + +def build_dynamic_tool_namespace( + tool_router: ToolRouter, + *, + local_mode: bool, +) -> tuple[dict[str, Any] | None, dict[str, str]]: + """Build the Codex dynamic-tool namespace and its dispatch name map.""" + tools: list[dict[str, Any]] = [] + dispatch: dict[str, str] = {} + + skipped = set(_UNSUPPORTED_CODEX_TOOLS) + if local_mode: + # Codex already provides sandboxed local shell and file tools. Keeping + # a second copy would create ambiguous or reserved tool names. + skipped.update(_CODEX_BUILTIN_LOCAL_TOOLS) + + for tool in tool_router.tools.values(): + if tool.name in skipped or not _DYNAMIC_TOOL_NAME.fullmatch(tool.name): + continue + tools.append( + { + "type": "function", + "name": tool.name, + "description": tool.description, + "inputSchema": tool.parameters, + } + ) + dispatch[tool.name] = tool.name + + if not tools: + return None, dispatch + + return ( + { + "type": "namespace", + "name": CODEX_TOOL_NAMESPACE, + "description": ( + "ML Intern tools for Hugging Face documentation, papers, " + "datasets, Hub repositories, Jobs, web research, notifications, " + "planning, and optional remote sandbox execution." + ), + "tools": tools, + }, + dispatch, + ) + + +class CodexAppServerRuntime: + """Manage one authenticated Codex app-server process and thread.""" + + def __init__( + self, + *, + config: Config, + tool_router: ToolRouter, + hf_token: str | None, + local_mode: bool, + cwd: str | Path, + autonomous_mode: bool, + approve_tool: ToolApprovalCallback | None = None, + on_delta: DeltaCallback | None = None, + on_tool: ToolCallback | None = None, + on_event: EventCallback | None = None, + codex_bin: str = "codex", + ) -> None: + self.config = config + self.tool_router = tool_router + self.hf_token = hf_token + self.local_mode = local_mode + self.cwd = str(Path(cwd).resolve()) + self.autonomous_mode = autonomous_mode + self.approve_tool = approve_tool + self.on_delta = on_delta + self.on_tool = on_tool + self.on_event = on_event + self.codex_bin = codex_bin + + self.auth_status: str | None = None + self.thread_id: str | None = None + self.active_turn_id: str | None = None + + self._process: asyncio.subprocess.Process | None = None + self._reader_task: asyncio.Task | None = None + self._stderr_task: asyncio.Task | None = None + self._event_task: asyncio.Task | None = None + self._write_lock = asyncio.Lock() + self._pending: dict[int, asyncio.Future] = {} + self._notifications: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._request_id = 0 + self._server_request_tasks: set[asyncio.Task] = set() + self._stderr_tail: deque[str] = deque(maxlen=20) + self._dispatch: dict[str, str] = {} + self._session_events: asyncio.Queue[Event] = asyncio.Queue() + self._tool_session: Session | None = None + self._tool_router_entered = False + + async def __aenter__(self) -> "CodexAppServerRuntime": + await self.start() + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + await self.close() + + async def start(self) -> None: + """Start Codex, initialize tools, and create an ephemeral thread.""" + if self._process is not None: + return + + self.auth_status = await codex_login_status(self.codex_bin) + resolved = shutil.which(self.codex_bin) + assert resolved is not None + + try: + await self.tool_router.__aenter__() + self._tool_router_entered = True + self._tool_session = Session( + self._session_events, + self.config, + tool_router=self.tool_router, + hf_token=self.hf_token, + hf_username="unknown", + local_mode=self.local_mode, + autonomous_mode=self.autonomous_mode, + stream=True, + ) + + namespace, self._dispatch = build_dynamic_tool_namespace( + self.tool_router, + local_mode=self.local_mode, + ) + + self._process = await asyncio.create_subprocess_exec( + resolved, + "app-server", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.cwd, + ) + self._reader_task = asyncio.create_task(self._read_loop()) + self._stderr_task = asyncio.create_task(self._read_stderr()) + self._event_task = asyncio.create_task(self._drain_session_events()) + + await self._request( + "initialize", + { + "clientInfo": { + "name": "ml_intern", + "title": "ML Intern", + "version": "0.1.0", + }, + "capabilities": {"experimentalApi": True}, + }, + ) + await self._notify("initialized", {}) + + params: dict[str, Any] = { + "cwd": self.cwd, + "ephemeral": True, + "approvalPolicy": "never", + "sandbox": "workspace-write" if self.local_mode else "read-only", + "serviceName": "ml-intern", + "developerInstructions": ( + "You are the OpenAI-authenticated Codex runtime inside ML " + "Intern. Use your built-in Codex tools for local repository " + "work. Use the ml_intern namespace for Hugging Face docs, " + "papers, datasets, Hub repositories, Jobs, web research, and " + "remote sandbox operations. Never claim that a ChatGPT login " + "is an OpenAI API key. If an ML Intern tool is denied or fails, " + "report that result instead of silently retrying a billable or " + "destructive operation." + ), + } + requested_model = codex_model_name(self.config.model_name) + if requested_model is not None: + params["model"] = requested_model + if namespace is not None: + params["dynamicTools"] = [namespace] + + response = await self._request("thread/start", params) + thread = response.get("thread") or {} + self.thread_id = thread.get("id") + if not self.thread_id: + raise CodexRuntimeError("Codex app-server did not return a thread id.") + except BaseException: + await self.close() + raise + + async def close(self) -> None: + """Close the Codex process and ML Intern tool resources.""" + for task in tuple(self._server_request_tasks): + task.cancel() + if self._server_request_tasks: + await asyncio.gather( + *self._server_request_tasks, + return_exceptions=True, + ) + self._server_request_tasks.clear() + + background_tasks = ( + self._event_task, + self._reader_task, + self._stderr_task, + ) + for task in background_tasks: + if task is not None: + task.cancel() + await asyncio.gather( + *(task for task in background_tasks if task is not None), + return_exceptions=True, + ) + self._event_task = None + self._reader_task = None + self._stderr_task = None + + process = self._process + self._process = None + if process is not None and process.returncode is None: + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=5) + except asyncio.TimeoutError: + process.kill() + await process.wait() + + if self._tool_router_entered: + await self.tool_router.__aexit__(None, None, None) + self._tool_router_entered = False + self._fail_pending(CodexRuntimeError("Codex app-server closed.")) + self.thread_id = None + self.active_turn_id = None + self._tool_session = None + + async def new_thread(self) -> None: + """Start a fresh ephemeral Codex thread with the same runtime.""" + if self._process is None: + raise CodexRuntimeError("Codex app-server is not running.") + # Recreate the process-level integration to ensure the same dynamic + # tools and model settings are applied consistently. + await self.close() + self.thread_id = None + self.active_turn_id = None + await self.start() + + async def interrupt(self) -> None: + """Interrupt the active Codex turn, if one is running.""" + if not self.thread_id or not self.active_turn_id: + return + try: + await self._request( + "turn/interrupt", + {"threadId": self.thread_id, "turnId": self.active_turn_id}, + ) + except Exception: + logger.debug("Failed to interrupt Codex turn", exc_info=True) + + async def run_turn(self, prompt: str, *, stream: bool = True) -> str: + """Run one user turn and return the final Codex answer.""" + if not self.thread_id: + raise CodexRuntimeError("Codex runtime has not been started.") + + response = await self._request( + "turn/start", + { + "threadId": self.thread_id, + "input": [{"type": "text", "text": prompt}], + }, + ) + turn = response.get("turn") or {} + self.active_turn_id = turn.get("id") + final_text = "" + + try: + while True: + message = await self._notifications.get() + method = message.get("method") + params = message.get("params") or {} + + if params.get("threadId") not in {None, self.thread_id}: + continue + notification_turn_id = params.get("turnId") + if ( + notification_turn_id + and self.active_turn_id + and notification_turn_id != self.active_turn_id + ): + continue + + if method == "item/agentMessage/delta" and stream: + delta = str(params.get("delta") or "") + if delta: + await _call_maybe_async(self.on_delta, delta) + elif method == "item/completed": + item = params.get("item") or {} + if item.get("type") == "agentMessage": + text = str(item.get("text") or "") + if item.get("phase") == "final_answer" or not final_text: + final_text = text + elif method == "turn/completed": + completed_turn = params.get("turn") or {} + if ( + self.active_turn_id + and completed_turn.get("id") != self.active_turn_id + ): + continue + if completed_turn.get("status") == "failed": + error = completed_turn.get("error") or {} + message_text = error.get("message") or "Codex turn failed." + raise CodexRuntimeError(str(message_text)) + return final_text + finally: + self.active_turn_id = None + + async def _request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: + process = self._process + if process is None or process.stdin is None: + raise CodexRuntimeError("Codex app-server is not running.") + self._request_id += 1 + request_id = self._request_id + future = asyncio.get_running_loop().create_future() + self._pending[request_id] = future + await self._write({"method": method, "id": request_id, "params": params}) + try: + return await future + finally: + self._pending.pop(request_id, None) + + async def _notify(self, method: str, params: dict[str, Any]) -> None: + await self._write({"method": method, "params": params}) + + async def _write(self, message: dict[str, Any]) -> None: + process = self._process + if process is None or process.stdin is None: + raise CodexRuntimeError("Codex app-server is not running.") + payload = (json.dumps(message, separators=(",", ":")) + "\n").encode() + async with self._write_lock: + process.stdin.write(payload) + await process.stdin.drain() + + async def _read_loop(self) -> None: + process = self._process + assert process is not None and process.stdout is not None + try: + while line := await process.stdout.readline(): + try: + message = json.loads(line) + except json.JSONDecodeError: + logger.warning( + "Ignoring non-JSON Codex app-server output: %r", + line[:200], + ) + continue + + request_id = message.get("id") + if request_id in self._pending and ( + "result" in message or "error" in message + ): + future = self._pending[request_id] + if "error" in message: + error = message.get("error") or {} + future.set_exception( + CodexRuntimeError( + str(error.get("message") or "Codex request failed.") + ) + ) + else: + future.set_result(message.get("result") or {}) + continue + + if message.get("method") and request_id is not None: + task = asyncio.create_task(self._handle_server_request(message)) + self._server_request_tasks.add(task) + task.add_done_callback(self._server_request_tasks.discard) + continue + + if message.get("method"): + await self._notifications.put(message) + except asyncio.CancelledError: + raise + except Exception as exc: + self._fail_pending(CodexRuntimeError(f"Codex read loop failed: {exc}")) + finally: + if self._process is not None and self._process.returncode is not None: + detail = "\n".join(self._stderr_tail) + suffix = f"\n{detail}" if detail else "" + self._fail_pending( + CodexRuntimeError(f"Codex app-server exited unexpectedly.{suffix}") + ) + + async def _read_stderr(self) -> None: + process = self._process + assert process is not None and process.stderr is not None + try: + while line := await process.stderr.readline(): + text = line.decode(errors="replace").rstrip() + if text: + self._stderr_tail.append(text) + logger.debug("codex app-server: %s", text) + except asyncio.CancelledError: + raise + + async def _drain_session_events(self) -> None: + try: + while True: + event = await self._session_events.get() + await _call_maybe_async(self.on_event, event) + except asyncio.CancelledError: + raise + + async def _handle_server_request(self, message: dict[str, Any]) -> None: + request_id = message["id"] + method = message.get("method") + if method != "item/tool/call": + await self._write( + { + "id": request_id, + "error": { + "code": -32601, + "message": f"Unsupported Codex server request: {method}", + }, + } + ) + return + + params = message.get("params") or {} + namespace = params.get("namespace") + codex_tool_name = str(params.get("tool") or "") + tool_name = self._dispatch.get(codex_tool_name) + arguments = params.get("arguments") or {} + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + arguments = {} + + if namespace != CODEX_TOOL_NAMESPACE or not tool_name: + await self._dynamic_tool_response( + request_id, + f"Unknown ML Intern tool: {namespace}.{codex_tool_name}", + False, + ) + return + + await _call_maybe_async(self.on_tool, tool_name, arguments, None, None) + + if _base_needs_approval(tool_name, arguments, self.config): + approved = bool( + await _call_maybe_async(self.approve_tool, tool_name, arguments) + ) + if not approved: + output = f"User denied ML Intern tool call: {tool_name}" + await _call_maybe_async( + self.on_tool, + tool_name, + arguments, + output, + False, + ) + await self._dynamic_tool_response(request_id, output, False) + return + + assert self._tool_session is not None + try: + output, success = await self.tool_router.call_tool( + tool_name, + arguments, + session=self._tool_session, + tool_call_id=str(params.get("callId") or request_id), + ) + except Exception as exc: + logger.exception("ML Intern tool failed in Codex runtime: %s", tool_name) + output = f"ML Intern tool failed: {exc}" + success = False + await _call_maybe_async( + self.on_tool, + tool_name, + arguments, + output, + success, + ) + await self._dynamic_tool_response(request_id, output, success) + + async def _dynamic_tool_response( + self, + request_id: int | str, + output: str, + success: bool, + ) -> None: + await self._write( + { + "id": request_id, + "result": { + "contentItems": [ + { + "type": "inputText", + "text": str(output), + } + ], + "success": bool(success), + }, + } + ) + + def _fail_pending(self, error: Exception) -> None: + for future in tuple(self._pending.values()): + if not future.done(): + future.set_exception(error) diff --git a/agent/core/llm_params.py b/agent/core/llm_params.py index d2f821c2..c6564ebe 100644 --- a/agent/core/llm_params.py +++ b/agent/core/llm_params.py @@ -7,6 +7,7 @@ import os +from agent.core.codex_models import is_codex_model_id from agent.core.hf_tokens import resolve_hf_router_token from agent.core.local_models import ( LOCAL_MODEL_API_KEY_DEFAULT, @@ -123,6 +124,13 @@ def _resolve_llm_params( """ normalized_model = strip_huggingface_model_prefix(model_name) or model_name + if is_codex_model_id(normalized_model): + raise ValueError( + "Codex model ids use the Codex app-server runtime and cannot be " + "sent through LiteLLM. Start ML Intern with " + "`--model codex/default`." + ) + if is_reserved_local_model_id(normalized_model): raise ValueError(f"Unsupported local model id: {normalized_model}") diff --git a/agent/core/model_switcher.py b/agent/core/model_switcher.py index 5ece764d..67d2ac3b 100644 --- a/agent/core/model_switcher.py +++ b/agent/core/model_switcher.py @@ -19,6 +19,7 @@ from litellm import acompletion +from agent.core.codex_models import CODEX_DEFAULT_MODEL_ID, is_codex_model_id from agent.core.effort_probe import ProbeInconclusive, probe_effort from agent.core.llm_params import _resolve_llm_params from agent.core.local_models import ( @@ -42,6 +43,7 @@ # ":cheapest", ":preferred", or ":" to override the default routing # policy (auto = fastest with failover). SUGGESTED_MODELS = [ + {"id": CODEX_DEFAULT_MODEL_ID, "label": "Codex (OpenAI auth)"}, {"id": CLAUDE_OPUS_48_MODEL_ID, "label": "Claude Opus 4.8"}, {"id": GPT_55_MODEL_ID, "label": "GPT-5.5"}, {"id": MINIMAX_M3_MODEL_ID, "label": "MiniMax M3"}, @@ -69,6 +71,8 @@ def is_valid_model_id(model_id: str) -> bool: if not model_id: return False normalized_model_id = strip_huggingface_model_prefix(model_id) or model_id + if is_codex_model_id(normalized_model_id): + return True if is_local_model_id(normalized_model_id): return True if is_reserved_local_model_id(normalized_model_id): @@ -92,7 +96,7 @@ def _print_hf_routing_info(model_id: str, console) -> bool: against the router catalog when possible; the probe below covers provider availability for uncataloged ids. """ - if is_local_model_id(model_id): + if is_codex_model_id(model_id) or is_local_model_id(model_id): return True from agent.core import hf_router_catalog as cat @@ -164,7 +168,9 @@ def print_model_listing(config, console) -> None: "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M3:novita').\n" "Add ':fastest', ':cheapest', ':preferred', or ':' to override routing.\n" "Use 'ollama/', 'vllm/', 'lm_studio/', or " - "'llamacpp/' for local OpenAI-compatible endpoints.[/dim]" + "'llamacpp/' for local OpenAI-compatible endpoints.\n" + "Use 'codex/default' at startup to reuse `codex login` " + "authentication.[/dim]" ) @@ -173,6 +179,7 @@ def print_invalid_id(arg: str, console) -> None: console.print( "[dim]Expected:\n" " • /[:tag] (HF router — paste from huggingface.co)\n" + " • codex/default | codex/ (Codex app-server)\n" " • ollama/ | vllm/ | lm_studio/ | llamacpp/[/dim]" ) @@ -228,6 +235,16 @@ async def probe_and_switch_model( ) return + if is_codex_model_id(model_id): + console.print( + "[yellow]Codex runtime selection takes effect at process startup.[/yellow]" + ) + console.print( + f"[dim]Restart with `ml-intern --model {model_id}`. " + "The current conversation keeps its existing model.[/dim]" + ) + return + preference = config.reasoning_effort if not _print_hf_routing_info(model_id, console): return diff --git a/agent/main.py b/agent/main.py index 6f29c06b..7b5bebf9 100644 --- a/agent/main.py +++ b/agent/main.py @@ -23,9 +23,12 @@ from prompt_toolkit import PromptSession from agent.config import load_config +from agent.codex_cli import run_codex_headless, run_codex_interactive from agent.core.approval_policy import is_scheduled_operation from agent.core.agent_loop import submission_loop from agent.core import model_switcher +from agent.core.codex_models import is_codex_model_id +from agent.core.codex_runtime import CodexRuntimeError from agent.core.hf_access import fetch_whoami_v2, normalize_hf_user_plan from agent.core.hf_tokens import resolve_hf_token from agent.core.local_models import is_local_model_id @@ -86,7 +89,8 @@ def _validate_cli_model_override(model: str) -> str: if not model_switcher.is_valid_model_id(model): raise ValueError( "Invalid model id. Use an HF Router id like " - "'zai-org/GLM-5.2:novita' or a supported local prefix." + "'zai-org/GLM-5.2:novita', 'codex/default' for Codex auth, " + "or a supported local prefix." ) return model.removeprefix("huggingface/") @@ -1192,6 +1196,23 @@ async def main(model: str | None = None, sandbox_tools: bool = False): _apply_tool_runtime_override(config, sandbox_tools=sandbox_tools) local_mode = _is_local_tool_runtime(config) + if is_codex_model_id(config.model_name): + # Codex owns OpenAI authentication. HF auth remains optional for ML + # Intern's Hub tools, except when the user explicitly selects the + # remote HF sandbox runtime. + hf_token = resolve_hf_token() + if not hf_token and not local_mode: + hf_token = await _prompt_and_save_hf_token(prompt_session) + hf_user, _hf_user_plan = await _get_hf_identity(hf_token) + await run_codex_interactive( + config=config, + prompt_session=prompt_session, + hf_token=hf_token, + hf_user=hf_user, + local_mode=local_mode, + ) + return + # HF token — required for Hub-backed models/tools and sandbox tools, but # not for local LLMs using only local filesystem tools. hf_token = resolve_hf_token() @@ -1447,6 +1468,33 @@ async def headless_main( local_mode = _is_local_tool_runtime(config) hf_token = resolve_hf_token() + if is_codex_model_id(config.model_name): + if not hf_token and not local_mode: + print( + "ERROR: HF sandbox tools require HF_TOKEN. Set it or use the " + "default local tool runtime.", + file=sys.stderr, + ) + sys.exit(1) + if max_iterations is not None: + config.max_iterations = max_iterations + print(f"Model: {config.model_name}", file=sys.stderr) + print(f"Tool runtime: {_tool_runtime_label(local_mode)}", file=sys.stderr) + print(f"Prompt: {prompt}", file=sys.stderr) + print("---", file=sys.stderr) + try: + await run_codex_headless( + prompt, + config=config, + hf_token=hf_token, + local_mode=local_mode, + stream=stream, + ) + except CodexRuntimeError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) + return + if not hf_token and (not is_local_model_id(config.model_name) or not local_mode): print( "ERROR: No HF token found. Set HF_TOKEN or run `hf auth login`.", diff --git a/tests/unit/test_cli_local_models.py b/tests/unit/test_cli_local_models.py index 7c675c59..c73647be 100644 --- a/tests/unit/test_cli_local_models.py +++ b/tests/unit/test_cli_local_models.py @@ -2,6 +2,12 @@ from agent.config import load_config from agent.core import model_switcher +from agent.core.codex_models import ( + CODEX_DEFAULT_MODEL_ID, + codex_model_name, + is_codex_model_id, +) +from agent.core.llm_params import _resolve_llm_params from agent.core.local_models import is_local_model_id from agent.main import CLI_CONFIG_PATH @@ -32,9 +38,23 @@ def test_openai_compat_prefix_is_not_supported(): assert not model_switcher.is_valid_model_id("openai-compat/custom-model") +def test_codex_model_ids_are_first_class_but_not_litellm_models(): + assert is_codex_model_id(CODEX_DEFAULT_MODEL_ID) + assert is_codex_model_id("codex/gpt-example") + assert not is_codex_model_id("codex/") + assert not is_codex_model_id("codex/model with spaces") + assert codex_model_name(CODEX_DEFAULT_MODEL_ID) is None + assert codex_model_name("codex/gpt-example") == "gpt-example" + assert model_switcher.is_valid_model_id(CODEX_DEFAULT_MODEL_ID) + + with pytest.raises(ValueError, match="Codex app-server runtime"): + _resolve_llm_params(CODEX_DEFAULT_MODEL_ID) + + def test_suggested_models_include_router_claude_models_and_no_native_ids(): ids = {m["id"] for m in model_switcher.SUGGESTED_MODELS} + assert CODEX_DEFAULT_MODEL_ID in ids assert "anthropic/claude-opus-4.8:fal-ai" in ids assert all(model_id.count("/") >= 1 for model_id in ids) diff --git a/tests/unit/test_cli_rendering.py b/tests/unit/test_cli_rendering.py index 285e943e..0c3d421e 100644 --- a/tests/unit/test_cli_rendering.py +++ b/tests/unit/test_cli_rendering.py @@ -227,6 +227,78 @@ def fake_banner(*, model=None, hf_user=None, tool_runtime=None): await main_mod.main(model="openai/gpt-5.5:fal-ai") +@pytest.mark.asyncio +async def test_interactive_codex_model_uses_codex_runtime_without_hf_prompt( + monkeypatch, +): + seen: dict[str, object] = {} + prompt_session = object() + + async def fail_hf_prompt(_prompt_session): + raise AssertionError("Codex local runtime must not require HF auth") + + async def fake_codex_runtime(**kwargs): + seen.update(kwargs) + + monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0) + monkeypatch.setattr(main_mod, "PromptSession", lambda: prompt_session) + monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: None) + monkeypatch.setattr(main_mod, "_prompt_and_save_hf_token", fail_hf_prompt) + monkeypatch.setattr(main_mod, "_get_hf_identity", _fake_hf_identity) + monkeypatch.setattr( + main_mod, + "load_config", + lambda _path, **_kwargs: SimpleNamespace( + model_name="zai-org/GLM-5.2:novita", + mcpServers={}, + tool_runtime="local", + ), + ) + monkeypatch.setattr(main_mod, "run_codex_interactive", fake_codex_runtime) + + await main_mod.main(model="codex/default") + + assert seen["prompt_session"] is prompt_session + assert seen["hf_token"] is None + assert seen["local_mode"] is True + assert seen["config"].model_name == "codex/default" + + +@pytest.mark.asyncio +async def test_headless_codex_model_uses_codex_runtime_without_hf_token(monkeypatch): + seen: dict[str, object] = {} + + async def fake_codex_runtime(prompt, **kwargs): + seen["prompt"] = prompt + seen.update(kwargs) + + monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: None) + monkeypatch.setattr( + main_mod, + "load_config", + lambda _path, **_kwargs: SimpleNamespace( + model_name="zai-org/GLM-5.2:novita", + mcpServers={}, + tool_runtime="local", + yolo_mode=False, + max_iterations=50, + ), + ) + monkeypatch.setattr(main_mod, "run_codex_headless", fake_codex_runtime) + + await main_mod.headless_main( + "inspect this repo", + model="codex/default", + stream=False, + ) + + assert seen["prompt"] == "inspect this repo" + assert seen["hf_token"] is None + assert seen["local_mode"] is True + assert seen["stream"] is False + assert seen["config"].model_name == "codex/default" + + @pytest.mark.asyncio async def test_local_model_local_runtime_skips_hf_token_prompt(monkeypatch): class StopAfterBanner(Exception): diff --git a/tests/unit/test_codex_runtime.py b/tests/unit/test_codex_runtime.py new file mode 100644 index 00000000..590b18a8 --- /dev/null +++ b/tests/unit/test_codex_runtime.py @@ -0,0 +1,145 @@ +from types import SimpleNamespace + +import pytest + +from agent.core.codex_runtime import ( + CODEX_TOOL_NAMESPACE, + CodexAppServerRuntime, + CodexRuntimeError, + build_dynamic_tool_namespace, + codex_login_status, +) +from agent.core.tools import ToolSpec + + +class StubRouter: + def __init__(self): + self.tools = { + "bash": ToolSpec( + name="bash", + description="shell", + parameters={"type": "object"}, + handler=None, + ), + "research": ToolSpec( + name="research", + description="nested research", + parameters={"type": "object"}, + handler=None, + ), + "hf_papers": ToolSpec( + name="hf_papers", + description="papers", + parameters={ + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + handler=None, + ), + "invalid.tool": ToolSpec( + name="invalid.tool", + description="invalid name", + parameters={"type": "object"}, + handler=None, + ), + } + + +def test_dynamic_namespace_keeps_ml_tools_and_skips_codex_duplicates(): + namespace, dispatch = build_dynamic_tool_namespace( + StubRouter(), + local_mode=True, + ) + + assert namespace is not None + assert namespace["name"] == CODEX_TOOL_NAMESPACE + assert [tool["name"] for tool in namespace["tools"]] == ["hf_papers"] + assert dispatch == {"hf_papers": "hf_papers"} + + +def test_remote_sandbox_tools_can_expose_namespaced_bash(): + namespace, dispatch = build_dynamic_tool_namespace( + StubRouter(), + local_mode=False, + ) + + assert namespace is not None + assert [tool["name"] for tool in namespace["tools"]] == [ + "bash", + "hf_papers", + ] + assert dispatch == {"bash": "bash", "hf_papers": "hf_papers"} + + +@pytest.mark.asyncio +async def test_codex_login_status_reports_missing_cli(monkeypatch): + monkeypatch.setattr( + "agent.core.codex_runtime.shutil.which", + lambda _binary: None, + ) + + with pytest.raises(CodexRuntimeError, match="Codex CLI is not installed"): + await codex_login_status() + + +@pytest.mark.asyncio +async def test_codex_login_status_uses_public_cli_status(monkeypatch): + class Process: + returncode = 0 + + async def communicate(self): + return b"Logged in using ChatGPT\n", b"" + + async def fake_subprocess(*_args, **_kwargs): + return Process() + + monkeypatch.setattr( + "agent.core.codex_runtime.shutil.which", + lambda _binary: "/usr/local/bin/codex", + ) + monkeypatch.setattr( + "agent.core.codex_runtime.asyncio.create_subprocess_exec", + fake_subprocess, + ) + + assert await codex_login_status() == "Logged in using ChatGPT" + + +@pytest.mark.asyncio +async def test_dynamic_tool_failure_is_returned_to_codex(tmp_path): + class FailingRouter: + async def call_tool(self, *_args, **_kwargs): + raise RuntimeError("probe failed") + + responses = [] + runtime = CodexAppServerRuntime( + config=SimpleNamespace(model_name="codex/default"), + tool_router=FailingRouter(), + hf_token=None, + local_mode=True, + cwd=tmp_path, + autonomous_mode=False, + ) + runtime._dispatch = {"hf_papers": "hf_papers"} + runtime._tool_session = object() + + async def capture_response(message): + responses.append(message) + + runtime._write = capture_response + + await runtime._handle_server_request( + { + "id": 7, + "method": "item/tool/call", + "params": { + "namespace": CODEX_TOOL_NAMESPACE, + "tool": "hf_papers", + "arguments": {"operation": "search", "query": "LoRA"}, + }, + } + ) + + assert responses[0]["id"] == 7 + assert responses[0]["result"]["success"] is False + assert "probe failed" in responses[0]["result"]["contentItems"][0]["text"] From 46fd4d812dae4202cbf6f484b910710aa47b30e3 Mon Sep 17 00:00:00 2001 From: tibetyalman <222756182+tibetyalman@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:50:18 +0300 Subject: [PATCH 2/5] feat: add Codex to local web UI --- README.md | 25 ++- agent/codex_cli.py | 14 +- agent/core/codex_runtime.py | 151 ++++++++++--- backend/codex_web.py | 201 ++++++++++++++++++ backend/routes/agent.py | 66 ++++-- backend/session_manager.py | 92 +++++++- frontend/src/components/Chat/ChatInput.tsx | 7 +- frontend/src/utils/model.ts | 1 + tests/unit/test_agent_model_gating.py | 78 +++++++ tests/unit/test_codex_runtime.py | 41 ++++ tests/unit/test_codex_web.py | 126 +++++++++++ .../unit/test_session_manager_persistence.py | 41 ++++ 12 files changed, 785 insertions(+), 58 deletions(-) create mode 100644 backend/codex_web.py create mode 100644 tests/unit/test_codex_web.py diff --git a/README.md b/README.md index f552b6f3..7ceaddf4 100644 --- a/README.md +++ b/README.md @@ -90,15 +90,22 @@ Codex's cached credentials. If `codex login` used **Sign in with ChatGPT**, Codex usage follows that ChatGPT plan's Codex allowance. If Codex was logged in with an OpenAI API key, standard API billing applies instead. -Codex supplies the main agent loop and local repository tools. ML Intern's -Hugging Face documentation, papers, datasets, Hub, Jobs, web research, and -optional sandbox tools are exposed to Codex under the `ml_intern` tool -namespace. `HF_TOKEN` is optional for the Codex model itself, but individual -Hub/Jobs tools and `--sandbox-tools` still require Hugging Face authentication. - -OpenAI-authenticated Codex mode currently targets the local CLI. The hosted web -app continues to use Hugging Face OAuth/Router because a local Codex login is a -device credential and must not be forwarded to the server. +To expose the same signed-in Codex runtime in the local Web UI: + +```bash +ML_INTERN_ENABLE_CODEX_WEB=1 uv run uvicorn backend.main:app --host ::1 --port 7860 +``` + +`Codex (ChatGPT)` then appears as the recommended/default Web model. This flag +is deliberately disabled when the backend detects Hugging Face Spaces or OAuth: +a developer's local ChatGPT allowance must never be shared with hosted visitors. + +Codex supplies the main agent loop. In the CLI it also has local repository +tools; in the Web UI the host repository stays read-only. ML Intern's Hugging +Face documentation, papers, datasets, Hub, Jobs, web research, and optional +sandbox tools are exposed to Codex under the `ml_intern` namespace. `HF_TOKEN` +is optional for the Codex model itself, but individual Hub/Jobs and sandbox +tools still require Hugging Face authentication. #### Local models diff --git a/agent/codex_cli.py b/agent/codex_cli.py index bb7a4e5d..f738c410 100644 --- a/agent/codex_cli.py +++ b/agent/codex_cli.py @@ -72,6 +72,7 @@ async def on_tool( arguments: dict[str, Any], output: str | None, success: bool | None, + _tool_call_id: str, ) -> None: if output is None: print_tool_call(name, json.dumps(arguments)[:120]) @@ -85,7 +86,11 @@ async def on_event(event: Event) -> None: str(event.data.get("log") or ""), ) - async def approve_tool(name: str, arguments: dict[str, Any]) -> bool: + async def approve_tool( + name: str, + arguments: dict[str, Any], + _tool_call_id: str, + ) -> bool: console.print(f"\n[bold yellow]Approval required:[/bold yellow] {name}") console.print_json(data=arguments) answer = await prompt_session.prompt_async("Approve this tool call? [y/N] ") @@ -177,13 +182,18 @@ async def on_tool( arguments: dict[str, Any], output: str | None, success: bool | None, + _tool_call_id: str, ) -> None: if output is None: print_tool_call(name, json.dumps(arguments)[:120]) else: print_tool_output(output, bool(success), truncate=True) - async def approve_tool(name: str, arguments: dict[str, Any]) -> bool: + async def approve_tool( + name: str, + arguments: dict[str, Any], + _tool_call_id: str, + ) -> bool: # Match the existing headless policy: scheduled Jobs never receive # automatic approval because they can create recurring spend. return not _is_scheduled_job(name, arguments) diff --git a/agent/core/codex_runtime.py b/agent/core/codex_runtime.py index bedf8f93..1bf6e1ec 100644 --- a/agent/core/codex_runtime.py +++ b/agent/core/codex_runtime.py @@ -41,10 +41,13 @@ "research", } -ToolApprovalCallback = Callable[[str, dict[str, Any]], Awaitable[bool] | bool] +ToolApprovalCallback = Callable[ + [str, dict[str, Any], str], + Awaitable[bool] | bool, +] DeltaCallback = Callable[[str], Awaitable[None] | None] ToolCallback = Callable[ - [str, dict[str, Any], str | None, bool | None], + [str, dict[str, Any], str | None, bool | None, str], Awaitable[None] | None, ] EventCallback = Callable[[Event], Awaitable[None] | None] @@ -151,6 +154,8 @@ def __init__( on_tool: ToolCallback | None = None, on_event: EventCallback | None = None, codex_bin: str = "codex", + manage_tool_router: bool = True, + tool_session: Session | None = None, ) -> None: self.config = config self.tool_router = tool_router @@ -163,6 +168,8 @@ def __init__( self.on_tool = on_tool self.on_event = on_event self.codex_bin = codex_bin + self.manage_tool_router = manage_tool_router + self._provided_tool_session = tool_session self.auth_status: str | None = None self.thread_id: str | None = None @@ -200,18 +207,22 @@ async def start(self) -> None: assert resolved is not None try: - await self.tool_router.__aenter__() - self._tool_router_entered = True - self._tool_session = Session( - self._session_events, - self.config, - tool_router=self.tool_router, - hf_token=self.hf_token, - hf_username="unknown", - local_mode=self.local_mode, - autonomous_mode=self.autonomous_mode, - stream=True, - ) + if self.manage_tool_router: + await self.tool_router.__aenter__() + self._tool_router_entered = True + if self._provided_tool_session is not None: + self._tool_session = self._provided_tool_session + else: + self._tool_session = Session( + self._session_events, + self.config, + tool_router=self.tool_router, + hf_token=self.hf_token, + hf_username="unknown", + local_mode=self.local_mode, + autonomous_mode=self.autonomous_mode, + stream=True, + ) namespace, self._dispatch = build_dynamic_tool_namespace( self.tool_router, @@ -228,7 +239,8 @@ async def start(self) -> None: ) self._reader_task = asyncio.create_task(self._read_loop()) self._stderr_task = asyncio.create_task(self._read_stderr()) - self._event_task = asyncio.create_task(self._drain_session_events()) + if self._provided_tool_session is None: + self._event_task = asyncio.create_task(self._drain_session_events()) await self._request( "initialize", @@ -249,16 +261,7 @@ async def start(self) -> None: "approvalPolicy": "never", "sandbox": "workspace-write" if self.local_mode else "read-only", "serviceName": "ml-intern", - "developerInstructions": ( - "You are the OpenAI-authenticated Codex runtime inside ML " - "Intern. Use your built-in Codex tools for local repository " - "work. Use the ml_intern namespace for Hugging Face docs, " - "papers, datasets, Hub repositories, Jobs, web research, and " - "remote sandbox operations. Never claim that a ChatGPT login " - "is an OpenAI API key. If an ML Intern tool is denied or fails, " - "report that result instead of silently retrying a billable or " - "destructive operation." - ), + "developerInstructions": self._developer_instructions(), } requested_model = codex_model_name(self.config.model_name) if requested_model is not None: @@ -320,6 +323,24 @@ async def close(self) -> None: self.active_turn_id = None self._tool_session = None + def _developer_instructions(self) -> str: + if self.local_mode: + tool_guidance = "Use your built-in Codex tools for local repository work. " + else: + tool_guidance = ( + "The host repository is read-only. Use the ml_intern namespace " + "for remote sandbox execution and ML research tools. " + ) + return ( + "You are the OpenAI-authenticated Codex runtime inside ML Intern. " + f"{tool_guidance}" + "Use the ml_intern namespace for Hugging Face docs, papers, datasets, " + "Hub repositories, Jobs, and web research. Never claim that a ChatGPT " + "login is an OpenAI API key. If an ML Intern tool is denied or fails, " + "report that result instead of silently retrying a billable or " + "destructive operation." + ) + async def new_thread(self) -> None: """Start a fresh ephemeral Codex thread with the same runtime.""" if self._process is None: @@ -379,12 +400,19 @@ async def run_turn(self, prompt: str, *, stream: bool = True) -> str: delta = str(params.get("delta") or "") if delta: await _call_maybe_async(self.on_delta, delta) + elif method == "item/started": + await self._emit_builtin_item( + params.get("item") or {}, + completed=False, + ) elif method == "item/completed": item = params.get("item") or {} if item.get("type") == "agentMessage": text = str(item.get("text") or "") if item.get("phase") == "final_answer" or not final_text: final_text = text + else: + await self._emit_builtin_item(item, completed=True) elif method == "turn/completed": completed_turn = params.get("turn") or {} if ( @@ -400,6 +428,58 @@ async def run_turn(self, prompt: str, *, stream: bool = True) -> str: finally: self.active_turn_id = None + async def _emit_builtin_item( + self, + item: dict[str, Any], + *, + completed: bool, + ) -> None: + """Surface Codex built-in tools through the same callback as ML tools.""" + item_type = str(item.get("type") or "") + if item_type == "commandExecution": + name = "codex_command" + arguments = { + "command": item.get("command"), + "cwd": item.get("cwd"), + } + output = str(item.get("aggregatedOutput") or "") if completed else None + success = ( + item.get("status") == "completed" and item.get("exitCode") in {None, 0} + if completed + else None + ) + elif item_type == "fileChange": + name = "codex_file_change" + arguments = {"changes": item.get("changes") or []} + output = json.dumps(item.get("changes") or [], ensure_ascii=False) + success = item.get("status") == "completed" if completed else None + if not completed: + output = None + elif item_type == "mcpToolCall": + name = f"mcp:{item.get('server')}.{item.get('tool')}" + arguments = item.get("arguments") or {} + result = item.get("result") + error = item.get("error") + output = ( + json.dumps(result if result is not None else error, ensure_ascii=False) + if completed + else None + ) + success = ( + item.get("status") == "completed" and not error if completed else None + ) + else: + return + + await _call_maybe_async( + self.on_tool, + name, + arguments, + output, + success, + str(item.get("id") or f"codex-{self._request_id}"), + ) + async def _request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: process = self._process if process is None or process.stdin is None: @@ -516,6 +596,7 @@ async def _handle_server_request(self, message: dict[str, Any]) -> None: codex_tool_name = str(params.get("tool") or "") tool_name = self._dispatch.get(codex_tool_name) arguments = params.get("arguments") or {} + tool_call_id = str(params.get("callId") or request_id) if isinstance(arguments, str): try: arguments = json.loads(arguments) @@ -530,11 +611,23 @@ async def _handle_server_request(self, message: dict[str, Any]) -> None: ) return - await _call_maybe_async(self.on_tool, tool_name, arguments, None, None) + await _call_maybe_async( + self.on_tool, + tool_name, + arguments, + None, + None, + tool_call_id, + ) if _base_needs_approval(tool_name, arguments, self.config): approved = bool( - await _call_maybe_async(self.approve_tool, tool_name, arguments) + await _call_maybe_async( + self.approve_tool, + tool_name, + arguments, + tool_call_id, + ) ) if not approved: output = f"User denied ML Intern tool call: {tool_name}" @@ -544,6 +637,7 @@ async def _handle_server_request(self, message: dict[str, Any]) -> None: arguments, output, False, + tool_call_id, ) await self._dynamic_tool_response(request_id, output, False) return @@ -554,7 +648,7 @@ async def _handle_server_request(self, message: dict[str, Any]) -> None: tool_name, arguments, session=self._tool_session, - tool_call_id=str(params.get("callId") or request_id), + tool_call_id=tool_call_id, ) except Exception as exc: logger.exception("ML Intern tool failed in Codex runtime: %s", tool_name) @@ -566,6 +660,7 @@ async def _handle_server_request(self, message: dict[str, Any]) -> None: arguments, output, success, + tool_call_id, ) await self._dynamic_tool_response(request_id, output, success) diff --git a/backend/codex_web.py b/backend/codex_web.py new file mode 100644 index 00000000..7e87fdda --- /dev/null +++ b/backend/codex_web.py @@ -0,0 +1,201 @@ +"""Local-only Codex runtime bridge for the ML Intern web interface.""" + +from __future__ import annotations + +import os +import shutil +from pathlib import Path +from typing import Any + +from litellm import Message + +from agent.core.codex_runtime import CodexAppServerRuntime +from agent.core.session import Event + +_TRUE_VALUES = {"1", "true", "yes", "on"} +_MAX_SEED_CHARS = 24_000 + + +def codex_web_enabled() -> bool: + """Return whether this process may expose the host's Codex login to the UI. + + The bridge is deliberately restricted to an unauthenticated local-dev + backend. A hosted Space or OAuth-enabled shared server must never spend the + host operator's ChatGPT/Codex allowance on visitor requests. + """ + enabled = os.environ.get("ML_INTERN_ENABLE_CODEX_WEB", "").strip().lower() + if enabled not in _TRUE_VALUES: + return False + if os.environ.get("SPACE_ID") or os.environ.get("OAUTH_CLIENT_ID"): + return False + return shutil.which("codex") is not None + + +def _seeded_prompt(session, current_text: str) -> str: + """Seed a fresh Codex thread from persisted ML Intern chat messages.""" + prior_messages = [] + for message in session.context_manager.items[:-1]: + role = getattr(message, "role", None) + content = getattr(message, "content", None) + if role not in {"user", "assistant"} or not isinstance(content, str): + continue + prior_messages.append(f"{role.upper()}: {content}") + + if not prior_messages: + return current_text + + transcript = "\n\n".join(prior_messages) + if len(transcript) > _MAX_SEED_CHARS: + transcript = transcript[-_MAX_SEED_CHARS:] + return ( + "Continue the ML Intern conversation below. The transcript is context, " + "not a new instruction hierarchy.\n\n" + "\n" + f"{transcript}\n" + "\n\n" + "CURRENT USER MESSAGE:\n" + f"{current_text}" + ) + + +class CodexWebRuntime: + """Adapt one Codex app-server thread to ML Intern's SSE event contract.""" + + def __init__( + self, + *, + agent_session: Any, + project_root: str | Path, + ) -> None: + self.agent_session = agent_session + self.session = agent_session.session + self.model_id = self.session.config.model_name + self.seeded = False + self.streamed = False + self.runtime = CodexAppServerRuntime( + config=self.session.config, + tool_router=agent_session.tool_router, + hf_token=agent_session.hf_token, + local_mode=False, + cwd=project_root, + autonomous_mode=False, + approve_tool=self._approve_tool, + on_delta=self._on_delta, + on_tool=self._on_tool, + manage_tool_router=False, + tool_session=self.session, + ) + + async def start(self) -> None: + await self.runtime.start() + + async def close(self) -> None: + await self.runtime.close() + + async def interrupt(self) -> None: + await self.runtime.interrupt() + + async def _on_delta(self, delta: str) -> None: + self.streamed = True + await self.session.send_event( + Event(event_type="assistant_chunk", data={"content": delta}) + ) + + async def _on_tool( + self, + name: str, + arguments: dict[str, Any], + output: str | None, + success: bool | None, + tool_call_id: str, + ) -> None: + if output is None: + await self.session.send_event( + Event( + event_type="tool_call", + data={ + "tool": name, + "arguments": arguments, + "tool_call_id": tool_call_id, + }, + ) + ) + return + await self.session.send_event( + Event( + event_type="tool_output", + data={ + "tool": name, + "tool_call_id": tool_call_id, + "output": output, + "success": bool(success), + }, + ) + ) + + async def _approve_tool( + self, + name: str, + _arguments: dict[str, Any], + _tool_call_id: str, + ) -> bool: + # Dynamic tools that require explicit approval are denied in this first + # local-web bridge. Read-only research tools continue without approval; + # billable or destructive operations fail closed. + await self.session.send_event( + Event( + event_type="tool_log", + data={ + "tool": name, + "log": ( + "Codex web mode denied an approval-required ML Intern " + "tool call." + ), + }, + ) + ) + return False + + async def run_user_input(self, text: str) -> str: + self.session.reset_cancel() + if text: + self.session.context_manager.add_message(Message(role="user", content=text)) + + await self.session.send_event( + Event(event_type="processing", data={"message": "Processing with Codex"}) + ) + + prompt = _seeded_prompt(self.session, text) if not self.seeded else text + self.seeded = True + self.streamed = False + final_text = await self.runtime.run_turn(prompt, stream=True) + + if self.streamed: + await self.session.send_event( + Event(event_type="assistant_stream_end", data={}) + ) + elif final_text: + await self.session.send_event( + Event(event_type="assistant_message", data={"content": final_text}) + ) + + if self.session.is_cancelled: + await self.session.send_event(Event(event_type="interrupted")) + return final_text + + if final_text: + self.session.context_manager.add_message( + Message(role="assistant", content=final_text) + ) + await self.session.send_event( + Event( + event_type="turn_complete", + data={ + "history_size": len(self.session.context_manager.items), + "final_response": final_text or None, + }, + ) + ) + self.session.increment_turn() + await self.session.auto_save_if_needed() + return final_text diff --git a/backend/routes/agent.py b/backend/routes/agent.py index 06cf1c38..03c8c6e7 100644 --- a/backend/routes/agent.py +++ b/backend/routes/agent.py @@ -50,7 +50,10 @@ SessionCapacityError, session_manager, ) +from codex_web import codex_web_enabled +from agent.core.codex_models import CODEX_DEFAULT_MODEL_ID, is_codex_model_id +from agent.core.codex_runtime import CodexRuntimeError, codex_login_status from agent.core.hf_access import get_jobs_access from agent.core.hf_tokens import resolve_hf_request_token from agent.core.local_models import local_model_provider @@ -116,7 +119,11 @@ def _schedule_usage_refresh_and_upload( task.add_done_callback(_background_route_tasks.discard) -def _available_models() -> list[dict[str, Any]]: +def _available_models( + *, + include_codex: bool | None = None, +) -> list[dict[str, Any]]: + include_codex = codex_web_enabled() if include_codex is None else include_codex models = [ { "id": CLAUDE_OPUS_48_MODEL_ID, @@ -144,14 +151,26 @@ def _available_models() -> list[dict[str, Any]]: "label": "DeepSeek V4 Pro", }, ] + if include_codex: + for model in models: + model.pop("recommended", None) + models.insert( + 0, + { + "id": CODEX_DEFAULT_MODEL_ID, + "label": "Codex (ChatGPT)", + "provider": "codex", + "recommended": True, + }, + ) return models -AVAILABLE_MODELS = _available_models() +AVAILABLE_MODELS = _available_models(include_codex=False) def _valid_model_ids() -> set[str]: - return {m["id"] for m in AVAILABLE_MODELS} + return {m["id"] for m in _available_models()} def _validate_model_id(model_id: str | None) -> None: @@ -161,13 +180,16 @@ def _validate_model_id(model_id: str | None) -> None: def _default_model() -> str: + if codex_web_enabled(): + return CODEX_DEFAULT_MODEL_ID return DEFAULT_MODEL_ID def _model_override_for_new_session(requested_model: str | None) -> str | None: """Return the model override to use when creating a new session. - Explicit model requests are honored. Empty web requests default to GLM 5.2. + Explicit model requests are honored. Empty requests use the active web + default (Codex for an enabled local Codex runtime, otherwise GLM 5.2). """ return requested_model or _default_model() @@ -180,7 +202,9 @@ def _user_hf_token(user: dict[str, Any] | None) -> str | None: def _model_requires_hf_router_token(model_id: str | None) -> bool: normalized = strip_huggingface_model_prefix(model_id) or model_id or "" - return local_model_provider(normalized) is None + return ( + not is_codex_model_id(normalized) and local_model_provider(normalized) is None + ) def _reject_oversize_dataset_upload(request: Request) -> None: @@ -297,6 +321,18 @@ async def llm_health_check( - timeout / network → provider unreachable """ model = _default_model() + if is_codex_model_id(model): + try: + await codex_login_status() + return LLMHealthResponse(status="ok", model=model) + except CodexRuntimeError as e: + return LLMHealthResponse( + status="error", + model=model, + error=str(e)[:500], + error_type="auth", + ) + hf_token = resolve_hf_request_token(request) if _model_requires_hf_router_token(model) and not hf_token: return LLMHealthResponse(status="skipped", model=model) @@ -351,8 +387,8 @@ async def llm_health_check( async def get_model() -> dict: """Get current model and available models. No auth required.""" return { - "current": session_manager.config.model_name, - "available": AVAILABLE_MODELS, + "current": _default_model(), + "available": _available_models(), } @@ -365,14 +401,18 @@ async def generate_title( ) -> dict: """Generate a short title for a chat session based on the first user message. - Always uses gpt-oss-120b via Cerebras on the HF router. The tab headline - renders as plain text, so the model is told to avoid markdown and any - stray formatting characters are stripped before returning. gpt-oss is a - reasoning model — reasoning_effort=low keeps the reasoning budget small - so the 60-token output budget isn't consumed before the title is written. + Codex sessions use a deterministic local title so they never invoke HF + Router. Other sessions use gpt-oss-120b via Cerebras on HF Router. The tab + headline renders as plain text, so formatting characters are stripped. """ try: - await _check_session_access(request.session_id, user) + agent_session = await _check_session_access(request.session_id, user) + if is_codex_model_id(agent_session.session.config.model_name): + fallback = request.text.strip() + title = fallback[:40].rstrip() + "…" if len(fallback) > 40 else fallback + await session_manager.update_session_title(request.session_id, title) + return {"title": title} + llm_params = _resolve_llm_params( "openai/gpt-oss-120b:cerebras", _user_hf_token(user), diff --git a/backend/session_manager.py b/backend/session_manager.py index 260eb8ff..9289f942 100644 --- a/backend/session_manager.py +++ b/backend/session_manager.py @@ -10,8 +10,11 @@ from pathlib import Path from typing import Any, Optional +from codex_web import CodexWebRuntime + from agent.config import load_config from agent.core.agent_loop import process_submission +from agent.core.codex_models import is_codex_model_id from agent.core.model_ids import ( GLM_52_MODEL_ID, is_known_router_model_id, @@ -135,6 +138,8 @@ class AgentSession: inference_billing_session_id: str | None = None usage_warning_next_threshold_usd: float = USAGE_WARNING_FIRST_THRESHOLD_USD usage_warning_spend_cache: dict[str, Any] = field(default_factory=dict) + codex_runtime: CodexWebRuntime | None = None + codex_runtime_model: str | None = None def __post_init__(self) -> None: if self.usage_window_started_at is None: @@ -266,7 +271,9 @@ def _model_from_saved_metadata( model: str | None, ) -> str: normalized = strip_huggingface_model_prefix(model) - if normalized and is_known_router_model_id(normalized): + if normalized and ( + is_known_router_model_id(normalized) or is_codex_model_id(normalized) + ): return normalized fallback_model = GLM_52_MODEL_ID @@ -592,9 +599,23 @@ async def refresh_session_usage_metrics( normalize_hf_billing_snapshot, summarize_usage_events, ) - from usage import build_hf_billing_snapshot session = agent_session.session + if is_codex_model_id(session.config.model_name): + hf_billing_snapshot = normalize_hf_billing_snapshot( + self._fallback_hf_billing_snapshot("codex_uses_chatgpt_allowance") + ) + session.usage_hf_billing_snapshot = hf_billing_snapshot + metrics = summarize_usage_events( + getattr(session, "logged_events", []) or [], + session_id=agent_session.session_id, + hf_billing_snapshot=hf_billing_snapshot, + ) + session.usage_metrics = metrics + return metrics + + from usage import build_hf_billing_snapshot + try: billing_snapshot = build_hf_billing_snapshot( self, @@ -829,6 +850,8 @@ async def _start_agent_session( @staticmethod def _start_cpu_sandbox_preload(agent_session: AgentSession) -> None: """Kick off a best-effort cpu-basic sandbox for the session.""" + if is_codex_model_id(agent_session.session.config.model_name): + return try: from agent.tools.sandbox_tool import start_cpu_sandbox_preload @@ -1656,6 +1679,41 @@ async def _reap_one(self, session_id: str, cutoff: datetime) -> bool: await self._cleanup_sandbox(session) return True + async def _close_codex_runtime(self, agent_session: AgentSession) -> None: + runtime = agent_session.codex_runtime + agent_session.codex_runtime = None + agent_session.codex_runtime_model = None + if runtime is not None: + try: + await runtime.close() + except Exception: + logger.warning( + "Failed to close Codex runtime for %s", + agent_session.session_id, + exc_info=True, + ) + + async def _ensure_codex_runtime( + self, + agent_session: AgentSession, + ) -> CodexWebRuntime: + model_id = agent_session.session.config.model_name + if ( + agent_session.codex_runtime is not None + and agent_session.codex_runtime_model == model_id + ): + return agent_session.codex_runtime + + await self._close_codex_runtime(agent_session) + runtime = CodexWebRuntime( + agent_session=agent_session, + project_root=PROJECT_ROOT, + ) + await runtime.start() + agent_session.codex_runtime = runtime + agent_session.codex_runtime_model = model_id + return runtime + async def _run_session( self, session_id: str, @@ -1692,9 +1750,24 @@ async def _run_session( agent_session.is_processing = True self._touch(agent_session) try: - should_continue = await process_submission( - session, submission - ) + if ( + submission.operation.op_type == OpType.USER_INPUT + and is_codex_model_id(session.config.model_name) + ): + codex_runtime = await self._ensure_codex_runtime( + agent_session + ) + text = ( + submission.operation.data.get("text", "") + if submission.operation.data + else "" + ) + await codex_runtime.run_user_input(text) + should_continue = True + else: + should_continue = await process_submission( + session, submission + ) finally: agent_session.is_processing = False # Stamp on turn finish too: a turn that ran longer @@ -1718,6 +1791,7 @@ async def _run_session( ) finally: + await self._close_codex_runtime(agent_session) broadcast_task.cancel() try: await broadcast_task @@ -1795,6 +1869,8 @@ async def interrupt(self, session_id: str) -> bool: if not agent_session or not agent_session.is_active: return False agent_session.session.cancel() + if agent_session.codex_runtime is not None: + await agent_session.codex_runtime.interrupt() return True async def undo(self, session_id: str) -> bool: @@ -1886,6 +1962,12 @@ async def update_session_model(self, session_id: str, model_id: str) -> bool: if not agent_session or not agent_session.is_active: return False agent_session.session.update_model(model_id) + if ( + not agent_session.is_processing + and agent_session.codex_runtime is not None + and agent_session.codex_runtime_model != model_id + ): + await self._close_codex_runtime(agent_session) self._touch(agent_session) await self.persist_session_snapshot(agent_session, runtime_state="idle") return True diff --git a/frontend/src/components/Chat/ChatInput.tsx b/frontend/src/components/Chat/ChatInput.tsx index 1cc64263..ffed5f27 100644 --- a/frontend/src/components/Chat/ChatInput.tsx +++ b/frontend/src/components/Chat/ChatInput.tsx @@ -25,6 +25,7 @@ import { useAgentStore } from '@/store/agentStore'; import { useSessionStore } from '@/store/sessionStore'; import { CLAUDE_OPUS_48_MODEL_PATH, + CODEX_DEFAULT_MODEL_PATH, DEEPSEEK_V4_PRO_MODEL_PATH, GLM_52_MODEL_PATH, GPT_55_MODEL_PATH, @@ -124,14 +125,18 @@ const modelOptionId = (modelPath: string) => ( const modelOptionFromApi = (model: { id?: string; label?: string; + provider?: string; recommended?: boolean; }): ModelOption | null => { if (!model.id) return null; + const avatarUrl = model.provider === 'codex' || model.id === CODEX_DEFAULT_MODEL_PATH + ? getHfAvatarUrl('openai') + : getHfAvatarUrl(model.id.replace(/^huggingface\//, '')); return { id: modelOptionId(model.id), name: model.label ?? model.id, modelPath: model.id, - avatarUrl: getHfAvatarUrl(model.id.replace(/^huggingface\//, '')), + avatarUrl, recommended: Boolean(model.recommended), }; }; diff --git a/frontend/src/utils/model.ts b/frontend/src/utils/model.ts index a1585f1e..3b157941 100644 --- a/frontend/src/utils/model.ts +++ b/frontend/src/utils/model.ts @@ -12,6 +12,7 @@ export const KIMI_K27_CODE_MODEL_PATH = 'moonshotai/Kimi-K2.7-Code:novita'; export const MINIMAX_M3_MODEL_PATH = 'MiniMaxAI/MiniMax-M3:novita'; export const GLM_52_MODEL_PATH = 'zai-org/GLM-5.2:novita'; export const DEEPSEEK_V4_PRO_MODEL_PATH = 'deepseek-ai/DeepSeek-V4-Pro:novita'; +export const CODEX_DEFAULT_MODEL_PATH = 'codex/default'; export function isClaudePath(modelPath: string | undefined): boolean { return !!modelPath && modelPath.includes('anthropic'); diff --git a/tests/unit/test_agent_model_gating.py b/tests/unit/test_agent_model_gating.py index 84621f84..26a5d77e 100644 --- a/tests/unit/test_agent_model_gating.py +++ b/tests/unit/test_agent_model_gating.py @@ -18,6 +18,11 @@ BILLING_SESSION_ID = "00000000-0000-4000-8000-000000000001" +@pytest.fixture(autouse=True) +def _disable_codex_web_by_default(monkeypatch): + monkeypatch.setattr(agent, "codex_web_enabled", lambda: False) + + def test_available_models_exclude_sonnet_and_have_no_pro_gate(): models = {model["id"]: model for model in agent.AVAILABLE_MODELS} @@ -36,6 +41,22 @@ def test_default_model_is_glm(): assert agent._default_model() == agent.DEFAULT_MODEL_ID +def test_local_codex_web_is_listed_and_becomes_default(monkeypatch): + monkeypatch.setattr(agent, "codex_web_enabled", lambda: True) + + models = {model["id"]: model for model in agent._available_models()} + + assert models[agent.CODEX_DEFAULT_MODEL_ID] == { + "id": agent.CODEX_DEFAULT_MODEL_ID, + "label": "Codex (ChatGPT)", + "provider": "codex", + "recommended": True, + } + assert "recommended" not in models[agent.DEFAULT_MODEL_ID] + assert agent._default_model() == agent.CODEX_DEFAULT_MODEL_ID + assert agent._model_override_for_new_session(None) == agent.CODEX_DEFAULT_MODEL_ID + + @pytest.mark.asyncio async def test_llm_health_uses_default_and_request_hf_token(monkeypatch): class Request: @@ -107,6 +128,28 @@ async def fail_acompletion(**kwargs): assert response.model == agent.DEFAULT_MODEL_ID +@pytest.mark.asyncio +async def test_codex_health_uses_codex_login_without_hf_completion(monkeypatch): + class Request: + headers = {} + cookies = {} + + async def fake_codex_status(): + return "Logged in using ChatGPT" + + async def fail_acompletion(**_kwargs): + raise AssertionError("Codex health must not use HF Router") + + monkeypatch.setattr(agent, "codex_web_enabled", lambda: True) + monkeypatch.setattr(agent, "codex_login_status", fake_codex_status) + monkeypatch.setattr(agent, "acompletion", fail_acompletion) + + response = await agent.llm_health_check(Request(), {"user_id": "dev"}) + + assert response.status == "ok" + assert response.model == agent.CODEX_DEFAULT_MODEL_ID + + @pytest.mark.asyncio async def test_generate_title_omits_session_id_from_hf_router(monkeypatch): completions = [] @@ -142,6 +185,7 @@ async def fake_check_session_access(session_id, user): session=SimpleNamespace( session_id="session-1", inference_billing_session_id=BILLING_SESSION_ID, + config=SimpleNamespace(model_name=agent.DEFAULT_MODEL_ID), ) ) @@ -168,6 +212,40 @@ async def fake_update_session_title(session_id, title): assert titles == [("session-1", "Clean title")] +@pytest.mark.asyncio +async def test_codex_title_avoids_hf_router(monkeypatch): + async def fake_check_session_access(_session_id, _user): + return SimpleNamespace( + session=SimpleNamespace( + config=SimpleNamespace(model_name=agent.CODEX_DEFAULT_MODEL_ID) + ) + ) + + async def fail_acompletion(**_kwargs): + raise AssertionError("Codex sessions must not use HF Router for titles") + + titles = [] + + async def fake_update_session_title(session_id, title): + titles.append((session_id, title)) + + monkeypatch.setattr(agent, "_check_session_access", fake_check_session_access) + monkeypatch.setattr(agent, "acompletion", fail_acompletion) + monkeypatch.setattr( + agent.session_manager, + "update_session_title", + fake_update_session_title, + ) + + response = await agent.generate_title( + agent.SubmitRequest(session_id="session-1", text="Research LoRA papers"), + {"user_id": "u1"}, + ) + + assert response == {"title": "Research LoRA papers"} + assert titles == [("session-1", "Research LoRA papers")] + + def test_empty_session_model_uses_glm_default(): assert agent._model_override_for_new_session(None) == agent.DEFAULT_MODEL_ID diff --git a/tests/unit/test_codex_runtime.py b/tests/unit/test_codex_runtime.py index 590b18a8..decbc9b3 100644 --- a/tests/unit/test_codex_runtime.py +++ b/tests/unit/test_codex_runtime.py @@ -143,3 +143,44 @@ async def capture_response(message): assert responses[0]["id"] == 7 assert responses[0]["result"]["success"] is False assert "probe failed" in responses[0]["result"]["contentItems"][0]["text"] + + +@pytest.mark.asyncio +async def test_builtin_command_item_uses_tool_callback(tmp_path): + calls = [] + + async def on_tool(name, arguments, output, success, tool_call_id): + calls.append((name, arguments, output, success, tool_call_id)) + + runtime = CodexAppServerRuntime( + config=SimpleNamespace(model_name="codex/default"), + tool_router=SimpleNamespace(), + hf_token=None, + local_mode=False, + cwd=tmp_path, + autonomous_mode=False, + on_tool=on_tool, + ) + + item = { + "id": "cmd-1", + "type": "commandExecution", + "command": "pwd", + "cwd": str(tmp_path), + "status": "completed", + "exitCode": 0, + "aggregatedOutput": str(tmp_path), + } + await runtime._emit_builtin_item(item, completed=False) + await runtime._emit_builtin_item(item, completed=True) + + assert calls[0] == ( + "codex_command", + {"command": "pwd", "cwd": str(tmp_path)}, + None, + None, + "cmd-1", + ) + assert calls[1][0] == "codex_command" + assert calls[1][2] == str(tmp_path) + assert calls[1][3] is True diff --git a/tests/unit/test_codex_web.py b/tests/unit/test_codex_web.py new file mode 100644 index 00000000..89db0905 --- /dev/null +++ b/tests/unit/test_codex_web.py @@ -0,0 +1,126 @@ +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend" +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + +from codex_web import CodexWebRuntime, codex_web_enabled # noqa: E402 + + +class StubContextManager: + def __init__(self): + self.items = [] + + def add_message(self, message): + self.items.append(message) + + +class StubSession: + def __init__(self): + self.config = SimpleNamespace(model_name="codex/default") + self.context_manager = StubContextManager() + self.events = [] + self.is_cancelled = False + self.turn_count = 0 + + def reset_cancel(self): + self.is_cancelled = False + + async def send_event(self, event): + self.events.append(event) + + def increment_turn(self): + self.turn_count += 1 + + async def auto_save_if_needed(self): + return None + + +def test_codex_web_requires_explicit_local_only_flag(monkeypatch): + monkeypatch.setenv("ML_INTERN_ENABLE_CODEX_WEB", "1") + monkeypatch.delenv("SPACE_ID", raising=False) + monkeypatch.delenv("OAUTH_CLIENT_ID", raising=False) + monkeypatch.setattr("codex_web.shutil.which", lambda _name: "/usr/bin/codex") + + assert codex_web_enabled() is True + + monkeypatch.setenv("SPACE_ID", "owner/space") + assert codex_web_enabled() is False + + +@pytest.mark.asyncio +async def test_web_runtime_emits_sse_contract_and_persists_messages(tmp_path): + session = StubSession() + agent_session = SimpleNamespace( + session=session, + tool_router=SimpleNamespace(), + hf_token=None, + ) + web_runtime = CodexWebRuntime( + agent_session=agent_session, + project_root=tmp_path, + ) + + class Runtime: + async def run_turn(self, prompt, *, stream): + assert prompt == "Find one LoRA paper" + assert stream is True + await web_runtime._on_delta("Found it") + return "Found it" + + web_runtime.runtime = Runtime() + + result = await web_runtime.run_user_input("Find one LoRA paper") + + assert result == "Found it" + assert [message.role for message in session.context_manager.items] == [ + "user", + "assistant", + ] + assert [event.event_type for event in session.events] == [ + "processing", + "assistant_chunk", + "assistant_stream_end", + "turn_complete", + ] + assert session.events[-1].data["final_response"] == "Found it" + assert session.turn_count == 1 + + +@pytest.mark.asyncio +async def test_web_runtime_surfaces_ml_tool_events(tmp_path): + session = StubSession() + web_runtime = CodexWebRuntime( + agent_session=SimpleNamespace( + session=session, + tool_router=SimpleNamespace(), + hf_token=None, + ), + project_root=tmp_path, + ) + + await web_runtime._on_tool( + "hf_papers", + {"query": "LoRA"}, + None, + None, + "call-1", + ) + await web_runtime._on_tool( + "hf_papers", + {"query": "LoRA"}, + "paper result", + True, + "call-1", + ) + + assert [event.event_type for event in session.events] == [ + "tool_call", + "tool_output", + ] + assert session.events[0].data["tool_call_id"] == "call-1" + assert session.events[1].data["success"] is True diff --git a/tests/unit/test_session_manager_persistence.py b/tests/unit/test_session_manager_persistence.py index d9c83286..86f16ae4 100644 --- a/tests/unit/test_session_manager_persistence.py +++ b/tests/unit/test_session_manager_persistence.py @@ -17,6 +17,7 @@ if str(_BACKEND_DIR) not in sys.path: sys.path.insert(0, str(_BACKEND_DIR)) +from agent.core.codex_models import CODEX_DEFAULT_MODEL_ID # noqa: E402 from agent.core.model_ids import GLM_52_MODEL_ID # noqa: E402 from agent.core.session_persistence import NoopSessionStore # noqa: E402 from agent.core.usage_thresholds import USAGE_THRESHOLD_TOOL_NAME # noqa: E402 @@ -457,6 +458,27 @@ async def test_refresh_usage_metrics_missing_token_falls_back_to_app_telemetry() } +@pytest.mark.asyncio +async def test_codex_usage_refresh_never_queries_hf_billing(monkeypatch): + manager = _manager_with_store(NoopSessionStore()) + agent_session = _runtime_agent_session("s1", hf_token="owner-token") + agent_session.session.config.model_name = CODEX_DEFAULT_MODEL_ID + + async def fail_billing_snapshot(*_args, **_kwargs): + raise AssertionError("Codex usage must not query Hugging Face billing") + + monkeypatch.setattr("usage.build_hf_billing_snapshot", fail_billing_snapshot) + + metrics = await manager.refresh_session_usage_metrics(agent_session) + + assert metrics["hf_billing"] == { + "source": "hf_billing_usage_v2", + "available": False, + "error": "codex_uses_chatgpt_allowance", + "current_session": None, + } + + @pytest.mark.asyncio async def test_refresh_usage_metrics_failure_records_error_code(monkeypatch): manager = _manager_with_store(NoopSessionStore()) @@ -859,6 +881,10 @@ def test_unknown_saved_model_defaults_to_glm(): assert model == GLM_52_MODEL_ID +def test_saved_codex_model_is_preserved(): + assert SessionManager._model_from_saved_metadata("codex/default") == "codex/default" + + @pytest.mark.asyncio async def test_update_session_auto_approval_defaults_to_five_dollars(): manager = _manager_with_store(NoopSessionStore()) @@ -1187,6 +1213,21 @@ def fake_start_cpu_sandbox_preload(agent_session: AgentSession) -> None: await _cancel_runtime_tasks(manager) +def test_codex_session_does_not_preload_hf_cpu_sandbox(monkeypatch): + agent_session = _runtime_agent_session("s1", hf_token="owner-token") + agent_session.session.config.model_name = CODEX_DEFAULT_MODEL_ID + + def fail_preload(_session): + raise AssertionError("Codex session must not preload an HF sandbox") + + monkeypatch.setattr( + "agent.tools.sandbox_tool.start_cpu_sandbox_preload", + fail_preload, + ) + + SessionManager._start_cpu_sandbox_preload(agent_session) + + @pytest.mark.asyncio async def test_lazy_restore_schedules_cpu_sandbox_preload(): manager = _manager_with_store(RestoreStore()) From ba1b7f46271a669631e06638a8d49c20d60ed301 Mon Sep 17 00:00:00 2001 From: tibetyalman <222756182+tibetyalman@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:32:26 +0300 Subject: [PATCH 3/5] feat: add Codex model and thinking controls --- README.md | 9 +- agent/core/codex_runtime.py | 109 +++++++- agent/core/session_persistence.py | 4 + backend/codex_web.py | 89 +++++- backend/models.py | 2 + backend/routes/agent.py | 133 +++++++-- backend/session_manager.py | 40 ++- frontend/src/components/Chat/ChatInput.tsx | 256 ++++++++++++++++-- tests/unit/test_agent_model_gating.py | 119 ++++++++ tests/unit/test_codex_runtime.py | 137 ++++++++++ tests/unit/test_codex_web.py | 47 +++- .../unit/test_session_manager_persistence.py | 7 +- 12 files changed, 902 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index 7ceaddf4..be2edbfa 100644 --- a/README.md +++ b/README.md @@ -96,9 +96,12 @@ To expose the same signed-in Codex runtime in the local Web UI: ML_INTERN_ENABLE_CODEX_WEB=1 uv run uvicorn backend.main:app --host ::1 --port 7860 ``` -`Codex (ChatGPT)` then appears as the recommended/default Web model. This flag -is deliberately disabled when the backend detects Hugging Face Spaces or OAuth: -a developer's local ChatGPT allowance must never be shared with hosted visitors. +The Web picker then loads the models available to the active Codex account and +their supported reasoning levels directly from `codex app-server`. `Codex Auto` +is the recommended/default choice; explicit Codex models and Thinking levels +can be changed per session. This flag is deliberately disabled when the backend +detects Hugging Face Spaces or OAuth: a developer's local ChatGPT allowance +must never be shared with hosted visitors. Codex supplies the main agent loop. In the CLI it also has local repository tools; in the Web UI the host repository stays read-only. ML Intern's Hugging diff --git a/agent/core/codex_runtime.py b/agent/core/codex_runtime.py index 1bf6e1ec..9eaa1065 100644 --- a/agent/core/codex_runtime.py +++ b/agent/core/codex_runtime.py @@ -91,6 +91,103 @@ async def codex_login_status(codex_bin: str = "codex") -> str: return text or "Codex authentication active" +async def codex_model_catalog( + codex_bin: str = "codex", + *, + timeout_s: float = 10.0, +) -> list[dict[str, Any]]: + """Return the picker-visible model catalog for the active Codex account.""" + await codex_login_status(codex_bin) + resolved = shutil.which(codex_bin) + assert resolved is not None + + process = await asyncio.create_subprocess_exec( + resolved, + "app-server", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + assert process.stdin is not None + assert process.stdout is not None + assert process.stderr is not None + stderr_task = asyncio.create_task(process.stderr.read()) + request_id = 0 + + async def write(message: dict[str, Any]) -> None: + payload = (json.dumps(message, separators=(",", ":")) + "\n").encode() + process.stdin.write(payload) + await process.stdin.drain() + + async def request(method: str, params: dict[str, Any]) -> dict[str, Any]: + nonlocal request_id + request_id += 1 + current_id = request_id + await write({"method": method, "id": current_id, "params": params}) + while True: + line = await asyncio.wait_for(process.stdout.readline(), timeout=timeout_s) + if not line: + raise CodexRuntimeError( + "Codex app-server closed while loading its model catalog." + ) + try: + message = json.loads(line) + except json.JSONDecodeError: + continue + if message.get("id") != current_id: + continue + if "error" in message: + error = message.get("error") or {} + raise CodexRuntimeError( + str(error.get("message") or f"Codex {method} request failed.") + ) + return message.get("result") or {} + + try: + await request( + "initialize", + { + "clientInfo": { + "name": "ml_intern", + "title": "ML Intern", + "version": "0.1.0", + }, + "capabilities": {"experimentalApi": True}, + }, + ) + await write({"method": "initialized", "params": {}}) + + models: list[dict[str, Any]] = [] + cursor: str | None = None + while True: + params: dict[str, Any] = { + "limit": 100, + "includeHidden": False, + } + if cursor: + params["cursor"] = cursor + result = await request("model/list", params) + models.extend( + model for model in result.get("data") or [] if isinstance(model, dict) + ) + cursor = result.get("nextCursor") + if not cursor: + return models + except TimeoutError as exc: + raise CodexRuntimeError("Timed out while loading Codex models.") from exc + finally: + if process.returncode is None: + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=5) + except asyncio.TimeoutError: + process.kill() + await process.wait() + stderr = (await stderr_task).decode(errors="replace").strip() + if process.returncode not in {None, 0, -15} and stderr: + logger.debug("Codex model catalog stderr: %s", stderr[-1000:]) + + def build_dynamic_tool_namespace( tool_router: ToolRouter, *, @@ -369,12 +466,16 @@ async def run_turn(self, prompt: str, *, stream: bool = True) -> str: if not self.thread_id: raise CodexRuntimeError("Codex runtime has not been started.") + params: dict[str, Any] = { + "threadId": self.thread_id, + "input": [{"type": "text", "text": prompt}], + } + reasoning_effort = getattr(self.config, "reasoning_effort", None) + if reasoning_effort: + params["effort"] = reasoning_effort response = await self._request( "turn/start", - { - "threadId": self.thread_id, - "input": [{"type": "text", "text": prompt}], - }, + params, ) turn = response.get("turn") or {} self.active_turn_id = turn.get("id") diff --git a/agent/core/session_persistence.py b/agent/core/session_persistence.py index 760b5934..c6e8afba 100644 --- a/agent/core/session_persistence.py +++ b/agent/core/session_persistence.py @@ -169,6 +169,7 @@ async def upsert_session( session_id: str, user_id: str, model: str, + reasoning_effort: str | None = None, title: str | None = None, surface: str = "frontend", created_at: datetime | None = None, @@ -203,6 +204,7 @@ async def upsert_session( "$set": { "title": title, "model": model, + "reasoning_effort": reasoning_effort, "usage_window_started_at": ( usage_window_started_at or created_at or now ), @@ -231,6 +233,7 @@ async def save_snapshot( user_id: str, model: str, messages: list[dict[str, Any]], + reasoning_effort: str | None = None, title: str | None = None, runtime_state: str = "idle", status: str = "active", @@ -255,6 +258,7 @@ async def save_snapshot( session_id=session_id, user_id=user_id, model=model, + reasoning_effort=reasoning_effort, title=title, created_at=created_at, runtime_state=runtime_state, diff --git a/backend/codex_web.py b/backend/codex_web.py index 7e87fdda..243c54c4 100644 --- a/backend/codex_web.py +++ b/backend/codex_web.py @@ -2,18 +2,24 @@ from __future__ import annotations +import asyncio import os import shutil +import time from pathlib import Path from typing import Any from litellm import Message -from agent.core.codex_runtime import CodexAppServerRuntime +from agent.core.codex_models import CODEX_DEFAULT_MODEL_ID +from agent.core.codex_runtime import CodexAppServerRuntime, codex_model_catalog from agent.core.session import Event _TRUE_VALUES = {"1", "true", "yes", "on"} _MAX_SEED_CHARS = 24_000 +_MODEL_CATALOG_TTL_S = 300.0 +_model_catalog_cache: tuple[float, list[dict[str, Any]]] | None = None +_model_catalog_lock = asyncio.Lock() def codex_web_enabled() -> bool: @@ -31,6 +37,87 @@ def codex_web_enabled() -> bool: return shutil.which("codex") is not None +def _reasoning_options(model: dict[str, Any]) -> list[dict[str, str]]: + options = [] + for item in model.get("supportedReasoningEfforts") or []: + effort = item.get("reasoningEffort") if isinstance(item, dict) else None + if not isinstance(effort, str) or not effort: + continue + options.append( + { + "id": effort, + "description": str(item.get("description") or ""), + } + ) + return options + + +def _catalog_to_web_models(catalog: list[dict[str, Any]]) -> list[dict[str, Any]]: + visible = [model for model in catalog if not model.get("hidden")] + default = next((model for model in visible if model.get("isDefault")), None) + models: list[dict[str, Any]] = [] + + auto: dict[str, Any] = { + "id": CODEX_DEFAULT_MODEL_ID, + "label": "Codex Auto", + "provider": "codex", + "recommended": True, + } + if default is not None: + display_name = str(default.get("displayName") or default.get("id") or "") + if display_name: + auto["label"] = f"Codex Auto ({display_name})" + auto["description"] = str(default.get("description") or "") + auto["default_reasoning_effort"] = default.get("defaultReasoningEffort") + auto["reasoning_efforts"] = _reasoning_options(default) + models.append(auto) + + for model in visible: + model_id = model.get("model") or model.get("id") + if not isinstance(model_id, str) or not model_id: + continue + models.append( + { + "id": f"codex/{model_id}", + "label": f"Codex · {model.get('displayName') or model_id}", + "provider": "codex", + "description": str(model.get("description") or ""), + "default_reasoning_effort": model.get("defaultReasoningEffort"), + "reasoning_efforts": _reasoning_options(model), + "upgrade": model.get("upgrade"), + } + ) + return models + + +async def codex_web_models(*, force_refresh: bool = False) -> list[dict[str, Any]]: + """Return a short-lived cached model catalog for the local Codex login.""" + global _model_catalog_cache + + if not codex_web_enabled(): + return [] + now = time.monotonic() + if ( + not force_refresh + and _model_catalog_cache is not None + and _model_catalog_cache[0] > now + ): + return _model_catalog_cache[1] + + async with _model_catalog_lock: + now = time.monotonic() + if ( + not force_refresh + and _model_catalog_cache is not None + and _model_catalog_cache[0] > now + ): + return _model_catalog_cache[1] + catalog = await codex_model_catalog() + models = _catalog_to_web_models(catalog) + _model_catalog_cache = (now + _MODEL_CATALOG_TTL_S, models) + return models + + def _seeded_prompt(session, current_text: str) -> str: """Seed a fresh Codex thread from persisted ML Intern chat messages.""" prior_messages = [] diff --git a/backend/models.py b/backend/models.py index 0c06cef3..c2b12873 100644 --- a/backend/models.py +++ b/backend/models.py @@ -69,6 +69,7 @@ class SessionResponse(BaseModel): session_id: str ready: bool = True model: str | None = None + reasoning_effort: str | None = None class PendingApprovalTool(BaseModel): @@ -100,6 +101,7 @@ class SessionInfo(BaseModel): user_id: str = "dev" pending_approval: list[PendingApprovalTool] | None = None model: str | None = None + reasoning_effort: str | None = None title: str | None = None notification_destinations: list[str] = Field(default_factory=list) auto_approval: SessionAutoApprovalInfo = Field( diff --git a/backend/routes/agent.py b/backend/routes/agent.py index 03c8c6e7..7d6be987 100644 --- a/backend/routes/agent.py +++ b/backend/routes/agent.py @@ -50,7 +50,7 @@ SessionCapacityError, session_manager, ) -from codex_web import codex_web_enabled +from codex_web import codex_web_enabled, codex_web_models from agent.core.codex_models import CODEX_DEFAULT_MODEL_ID, is_codex_model_id from agent.core.codex_runtime import CodexRuntimeError, codex_login_status @@ -122,6 +122,7 @@ def _schedule_usage_refresh_and_upload( def _available_models( *, include_codex: bool | None = None, + codex_models: list[dict[str, Any]] | None = None, ) -> list[dict[str, Any]]: include_codex = codex_web_enabled() if include_codex is None else include_codex models = [ @@ -154,31 +155,84 @@ def _available_models( if include_codex: for model in models: model.pop("recommended", None) - models.insert( - 0, + codex_options = codex_models or [ { "id": CODEX_DEFAULT_MODEL_ID, "label": "Codex (ChatGPT)", "provider": "codex", "recommended": True, - }, - ) + } + ] + models[0:0] = [dict(model) for model in codex_options] return models AVAILABLE_MODELS = _available_models(include_codex=False) -def _valid_model_ids() -> set[str]: - return {m["id"] for m in _available_models()} +async def _live_available_models() -> list[dict[str, Any]]: + codex_models: list[dict[str, Any]] | None = None + if codex_web_enabled(): + try: + codex_models = await codex_web_models() + except CodexRuntimeError as e: + logger.warning("Could not load Codex model catalog: %s", e) + return _available_models(codex_models=codex_models) -def _validate_model_id(model_id: str | None) -> None: - if not model_id or model_id in _valid_model_ids(): +def _valid_model_ids(models: list[dict[str, Any]]) -> set[str]: + return {str(model["id"]) for model in models} + + +def _validate_model_id( + model_id: str | None, + models: list[dict[str, Any]], +) -> None: + if not model_id or model_id in _valid_model_ids(models): return raise HTTPException(status_code=400, detail=f"Unknown model: {model_id}") +def _resolve_codex_reasoning_effort( + model_id: str, + requested_effort: Any, + models: list[dict[str, Any]], + *, + current_effort: str | None = None, +) -> str | None: + if not is_codex_model_id(model_id): + if requested_effort is not None: + raise HTTPException( + status_code=400, + detail="Reasoning effort is only configurable for Codex models.", + ) + return None + + option = next((model for model in models if model.get("id") == model_id), None) + if option is None: + raise HTTPException(status_code=400, detail=f"Unknown model: {model_id}") + supported = [ + str(item["id"]) + for item in option.get("reasoning_efforts") or [] + if isinstance(item, dict) and item.get("id") + ] + if requested_effort is not None: + if not isinstance(requested_effort, str) or requested_effort not in supported: + raise HTTPException( + status_code=400, + detail=( + f"Unsupported reasoning effort for {model_id}: {requested_effort}" + ), + ) + return requested_effort + if current_effort in supported: + return current_effort + default = option.get("default_reasoning_effort") + if isinstance(default, str) and default in supported: + return default + return supported[0] if supported else None + + def _default_model() -> str: if codex_web_enabled(): return CODEX_DEFAULT_MODEL_ID @@ -388,7 +442,7 @@ async def get_model() -> dict: """Get current model and available models. No auth required.""" return { "current": _default_model(), - "available": _available_models(), + "available": await _live_available_models(), } @@ -488,13 +542,21 @@ async def create_session( body = await request.json() except Exception: body = None + reasoning_effort = None if isinstance(body, dict): model = body.get("model") + reasoning_effort = body.get("reasoning_effort") - _validate_model_id(model) + models = await _live_available_models() + _validate_model_id(model, models) # Empty requests use the web default. model = _model_override_for_new_session(model) + reasoning_effort = _resolve_codex_reasoning_effort( + model, + reasoning_effort, + models, + ) try: session_id = await session_manager.create_session( @@ -503,6 +565,7 @@ async def create_session( hf_token=hf_token, user_plan=user.get("plan"), model=model, + reasoning_effort=reasoning_effort, is_pro=user.get("plan") == "pro", ) except SessionCapacityError as e: @@ -514,6 +577,7 @@ async def create_session( session_id=session_id, ready=True, model=model, + reasoning_effort=reasoning_effort, ) @@ -536,9 +600,15 @@ async def restore_session_summary( hf_token = resolve_hf_request_token(request) model = body.get("model") - _validate_model_id(model) + models = await _live_available_models() + _validate_model_id(model, models) model = _model_override_for_new_session(model) + reasoning_effort = _resolve_codex_reasoning_effort( + model, + body.get("reasoning_effort"), + models, + ) try: session_id = await session_manager.create_session( @@ -547,6 +617,7 @@ async def restore_session_summary( hf_token=hf_token, user_plan=user.get("plan"), model=model, + reasoning_effort=reasoning_effort, is_pro=user.get("plan") == "pro", ) except SessionCapacityError as e: @@ -576,6 +647,7 @@ async def restore_session_summary( session_id=session_id, ready=True, model=model, + reasoning_effort=reasoning_effort, ) @@ -619,15 +691,44 @@ async def set_session_model( model_id = body.get("model") if not model_id: raise HTTPException(status_code=400, detail="Missing 'model' field") - _validate_model_id(model_id) + models = await _live_available_models() + _validate_model_id(model_id, models) if not agent_session: raise HTTPException(status_code=404, detail="Session not found") - await session_manager.update_session_model(session_id, model_id) + current_session = getattr(agent_session, "session", None) + current_config = getattr(current_session, "config", None) + current_model = getattr(current_config, "model_name", None) + current_effort = ( + getattr(current_config, "reasoning_effort", None) + if is_codex_model_id(current_model) + else None + ) + reasoning_effort = _resolve_codex_reasoning_effort( + model_id, + body.get("reasoning_effort"), + models, + current_effort=current_effort, + ) + if is_codex_model_id(model_id): + await session_manager.update_session_model( + session_id, + model_id, + reasoning_effort=reasoning_effort, + ) + else: + await session_manager.update_session_model(session_id, model_id) logger.info( f"Session {session_id} model → {model_id} " - f"(by {user.get('username', 'unknown')})" + f"(reasoning={reasoning_effort or 'default'}, " + f"by {user.get('username', 'unknown')})" ) - return {"session_id": session_id, "model": model_id} + response = { + "session_id": session_id, + "model": model_id, + } + if is_codex_model_id(model_id): + response["reasoning_effort"] = reasoning_effort + return response @router.post("/session/{session_id}/notifications") diff --git a/backend/session_manager.py b/backend/session_manager.py index 9289f942..50133ef0 100644 --- a/backend/session_manager.py +++ b/backend/session_manager.py @@ -140,6 +140,7 @@ class AgentSession: usage_warning_spend_cache: dict[str, Any] = field(default_factory=dict) codex_runtime: CodexWebRuntime | None = None codex_runtime_model: str | None = None + codex_runtime_effort: str | None = None def __post_init__(self) -> None: if self.usage_window_started_at is None: @@ -293,6 +294,7 @@ def _create_session_sync( hf_token: str | None, user_plan: str | None, model: str | None, + reasoning_effort: str | None, event_queue: asyncio.Queue, notification_destinations: list[str] | None = None, ) -> tuple[ToolRouter, Session]: @@ -307,6 +309,10 @@ def _create_session_sync( normalized_model = strip_huggingface_model_prefix(model) if normalized_model: session_config.model_name = normalized_model + if is_codex_model_id(session_config.model_name): + session_config.reasoning_effort = reasoning_effort + elif reasoning_effort is not None: + session_config.reasoning_effort = reasoning_effort session = Session( event_queue=event_queue, config=session_config, @@ -1032,6 +1038,11 @@ async def persist_session_snapshot( session_id=agent_session.session_id, user_id=agent_session.user_id, model=agent_session.session.config.model_name, + reasoning_effort=getattr( + agent_session.session.config, + "reasoning_effort", + None, + ), title=agent_session.title, messages=self._serialize_messages(agent_session.session), runtime_state=runtime_state or self._runtime_state(agent_session), @@ -1159,6 +1170,7 @@ async def ensure_session_loaded( hf_token=hf_token, user_plan=user_plan, model=model, + reasoning_effort=meta.get("reasoning_effort"), event_queue=event_queue, notification_destinations=meta.get("notification_destinations") or [], ) @@ -1278,6 +1290,7 @@ async def create_session( hf_token: str | None = None, user_plan: str | None = None, model: str | None = None, + reasoning_effort: str | None = None, is_pro: bool | None = None, ) -> str: """Create a new agent session and return its ID. @@ -1294,6 +1307,8 @@ async def create_session( model: Optional model override. When set, replaces ``model_name`` on the per-session config clone. None falls back to the config default. + reasoning_effort: Optional per-session reasoning effort. Codex + sessions leave this unset to use the selected model's default. Raises: SessionCapacityError: If the server or user has reached the @@ -1343,6 +1358,7 @@ async def create_session( hf_token=hf_token, user_plan=user_plan, model=model, + reasoning_effort=reasoning_effort, event_queue=event_queue, ) @@ -1683,6 +1699,7 @@ async def _close_codex_runtime(self, agent_session: AgentSession) -> None: runtime = agent_session.codex_runtime agent_session.codex_runtime = None agent_session.codex_runtime_model = None + agent_session.codex_runtime_effort = None if runtime is not None: try: await runtime.close() @@ -1698,9 +1715,11 @@ async def _ensure_codex_runtime( agent_session: AgentSession, ) -> CodexWebRuntime: model_id = agent_session.session.config.model_name + reasoning_effort = agent_session.session.config.reasoning_effort if ( agent_session.codex_runtime is not None and agent_session.codex_runtime_model == model_id + and agent_session.codex_runtime_effort == reasoning_effort ): return agent_session.codex_runtime @@ -1712,6 +1731,7 @@ async def _ensure_codex_runtime( await runtime.start() agent_session.codex_runtime = runtime agent_session.codex_runtime_model = model_id + agent_session.codex_runtime_effort = reasoning_effort return runtime async def _run_session( @@ -1957,15 +1977,26 @@ async def update_session_title(self, session_id: str, title: str | None) -> None agent_session.title = title await self._store().update_session_fields(session_id, title=title) - async def update_session_model(self, session_id: str, model_id: str) -> bool: + async def update_session_model( + self, + session_id: str, + model_id: str, + *, + reasoning_effort: str | None = None, + ) -> bool: agent_session = self.sessions.get(session_id) if not agent_session or not agent_session.is_active: return False agent_session.session.update_model(model_id) + if is_codex_model_id(model_id): + agent_session.session.config.reasoning_effort = reasoning_effort if ( not agent_session.is_processing and agent_session.codex_runtime is not None - and agent_session.codex_runtime_model != model_id + and ( + agent_session.codex_runtime_model != model_id + or agent_session.codex_runtime_effort != reasoning_effort + ) ): await self._close_codex_runtime(agent_session) self._touch(agent_session) @@ -2059,6 +2090,11 @@ def get_session_info(self, session_id: str) -> dict[str, Any] | None: "user_id": agent_session.user_id, "pending_approval": pending_approval, "model": agent_session.session.config.model_name, + "reasoning_effort": getattr( + agent_session.session.config, + "reasoning_effort", + None, + ), "title": agent_session.title, "notification_destinations": list( agent_session.session.notification_destinations diff --git a/frontend/src/components/Chat/ChatInput.tsx b/frontend/src/components/Chat/ChatInput.tsx index ffed5f27..0b463166 100644 --- a/frontend/src/components/Chat/ChatInput.tsx +++ b/frontend/src/components/Chat/ChatInput.tsx @@ -19,6 +19,7 @@ import ArrowUpwardIcon from '@mui/icons-material/ArrowUpward'; import ArrowDropDownIcon from '@mui/icons-material/ArrowDropDown'; import StopIcon from '@mui/icons-material/Stop'; import AddIcon from '@mui/icons-material/Add'; +import PsychologyAltIcon from '@mui/icons-material/PsychologyAlt'; import { apiFetch, apiUpload } from '@/utils/api'; import JobsUpgradeDialog from '@/components/JobsUpgradeDialog'; import { useAgentStore } from '@/store/agentStore'; @@ -40,9 +41,18 @@ interface ModelOption { name: string; modelPath: string; avatarUrl: string; + provider?: string; + description?: string; + defaultReasoningEffort?: string; + reasoningEfforts?: ReasoningEffortOption[]; recommended?: boolean; } +interface ReasoningEffortOption { + id: string; + description?: string; +} + const getHfAvatarUrl = (modelId: string) => { const org = modelId.split('/')[0]; return `https://huggingface.co/api/avatars/${org}`; @@ -126,6 +136,9 @@ const modelOptionFromApi = (model: { id?: string; label?: string; provider?: string; + description?: string; + default_reasoning_effort?: string; + reasoning_efforts?: ReasoningEffortOption[]; recommended?: boolean; }): ModelOption | null => { if (!model.id) return null; @@ -137,10 +150,31 @@ const modelOptionFromApi = (model: { name: model.label ?? model.id, modelPath: model.id, avatarUrl, + provider: model.provider, + description: model.description, + defaultReasoningEffort: model.default_reasoning_effort, + reasoningEfforts: model.reasoning_efforts ?? [], recommended: Boolean(model.recommended), }; }; +const isCodexModel = (model: ModelOption | undefined) => ( + model?.provider === 'codex' || model?.modelPath.startsWith('codex/') +); + +const reasoningEffortLabel = (effort: string | undefined) => { + switch (effort) { + case 'minimal': return 'Minimal'; + case 'low': return 'Low'; + case 'medium': return 'Medium'; + case 'high': return 'High'; + case 'xhigh': return 'Extra High'; + case 'max': return 'Max'; + case 'ultra': return 'Ultra'; + default: return effort || 'Default'; + } +}; + const readApiErrorMessage = async (res: Response, fallback: string): Promise => { try { const data = await res.json(); @@ -210,6 +244,8 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, ), ); const [modelAnchorEl, setModelAnchorEl] = useState(null); + const [reasoningAnchorEl, setReasoningAnchorEl] = useState(null); + const [selectedReasoningEffort, setSelectedReasoningEffort] = useState(null); const jobsUpgradeRequired = useAgentStore((s) => s.jobsUpgradeRequired); const setJobsUpgradeRequired = useAgentStore((s) => s.setJobsUpgradeRequired); const updateSessionModel = useSessionStore((s) => s.updateSessionModel); @@ -243,7 +279,10 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, setModelOptions(next); if (!sessionIdRef.current) { const current = data.current ? findModelByPath(data.current, next) : null; - if (current) setSelectedModelPath(current.modelPath); + if (current) { + setSelectedModelPath(current.modelPath); + setSelectedReasoningEffort(current.defaultReasoningEffort ?? null); + } } }) .catch(() => { /* ignore */ }); @@ -262,6 +301,9 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, if (data?.model) { const model = findModelByPath(data.model, modelOptionsRef.current); setSelectedModelPath(model?.modelPath ?? data.model); + setSelectedReasoningEffort( + data.reasoning_effort ?? model?.defaultReasoningEffort ?? null, + ); updateSessionModel(sessionId, data.model); } }) @@ -278,6 +320,24 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, || visibleModelOptions[0] || modelOptions[0] ); + const selectedReasoningOptions = selectedModel.reasoningEfforts ?? []; + const selectedReasoningLabel = reasoningEffortLabel( + selectedReasoningEffort ?? selectedModel.defaultReasoningEffort, + ); + + useEffect(() => { + if (!selectedModelPath.startsWith('codex/')) { + setSelectedReasoningEffort(null); + return; + } + const model = findModelByPath(selectedModelPath, modelOptions); + if (!model) return; + const supported = new Set((model.reasoningEfforts ?? []).map((effort) => effort.id)); + setSelectedReasoningEffort((current) => { + if (current && supported.has(current)) return current; + return model.defaultReasoningEffort ?? model.reasoningEfforts?.[0]?.id ?? null; + }); + }, [modelOptions, selectedModelPath]); // Auto-focus the textarea when the session becomes ready useEffect(() => { @@ -386,16 +446,40 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, setModelAnchorEl(null); }; + const handleReasoningClick = (event: React.MouseEvent) => { + setReasoningAnchorEl(event.currentTarget); + }; + + const handleReasoningClose = () => { + setReasoningAnchorEl(null); + }; + const handleSelectModel = async (model: ModelOption) => { handleModelClose(); if (!sessionId) return; + const supported = new Set((model.reasoningEfforts ?? []).map((effort) => effort.id)); + const nextReasoningEffort = isCodexModel(model) + ? ( + (selectedReasoningEffort && supported.has(selectedReasoningEffort) + ? selectedReasoningEffort + : null) + ?? model.defaultReasoningEffort + ?? model.reasoningEfforts?.[0]?.id + ?? null + ) + : null; try { const res = await apiFetch(`/api/session/${sessionId}/model`, { method: 'POST', - body: JSON.stringify({ model: model.modelPath }), + body: JSON.stringify({ + model: model.modelPath, + ...(nextReasoningEffort ? { reasoning_effort: nextReasoningEffort } : {}), + }), }); if (res.ok) { + const data = await res.json(); setSelectedModelPath(model.modelPath); + setSelectedReasoningEffort(data.reasoning_effort ?? nextReasoningEffort); updateSessionModel(sessionId, model.modelPath); setModelSwitchError(null); return; @@ -406,6 +490,31 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, } }; + const handleSelectReasoningEffort = async (effort: ReasoningEffortOption) => { + handleReasoningClose(); + if (!sessionId || !isCodexModel(selectedModel)) return; + try { + const res = await apiFetch(`/api/session/${sessionId}/model`, { + method: 'POST', + body: JSON.stringify({ + model: selectedModel.modelPath, + reasoning_effort: effort.id, + }), + }); + if (res.ok) { + const data = await res.json(); + setSelectedReasoningEffort(data.reasoning_effort ?? effort.id); + setModelSwitchError(null); + return; + } + setModelSwitchError(await readApiErrorMessage(res, 'Could not change thinking level.')); + } catch (error) { + setModelSwitchError( + error instanceof Error ? error.message : 'Could not change thinking level.', + ); + } + }; + const handleJobsUpgradeClose = useCallback(() => { setJobsUpgradeRequired(null); setAwaitingTopUp(false); @@ -655,33 +764,70 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, {/* Powered By Badge */} - - powered by - - {selectedModel.name} - - {selectedModel.name} - - + + + powered by + + {selectedModel.name} + + {selectedModel.name} + + + + {isCodexModel(selectedModel) && selectedReasoningOptions.length > 0 && ( + <> + · + + + + + Thinking: {selectedReasoningLabel} + + + + + + )} {/* Model Selection Menu */} @@ -746,6 +892,72 @@ export default function ChatInput({ sessionId, initialModelPath, onSend, onStop, )} } + secondary={model.description || undefined} + secondaryTypographyProps={{ + sx: { + color: 'var(--muted-text)', + fontSize: '11px', + maxWidth: 360, + whiteSpace: 'normal', + lineHeight: 1.3, + }, + }} + /> + + ))} + + + + {selectedReasoningOptions.map((effort) => ( + handleSelectReasoningEffort(effort)} + selected={selectedReasoningEffort === effort.id} + sx={{ + py: 1.2, + '&.Mui-selected': { + bgcolor: 'rgba(255,255,255,0.05)', + }, + }} + > + + + + ))} diff --git a/tests/unit/test_agent_model_gating.py b/tests/unit/test_agent_model_gating.py index 26a5d77e..71903eeb 100644 --- a/tests/unit/test_agent_model_gating.py +++ b/tests/unit/test_agent_model_gating.py @@ -57,6 +57,43 @@ def test_local_codex_web_is_listed_and_becomes_default(monkeypatch): assert agent._model_override_for_new_session(None) == agent.CODEX_DEFAULT_MODEL_ID +@pytest.mark.asyncio +async def test_model_config_uses_live_codex_catalog(monkeypatch): + async def fake_codex_models(): + return [ + { + "id": agent.CODEX_DEFAULT_MODEL_ID, + "label": "Codex Auto (GPT-5.6-Sol)", + "provider": "codex", + "recommended": True, + "default_reasoning_effort": "low", + "reasoning_efforts": [ + {"id": "low", "description": "Fast"}, + {"id": "max", "description": "Deep"}, + ], + }, + { + "id": "codex/gpt-5.6-terra", + "label": "Codex · GPT-5.6-Terra", + "provider": "codex", + "default_reasoning_effort": "medium", + "reasoning_efforts": [{"id": "medium", "description": "Balanced"}], + }, + ] + + monkeypatch.setattr(agent, "codex_web_enabled", lambda: True) + monkeypatch.setattr(agent, "codex_web_models", fake_codex_models) + + payload = await agent.get_model() + + assert payload["current"] == agent.CODEX_DEFAULT_MODEL_ID + assert [model["id"] for model in payload["available"][:2]] == [ + agent.CODEX_DEFAULT_MODEL_ID, + "codex/gpt-5.6-terra", + ] + assert payload["available"][0]["reasoning_efforts"][1]["id"] == "max" + + @pytest.mark.asyncio async def test_llm_health_uses_default_and_request_hf_token(monkeypatch): class Request: @@ -314,6 +351,88 @@ async def fake_update_session_model(session_id, model_id): assert updated == [("s1", agent.DEFAULT_GPT_MODEL_ID)] +@pytest.mark.asyncio +async def test_switching_codex_model_and_reasoning_effort(monkeypatch): + updated = [] + codex_models = [ + { + "id": "codex/gpt-5.6-sol", + "label": "Codex · GPT-5.6-Sol", + "provider": "codex", + "default_reasoning_effort": "low", + "reasoning_efforts": [ + {"id": "low", "description": "Fast"}, + {"id": "max", "description": "Deep"}, + ], + } + ] + + async def fake_live_models(): + return codex_models + + async def fake_check_session_access(_session_id, _user, _request=None): + return SimpleNamespace( + session=SimpleNamespace( + config=SimpleNamespace( + model_name=agent.CODEX_DEFAULT_MODEL_ID, + reasoning_effort="low", + ) + ) + ) + + async def fake_update_session_model(session_id, model_id, **kwargs): + updated.append((session_id, model_id, kwargs)) + + monkeypatch.setattr(agent, "_live_available_models", fake_live_models) + monkeypatch.setattr(agent, "_check_session_access", fake_check_session_access) + monkeypatch.setattr( + agent.session_manager, + "update_session_model", + fake_update_session_model, + ) + + response = await agent.set_session_model( + "s1", + { + "model": "codex/gpt-5.6-sol", + "reasoning_effort": "max", + }, + request=None, + user={"user_id": "u1"}, + ) + + assert response == { + "session_id": "s1", + "model": "codex/gpt-5.6-sol", + "reasoning_effort": "max", + } + assert updated == [ + ( + "s1", + "codex/gpt-5.6-sol", + {"reasoning_effort": "max"}, + ) + ] + + +def test_codex_rejects_unsupported_reasoning_effort(): + with pytest.raises(HTTPException) as exc_info: + agent._resolve_codex_reasoning_effort( + "codex/gpt-5.6-sol", + "ultra", + [ + { + "id": "codex/gpt-5.6-sol", + "reasoning_efforts": [{"id": "low"}], + "default_reasoning_effort": "low", + } + ], + ) + + assert exc_info.value.status_code == 400 + assert "Unsupported reasoning effort" in exc_info.value.detail + + @pytest.mark.asyncio async def test_switching_to_unknown_model_id_is_rejected(monkeypatch): async def fake_check_session_access(session_id, user, request=None): diff --git a/tests/unit/test_codex_runtime.py b/tests/unit/test_codex_runtime.py index decbc9b3..46351364 100644 --- a/tests/unit/test_codex_runtime.py +++ b/tests/unit/test_codex_runtime.py @@ -8,6 +8,7 @@ CodexRuntimeError, build_dynamic_tool_namespace, codex_login_status, + codex_model_catalog, ) from agent.core.tools import ToolSpec @@ -105,6 +106,142 @@ async def fake_subprocess(*_args, **_kwargs): assert await codex_login_status() == "Logged in using ChatGPT" +@pytest.mark.asyncio +async def test_codex_model_catalog_uses_account_visible_models(monkeypatch): + writes = [] + + class Stdin: + def write(self, payload): + writes.append(payload.decode()) + + async def drain(self): + return None + + class Stdout: + def __init__(self): + self.lines = [ + b'{"id":1,"result":{"userAgent":"Codex"}}\n', + b'{"method":"remoteControl/status/changed","params":{}}\n', + ( + b'{"id":2,"result":{"data":[{"id":"gpt-5.6-sol",' + b'"model":"gpt-5.6-sol","displayName":"GPT-5.6-Sol",' + b'"defaultReasoningEffort":"low",' + b'"supportedReasoningEfforts":[{"reasoningEffort":"low"}],' + b'"isDefault":true}],"nextCursor":null}}\n' + ), + ] + + async def readline(self): + return self.lines.pop(0) + + class Stderr: + async def read(self): + return b"" + + class Process: + def __init__(self): + self.stdin = Stdin() + self.stdout = Stdout() + self.stderr = Stderr() + self.returncode = None + + def terminate(self): + self.returncode = -15 + + def kill(self): + self.returncode = -9 + + async def wait(self): + return self.returncode + + async def fake_login_status(_codex_bin): + return "Logged in using ChatGPT" + + async def fake_subprocess(*_args, **_kwargs): + return Process() + + monkeypatch.setattr( + "agent.core.codex_runtime.codex_login_status", + fake_login_status, + ) + monkeypatch.setattr( + "agent.core.codex_runtime.shutil.which", + lambda _binary: "/usr/local/bin/codex", + ) + monkeypatch.setattr( + "agent.core.codex_runtime.asyncio.create_subprocess_exec", + fake_subprocess, + ) + + models = await codex_model_catalog() + + assert models[0]["model"] == "gpt-5.6-sol" + assert models[0]["defaultReasoningEffort"] == "low" + assert '"method":"model/list"' in writes[-1] + assert '"includeHidden":false' in writes[-1] + + +@pytest.mark.asyncio +async def test_run_turn_sends_selected_reasoning_effort(tmp_path): + requests = [] + runtime = CodexAppServerRuntime( + config=SimpleNamespace( + model_name="codex/gpt-5.6-sol", + reasoning_effort="max", + ), + tool_router=SimpleNamespace(), + hf_token=None, + local_mode=False, + cwd=tmp_path, + autonomous_mode=False, + ) + runtime.thread_id = "thread-1" + + async def fake_request(method, params): + requests.append((method, params)) + return {"turn": {"id": "turn-1"}} + + runtime._request = fake_request + await runtime._notifications.put( + { + "method": "item/completed", + "params": { + "threadId": "thread-1", + "turnId": "turn-1", + "item": { + "type": "agentMessage", + "phase": "final_answer", + "text": "done", + }, + }, + } + ) + await runtime._notifications.put( + { + "method": "turn/completed", + "params": { + "threadId": "thread-1", + "turnId": "turn-1", + "turn": {"id": "turn-1", "status": "completed"}, + }, + } + ) + + result = await runtime.run_turn("hello") + + assert result == "done" + assert requests == [ + ( + "turn/start", + { + "threadId": "thread-1", + "input": [{"type": "text", "text": "hello"}], + "effort": "max", + }, + ) + ] + + @pytest.mark.asyncio async def test_dynamic_tool_failure_is_returned_to_codex(tmp_path): class FailingRouter: diff --git a/tests/unit/test_codex_web.py b/tests/unit/test_codex_web.py index 89db0905..566882ae 100644 --- a/tests/unit/test_codex_web.py +++ b/tests/unit/test_codex_web.py @@ -8,7 +8,11 @@ if str(_BACKEND_DIR) not in sys.path: sys.path.insert(0, str(_BACKEND_DIR)) -from codex_web import CodexWebRuntime, codex_web_enabled # noqa: E402 +from codex_web import ( # noqa: E402 + CodexWebRuntime, + _catalog_to_web_models, + codex_web_enabled, +) class StubContextManager: @@ -52,6 +56,47 @@ def test_codex_web_requires_explicit_local_only_flag(monkeypatch): assert codex_web_enabled() is False +def test_codex_catalog_maps_models_and_reasoning_options(): + models = _catalog_to_web_models( + [ + { + "id": "gpt-5.6-sol", + "model": "gpt-5.6-sol", + "displayName": "GPT-5.6-Sol", + "description": "Latest frontier model", + "isDefault": True, + "defaultReasoningEffort": "low", + "supportedReasoningEfforts": [ + { + "reasoningEffort": "low", + "description": "Fast responses", + }, + { + "reasoningEffort": "max", + "description": "Maximum reasoning", + }, + ], + }, + { + "id": "hidden-model", + "model": "hidden-model", + "hidden": True, + }, + ] + ) + + assert [model["id"] for model in models] == [ + "codex/default", + "codex/gpt-5.6-sol", + ] + assert models[0]["label"] == "Codex Auto (GPT-5.6-Sol)" + assert models[0]["default_reasoning_effort"] == "low" + assert models[1]["reasoning_efforts"] == [ + {"id": "low", "description": "Fast responses"}, + {"id": "max", "description": "Maximum reasoning"}, + ] + + @pytest.mark.asyncio async def test_web_runtime_emits_sse_contract_and_persists_messages(tmp_path): session = StubSession() diff --git a/tests/unit/test_session_manager_persistence.py b/tests/unit/test_session_manager_persistence.py index 86f16ae4..09f3b036 100644 --- a/tests/unit/test_session_manager_persistence.py +++ b/tests/unit/test_session_manager_persistence.py @@ -36,13 +36,17 @@ def __init__( hf_token: str | None = None, user_plan: str | None = None, model: str = "test-model", + reasoning_effort: str | None = None, ): self.hf_token = hf_token self.user_plan = user_plan self.context_manager = SimpleNamespace(items=[]) self.pending_approval = None self.turn_count = 0 - self.config = SimpleNamespace(model_name=model) + self.config = SimpleNamespace( + model_name=model, + reasoning_effort=reasoning_effort, + ) self.notification_destinations = [] self.auto_approval_enabled = False self.auto_approval_cost_cap_usd = None @@ -912,6 +916,7 @@ def fake_create_session_sync(**kwargs: Any): hf_token=kwargs.get("hf_token"), user_plan=kwargs.get("user_plan"), model=kwargs.get("model") or "test-model", + reasoning_effort=kwargs.get("reasoning_effort"), ) async def fake_run_session(*_: Any) -> None: From c3322f80a5c3bb95c6b75b991be43569c79501ce Mon Sep 17 00:00:00 2001 From: tibetyalman <222756182+tibetyalman@users.noreply.github.com> Date: Sat, 1 Aug 2026 04:37:30 +0300 Subject: [PATCH 4/5] fix: avoid local Hugging Face sign-in flash --- .../WelcomeScreen/WelcomeScreen.tsx | 14 +++- frontend/src/hooks/useAuth.ts | 65 ++++++++++++------- frontend/src/store/agentStore.ts | 4 ++ 3 files changed, 58 insertions(+), 25 deletions(-) diff --git a/frontend/src/components/WelcomeScreen/WelcomeScreen.tsx b/frontend/src/components/WelcomeScreen/WelcomeScreen.tsx index 2e6d9578..a8b00348 100644 --- a/frontend/src/components/WelcomeScreen/WelcomeScreen.tsx +++ b/frontend/src/components/WelcomeScreen/WelcomeScreen.tsx @@ -181,7 +181,7 @@ function ChecklistStep({ export default function WelcomeScreen() { const { createSession } = useSessionStore(); - const { setPlan, clearPanel, user } = useAgentStore(); + const { setPlan, clearPanel, authChecking, user } = useAgentStore(); const [isCreating, setIsCreating] = useState(false); const [error, setError] = useState(null); @@ -313,7 +313,17 @@ export default function WelcomeScreen() { mx: 2, }} > - {isDevUser ? ( + {authChecking ? ( + + ) : isDevUser ? ( /* Dev mode: single step */ s.setUser); + const setAuthChecking = useAgentStore((s) => s.setAuthChecking); useEffect(() => { let cancelled = false; + async function hydrateCurrentUser(): Promise { + const response = await fetch('/auth/me', { credentials: 'include' }); + if (!response.ok) return false; + + const data = await response.json(); + if (!data.authenticated) return false; + if (!cancelled) { + setUser({ + authenticated: true, + username: data.username, + name: data.name, + picture: data.picture, + plan: data.plan === 'pro' ? 'pro' : 'free', + }); + logger.log('Authenticated as', data.username); + } + return true; + } + async function checkAuth() { try { - // Check if user is already authenticated (cookie-based) - const response = await fetch('/auth/me', { credentials: 'include' }); - if (response.ok) { - const data = await response.json(); - if (!cancelled && data.authenticated) { - setUser({ - authenticated: true, - username: data.username, - name: data.name, - picture: data.picture, - plan: data.plan === 'pro' ? 'pro' : 'free', - }); - logger.log('Authenticated as', data.username); - return; - } - } - - // Not authenticated — check if auth is enabled + // Check the cheap instance status first. Local development bypasses + // OAuth, so mark it authenticated immediately instead of flashing a + // misleading "Sign in" gate while /auth/me validates HF_TOKEN. const statusRes = await fetch('/auth/status', { credentials: 'include' }); - const statusData = await statusRes.json(); - if (!statusData.auth_enabled) { - // Dev mode — no OAuth configured - if (!cancelled) setUser({ authenticated: true, username: 'dev', plan: 'pro' }); + const statusData = statusRes.ok ? await statusRes.json() : null; + if (statusData && !statusData.auth_enabled) { + if (!cancelled) { + setUser({ authenticated: true, username: 'dev', plan: 'pro' }); + setAuthChecking(false); + } + // Resolve the real HF identity in the background when HF_TOKEN is + // configured. Failure is harmless because dev auth is already valid. + try { + await hydrateCurrentUser(); + } catch { + // Keep the local dev identity. + } return; } + // Hosted/OAuth mode: restore the cookie-backed user when available. + if (await hydrateCurrentUser()) return; + // Auth enabled but not logged in — welcome screen will handle it if (!cancelled) setUser(null); } catch { // Backend unreachable — assume dev mode if (!cancelled) setUser({ authenticated: true, username: 'dev', plan: 'pro' }); + } finally { + if (!cancelled) setAuthChecking(false); } } checkAuth(); return () => { cancelled = true; }; - }, [setUser]); + }, [setAuthChecking, setUser]); } diff --git a/frontend/src/store/agentStore.ts b/frontend/src/store/agentStore.ts index a09e4492..b336980b 100644 --- a/frontend/src/store/agentStore.ts +++ b/frontend/src/store/agentStore.ts @@ -116,6 +116,7 @@ interface AgentStore { isProcessing: boolean; isConnected: boolean; activityStatus: ActivityStatus; + authChecking: boolean; user: User | null; llmHealthError: LLMHealthError | null; jobsUpgradeRequired: JobsUpgradeState | null; @@ -169,6 +170,7 @@ interface AgentStore { setProcessing: (isProcessing: boolean) => void; setConnected: (isConnected: boolean) => void; setActivityStatus: (status: ActivityStatus) => void; + setAuthChecking: (checking: boolean) => void; setUser: (user: User | null) => void; setLlmHealthError: (error: LLMHealthError | null) => void; setJobsUpgradeRequired: (state: JobsUpgradeState | null) => void; @@ -289,6 +291,7 @@ export const useAgentStore = create()((set, get) => ({ isProcessing: false, isConnected: false, activityStatus: { type: 'idle' }, + authChecking: true, user: null, llmHealthError: null, jobsUpgradeRequired: null, @@ -401,6 +404,7 @@ export const useAgentStore = create()((set, get) => ({ }, setConnected: (isConnected) => set({ isConnected }), setActivityStatus: (status) => set({ activityStatus: status }), + setAuthChecking: (authChecking) => set({ authChecking }), setUser: (user) => set({ user }), setLlmHealthError: (error) => set({ llmHealthError: error }), setJobsUpgradeRequired: (state) => set({ jobsUpgradeRequired: state }), From fb6851e5bdd1cc163438c453b541c3924c9f1ee4 Mon Sep 17 00:00:00 2001 From: tibetyalman <222756182+tibetyalman@users.noreply.github.com> Date: Sat, 1 Aug 2026 05:10:21 +0300 Subject: [PATCH 5/5] fix: recover stalled Codex turns --- agent/core/codex_runtime.py | 41 +++++++++- backend/session_manager.py | 46 +++++++++-- tests/unit/test_codex_runtime.py | 80 +++++++++++++++++++ .../unit/test_session_manager_persistence.py | 42 ++++++++++ 4 files changed, 200 insertions(+), 9 deletions(-) diff --git a/agent/core/codex_runtime.py b/agent/core/codex_runtime.py index 9eaa1065..efaec0a6 100644 --- a/agent/core/codex_runtime.py +++ b/agent/core/codex_runtime.py @@ -34,6 +34,9 @@ CODEX_TOOL_NAMESPACE = "ml_intern" _DYNAMIC_TOOL_NAME = re.compile(r"^[A-Za-z0-9_-]+$") _CODEX_BUILTIN_LOCAL_TOOLS = {"bash", "read", "write", "edit"} +_CODEX_INTERRUPT_TIMEOUT_S = 3.0 +_CODEX_NOTIFICATION_POLL_S = 1.0 +_CODEX_TURN_IDLE_TIMEOUT_S = 300.0 _UNSUPPORTED_CODEX_TOOLS = { # This tool creates a nested LiteLLM research loop using the active model. # Codex already has its own independent context and can call the underlying @@ -454,9 +457,17 @@ async def interrupt(self) -> None: if not self.thread_id or not self.active_turn_id: return try: - await self._request( - "turn/interrupt", - {"threadId": self.thread_id, "turnId": self.active_turn_id}, + await asyncio.wait_for( + self._request( + "turn/interrupt", + {"threadId": self.thread_id, "turnId": self.active_turn_id}, + ), + timeout=_CODEX_INTERRUPT_TIMEOUT_S, + ) + except TimeoutError: + logger.warning( + "Codex app-server did not acknowledge interrupt within %.0fs", + _CODEX_INTERRUPT_TIMEOUT_S, ) except Exception: logger.debug("Failed to interrupt Codex turn", exc_info=True) @@ -480,10 +491,32 @@ async def run_turn(self, prompt: str, *, stream: bool = True) -> str: turn = response.get("turn") or {} self.active_turn_id = turn.get("id") final_text = "" + loop = asyncio.get_running_loop() + last_activity_at = loop.time() try: while True: - message = await self._notifications.get() + if getattr(self._tool_session, "is_cancelled", False): + return final_text + + idle_for = loop.time() - last_activity_at + remaining_idle_s = _CODEX_TURN_IDLE_TIMEOUT_S - idle_for + if remaining_idle_s <= 0: + raise CodexRuntimeError( + "Codex produced no activity for 5 minutes, so the stuck " + "turn was stopped. Retry or choose a lower Thinking level." + ) + try: + message = await asyncio.wait_for( + self._notifications.get(), + timeout=min( + _CODEX_NOTIFICATION_POLL_S, + remaining_idle_s, + ), + ) + except TimeoutError: + continue + last_activity_at = loop.time() method = message.get("method") params = message.get("params") or {} diff --git a/backend/session_manager.py b/backend/session_manager.py index 50133ef0..c2a541f5 100644 --- a/backend/session_manager.py +++ b/backend/session_manager.py @@ -141,6 +141,7 @@ class AgentSession: codex_runtime: CodexWebRuntime | None = None codex_runtime_model: str | None = None codex_runtime_effort: str | None = None + codex_turn_task: asyncio.Task | None = None def __post_init__(self) -> None: if self.usage_window_started_at is None: @@ -1774,15 +1775,42 @@ async def _run_session( submission.operation.op_type == OpType.USER_INPUT and is_codex_model_id(session.config.model_name) ): - codex_runtime = await self._ensure_codex_runtime( - agent_session - ) text = ( submission.operation.data.get("text", "") if submission.operation.data else "" ) - await codex_runtime.run_user_input(text) + + async def run_codex_turn() -> None: + codex_runtime = await self._ensure_codex_runtime( + agent_session + ) + try: + await codex_runtime.run_user_input(text) + except Exception: + await self._close_codex_runtime(agent_session) + raise + + codex_turn_task = asyncio.create_task(run_codex_turn()) + agent_session.codex_turn_task = codex_turn_task + try: + await codex_turn_task + except asyncio.CancelledError: + # A Stop request cancels only the active Codex + # turn. Cancellation of the outer session task + # must still shut the whole session down. + current_task = asyncio.current_task() + if ( + current_task is not None + and current_task.cancelling() + ): + raise + await session.send_event( + Event(event_type="interrupted") + ) + finally: + if agent_session.codex_turn_task is codex_turn_task: + agent_session.codex_turn_task = None should_continue = True else: should_continue = await process_submission( @@ -1889,8 +1917,16 @@ async def interrupt(self, session_id: str) -> bool: if not agent_session or not agent_session.is_active: return False agent_session.session.cancel() + codex_turn_task = agent_session.codex_turn_task + if codex_turn_task is not None and not codex_turn_task.done(): + codex_turn_task.cancel() if agent_session.codex_runtime is not None: - await agent_session.codex_runtime.interrupt() + try: + await agent_session.codex_runtime.interrupt() + finally: + # A cancelled app-server turn can remain wedged even after the + # interrupt request. Recreate it for the next user message. + await self._close_codex_runtime(agent_session) return True async def undo(self, session_id: str) -> bool: diff --git a/tests/unit/test_codex_runtime.py b/tests/unit/test_codex_runtime.py index 46351364..b8811878 100644 --- a/tests/unit/test_codex_runtime.py +++ b/tests/unit/test_codex_runtime.py @@ -1,3 +1,4 @@ +import asyncio from types import SimpleNamespace import pytest @@ -242,6 +243,85 @@ async def fake_request(method, params): ] +@pytest.mark.asyncio +async def test_run_turn_exits_when_session_is_cancelled(tmp_path): + runtime = CodexAppServerRuntime( + config=SimpleNamespace(model_name="codex/default"), + tool_router=SimpleNamespace(), + hf_token=None, + local_mode=False, + cwd=tmp_path, + autonomous_mode=False, + ) + runtime.thread_id = "thread-1" + runtime._tool_session = SimpleNamespace(is_cancelled=True) + + async def fake_request(_method, _params): + return {"turn": {"id": "turn-1"}} + + runtime._request = fake_request + + assert await runtime.run_turn("hello") == "" + assert runtime.active_turn_id is None + + +@pytest.mark.asyncio +async def test_run_turn_times_out_when_codex_stops_emitting_activity( + monkeypatch, + tmp_path, +): + monkeypatch.setattr( + "agent.core.codex_runtime._CODEX_NOTIFICATION_POLL_S", + 0.005, + ) + monkeypatch.setattr( + "agent.core.codex_runtime._CODEX_TURN_IDLE_TIMEOUT_S", + 0.01, + ) + runtime = CodexAppServerRuntime( + config=SimpleNamespace(model_name="codex/default"), + tool_router=SimpleNamespace(), + hf_token=None, + local_mode=False, + cwd=tmp_path, + autonomous_mode=False, + ) + runtime.thread_id = "thread-1" + + async def fake_request(_method, _params): + return {"turn": {"id": "turn-1"}} + + runtime._request = fake_request + + with pytest.raises(CodexRuntimeError, match="no activity"): + await runtime.run_turn("hello") + + +@pytest.mark.asyncio +async def test_interrupt_does_not_wait_forever_for_app_server(monkeypatch, tmp_path): + monkeypatch.setattr( + "agent.core.codex_runtime._CODEX_INTERRUPT_TIMEOUT_S", + 0.01, + ) + runtime = CodexAppServerRuntime( + config=SimpleNamespace(model_name="codex/default"), + tool_router=SimpleNamespace(), + hf_token=None, + local_mode=False, + cwd=tmp_path, + autonomous_mode=False, + ) + runtime.thread_id = "thread-1" + runtime.active_turn_id = "turn-1" + + async def hanging_request(_method, _params): + await asyncio.Event().wait() + + runtime._request = hanging_request + + await runtime.interrupt() + + @pytest.mark.asyncio async def test_dynamic_tool_failure_is_returned_to_codex(tmp_path): class FailingRouter: diff --git a/tests/unit/test_session_manager_persistence.py b/tests/unit/test_session_manager_persistence.py index 09f3b036..245d660d 100644 --- a/tests/unit/test_session_manager_persistence.py +++ b/tests/unit/test_session_manager_persistence.py @@ -61,10 +61,21 @@ def __init__( self.sandbox_preload_cancel_event = None self.events = [] self.session_id = "s1" + self._cancelled = False async def send_event(self, event): self.events.append(event) + def cancel(self): + self._cancelled = True + + def reset_cancel(self): + self._cancelled = False + + @property + def is_cancelled(self): + return self._cancelled + def auto_approval_policy_summary(self): cap = self.auto_approval_cost_cap_usd remaining = ( @@ -185,6 +196,37 @@ def test_agent_session_replaces_non_uuid_inference_billing_session_id(): ) +@pytest.mark.asyncio +async def test_interrupt_cancels_active_codex_turn_and_closes_runtime(): + class FakeCodexRuntime: + def __init__(self) -> None: + self.interrupted = False + self.closed = False + + async def interrupt(self) -> None: + self.interrupted = True + + async def close(self) -> None: + self.closed = True + + manager = _manager_with_store(NoopSessionStore()) + agent_session = _runtime_agent_session("codex-stop") + runtime = FakeCodexRuntime() + turn_task = asyncio.create_task(asyncio.Event().wait()) + agent_session.codex_runtime = runtime # type: ignore[assignment] + agent_session.codex_turn_task = turn_task + manager.sessions[agent_session.session_id] = agent_session + + assert await manager.interrupt(agent_session.session_id) is True + await asyncio.gather(turn_task, return_exceptions=True) + + assert agent_session.session.is_cancelled is True + assert turn_task.cancelled() + assert runtime.interrupted is True + assert runtime.closed is True + assert agent_session.codex_runtime is None + + @pytest.mark.asyncio async def test_reset_session_usage_window_updates_runtime_and_store(): store = RestoreStore()