From c6049c3fc18da6c8f52e28ebd7211002b81d8044 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:23:13 +0800 Subject: [PATCH] refactor(codex)!: share app-server processes across agent wrappers (#709) * refactor(codex)!: share app-server and reuse application job tools * test(codex): use FastMCP input_schema without compatibility aliases * feat(codex): configure native model catalog and preserve streamed metadata * fix(codex): default to unattended tools and preserve retry streams * fix(codex): use an explicit Responses provider for custom endpoints * fix(codex): store runtime state in workspace .codex * feat(config): add minimal Studio workspace config with Codex chat agent - Introduce proactive.yaml defining a workspace with HTTP backend and web enabled - Add background job for Markdown index update watching daily and digest directories - Define multiple base and stream backend jobs for chat, config, version, status, health, graph, reindex, and file operations - Configure Codex app server with API key auth and detailed model catalog including qwen3.8-max - Set up default agent wrapper using Codex backend with high effort and read-only sandbox - Establish default tokenizer (regex), keyword index (bm25), file graph and chunker for markdown files - Configure file storage backend with local implementation and indexing components * fix(codex): settle cancelled dispatch and isolate TLS callbacks * test(codex): check cancellation barriers across Python versions --- docs/en/framework.md | 39 +- docs/zh/framework.md | 34 +- reme/components/__init__.py | 2 + .../agent_wrapper/codex_agent_wrapper.py | 551 +--- .../agent_wrapper/codex_mcp_server.py | 85 - reme/components/client/http_client.py | 5 +- reme/components/codex_app_server.py | 421 +++ reme/components/service/http_service.py | 14 +- reme/components/service/mcp_tools.py | 24 +- reme/components/service/scoped_mcp.py | 99 + reme/config/default.yaml | 68 +- reme/config/demo.yaml | 56 +- reme/config/proactive.yaml | 266 ++ reme/enumeration/component_enum.py | 2 + reme/steps/common/stream_llm_demo.py | 4 +- tests/integration/test_codex_agent_wrapper.py | 284 +- tests/integration/test_codex_demo.py | 176 ++ .../test_codex_wrapper_capabilities.py | 139 + tests/unit/test_codex_agent_wrapper.py | 2453 +++++++---------- tests/unit/test_common_steps.py | 32 +- tests/unit/test_http_client.py | 33 + 21 files changed, 2622 insertions(+), 2165 deletions(-) delete mode 100644 reme/components/agent_wrapper/codex_mcp_server.py create mode 100644 reme/components/codex_app_server.py create mode 100644 reme/components/service/scoped_mcp.py create mode 100644 reme/config/proactive.yaml create mode 100644 tests/integration/test_codex_demo.py create mode 100644 tests/integration/test_codex_wrapper_capabilities.py diff --git a/docs/en/framework.md b/docs/en/framework.md index 8b6fc4ca..fb6757fa 100644 --- a/docs/en/framework.md +++ b/docs/en/framework.md @@ -75,6 +75,7 @@ reme/ as_llm/, as_embedding/ # model wrappers as_decision_model/ # optional decision model adapters agent_wrapper/ # AgentScope / Claude Code / Codex wrappers + codex_app_server.py # shared Codex process and conversation locks steps/ base_step.py # BaseStep, Ref, dispatch_steps common/ # version, help, health_check, status, chat @@ -547,7 +548,8 @@ Current default components in `reme/config/default.yaml`: | `as_llm` | `default` | `${LLM_BACKEND:-openai}` | LLM model wrapper. | | `agent_wrapper` | `default` | `agentscope` | AgentScope wrapper. | | `agent_wrapper` | `claude_code` | `claude_code` | Claude Code wrapper. | -| `agent_wrapper` | `codex/codex_oauth` | `codex` | Codex wrappers for API-key and OAuth authentication. | +| `codex_app_server` | `default/oauth` | `sdk` | Shared Codex processes for API-key and OAuth authentication. | +| `agent_wrapper` | `codex/codex_oauth` | `codex` | Agent configuration referencing a named Codex app-server. | | `file_graph` | `default` | `local` | Wikilink graph. | | `file_catalog` | `default/resource/digest/dream` | `local` | File-change checkpoints. | | `file_chunker` | `markdown` | `markdown` | Markdown AST chunking. | @@ -555,6 +557,41 @@ Current default components in `reme/config/default.yaml`: | `keyword_index` | `default` | `bm25` | BM25 keyword index. | | `file_store` | `default` | `local` | Combines file_graph and keyword_index; defaults to `embedding_store: ""`. | +Codex wrappers reference `app_server: default` or `app_server: oauth`. Multiple wrappers can share one server while +keeping their own models, instructions and conversations. Omit `session_id` to start a new thread; pass a returned +`session_id` to continue it. Different threads can execute concurrently; operations on one thread are serialized. + +The default is `approval_mode: deny_all` with `sandbox: full-access`: the SDK maps this mode to +`approval_policy: never`, so available tools run without human or model approval. Requests that need additional +permissions are rejected rather than prompting. Set `approval_mode: auto_review` to opt into the configured +automatic reviewer. Tool allowlists and injected argument constraints continue to apply. + +With an API key and a custom `base_url`, the shared server configures a dedicated Responses provider that uses +HTTP/SSE by default. Set `supports_websockets: true` on `codex_app_server` only for endpoints supporting the Responses +WebSocket transport. Explicit `config_overrides` retain precedence over generated provider defaults. + +API-key defaults use `qwen3.8-max` for main conversations, code reviews, subagents and native memory consolidation, +and `qwen3.8-flash` for automatic approval review and native memory extraction. The model catalog, main requests and +default subagents select `xhigh` reasoning. The catalog exposes only `xhigh`, so native approval reviews also use +that level rather than the SDK's internal `low` default. Main-model overrides use `model` +on the wrapper; the remaining native settings are explicit in `codex_app_server.default.config_overrides` and +`model_catalog`. ReMe writes this native Codex catalog as rebuildable metadata in the server's private home. +`auto_review_model_override` selects the approval model separately from the main model; changing only `model` does +not change that reviewer. Native conversation compaction uses the thread's model. OAuth uses account-supported +models and the separate `CODEX_OAUTH_MODEL_NAME` override. See [demo configuration](../../reme/config/demo.yaml) for +the complete API-key setup using `LLM_API_KEY` and `LLM_BASE_URL`. + +Migration: move `auth_mode`, `api_key`, `base_url`, `codex_home`, `codex_bin`, `config_overrides` and `experimental_api` +from wrappers to `components.codex_app_server.` (`backend: sdk`). The default Codex home is +`{workspace_dir}/.codex`; set `codex_home` to the previous directory to retain its login and saved threads. +Configure separate `codex_home` values when servers need independent account or session storage. +Remove `mcp_config` and `release_threads`. Job tools run against the current Application through a +session-scoped MCP endpoint, reusing its plain HTTP listener when available. TLS, CLI and embedded applications use +a lifecycle-managed loopback listener. Tool selection and injected arguments remain bound to each loaded thread; an unloaded +thread can be resumed with a new binding. Internal jobs are available only through the explicitly selected tools, +without changing their public service exposure. Closing a wrapper leaves the shared server running; application +shutdown closes it. A disconnected request fails without replaying tools; later requests can reconnect and resume. + Note that the `search` Step configuration contains `vector_weight`, but `file_store.default.embedding_store` is empty by default. Vector retrieval is available only when the runtime configuration enables an embedding store. diff --git a/docs/zh/framework.md b/docs/zh/framework.md index ad11342e..1910b1b1 100644 --- a/docs/zh/framework.md +++ b/docs/zh/framework.md @@ -70,6 +70,7 @@ reme/ as_llm/, as_embedding/ # 模型封装 as_decision_model/ # 可选的决策模型适配器 agent_wrapper/ # AgentScope / Claude Code / Codex wrapper + codex_app_server.py # 共享 Codex 进程与会话锁 steps/ base_step.py # BaseStep、Ref、dispatch_steps common/ # version、help、health_check、status、chat @@ -528,7 +529,8 @@ flowchart LR | `as_decision_model` | `default` | 默认未配置；示例为 `agentscope` | 可选概率决策；通过继承 default 的 decision_demo 配置演示 | | `agent_wrapper` | `default` | `agentscope` | AgentScope wrapper | | `agent_wrapper` | `claude_code` | `claude_code` | Claude Code wrapper | -| `agent_wrapper` | `codex/codex_oauth` | `codex` | 分别使用 API key 与 OAuth 的 Codex wrapper | +| `codex_app_server` | `default/oauth` | `sdk` | 分别使用 API key 与 OAuth 的共享 Codex 进程 | +| `agent_wrapper` | `codex/codex_oauth` | `codex` | 引用具名 Codex app-server 的 agent 配置 | | `file_graph` | `default` | `local` | wikilink 图谱 | | `file_catalog` | `default/resource/digest/dream` | `local` | 文件变更 checkpoint | | `file_chunker` | `markdown` | `markdown` | Markdown AST 分块 | @@ -536,6 +538,36 @@ flowchart LR | `keyword_index` | `default` | `bm25` | BM25 关键词索引 | | `file_store` | `default` | `local` | 组合 file_graph、keyword_index；默认 `embedding_store: ""` | +Codex wrapper 通过 `app_server: default` 或 `app_server: oauth` 引用共享组件。多个 wrapper 可以共享同一个进程， +并分别配置模型、指令和会话。不传 `session_id` 时创建新 thread；传回已有 `session_id` 时续聊。 +不同 thread 可以并发，同一 thread 的操作串行。 + +默认使用 `approval_mode: deny_all` 和 `sandbox: full-access`；SDK 将该模式映射为 `approval_policy: never`， +可用工具直接执行，不经过人工或模型审批。需要新增权限的请求直接拒绝，不弹出审批。 +如需自动审批，可显式设置 `approval_mode: auto_review`。工具白名单与注入参数约束继续生效。 + +使用 API key 和自定义 `base_url` 时，共享组件会配置独立的 Responses provider，默认直接使用 HTTP/SSE。 +仅当端点支持 Responses WebSocket 传输时，才在 `codex_app_server` 上设置 `supports_websockets: true`。 +显式 `config_overrides` 优先于自动生成的 provider 默认值。 + +API key 模式默认使用 `qwen3.8-max` 处理主对话、代码审查、子任务和原生记忆整合，使用 `qwen3.8-flash` +处理自动审批和原生记忆提取。模型目录、主请求和默认子任务的 reasoning 均设为 `xhigh`。 +模型目录仅开放 `xhigh` 档位，使原生自动审批也使用该档位，避免继续采用 SDK 内部的 `low` 默认值。 +主模型由 wrapper 的 `model` 配置，其余原生配置显式写在 +`codex_app_server.default.config_overrides` 和 `model_catalog` 中；ReMe 将原生 Codex 模型目录写为共享组件私有目录内的 +可重建元数据。`auto_review_model_override` 单独选择审批模型，仅修改主模型不会改变审批模型。原生会话压缩沿用 +thread 的模型。OAuth 使用账号支持的模型，通过独立的 `CODEX_OAUTH_MODEL_NAME` 覆盖。 +使用 `LLM_API_KEY`、`LLM_BASE_URL` 的完整 API key 配置见 [demo 配置](../../reme/config/demo.yaml)。 + +迁移：将 wrapper 上的 `auth_mode`、`api_key`、`base_url`、`codex_home`、`codex_bin`、`config_overrides` 和 +`experimental_api` 移到 `components.codex_app_server.<名称>`，并使用 `backend: sdk`。默认 Codex 目录为 +`{workspace_dir}/.codex`；如需继续使用原有登录和会话，请将 `codex_home` 指向旧目录。 +多个服务组件需要独立账号或会话存储时，分别显式配置 `codex_home`。 +删除 `mcp_config` 和 `release_threads`。job 工具通过受限 MCP 入口调用当前 Application 中的原有 job，明文 HTTP 模式复用现有 +监听器，TLS、CLI 和嵌入式模式使用随组件关闭的本地监听器。工具白名单及隐藏参数绑定到已加载的 thread；thread 卸载后可在 +恢复时重新绑定。内部 job 仅对显式选择的工具开放，不改变其公共服务可见性。关闭 wrapper 不关闭共享进程，由 Application +统一关闭。连接断开时当前请求报错且不重放工具，后续请求可以重连并恢复会话。 + 注意：`search` step 的配置含 `vector_weight`，但默认 `file_store.default.embedding_store` 为空，因此实际是否有向量检索取决于运行配置是否启用 embedding store。 diff --git a/reme/components/__init__.py b/reme/components/__init__.py index 7ffbda4c..a4620583 100644 --- a/reme/components/__init__.py +++ b/reme/components/__init__.py @@ -4,6 +4,7 @@ from . import agent_wrapper from . import as_llm from . import as_decision_model from . import client +from .codex_app_server import CodexAppServer from . import as_embedding from . import embedding_store from . import file_catalog @@ -27,6 +28,7 @@ __all__ = [ "BaseComponent", "ComponentMixin", "ComponentRegistry", + "CodexAppServer", "R", "create_application_registry", "PromptHandler", diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index b744452d..d7d5afb5 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -1,140 +1,74 @@ -"""Codex Python SDK backend for the unified agent wrapper.""" +"""Codex agent sessions backed by an application-owned app-server.""" from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator -from contextlib import suppress -from dataclasses import dataclass +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager from functools import partial -import hashlib import json import os -from pathlib import Path -import sys -import tempfile from typing import Any, TYPE_CHECKING from .base_agent_wrapper import BaseAgentWrapper +from ..codex_app_server import CodexAppServer from ..component_registry import R +from ..job import BackgroundJob, StreamJob from ...enumeration import ChunkEnum from ...schema import StreamChunk, TokenUsage if TYPE_CHECKING: - from openai_codex import AsyncCodex, AsyncThread, CodexConfig, RunInput + from openai_codex import AsyncThread, RunInput from openai_codex.types import Notification -else: - # Keep the optional Codex SDK out of ReMe's package import path. Tests may - # also replace this value before the SDK is loaded. - AsyncCodex: Any = None - - -def _get_async_codex_class(): - """Return the optional Codex client class, importing its SDK on first use.""" - global AsyncCodex # pylint: disable=global-statement - if AsyncCodex is None: - from openai_codex import AsyncCodex as AsyncCodexClass - - AsyncCodex = AsyncCodexClass - return AsyncCodex - - -@dataclass(frozen=True) -class _CodexAuthConfig: - """Resolved authentication settings for one Codex app-server.""" - - mode: str - api_key: str = "" - base_url: str = "" - - -@dataclass(frozen=True) -class _CodexLaunchConfig: - """Options fixed for the lifetime of one Codex app-server.""" - - auth_mode: str - api_key: str - base_url: str - codex_bin: str | None - config_overrides: tuple[str, ...] - experimental_api: bool @R.register("codex") class CodexAgentWrapper(BaseAgentWrapper): - """Agent wrapper backed by the Codex Python SDK.""" + """Configure agents while sharing transport, authentication and session locks.""" SDK_PACKAGE = "openai-codex" - _CLIENT_OPTION_NAMES = frozenset( + _SERVER_OPTIONS = frozenset( { - "api_key", "auth_mode", + "api_key", "base_url", + "supports_websockets", "codex_bin", "codex_home", "config_overrides", - "cwd", + "model_catalog", "experimental_api", "launch_args_override", + "mcp_config", + "release_threads", }, ) + def __init__(self, app_server: str = "default", **kwargs) -> None: + self._reject_server_options(kwargs) + super().__init__(**kwargs) + self.app_server = self.bind(app_server, CodexAppServer, default_factory=CodexAppServer, optional=False) + self._active = 0 + self._idle = asyncio.Event() + self._idle.set() + + @classmethod + def _reject_server_options(cls, kwargs: dict[str, Any]) -> None: + invalid = sorted(cls._SERVER_OPTIONS.intersection(kwargs)) + if invalid: + raise TypeError(f"Configure Codex process options on codex_app_server: {', '.join(invalid)}") + + def _reply_kwargs(self, kwargs: dict[str, Any], *, stream: bool = False) -> dict[str, Any]: + self._reject_server_options(kwargs) + invalid = sorted({"app_server", "cwd"}.intersection(kwargs)) + if invalid: + raise TypeError(f"Configure Codex wrapper options on the component: {', '.join(invalid)}") + return self._merged_stream_kwargs(kwargs) if stream else self._merged_kwargs(kwargs) + @staticmethod def _codex_usage(usage: Any) -> TokenUsage: - """Normalize Codex's full-turn input/output usage snapshot.""" return TokenUsage.from_provider(usage) - # pylint: disable=too-many-arguments - def __init__( - self, - mcp_config: str | None = None, - codex_home: str | Path | None = None, - *, - auth_mode: str = "auto", - api_key: str = "", - base_url: str = "", - codex_bin: str | None = None, - config_overrides: list[str] | tuple[str, ...] | None = None, - experimental_api: bool = True, - release_threads: bool = True, - **kwargs, - ) -> None: - if "launch_args_override" in kwargs: - raise TypeError("launch_args_override is not supported; configure codex_bin instead") - super().__init__(**kwargs) - self.mcp_config = mcp_config - self._codex_home = codex_home - self._launch_config = _CodexLaunchConfig( - auth_mode=auth_mode, - api_key=api_key, - base_url=base_url, - codex_bin=codex_bin, - config_overrides=tuple(config_overrides or ()), - experimental_api=experimental_api, - ) - self._codex: AsyncCodex | None = None - self._release_threads = release_threads - self._turn_lock = asyncio.Lock() - self._mcp_snapshot_path: Path | None = None - self._thread_tool_contexts: dict[str, str] = {} - - @property - def session_path(self) -> Path: - """Directory used for Codex state and persisted threads.""" - if self._codex_home: - path = Path(self._codex_home).expanduser() - return path if path.is_absolute() else self.workspace_path / path - if self.app_context is None: - return self.workspace_path / "mem_session" / "codex" - return self.workspace_path / self.app_context.app_config.mem_session_dir / "codex" - - @property - def wrapper_session_path(self) -> Path: - """ReMe-owned session data, kept separate from a shared OAuth CODEX_HOME.""" - if self.app_context is None: - return self.workspace_path / "mem_session" / "codex" - return self.workspace_path / self.app_context.app_config.mem_session_dir / "codex" - def _ensure_skills(self, skills: list[str] | str | None) -> None: """Expose selected project skills through Codex's repo-level directory.""" sources = self._resolve_project_skills(skills) @@ -154,178 +88,25 @@ class CodexAgentWrapper(BaseAgentWrapper): relative_source = os.path.relpath(source, target.parent) target.symlink_to(relative_source, target_is_directory=True) - def _explicit_mcp_config(self, kwargs: dict[str, Any]) -> str | None: - value = kwargs.get("mcp_config") if "mcp_config" in kwargs else self.mcp_config - if value is None: - return None - source = Path(str(value)).expanduser() - if source.suffix in {".yaml", ".yml", ".json"}: - if not source.is_absolute(): - source = self.workspace_path / source - return str(source.absolute()) - return str(value) - - def _effective_config_snapshot(self) -> Path: - """Create one private snapshot that remains valid for the client lifetime.""" - if self._mcp_snapshot_path is not None: - return self._mcp_snapshot_path - if self.app_context is None: - raise RuntimeError("Cannot snapshot MCP config without an app_context") - snapshot_dir = self.wrapper_session_path / "reme-mcp" - snapshot_dir.mkdir(parents=True, exist_ok=True) - fd, raw_path = tempfile.mkstemp(prefix="config-", suffix=".json", dir=snapshot_dir) - try: - if hasattr(os, "fchmod"): - os.fchmod(fd, 0o600) - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump(self.app_context.app_config.model_dump(mode="json"), stream) - except BaseException: - with suppress(OSError): - os.close(fd) - Path(raw_path).unlink(missing_ok=True) - raise - self._mcp_snapshot_path = Path(raw_path) - return self._mcp_snapshot_path - - def _mcp_config_source(self, kwargs: dict[str, Any]) -> str: - return self._explicit_mcp_config(kwargs) or str(self._effective_config_snapshot()) - - @staticmethod - def _resolve_auth_config(auth_mode: str, api_key: str = "", base_url: str = "") -> _CodexAuthConfig: - """Resolve login from wrapper options. - - The Codex child process still inherits the parent process environment. - """ - requested_mode = str(auth_mode or "auto").lower() - if requested_mode not in {"auto", "api_key", "oauth"}: - raise ValueError("auth_mode must be one of: auto, api_key, oauth") - if requested_mode == "oauth": - return _CodexAuthConfig(mode="oauth") - - api_key = api_key if isinstance(api_key, str) else "" - if requested_mode == "api_key" and not api_key: - raise ValueError("auth_mode='api_key' requires a non-empty API key") - if not api_key: - return _CodexAuthConfig(mode="oauth") - - base_url = base_url if isinstance(base_url, str) else "" - return _CodexAuthConfig(mode="api_key", api_key=api_key, base_url=base_url) - - def _build_client_config(self, auth: _CodexAuthConfig) -> CodexConfig: - from openai_codex import CodexConfig - - env = dict(self.subprocess_environment) - self.session_path.mkdir(parents=True, exist_ok=True) - env["CODEX_HOME"] = str(self.session_path) - - overrides = list(self._launch_config.config_overrides) - if auth.base_url: - overrides.append(f"openai_base_url={json.dumps(auth.base_url)}") - login_method = "api" if auth.mode == "api_key" else "chatgpt" - overrides.append(f"forced_login_method={json.dumps(login_method)}") - return CodexConfig( - codex_bin=self._launch_config.codex_bin, - config_overrides=tuple(overrides), - cwd=str(self.cwd), - env=env, - client_name="reme", - client_title="ReMe", - experimental_api=self._launch_config.experimental_api, - ) - - @classmethod - def _reject_client_options(cls, kwargs: dict[str, Any]) -> None: - invalid = sorted(cls._CLIENT_OPTION_NAMES.intersection(kwargs)) - if invalid: - names = ", ".join(invalid) - raise TypeError(f"Codex client options must be configured on the wrapper: {names}") - - def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: - from ..job import BackgroundJob, StreamJob - - job_names = list(dict.fromkeys(kwargs.get("job_tools") or [])) - if not job_names: - return None - jobs = self._resolve_job_tools(job_names) - unsupported = [job.name for job in jobs if isinstance(job, (BackgroundJob, StreamJob))] - if unsupported: - raise TypeError(f"Codex job_tools must be non-stream request jobs: {', '.join(unsupported)}") - - config_source = self._mcp_config_source(kwargs) - args = [ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - config_source, - "--workspace", - str(self.workspace_path), - ] - for name in job_names: - args.extend(["--job", name]) - args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) - if injected_job_kwargs := dict(kwargs.get("injected_job_kwargs") or {}): - args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) - return { - "command": sys.executable, - "args": args, - "cwd": str(self.project_path), - "required": True, - "enabled_tools": job_names, - "startup_timeout_sec": kwargs.get("mcp_startup_timeout", 30), - "tool_timeout_sec": kwargs.get("mcp_tool_timeout", 300), - } - - def _thread_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: - config = dict(kwargs.get("config") or {}) - if proxy_environment := self.command_proxy_environment: - shell_environment_policy = dict(config.get("shell_environment_policy") or {}) - environment = dict(shell_environment_policy.get("set") or {}) - environment.update(proxy_environment) - shell_environment_policy["set"] = environment - config["shell_environment_policy"] = shell_environment_policy - if server := self._mcp_server_config(kwargs): - servers = dict(config.get("mcp_servers") or {}) - server_key = hashlib.sha256(json.dumps(server, sort_keys=True).encode()).hexdigest()[:12] - servers[f"reme_jobs_{server_key}"] = server - config["mcp_servers"] = servers - return config or None - @staticmethod def _enum(enum_cls: Any, value: Any, default: Any = None) -> Any: if value is None: return default return value if isinstance(value, enum_cls) else enum_cls(value) - async def _open_thread(self, codex: AsyncCodex, kwargs: dict[str, Any]) -> AsyncThread: + def _thread_options(self, kwargs: dict[str, Any]) -> dict[str, Any]: from openai_codex import ApprovalMode, Sandbox from openai_codex.types import Personality, ThreadSource, ThreadStartSource - resume = kwargs.get("resume") or "" - session_id = kwargs.get("session_id") or "" - if resume and session_id and resume != session_id: - raise ValueError("resume and session_id must identify the same Codex thread") - thread_id = resume or session_id - fork_session = bool(kwargs.get("fork_session", False)) - if fork_session and not thread_id: - raise ValueError("fork_session=True requires resume or session_id") - if self._release_threads and self._thread_tool_contexts: - await self._prune_thread_tool_contexts(codex) - requested_tool_context = json.dumps( - { - "tool_context_id": str(kwargs.get("tool_context_id") or ""), - "injected_job_kwargs": dict(kwargs.get("injected_job_kwargs") or {}), - }, - sort_keys=True, - default=str, - ) - if not fork_session and thread_id in self._thread_tool_contexts: - if requested_tool_context != self._thread_tool_contexts[thread_id]: - raise ValueError("tool_context_id and injected_job_kwargs cannot change when resuming a Codex thread") - - common = { - "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.auto_review), + config = dict(kwargs.get("config") or {}) + if proxy_environment := self.command_proxy_environment: + policy = dict(config.get("shell_environment_policy") or {}) + policy["set"] = {**policy.get("set", {}), **proxy_environment} + config["shell_environment_policy"] = policy + options = { + "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.deny_all), "base_instructions": kwargs.get("base_instructions"), - "config": self._thread_config(kwargs), + "config": config, "cwd": str(self.cwd), "developer_instructions": kwargs.get("system_prompt"), "model": kwargs.get("model"), @@ -333,122 +114,64 @@ class CodexAgentWrapper(BaseAgentWrapper): "sandbox": self._enum(Sandbox, kwargs.get("sandbox"), Sandbox.full_access), "service_tier": kwargs.get("service_tier"), } - personality = self._enum(Personality, kwargs.get("personality")) - thread_source = self._enum(ThreadSource, kwargs.get("thread_source")) - if fork_session: - thread = await codex.thread_fork( - thread_id, + if kwargs.get("fork_session"): + options.update( ephemeral=kwargs.get("ephemeral"), - thread_source=thread_source, - **common, + thread_source=self._enum(ThreadSource, kwargs.get("thread_source")), ) - elif thread_id: - thread = await codex.thread_resume(thread_id, personality=personality, **common) else: - thread = await codex.thread_start( - ephemeral=kwargs.get("ephemeral"), - personality=personality, - service_name=kwargs.get("service_name"), - session_start_source=self._enum(ThreadStartSource, kwargs.get("session_start_source")), - thread_source=thread_source, - **common, - ) - self._thread_tool_contexts[thread.id] = requested_tool_context - return thread + options["personality"] = self._enum(Personality, kwargs.get("personality")) + if not (kwargs.get("resume") or kwargs.get("session_id")): + options.update( + ephemeral=kwargs.get("ephemeral"), + service_name=kwargs.get("service_name"), + session_start_source=self._enum(ThreadStartSource, kwargs.get("session_start_source")), + thread_source=self._enum(ThreadSource, kwargs.get("thread_source")), + ) + return options - def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: - from openai_codex import ApprovalMode, Sandbox - from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary - - return { - "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), - "cwd": str(self.cwd), - "effort": self._enum(ReasoningEffort, kwargs.get("effort")), - "model": kwargs.get("model"), - "output_schema": kwargs.get("output_schema"), - "personality": self._enum(Personality, kwargs.get("personality")), - "sandbox": self._enum(Sandbox, kwargs.get("sandbox")), - "service_tier": kwargs.get("service_tier"), - "summary": self._enum(ReasoningSummary, kwargs.get("summary")), - } - - async def _get_codex(self) -> AsyncCodex: - """Lazily start the component-owned client from its fixed launch configuration.""" - if self._codex is not None: - return self._codex - auth = self._resolve_auth_config( - self._launch_config.auth_mode, - self._launch_config.api_key, - self._launch_config.base_url, - ) - codex = _get_async_codex_class()(self._build_client_config(auth)) + @asynccontextmanager + async def _session(self, kwargs: dict[str, Any], *, compact: bool = False) -> AsyncIterator[AsyncThread]: + resume, session_id = kwargs.get("resume") or "", kwargs.get("session_id") or "" + if resume and session_id and resume != session_id: + raise ValueError("resume and session_id must identify the same Codex thread") + thread_id = resume or session_id + fork = bool(kwargs.get("fork_session")) + if fork and not thread_id: + raise ValueError("fork_session=True requires resume or session_id") + jobs = None if compact else self._resolve_job_tools(list(dict.fromkeys(kwargs.get("job_tools") or []))) + if any(isinstance(job, (BackgroundJob, StreamJob)) for job in jobs or []): + raise TypeError("Codex job_tools must be non-stream request jobs") + if any(not getattr(job, "enabled", True) for job in jobs or []): + raise ValueError("Codex job_tools must be enabled") + self._ensure_skills(kwargs.get("skills")) + await self.start() + self._active += 1 + self._idle.clear() try: - if auth.mode == "api_key": - await codex.login_api_key(auth.api_key) - else: - account = await codex.account() - if account.account is None: - raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") - except BaseException: - await codex.close() - raise - self._codex = codex - return codex + async with self.app_server.thread( + thread_id, + fork=fork, + options={} if compact else self._thread_options(kwargs), + jobs=jobs, + injected=self._resolve_injected_job_kwargs(kwargs), + mcp_startup_timeout=kwargs.get("mcp_startup_timeout", 30), + mcp_tool_timeout=kwargs.get("mcp_tool_timeout", 300), + ) as thread: + yield thread + finally: + self._active -= 1 + if not self._active: + self._idle.set() async def _close(self) -> None: - """Close the persistent app-server and remove its private config snapshot.""" - async with self._turn_lock: - codex, self._codex = self._codex, None - self._thread_tool_contexts.clear() - try: - if codex is not None: - await codex.close() - finally: - if self._mcp_snapshot_path is not None: - self._mcp_snapshot_path.unlink(missing_ok=True) - self._mcp_snapshot_path = None + await self._idle.wait() - async def _prune_thread_tool_contexts(self, codex: AsyncCodex) -> None: - """Forget tool constraints only after the server confirms a thread is unloaded.""" + async def _interrupt(self, turn) -> None: try: - from openai_codex.generated.v2_all import ThreadLoadedListResponse - - loaded_ids: set[str] = set() - cursor = None - while True: - page = await codex._client.request( # pylint: disable=protected-access - "thread/loaded/list", - {"cursor": cursor}, - response_model=ThreadLoadedListResponse, - ) - loaded_ids.update(page.data) - cursor = page.next_cursor - if not cursor: - break - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to list loaded Codex threads: {exc}") - return - self._thread_tool_contexts = { - thread_id: context for thread_id, context in self._thread_tool_contexts.items() if thread_id in loaded_ids - } - - async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None: - """Unsubscribe so the app-server can unload the thread once it is idle.""" - try: - from openai_codex.generated.v2_all import ThreadUnsubscribeResponse - - # The SDK exposes no public unsubscribe; request() is public but - # reached through the private client attribute on SDK 0.162.x. - await codex._client.request( # pylint: disable=protected-access - "thread/unsubscribe", - {"threadId": thread_id}, - response_model=ThreadUnsubscribeResponse, - ) - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to unsubscribe Codex thread {thread_id}: {exc}") - # Unsubscribe is not an unload acknowledgement. In particular, a - # cancelled turn can still be running with its original MCP constraints. - # Keep the guard until thread/loaded/list confirms the thread is gone. + await turn.interrupt() + except Exception as exc: + self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") @staticmethod def _serialize(value: Any) -> Any: @@ -457,20 +180,15 @@ class CodexAgentWrapper(BaseAgentWrapper): return to_jsonable_python(value, by_alias=True) async def reply(self, inputs: RunInput, **kwargs) -> dict: - """Run one Codex turn and return its final response.""" - self._reject_client_options(kwargs) - kwargs = self._merged_kwargs(kwargs) - self._ensure_skills(kwargs.get("skills")) - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await self._open_thread(codex, kwargs) + """Run an isolated or explicitly resumed conversation on the shared server.""" + kwargs = self._reply_kwargs(kwargs) + async with self._session(kwargs) as thread: + turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) try: - result = await thread.run(inputs, **self._turn_kwargs(kwargs)) - finally: - if self._release_threads: - await self._release_thread(codex, thread.id) - + result = await turn.run() + except BaseException: + await self._interrupt(turn) + raise final_response = result.final_response or "" response = { "session_id": thread.id, @@ -491,16 +209,25 @@ class CodexAgentWrapper(BaseAgentWrapper): return response async def compact_session(self, session_id: str) -> None: - """Start native compaction of a Codex thread.""" - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await codex.thread_resume(session_id) - try: - await thread.compact() - finally: - if self._release_threads: - await self._release_thread(codex, thread.id) + """Hold the conversation lock until native compaction has finished.""" + async with self._session({"session_id": session_id}, compact=True) as thread: + await self.app_server.compact(thread) + + def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + from openai_codex import ApprovalMode, Sandbox + from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary + + return { + "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), + "cwd": str(self.cwd), + "effort": self._enum(ReasoningEffort, kwargs.get("effort")), + "model": kwargs.get("model"), + "output_schema": kwargs.get("output_schema"), + "personality": self._enum(Personality, kwargs.get("personality")), + "sandbox": self._enum(Sandbox, kwargs.get("sandbox")), + "service_tier": kwargs.get("service_tier"), + "summary": self._enum(ReasoningSummary, kwargs.get("summary")), + } @classmethod # pylint: disable=too-many-return-statements @@ -619,31 +346,19 @@ class CodexAgentWrapper(BaseAgentWrapper): return [] async def reply_stream(self, inputs: RunInput, **kwargs) -> AsyncGenerator[StreamChunk, None]: - """Stream Codex app-server notifications as unified chunks.""" - self._reject_client_options(kwargs) - kwargs = self._merged_stream_kwargs(kwargs) - self._ensure_skills(kwargs.get("skills")) - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await self._open_thread(codex, kwargs) + """Stream notifications routed by the SDK to this conversation's turn.""" + kwargs = self._reply_kwargs(kwargs, stream=True) + async with self._session(kwargs) as thread: + turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) + stream = turn.stream() + completed = False try: - turn = await thread.turn(inputs, **self._turn_kwargs(kwargs)) - stream = turn.stream() - completed = False - try: - async for event in stream: - if event.method == "turn/completed": - completed = True - for chunk in self._event_to_chunks(event, thread.id): - yield chunk - finally: - if not completed: - try: - await turn.interrupt() - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") - await stream.aclose() + async for event in stream: + if event.method == "turn/completed": + completed = True + for chunk in self._event_to_chunks(event, thread.id): + yield chunk finally: - if self._release_threads: - await self._release_thread(codex, thread.id) + if not completed: + await self._interrupt(turn) + await stream.aclose() diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py deleted file mode 100644 index 9ed86f5b..00000000 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ /dev/null @@ -1,85 +0,0 @@ -"""FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" - -import argparse -import json -from pathlib import Path -from typing import Any - -from ...config import resolve_app_config -from ...reme import ReMe - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Expose selected ReMe jobs over FastMCP STDIO") - parser.add_argument("--config", default="default", help="ReMe config name or file path") - parser.add_argument("--workspace", required=True, help="ReMe workspace directory") - parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") - parser.add_argument("--tool-context-id", default="", help="Context id injected into every job call") - parser.add_argument( - "--injected-job-kwargs", - default="", - help="JSON object of server-owned kwargs injected into every job call", - ) - return parser.parse_args() - - -def _prepare_config( - config: dict[str, Any], - job_names: list[str], - tool_context_id: str = "", - injected_job_kwargs: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Configure the dedicated child Application to serve selected jobs over MCP STDIO.""" - selected = set(job_names) - jobs: dict[str, dict[str, Any]] = {} - for name, raw_job_config in (config.get("jobs") or {}).items(): - job_config = dict(raw_job_config) - if job_config.get("backend") in {"background", "cron"}: - continue - if name in selected: - # An explicit Codex job_tools selection has always overridden enable_serve. - job_config["enable_serve"] = True - jobs[name] = job_config - - missing = sorted(selected.difference(jobs)) - if missing: - raise KeyError(f"Codex job tools not found or not request jobs: {', '.join(missing)}") - - service: dict[str, Any] = { - "backend": "mcp", - "transport": "stdio", - "jobs": job_names, - "tool_error_on_failure": True, - } - injected = dict(injected_job_kwargs or {}) - if tool_context_id: - injected["tool_context_id"] = tool_context_id - if injected: - service["injected_job_kwargs"] = injected - - prepared = dict(config) - prepared["jobs"] = jobs - prepared["service"] = service - return prepared - - -def main() -> None: - """Load ReMe and serve the requested jobs over STDIO.""" - args = _parse_args() - config = resolve_app_config( - config=args.config, - workspace_dir=str(Path(args.workspace).absolute()), - enable_logo=False, - log_to_console=False, - log_to_file=False, - log_config=False, - ) - injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else None - if injected_job_kwargs is not None and not isinstance(injected_job_kwargs, dict): - raise TypeError("--injected-job-kwargs must be a JSON object") - config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) - ReMe(**config).run_app() - - -if __name__ == "__main__": - main() diff --git a/reme/components/client/http_client.py b/reme/components/client/http_client.py index ef4da0a9..b304e1bf 100644 --- a/reme/components/client/http_client.py +++ b/reme/components/client/http_client.py @@ -60,9 +60,8 @@ class HttpClient(BaseClient): except json.JSONDecodeError: continue chunk = StreamChunk(**data) - if chunk.chunk_type == ChunkEnum.ERROR: - # Surface server-side errors as exceptions so callers don't - # mistake error chunks for valid content. + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True: + # Recoverable notifications remain visible; terminal errors stop the request. raise RuntimeError(str(chunk.chunk)) if chunk.done: return diff --git a/reme/components/codex_app_server.py b/reme/components/codex_app_server.py new file mode 100644 index 00000000..da465865 --- /dev/null +++ b/reme/components/codex_app_server.py @@ -0,0 +1,421 @@ +"""Application-owned Codex process, concurrent sessions, and local job tools.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field +import json +from pathlib import Path +import sys +from typing import Any, TYPE_CHECKING + +from .base_component import BaseComponent +from .component_registry import R +from .service.scoped_mcp import ScopedMCP +from ..constants import normalize_connect_host +from ..enumeration import ComponentEnum +from ..utils.async_utils import finish_cleanup + +if TYPE_CHECKING: + from openai_codex import AsyncCodex, AsyncThread, CodexConfig + + +@dataclass +class _ThreadState: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + users: int = 0 + revision: int = 0 + binding: str | None = None + token: str | None = None + completion: asyncio.Future[None] | None = None + operation_started: bool = False + + +@R.register("sdk") +class CodexAppServer(BaseComponent): + """Share one SDK client across wrappers with the same account and launch settings.""" + + component_type = ComponentEnum.CODEX_APP_SERVER + + def __init__( + self, + auth_mode: str = "auto", + api_key: str = "", + base_url: str = "", + supports_websockets: bool = False, + codex_home: str | Path | None = None, + codex_bin: str | None = None, + config_overrides: list[str] | tuple[str, ...] = (), + model_catalog: dict[str, Any] | None = None, + cwd: str | Path | None = None, + experimental_api: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if auth_mode not in {"auto", "api_key", "oauth"}: + raise ValueError("auth_mode must be one of: auto, api_key, oauth") + self._auth_mode = auth_mode + self._api_key = api_key + self._base_url = base_url + self._supports_websockets = supports_websockets + self._codex_bin = codex_bin + self._config_overrides = tuple(config_overrides) + self._model_catalog = model_catalog + self._cwd = cwd + self._experimental_api = experimental_api + self._codex_home = codex_home + self._codex: AsyncCodex | None = None + self._notifications: asyncio.Task | None = None + self._client_lock = asyncio.Lock() + self._threads: dict[str, _ThreadState] = {} + self._active = 0 + self._idle = asyncio.Event() + self._idle.set() + self._failed = False + self._http_url: str | None = None + self.jobs_mcp = ScopedMCP() + + @property + def session_path(self) -> Path: + """Persistent Codex state under the workspace, unless an explicit home is configured.""" + if self._codex_home: + path = Path(self._codex_home).expanduser() + return path if path.is_absolute() else self.workspace_path / path + return self.workspace_path / ".codex" + + def _client_config(self) -> CodexConfig: + from openai_codex import CodexConfig + + environment = dict(self.app_context.app_config.environment) if self.app_context else {} + environment["CODEX_HOME"] = str(self.session_path) + overrides = list(self._config_overrides) + api_login = self._auth_mode != "oauth" and bool(self._api_key) + if self._auth_mode == "api_key" and not api_login: + raise ValueError("auth_mode='api_key' requires a non-empty API key") + self.session_path.mkdir(parents=True, exist_ok=True) + if self._model_catalog is not None: + if any(option.partition("=")[0].strip() == "model_catalog_json" for option in overrides): + raise ValueError("Configure either model_catalog or config_overrides.model_catalog_json, not both") + catalog_path = self.session_path / "reme-model-catalog.json" + catalog_path.write_text(json.dumps(self._model_catalog, ensure_ascii=False), encoding="utf-8") + overrides.append(f"model_catalog_json={json.dumps(str(catalog_path))}") + if api_login and self._base_url: + # Custom endpoints must not inherit the built-in OpenAI provider's WebSocket capability. + overrides = [ + 'model_provider="reme"', + 'model_providers.reme.name="ReMe API"', + f"model_providers.reme.base_url={json.dumps(self._base_url)}", + 'model_providers.reme.wire_api="responses"', + "model_providers.reme.requires_openai_auth=true", + f"model_providers.reme.supports_websockets={json.dumps(self._supports_websockets)}", + *overrides, + ] + overrides.append(f"forced_login_method={json.dumps('api' if api_login else 'chatgpt')}") + cwd = Path(self._cwd).expanduser() if self._cwd else self.workspace_path + if not cwd.is_absolute(): + cwd = self.workspace_path / cwd + return CodexConfig( + codex_bin=self._codex_bin, + config_overrides=tuple(overrides), + cwd=str(cwd), + env=environment, + client_name="reme", + client_title="ReMe", + experimental_api=self._experimental_api, + ) + + async def _get_client(self) -> AsyncCodex: + from openai_codex import AsyncCodex + + async with self._client_lock: + if self._failed: + if self._active > 1: + raise RuntimeError("Codex disconnected; wait for its active requests to finish before reconnecting") + await self._close_client() + self._threads.clear() + self.jobs_mcp.scopes.clear() + self._failed = False + if self._codex is None: + codex = AsyncCodex(self._client_config()) + try: + if self._auth_mode != "oauth" and self._api_key: + await codex.login_api_key(self._api_key) + elif (await codex.account()).account is None: + raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") + except BaseException: + await codex.close() + raise + self._codex = codex + self._notifications = asyncio.create_task(self._watch_notifications(codex)) + return self._codex + + async def _watch_notifications(self, codex: AsyncCodex) -> None: + """Route SDK status events to the operation that owns each conversation.""" + try: + while self._codex is codex: + event = await codex._client.next_notification() # pylint: disable=protected-access + if event.method != "thread/status/changed": + continue + state = self._threads.get(event.payload.thread_id) + if state is None or state.completion is None or state.completion.done(): + continue + if event.payload.status.root.type == "active": + state.operation_started = True + elif state.operation_started: + state.completion.set_result(None) + except Exception as exc: + if self._codex is codex: + self._failed = True + for state in self._threads.values(): + if state.completion is not None and not state.completion.done(): + state.completion.set_exception(exc) + + async def _close_client(self) -> None: + codex, self._codex = self._codex, None + try: + if codex is not None: + await codex.close() + finally: + if self._notifications is not None: + self._notifications.cancel() + with suppress(asyncio.CancelledError): + await self._notifications + self._notifications = None + + async def _mcp_url(self) -> str: + from .service.http_service import HttpService + + service = self.app_context.service if self.app_context else None + if ( + isinstance(service, HttpService) + and service.service is not None + and any(getattr(route, "endpoint", None) is self.jobs_mcp for route in service.service.routes) + ): + if self._http_url is None: + host = normalize_connect_host(service.host) + host = "::1" if host == "::" else host + # Background jobs can run before uvicorn finishes binding the main listener. + async with asyncio.timeout(30): + while True: + try: + _, writer = await asyncio.open_connection(host, service.port) + except OSError: + await asyncio.sleep(0.01) + else: + writer.close() + await writer.wait_closed() + break + authority = f"[{host}]" if ":" in host else host + self._http_url = f"http://{authority}:{service.port}{self.jobs_mcp.path}" + return self._http_url + return await self.jobs_mcp.listen() + + async def _prune_threads(self, codex: AsyncCodex) -> None: + """Revoke idle tool bindings only after the server confirms their threads unloaded.""" + from openai_codex.generated.v2_all import ThreadLoadedListResponse + + candidates = {key: (state, state.revision) for key, state in self._threads.items() if not state.users} + if not candidates: + return + try: + loaded: set[str] = set() + cursor = None + while True: + page = await codex._client.request( # pylint: disable=protected-access + "thread/loaded/list", + {"cursor": cursor}, + response_model=ThreadLoadedListResponse, + ) + loaded.update(page.data) + cursor = page.next_cursor + if not cursor: + break + except Exception as exc: + self.logger.warning(f"Failed to list loaded Codex threads: {exc}") + return + for key, (state, revision) in candidates.items(): + if key not in loaded and not state.users and state.revision == revision and self._threads.get(key) is state: + self._threads.pop(key) + self.jobs_mcp.scopes.pop(state.token, None) + + async def wait_idle(self, thread: AsyncThread) -> None: + """Keep the thread lock through cancellation or asynchronous compaction.""" + state = self._threads[thread.id] + if state.completion is not None: + await asyncio.shield(state.completion) + state.completion = None + async with asyncio.timeout(30): + while (await thread.read()).thread.status.root.type == "active": + await asyncio.sleep(0.05) + + @asynccontextmanager + async def _operation(self, thread: AsyncThread) -> AsyncIterator[None]: + """Track a shielded dispatch; clear its barrier if the SDK rejects the request.""" + state = self._threads[thread.id] + state.operation_started = False + state.completion = asyncio.get_running_loop().create_future() + state.completion.add_done_callback(lambda done: done.exception() if not done.cancelled() else None) + try: + yield + except Exception: + state.completion.cancel() + state.completion = None + raise + + async def turn(self, thread: AsyncThread, inputs: Any, **kwargs): + """Start a normal SDK turn with completion tracking owned by this server.""" + turn = None + + async def dispatch(): + nonlocal turn + async with self._operation(thread): + turn = await thread.turn(inputs, **kwargs) + + try: + await finish_cleanup(dispatch()) + except asyncio.CancelledError: + if turn is not None: + try: + await finish_cleanup(turn.interrupt()) + except asyncio.CancelledError: + pass + except Exception as exc: + self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") + raise + return turn + + async def compact(self, thread: AsyncThread) -> None: + """Enqueue native compaction; the thread lease waits for its completion event.""" + + async def dispatch(): + async with self._operation(thread): + await thread.compact() + + await finish_cleanup(dispatch()) + + @asynccontextmanager + async def thread( + self, + thread_id: str = "", + *, + fork: bool = False, + options: dict[str, Any] | None = None, + jobs: list[Any] | None = None, + injected: dict[str, Any] | None = None, + mcp_startup_timeout: float = 30, + mcp_tool_timeout: float = 300, + ) -> AsyncIterator[AsyncThread]: + """Serialize one conversation while allowing independent conversations to run concurrently.""" + from openai_codex import TransportClosedError + + await self.start() + self._active += 1 + self._idle.clear() + state = None + created_token = None + thread = None + codex = None + try: + codex = await self._get_client() + await self._prune_threads(codex) + state = self._threads.setdefault(thread_id, _ThreadState()) if thread_id and not fork else _ThreadState() + state.users += 1 + state.revision += 1 + async with state.lock: + options = dict(options or {}) + if jobs is not None: + injected = json.loads(json.dumps(injected or {})) + binding = json.dumps({"jobs": [job.name for job in jobs], "injected": injected}, sort_keys=True) + if state.binding is not None and state.binding != binding: + raise ValueError( + "job_tools, tool_context_id and injected_job_kwargs cannot change on a loaded thread", + ) + state.binding = binding + config = dict(options.get("config") or {}) + servers = dict(config.get("mcp_servers") or {}) + if "reme_jobs" in servers: + raise ValueError("mcp_servers.reme_jobs is reserved for ReMe job tools") + # Codex validates the transport even when a server is disabled. + servers["reme_jobs"] = {"url": "http://127.0.0.1", "enabled": False} + if jobs: + url = await self._mcp_url() + if state.token is None: + state.token = self.jobs_mcp.register(jobs, injected) + created_token = state.token + servers["reme_jobs"] = { + "url": url, + "http_headers": {"Authorization": f"Bearer {state.token}"}, + "required": True, + "enabled_tools": [job.name for job in jobs], + "startup_timeout_sec": mcp_startup_timeout, + "tool_timeout_sec": mcp_tool_timeout, + } + config["mcp_servers"] = servers + options["config"] = config + if fork: + source = self._threads.setdefault(thread_id, _ThreadState()) + source.users += 1 + source.revision += 1 + try: + async with source.lock: + thread = await codex.thread_fork(thread_id, **options) + finally: + source.users -= 1 + elif thread_id: + thread = await codex.thread_resume(thread_id, **options) + else: + thread = await codex.thread_start(**options) + self._threads[thread.id] = state + if thread_id and not fork: + await self.wait_idle(thread) + try: + yield thread + finally: + # Interrupt acknowledgement alone does not mean that tool execution stopped. + error = sys.exception() + try: + async with asyncio.timeout(30 if error is not None else None): + await self.wait_idle(thread) + except Exception as exc: + if isinstance(exc, TransportClosedError) and self._codex is codex: + self._failed = True + if error is None: + raise + self.logger.warning(f"Failed to wait for Codex thread {thread.id}: {exc}") + else: + from openai_codex.generated.v2_all import ThreadUnsubscribeResponse + + try: + await codex._client.request( # pylint: disable=protected-access + "thread/unsubscribe", + {"threadId": thread.id}, + response_model=ThreadUnsubscribeResponse, + ) + except Exception as exc: + self.logger.warning(f"Failed to unsubscribe Codex thread {thread.id}: {exc}") + except TransportClosedError: + if codex is not None and self._codex is codex: + self._failed = True + raise + finally: + if state is not None: + state.users -= 1 + if thread is None and created_token is not None: + self.jobs_mcp.scopes.pop(created_token, None) + state.token = None + self._active -= 1 + if not self._active: + self._idle.set() + + async def _close(self) -> None: + await self._idle.wait() + try: + await self._close_client() + finally: + self._codex = None + self._failed = False + self._http_url = None + self._threads.clear() + await self.jobs_mcp.close() diff --git a/reme/components/service/http_service.py b/reme/components/service/http_service.py index 1d90ab2f..5863852a 100644 --- a/reme/components/service/http_service.py +++ b/reme/components/service/http_service.py @@ -67,10 +67,15 @@ class HttpService(BaseService): def build_service(self, app: "Application") -> None: """Create one FastAPI app containing JSON/SSE and optional MCP routes.""" + from fastmcp.utilities.lifespan import combine_lifespans + lifespan = self._lifespan(app, self.host, self.port) + servers = self.app_context.components.get("codex_app_server", {}) if self.app_context else {} + # TLS listeners may use certificates that cannot be verified through a loopback address. + if self.kwargs.get("ssl_certfile") or self.kwargs.get("ssl_keyfile"): + servers = {} if self.mcp_enabled: from fastmcp import FastMCP - from fastmcp.utilities.lifespan import combine_lifespans self.mcp_server = FastMCP(name=app.config.app_name) self.mcp_app = self.mcp_server.http_app( @@ -84,6 +89,13 @@ class HttpService(BaseService): title=app.config.app_name, lifespan=lifespan, ) + for server in servers.values(): + gateway = server.jobs_mcp + self.service.router.lifespan_context = combine_lifespans( + gateway.app.lifespan, + self.service.router.lifespan_context, + ) + self.service.router.routes.append(Route(gateway.path, endpoint=gateway, methods=["GET", "POST", "DELETE"])) cors_origins = ["*"] self.service.add_middleware( CORSMiddleware, # type: ignore[arg-type] diff --git a/reme/components/service/mcp_tools.py b/reme/components/service/mcp_tools.py index 77209f2c..e29bd323 100644 --- a/reme/components/service/mcp_tools.py +++ b/reme/components/service/mcp_tools.py @@ -13,11 +13,20 @@ def add_mcp_job( tool_error_on_failure: bool, ) -> bool: """Register a non-stream job as an MCP tool on ``server``.""" + tool = make_mcp_job_tool(job, injected_job_kwargs=injected_job_kwargs, tool_error_on_failure=tool_error_on_failure) + if tool is None: + return False + server.add_tool(tool) + return True + + +def make_mcp_job_tool(job: BaseJob, *, injected_job_kwargs: dict[str, Any], tool_error_on_failure: bool) -> Any: + """Build a tool bound to an existing job and server-owned arguments.""" from fastmcp.exceptions import ToolError from fastmcp.tools import FunctionTool if isinstance(job, StreamJob): - return False + return None async def execute_tool(**kwargs): conflicts = sorted(injected_job_kwargs.keys() & kwargs.keys()) @@ -43,12 +52,9 @@ def add_mcp_job( if "required" in parameters: parameters["required"] = [name for name in parameters["required"] if name not in injected_names] - server.add_tool( - FunctionTool( - name=job.name, - description=job.description, - fn=execute_tool, - parameters=parameters, - ), + return FunctionTool( + name=job.name, + description=job.description, + fn=execute_tool, + parameters=parameters, ) - return True diff --git a/reme/components/service/scoped_mcp.py b/reme/components/service/scoped_mcp.py new file mode 100644 index 00000000..c5c1e681 --- /dev/null +++ b/reme/components/service/scoped_mcp.py @@ -0,0 +1,99 @@ +"""Session-scoped MCP tools backed by the current Application's jobs.""" + +import asyncio +from contextlib import contextmanager +import secrets +import socket +from typing import Any + +from .mcp_tools import make_mcp_job_tool + + +class ScopedMCP: + """Serve isolated tool catalogs on one authenticated, application-owned endpoint.""" + + def __init__(self) -> None: + self.path = f"/_reme/codex/{secrets.token_hex(12)}" + self.scopes: dict[str, list[Any]] = {} + self._app: Any = None + self._server: Any = None + self._task: asyncio.Task | None = None + self._socket: socket.socket | None = None + + @property + def app(self) -> Any: + """Build one MCP transport; resolve tools separately for every HTTP request.""" + if self._app is None: + from fastmcp import FastMCP + from fastmcp.server.dependencies import get_http_headers + from fastmcp.server.providers import Provider + + scopes = self.scopes + + class JobProvider(Provider): + """Select the immutable tool binding carried by the request's bearer token.""" + + async def _list_tools(self): + authorization = get_http_headers(include={"authorization"}).get("authorization", "") + return scopes.get(authorization.removeprefix("Bearer "), []) + + self._app = FastMCP("ReMe agent tools", providers=[JobProvider()], tasks=False).http_app( + path=self.path, + stateless_http=True, + ) + return self._app + + def register(self, jobs: list[Any], injected: dict[str, Any]) -> str: + """Bind selected jobs and hidden arguments without cloning the Application.""" + token = secrets.token_urlsafe(32) + self.scopes[token] = [ + make_mcp_job_tool(job, injected_job_kwargs=injected, tool_error_on_failure=True) for job in jobs + ] + return token + + async def __call__(self, scope, receive, send): + if scope["type"] == "http": + authorization = dict(scope["headers"]).get(b"authorization", b"").decode("latin-1") + if not authorization.startswith("Bearer ") or authorization[7:] not in self.scopes: + from starlette.responses import Response + + await Response(status_code=403)(scope, receive, send) + return + await self.app(scope, receive, send) # pylint: disable=not-callable + + async def listen(self) -> str: + """Start a loopback listener for embedded/CLI hosts that have no HTTP service.""" + if self._server is None: + import uvicorn + + class LocalServer(uvicorn.Server): + """Leave signal handling to the containing ReMe application.""" + + @contextmanager + def capture_signals(self): + """Keep the caller's signal handlers intact.""" + yield + + self._socket = socket.socket() + self._socket.bind(("127.0.0.1", 0)) + self._server = LocalServer(uvicorn.Config(self, log_config=None, access_log=False, ws="none")) + self._task = asyncio.create_task(self._server.serve(sockets=[self._socket])) + while not self._server.started: + if self._task.done(): + await self._task + raise RuntimeError("Codex tool endpoint stopped before startup") + await asyncio.sleep(0.01) + return f"http://127.0.0.1:{self._socket.getsockname()[1]}{self.path}" + + async def close(self) -> None: + """Stop the optional listener and revoke all tool bindings.""" + try: + if self._server is not None: + self._server.should_exit = True # pylint: disable=attribute-defined-outside-init + try: + await self._task + finally: + self._socket.close() + self._server = self._task = self._socket = None + finally: + self.scopes.clear() diff --git a/reme/config/default.yaml b/reme/config/default.yaml index d2d481ba..677adf8c 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -930,6 +930,58 @@ components: max_tokens: 65536 thinking_enable: false + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${CODEX_API_KEY:-} + base_url: ${CODEX_BASE_URL:-} + codex_home: .codex + supports_websockets: false + config_overrides: + - 'model_reasoning_effort="xhigh"' + - 'review_model="qwen3.8-max"' + - 'agents.default_subagent_model="qwen3.8-max"' + - 'agents.default_subagent_reasoning_effort="xhigh"' + - 'memories.extract_model="qwen3.8-flash"' + - 'memories.consolidation_model="qwen3.8-max"' + model_catalog: + models: + - &codex_model_defaults + slug: &codex_simple_model qwen3.8-flash + display_name: Qwen3.8 Flash + description: Simple tasks and automatic approval review + auto_review_model_override: *codex_simple_model + default_reasoning_level: xhigh + multi_agent_reasoning_effort: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 1 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + - <<: *codex_model_defaults + slug: &codex_complex_model qwen3.8-max + display_name: Qwen3.8 Max + description: Main conversations and complex tasks + default_reasoning_level: xhigh + priority: 0 + oauth: + backend: sdk + auth_mode: oauth + codex_home: ${CODEX_HOME:-~/.codex} + agent_wrapper: default: backend: agentscope @@ -951,18 +1003,16 @@ components: permission_mode: bypassPermissions codex: backend: codex - auth_mode: api_key - model: ${CODEX_MODEL_NAME:-} - api_key: ${CODEX_API_KEY:-} - base_url: ${CODEX_BASE_URL:-} - approval_mode: auto_review + app_server: default + model: ${CODEX_MODEL_NAME:-qwen3.8-max} + effort: xhigh + approval_mode: deny_all sandbox: full-access codex_oauth: backend: codex - auth_mode: oauth - model: ${CODEX_MODEL_NAME:-} - codex_home: ${CODEX_HOME:-~/.codex} - approval_mode: auto_review + app_server: oauth + model: ${CODEX_OAUTH_MODEL_NAME:-} + approval_mode: deny_all sandbox: full-access file_graph: diff --git a/reme/config/demo.yaml b/reme/config/demo.yaml index b883ce83..dd66cafb 100644 --- a/reme/config/demo.yaml +++ b/reme/config/demo.yaml @@ -115,6 +115,54 @@ components: max_tokens: 65536 thinking_enable: false + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + codex_home: .codex + supports_websockets: false + config_overrides: + - 'model_reasoning_effort="xhigh"' + - 'review_model="qwen3.8-max"' + - 'agents.default_subagent_model="qwen3.8-max"' + - 'agents.default_subagent_reasoning_effort="xhigh"' + - 'memories.extract_model="qwen3.8-flash"' + - 'memories.consolidation_model="qwen3.8-max"' + model_catalog: + models: + - &codex_model_defaults + slug: &codex_simple_model qwen3.8-flash + display_name: Qwen3.8 Flash + description: Simple tasks and automatic approval review + auto_review_model_override: *codex_simple_model + default_reasoning_level: xhigh + multi_agent_reasoning_effort: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 1 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + - <<: *codex_model_defaults + slug: &codex_complex_model qwen3.8-max + display_name: Qwen3.8 Max + description: Main conversations and complex tasks + default_reasoning_level: xhigh + priority: 0 + agent_wrapper: as: backend: agentscope @@ -140,10 +188,10 @@ components: permission_mode: bypassPermissions codex: backend: codex - auth_mode: oauth - model: ${CODEX_MODEL_NAME:-} - codex_home: ${CODEX_HOME:-~/.codex} + app_server: default + model: *codex_complex_model + effort: xhigh job_tools: - add - approval_mode: auto_review + approval_mode: deny_all sandbox: full-access diff --git a/reme/config/proactive.yaml b/reme/config/proactive.yaml new file mode 100644 index 00000000..7f189229 --- /dev/null +++ b/reme/config/proactive.yaml @@ -0,0 +1,266 @@ +# Minimal Studio workspace with one Codex chat agent using API-key authentication. +# Set LLM_API_KEY / LLM_BASE_URL (and optionally CODEX_MODEL_NAME) for chat. +# Custom LLM_BASE_URL endpoints must support the Responses API. +# Start: reme start config=proactive +# Open: http://127.0.0.1:2333 + +service: + backend: http + web_enabled: true + mcp_enabled: false + +jobs: + # Keep Markdown search and graph data in sync with the workspace files. + index_update_loop: + backend: background + watch_dirs: [daily_dir, digest_dir] + watch_suffixes: [md] + steps: + - backend: init_changes_step + monitor_type: file_store + monitor_name: default + dispatch_steps: [update_index_step] + - backend: watch_changes_step + monitor_type: file_store + dispatch_steps: + - backend: update_index_step + persist: false + + chat: + backend: stream + description: "Stream a read-only agent conversation over the ReMe workspace." + parameters: + type: object + properties: + query: {type: string, description: "user message"} + session_id: {type: string, description: "agent session to resume"} + system_prompt: {type: string, description: "optional system prompt override"} + required: [query] + steps: + - backend: chat_step + agent_wrapper: default + + app_config: + backend: base + description: "Return the effective application config with secrets redacted." + steps: + - backend: app_config_step + + version: + backend: base + description: "Return the ReMe version." + steps: + - backend: version_step + + status: + backend: base + description: "Report component memory estimates and process RSS." + steps: + - backend: status_step + + health_check: + backend: base + description: "Return a concise health-check snapshot of ReMe components." + steps: + - backend: health_check_step + + graph_snapshot: + backend: base + description: "Return the category-rooted digest wikilink graph with daily-note leaves." + steps: + - backend: graph_snapshot_step + + reindex: + backend: base + description: "Rebuild the local BM25 index from already-ingested chunks." + parameters: + type: object + properties: + scope: {type: string, enum: [all, bm25], default: all} + steps: + - backend: reindex_step + + list: + backend: base + description: "List files under a workspace path." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative directory; empty = root", default: ""} + recursive: {type: boolean, default: false} + sort_by: {type: string, enum: [mtime]} + extensions: {type: array, items: {type: string}} + limit: {type: integer, default: 100} + steps: + - backend: list_step + + load: + backend: base + description: "Load a complete text file for the local editor without agent-output truncation." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + required: [path] + steps: + - backend: load_step + max_bytes: 5242880 + + save: + backend: base + description: "Save text verbatim, optionally rejecting external changes since the file was opened." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + content: {type: string, description: "complete file content, including frontmatter"} + expected_mtime: {type: string, description: "optional ISO mtime returned by load/stat"} + required: [path, content] + steps: + - backend: save_step + + # ChatStep exposes these read-only jobs as the agent's tools. + search: + backend: base + description: "Search indexed daily and digest Markdown files with BM25." + parameters: + type: object + properties: + query: {type: string, description: "search query"} + limit: {type: integer, default: 5} + min_score: {type: number, default: 0.0} + start_date: {type: string, description: "optional inclusive YYYY-MM-DD start date"} + end_date: {type: string, description: "optional inclusive YYYY-MM-DD end date"} + required: [query] + steps: + - backend: search_step + + read: + backend: base + description: "Read a Markdown file under the workspace." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative Markdown path"} + start_line: {type: integer, description: "first line (1-based, inclusive)"} + end_line: {type: integer, description: "last line (1-based, inclusive)"} + required: [path] + steps: + - backend: read_step + + read_image: + backend: base + description: "Read an image file as base64." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative image path"} + required: [path] + steps: + - backend: read_image_step + max_bytes: 5242880 + + frontmatter_read: + backend: base + description: "Read a file's frontmatter." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative Markdown path"} + required: [path] + steps: + - backend: frontmatter_read_step + + stat: + backend: base + description: "Stat a workspace path (size, mtime, exists, is_dir, is_file)." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + required: [path] + steps: + - backend: stat_step + + traverse: + backend: base + description: "Return a bounded wikilink graph with frontend-ready nodes and directed edges." + parameters: + type: object + properties: + path: + oneOf: + - {type: string} + - {type: array, items: {type: string}, minItems: 1} + description: "workspace-relative path or paths used as graph seeds" + depth: {type: integer, minimum: 0, default: 1} + direction: {type: string, enum: [forward, backward, both], default: both} + required: [path] + steps: + - backend: traverse_step + +components: + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + codex_home: .codex + model_catalog: + models: + - slug: qwen3.8-max + display_name: Qwen3.8 Max + description: Studio chat + default_reasoning_level: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 0 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + + agent_wrapper: + default: + backend: codex + app_server: default + model: ${CODEX_MODEL_NAME:-qwen3.8-max} + effort: xhigh + approval_mode: deny_all + sandbox: read-only + + tokenizer: + default: + backend: regex + + keyword_index: + default: + backend: bm25 + tokenizer: default + + file_graph: + default: + backend: local + + file_chunker: + markdown: + backend: markdown + supported_extensions: [md] + + file_store: + default: + backend: local + embedding_store: "" + keyword_index: default + file_graph: default diff --git a/reme/enumeration/component_enum.py b/reme/enumeration/component_enum.py index 8b317d61..37e03962 100644 --- a/reme/enumeration/component_enum.py +++ b/reme/enumeration/component_enum.py @@ -40,6 +40,8 @@ class ComponentEnum(str, Enum): AGENT_WRAPPER = "agent_wrapper" + CODEX_APP_SERVER = "codex_app_server" + MAIL = "mail" OUTBOUND_PROXY = "outbound_proxy" diff --git a/reme/steps/common/stream_llm_demo.py b/reme/steps/common/stream_llm_demo.py index 4fca8e71..aa554acf 100644 --- a/reme/steps/common/stream_llm_demo.py +++ b/reme/steps/common/stream_llm_demo.py @@ -10,7 +10,7 @@ class StreamLLMDemoStep(BaseStep): """Drive an Agent powered by the ``agent_wrapper`` component with streaming output. When streaming is enabled on the context, text/thinking/tool events are - pushed chunk-by-chunk via ``self.context.add_stream_string``. + pushed chunk-by-chunk via ``self.context.add_stream_chunk``. When streaming is not enabled, falls back to non-streaming reply. Inputs (from RuntimeContext): @@ -63,7 +63,7 @@ class StreamLLMDemoStep(BaseStep): text_parts: list[str] = [] async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - await self.context.add_stream_string(chunk.chunk, chunk.chunk_type) + await self.context.add_stream_chunk(chunk) if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): text_parts.append(chunk.chunk) diff --git a/tests/integration/test_codex_agent_wrapper.py b/tests/integration/test_codex_agent_wrapper.py index 2cc45a11..3cef6b4d 100644 --- a/tests/integration/test_codex_agent_wrapper.py +++ b/tests/integration/test_codex_agent_wrapper.py @@ -1,23 +1,17 @@ -"""Opt-in live coverage for the Codex app-server wrapper. - -Run with ``REME_CODEX_INTEGRATION=1``. These tests consume the caller's active -API-key or Codex OAuth account and are intentionally excluded from normal CI. -""" +"""Opt-in live Codex SDK coverage; set REME_CODEX_INTEGRATION=1 to use credentials.""" # pylint: disable=protected-access import asyncio import os from pathlib import Path -import subprocess -from types import SimpleNamespace import pytest from pydantic import BaseModel -from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper -from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import ApplicationConfig, Response +from reme.application import Application +from reme.enumeration import ChunkEnum +from reme.steps.base_step import BaseStep pytestmark = pytest.mark.skipif( os.getenv("REME_CODEX_INTEGRATION") != "1", @@ -29,213 +23,123 @@ class _StructuredResult(BaseModel): marker: str -class _CustomJob: - name = "only_custom" - description = "Return the fixed marker CUSTOM_JOB_OK." - parameters = {"type": "object", "properties": {}, "additionalProperties": False} +class _CustomStep(BaseStep): + """Exercise a parent-only registration and the actual parent metadata.""" - async def __call__(self, **_kwargs): - return Response(answer="CUSTOM_JOB_OK") - - -class _DraftJob: - def __init__(self, name: str, parameters: dict): - self.name = name - self.description = f"Live tool-context contract job: {name}" - self.parameters = parameters - - -def _child_pids(root_pid: int) -> set[int]: - result = subprocess.run( - ["ps", "-axo", "pid=,ppid="], - check=True, - capture_output=True, - text=True, - ) - children: dict[int, list[int]] = {} - for line in result.stdout.splitlines(): - pid, parent = (int(value) for value in line.split()) - children.setdefault(parent, []).append(pid) - found: set[int] = set() - pending = list(children.get(root_pid, [])) - while pending: - pid = pending.pop() - if pid not in found: - found.add(pid) - pending.extend(children.get(pid, [])) - return found + async def execute(self): + self.app_context.metadata["live_custom_calls"] = self.app_context.metadata.get("live_custom_calls", 0) + 1 + self.context.response.answer = "CUSTOM_JOB_OK" @pytest.mark.asyncio -async def test_live_codex_reply_stream_tools_skills_resume_fork_approval_and_close(tmp_path): - """Exercise the live Codex wrapper contract when explicitly enabled.""" +async def test_live_shared_codex_sessions_tools_streaming_and_close(tmp_path): + """Use one real SDK process and execute every job inside the containing Application.""" workspace = tmp_path / "workspace" workspace.mkdir() - project_skill = Path(__file__).resolve().parents[2] / "skills" / "reme_memory" skills_root = workspace / "skills" skills_root.mkdir() + project_skill = Path(__file__).resolve().parents[2] / "skills" / "reme_memory" (skills_root / "reme_memory").symlink_to(project_skill, target_is_directory=True) - - app_config = ApplicationConfig( + api_key = base_url = "" + for prefix in ("CODEX", "OPENAI", "LLM"): + if api_key := os.getenv(f"{prefix}_API_KEY", ""): + base_url = os.getenv(f"{prefix}_BASE_URL", "") + break + codex_home = os.getenv("REME_CODEX_HOME") + if not api_key and codex_home is None: + codex_home = str(Path.home() / ".codex") + app = Application( workspace_dir=str(workspace), - mem_session_dir="sessions", enable_logo=False, log_to_console=False, log_to_file=False, service={"backend": "mcp"}, + components={ + "codex_app_server": { + "default": { + "backend": "sdk", + "auth_mode": "api_key" if api_key else "oauth", + "api_key": api_key, + "base_url": base_url, + "codex_home": codex_home, + }, + }, + "agent_wrapper": { + name: {"backend": "codex", "model": os.getenv("CODEX_MODEL_NAME")} for name in ("first", "second") + }, + }, jobs={ "only_custom": { "backend": "base", - "description": _CustomJob.description, - "parameters": _CustomJob.parameters, - "steps": [], + "description": "Return CUSTOM_JOB_OK.", + "enable_serve": False, + "steps": [{"backend": "live_custom_step"}], }, - "add_live_draft": { + "add_draft": { "backend": "base", "description": "Append text to the current tool context.", - "parameters": { - "type": "object", - "properties": {"text": {"type": "string"}}, - "required": ["text"], - }, + "parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}, "steps": [{"backend": "add_draft_step"}], }, - "read_live_draft": { + "read_draft": { "backend": "base", "description": "Read text accumulated in the current tool context.", "parameters": {"type": "object", "properties": {}}, "steps": [{"backend": "read_all_draft_step"}], }, }, - components={ComponentEnum.AS_LLM: {}}, ) - context = SimpleNamespace( - app_config=app_config, - jobs={ - "only_custom": _CustomJob(), - "add_live_draft": _DraftJob( - "add_live_draft", - { - "type": "object", - "properties": {"text": {"type": "string"}}, - "required": ["text"], - }, - ), - "read_live_draft": _DraftJob( - "read_live_draft", - {"type": "object", "properties": {}}, - ), - }, - ) - codex_home = os.getenv("REME_CODEX_HOME") - if codex_home is None and not any(os.getenv(name) for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY")): - codex_home = str(Path.home() / ".codex") - wrapper = CodexAgentWrapper(app_context=context, codex_home=codex_home) - - await wrapper.start() - first = await wrapper.reply( - "Call add_live_draft with text STATE_OK, then reply FIRST_TURN_OK.", - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert "FIRST_TURN_OK" in first["last_message"] - - resumed = await wrapper.reply( - "Call read_live_draft. Include its exact result and RESUME_OK in your answer.", - resume=first["session_id"], - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert resumed["session_id"] == first["session_id"] - assert "RESUME_OK" in resumed["last_message"] - assert "STATE_OK" in resumed["last_message"] - - isolated = await wrapper.reply( - "Call read_live_draft. If it is empty, reply ISOLATED_OK.", - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="other-context", - ) - assert "ISOLATED_OK" in isolated["last_message"] - assert "STATE_OK" not in isolated["last_message"] - - resumed_after_isolated_context = await wrapper.reply( - "Call read_live_draft again and include its exact result.", - resume=first["session_id"], - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert "STATE_OK" in resumed_after_isolated_context["last_message"] - - forked = await wrapper.reply( - "Reply with exactly FORK_OK.", - resume=first["session_id"], - fork_session=True, - tool_context_id="fork-context", - ) - assert forked["session_id"] != first["session_id"] - assert "FORK_OK" in forked["last_message"] - - structured = await wrapper.reply( - "Return marker STRUCTURED_OK.", - output_schema=_StructuredResult, - ) - assert structured["structured_output"] == {"marker": "STRUCTURED_OK"} - - tool_result = await wrapper.reply( - "Call the only_custom tool once, then include its result in your answer.", - job_tools=["only_custom"], - tool_context_id="tool-context", - ) - assert "CUSTOM_JOB_OK" in tool_result["last_message"] - - skill_result = await wrapper.reply( - "Use the reme_memory skill. Read only its instructions and reply SKILL_OK; do not run its scripts or CLI.", - skills=["reme_memory"], - ) - assert "SKILL_OK" in skill_result["last_message"] - assert (workspace / ".agents" / "skills" / "reme_memory").is_symlink() - - outside_target = tmp_path / "approval-target.txt" - approval_chunks = [ - chunk - async for chunk in wrapper.reply_stream( - f"Try to write the word approved to {outside_target} using a shell command.", - approval_mode="auto_review", - sandbox="workspace-write", + app.context.registry.add("live_custom_step", _CustomStep, owner=__name__) + server = app.context.components["codex_app_server"]["default"] + async with app: + first, second = (app.context.components["agent_wrapper"][name] for name in ("first", "second")) + original = await first.reply( + "Call add_draft with text STATE_OK, then reply FIRST_OK.", + job_tools=["add_draft", "read_draft"], + tool_context_id="one", ) - ] - assert approval_chunks[0].chunk_type == ChunkEnum.REPLY_START - assert approval_chunks[-1].chunk_type == ChunkEnum.REPLY_END - assert any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in approval_chunks) - - codex_proc = wrapper._codex._client._sync._proc - assert codex_proc is not None and codex_proc.poll() is None - child_pids = _child_pids(codex_proc.pid) - await wrapper.close() - await asyncio.sleep(0.2) - - assert codex_proc.poll() is not None - running_pids = { - int(line) - for line in subprocess.run( - ["ps", "-axo", "pid="], - check=True, - capture_output=True, - text=True, - ).stdout.split() - } - assert child_pids.isdisjoint(running_pids) - assert wrapper._codex is None - assert wrapper._mcp_snapshot_path is None - - second_context = SimpleNamespace(app_config=app_config, jobs=context.jobs.copy()) - second_wrapper = CodexAgentWrapper(app_context=second_context, codex_home=codex_home) - await second_wrapper.start() - isolated_application = await second_wrapper.reply( - "Call read_live_draft. If it is empty, reply NEW_APPLICATION_OK.", - job_tools=["read_live_draft"], - tool_context_id="live-context", - ) - await second_wrapper.close() - assert "NEW_APPLICATION_OK" in isolated_application["last_message"] - assert "STATE_OK" not in isolated_application["last_message"] + assert "FIRST_OK" in original["last_message"] + resumed = await second.reply( + "Call read_draft and include its exact result.", + session_id=original["session_id"], + job_tools=["add_draft", "read_draft"], + tool_context_id="one", + ) + assert resumed["session_id"] == original["session_id"] + assert "STATE_OK" in resumed["last_message"] + isolated = await second.reply( + "Call read_draft. If it is empty, reply ISOLATED_OK.", + job_tools=["read_draft"], + tool_context_id="two", + ) + assert isolated["session_id"] != original["session_id"] + assert "ISOLATED_OK" in isolated["last_message"] + assert "STATE_OK" not in isolated["last_message"] + forked = await first.reply("Reply FORK_OK.", resume=original["session_id"], fork_session=True) + assert forked["session_id"] != original["session_id"] + assert "FORK_OK" in forked["last_message"] + custom = await first.reply("Call only_custom and include its result.", job_tools=["only_custom"]) + assert "CUSTOM_JOB_OK" in custom["last_message"] + assert app.context.metadata["live_custom_calls"] >= 1 + structured = await first.reply("Return marker STRUCTURED_OK.", output_schema=_StructuredResult) + assert structured["structured_output"] == {"marker": "STRUCTURED_OK"} + await first.compact_session(structured["session_id"]) + skill = await first.reply( + "Read the reme_memory skill instructions and reply SKILL_OK. Do not execute its scripts.", + skills=["reme_memory"], + ) + assert "SKILL_OK" in skill["last_message"] + assert (workspace / ".agents/skills/reme_memory").is_symlink() + concurrent = await asyncio.gather(first.reply("Reply ONE."), second.reply("Reply TWO.")) + assert concurrent[0]["session_id"] != concurrent[1]["session_id"] + chunks = [chunk async for chunk in second.reply_stream("Reply STREAM_OK.")] + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.REPLY_END + process = server._codex._client._sync._proc + assert process.poll() is None + await first.close() + assert process.poll() is None + assert (await second.reply("Reply STILL_OK."))["result"] + assert process.poll() is not None + assert server._codex is None + assert server.jobs_mcp._task is None diff --git a/tests/integration/test_codex_demo.py b/tests/integration/test_codex_demo.py new file mode 100644 index 00000000..2755ca62 --- /dev/null +++ b/tests/integration/test_codex_demo.py @@ -0,0 +1,176 @@ +"""Opt-in live coverage of both demo Steps using Codex API-key authentication.""" + +# pylint: disable=redefined-outer-name + +import asyncio +from collections import Counter +import json +import os +from pathlib import Path + +import pytest +import pytest_asyncio + +from reme.application import Application +from reme.config import resolve_app_config +from reme.enumeration import ChunkEnum +from reme.steps.common.stream_llm_demo import StreamLLMDemoStep +from reme.utils import global_counter_get, parse_env_file + +pytestmark = pytest.mark.skipif( + os.getenv("REME_CODEX_INTEGRATION") != "1", + reason="set REME_CODEX_INTEGRATION=1 to run live Codex tests", +) + + +@pytest_asyncio.fixture +async def demo_app(tmp_path, monkeypatch, request): + """Load the requested LLM credentials while keeping all runtime state temporary.""" + env_path = Path(__file__).resolve().parents[2] / ".env" + environment = {**(parse_env_file(env_path) if env_path.is_file() else {}), **os.environ} + for name in ("LLM_API_KEY", "LLM_BASE_URL"): + assert environment.get(name), f"{name} is required for the live Codex demo tests" + monkeypatch.setenv(name, environment[name]) + config = resolve_app_config( + config="demo", + log_config=False, + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + ) + approval_mode = getattr(request, "param", "deny_all") + config["components"]["agent_wrapper"]["codex"]["approval_mode"] = approval_mode + server_config = config["components"]["codex_app_server"]["default"] + assert server_config["auth_mode"] == "api_key" + assert server_config["api_key"] == environment["LLM_API_KEY"] + assert server_config["base_url"] == environment["LLM_BASE_URL"] + app = Application(**config) + server = app.context.components["codex_app_server"]["default"] + async with app: + yield app + models = set() + reasoning = set() + for path in (server.session_path / "sessions").rglob("*.jsonl"): + for line in path.read_text(encoding="utf-8").splitlines(): + record = json.loads(line) + payload = record.get("payload", {}) + if record.get("type") == "event_msg" and payload.get("type") == "thread_settings_applied": + settings = payload["thread_settings"] + models.add(settings["model"]) + reasoning.add((settings["model"], settings.get("reasoning_effort"))) + print(f"\nCodex models observed: {sorted(models)}") + print(f"Codex reasoning observed: {sorted(reasoning)}") + assert ("qwen3.8-max", "xhigh") in reasoning + assert {effort for _, effort in reasoning} == {"xhigh"} + assert models == ( + {"qwen3.8-max", "qwen3.8-flash"} + if approval_mode == "auto_review" and global_counter_get(app.context.metadata, ["__job_counter", "add"]) + else {"qwen3.8-max"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True], ids=["non-stream", "stream"]) +@pytest.mark.parametrize("use_tool", [False, True], ids=["chat", "add-tool"]) +async def test_live_codex_demo_chat_and_add_tool(demo_app, streaming, use_tool): + """Verify text delivery and actual ReMe add execution through each demo Job.""" + query = ( + "Use the add tool to calculate 21 + 21 and return only its result." + if use_tool + else "Write the integers from 1 through 20 in order, separated by spaces." + ) + sys_prompt = "Use the add tool for addition. Otherwise follow the requested output format." + previous_calls = global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) + async with asyncio.timeout(90): + if streaming: + chunks = [ + chunk + async for chunk in demo_app.run_stream_job( + "stream_llm_demo", + query=query, + sys_prompt=sys_prompt, + ) + ] + errors = [ + chunk.chunk + for chunk in chunks + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True + ] + assert not errors, errors + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.DONE + assert chunks[-1].done + assert not any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert any(chunk.chunk_type == ChunkEnum.REPLY_END for chunk in chunks) + assert chunks[-2].metadata.get("status") == "completed" + content = [chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk] + text = "".join(content).strip() + counts = Counter(chunk.chunk_type.value for chunk in chunks) + print(f"\nstream tool={use_tool}: chunks={dict(counts)}, answer={text!r}") + if use_tool: + assert counts[ChunkEnum.TOOL_CALL.value] > 0 + assert counts[ChunkEnum.TOOL_RESULT.value] > 0 + else: + response = await demo_app.run_job("llm_demo", query=query, sys_prompt=sys_prompt) + assert response.success, response.answer + text = response.answer.strip() + assert response.metadata["response"] == text + print(f"\nnon-stream tool={use_tool}: answer={text!r}") + assert text + if use_tool: + assert "42" in text + assert global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) > previous_calls + else: + assert text.split() == [str(number) for number in range(1, 21)] + + +@pytest.mark.asyncio +async def test_live_codex_stream_demo_without_stream_queue(demo_app): + """Without a queue, StreamLLMDemoStep uses the non-streaming wrapper reply.""" + step = StreamLLMDemoStep(app_context=demo_app.context, agent_wrapper="codex") + async with asyncio.timeout(90): + response = await step(query="Reply exactly FALLBACK_OK.") + assert response.success, response.answer + assert response.answer.strip() == "FALLBACK_OK" + + +@pytest.mark.asyncio +async def test_live_codex_native_shell_tool_needs_no_approval(demo_app): + """The default policy also executes a native shell tool without an approval request.""" + wrapper = demo_app.context.components["agent_wrapper"]["codex"] + async with asyncio.timeout(90): + chunks = [ + chunk + async for chunk in wrapper.reply_stream( + "Run the shell command printf NATIVE_TOOL_OK, then return only the command output.", + ) + ] + assert any( + chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_name == "commandExecution" for chunk in chunks + ) + assert any( + chunk.chunk_type == ChunkEnum.TOOL_RESULT and chunk.tool_call_name == "commandExecution" for chunk in chunks + ) + assert not any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert "NATIVE_TOOL_OK" in "".join( + chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("demo_app", ["auto_review"], indirect=True) +async def test_live_codex_automatic_review_remains_available(demo_app): + """Explicitly selecting auto_review still invokes the configured Flash reviewer.""" + async with asyncio.timeout(90): + chunks = [ + chunk + async for chunk in demo_app.run_stream_job( + "stream_llm_demo", + query="Call the add tool to calculate 21 + 21, then return only its result.", + ) + ] + assert any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) > 0 + assert "42" in "".join(chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT) diff --git a/tests/integration/test_codex_wrapper_capabilities.py b/tests/integration/test_codex_wrapper_capabilities.py new file mode 100644 index 00000000..f529c76a --- /dev/null +++ b/tests/integration/test_codex_wrapper_capabilities.py @@ -0,0 +1,139 @@ +"""Opt-in live checks for skills, additional instructions, and ReMe job tools.""" + +# pylint: disable=redefined-outer-name + +import os +from pathlib import Path +import secrets + +import pytest +import pytest_asyncio + +from reme.application import Application +from reme.steps.base_step import BaseStep + +pytestmark = pytest.mark.skipif( + os.getenv("REME_CODEX_INTEGRATION") != "1", + reason="set REME_CODEX_INTEGRATION=1 to run live Codex tests", +) + + +class _ProbeStep(BaseStep): + """Record execution in the containing Application rather than a child process.""" + + async def execute(self): + """Capture the tool input and return a deterministic result.""" + text = self.context.data["text"] + self.app_context.metadata.setdefault("codex_probe_calls", []).append( + {"text": text, "tool_context_id": self.context.data.get("tool_context_id")}, + ) + self.context.response.answer = f"JOB_OK:{text}" + + +@pytest_asyncio.fixture +async def wrapper(tmp_path): + """Use configured credentials with an isolated workspace and real SDK transport.""" + api_key = base_url = "" + for prefix in ("CODEX", "OPENAI", "LLM"): + if api_key := os.getenv(f"{prefix}_API_KEY", ""): + base_url = os.getenv(f"{prefix}_BASE_URL", "") + break + codex_home = os.getenv("REME_CODEX_HOME") + if not api_key and codex_home is None: + codex_home = str(Path.home() / ".codex") + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "mcp"}, + components={ + "codex_app_server": { + "default": { + "backend": "sdk", + "auth_mode": "api_key" if api_key else "oauth", + "api_key": api_key, + "base_url": base_url, + "codex_home": codex_home, + }, + }, + "agent_wrapper": {"default": {"backend": "codex", "model": os.getenv("CODEX_MODEL_NAME")}}, + }, + jobs={ + "capability_probe_job": { + "backend": "base", + "description": "Return the supplied text prefixed with JOB_OK:.", + "enable_serve": False, + "parameters": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + }, + "steps": [{"backend": "codex_capability_probe_step"}], + }, + }, + ) + app.context.registry.add("codex_capability_probe_step", _ProbeStep, owner=__name__) + async with app: + yield app.context.components["agent_wrapper"]["default"] + + +@pytest.mark.asyncio +async def test_live_codex_add_skill_reads_instructions(wrapper, tmp_path): + """The expected answer exists only inside the newly added skill file.""" + marker = secrets.token_hex(16) + source = tmp_path / "skills" / "capability_probe" + source.mkdir(parents=True) + (source / "SKILL.md").write_text( + "---\nname: capability_probe\ndescription: Return the verification token when explicitly invoked.\n---\n\n" + f"When invoked, reply with exactly this token and no other text: {marker}\n", + encoding="utf-8", + ) + wrapper.add_skills(["capability_probe"]) + result = await wrapper.reply("$capability_probe") + assert result["result"].strip() == marker + assert (tmp_path / ".agents/skills/capability_probe").resolve() == source.resolve() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resume", + [ + False, + pytest.param( + True, + marks=pytest.mark.xfail( + reason="Codex ignores developer instruction overrides on an already loaded thread", + strict=True, + raises=AssertionError, + ), + ), + ], +) +async def test_live_codex_extend_system_prompt(wrapper, resume): + """Check additional instructions on new sessions and expose the loaded-session limitation.""" + first_marker, added_marker = secrets.token_hex(16), secrets.token_hex(16) + wrapper.set_system_prompt(f"For input ORIGINAL_PROMPT_PROBE, reply exactly {first_marker}.") + original = await wrapper.reply("ORIGINAL_PROMPT_PROBE") + assert original["result"].strip() == first_marker + wrapper.set_system_prompt( + wrapper.kwargs["system_prompt"] + f"\nFor input ADDED_PROMPT_PROBE, reply exactly {added_marker}.", + ) + options = {"session_id": original["session_id"]} if resume else {} + extended = await wrapper.reply("ADDED_PROMPT_PROBE", **options) + assert (extended["session_id"] == original["session_id"]) == resume + assert extended["result"].strip() == added_marker + + +@pytest.mark.asyncio +async def test_live_codex_add_job_tool_calls_current_application(wrapper): + """A private ReMe job executes through MCP with its hidden context binding.""" + marker = secrets.token_hex(16) + wrapper.add_job_tools(["capability_probe_job"]) + result = await wrapper.reply( + f"Call capability_probe_job with text {marker}, then return only its exact result.", + tool_context_id="codex_capability_probe", + ) + calls = wrapper.app_context.metadata["codex_probe_calls"] + assert calls == [{"text": marker, "tool_context_id": "codex_capability_probe"}] + assert result["result"].strip() == f"JOB_OK:{marker}" diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index 6f436d52..f4f86937 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -1,340 +1,1052 @@ -"""Unit tests for the Codex agent wrapper and its FastMCP bridge.""" +"""Shared Codex runtime, isolated conversations, and same-Application MCP tools.""" -# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,too-many-lines +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,redefined-outer-name import asyncio from dataclasses import dataclass import json -import os from pathlib import Path -import stat -import sys +import socket from types import SimpleNamespace +from unittest.mock import AsyncMock +import httpx import pytest +import pytest_asyncio +from fastmcp import Client +from fastmcp.client import StreamableHttpTransport +from fastmcp.exceptions import ToolError from openai_codex.generated.v2_all import TokenUsageBreakdown from pydantic import BaseModel +from reme.application import Application +from reme.components.application_context import ApplicationContext from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper -from reme.components.agent_wrapper.codex_mcp_server import _prepare_config -from reme.components.job import BackgroundJob -from reme.components.outbound_proxy import FixedHttpOutboundProxy +from reme.components.codex_app_server import CodexAppServer +from reme.components.job import BackgroundJob, StreamJob +from reme.components.service.scoped_mcp import ScopedMCP from reme.config import resolve_app_config from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import ApplicationConfig, Response - -# A fresh bridge imports ReMe before replying to initialize; keep this separate -# from the 10-second budget for requests to an already initialized server. -STDIO_INIT_TIMEOUT = 30 +from reme.schema import Response class _Job: - def __init__(self, name="search"): - self.name = name - self.description = "Search memory" - self.parameters = { - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - } + name = "search" + description = "Search memory" + enabled = True + enable_serve = False + parameters = { + "type": "object", + "properties": {"query": {"type": "string"}, "_allowed_paths": {"type": "array"}}, + "required": ["query"], + } + + def __init__(self): self.calls = [] async def __call__(self, **kwargs): self.calls.append(kwargs) - return Response(answer=f"found:{kwargs['query']}") + return Response(answer={"query": kwargs["query"], "calls": len(self.calls)}) + + +@dataclass +class _Result: + final_response: str + usage: object = None + + +class _Turn: + def __init__(self, thread, inputs): + self.thread = thread + self.inputs = inputs + self.id = f"turn-{len(thread.inputs)}" + + async def run(self): + try: + if self.thread.codex.run_hook: + await self.thread.codex.run_hook(self.thread, self.inputs) + return _Result(self.inputs) + finally: + self.thread.busy = False + + async def interrupt(self): + self.thread.interrupts += 1 + self.thread.busy = False + + async def stream(self): + try: + yield SimpleNamespace(method="turn/started", payload=SimpleNamespace(turn=SimpleNamespace(id=self.id))) + yield SimpleNamespace( + method="item/agentMessage/delta", + payload=SimpleNamespace(item_id="item", delta=self.inputs), + ) + self.thread.busy = False + yield SimpleNamespace( + method="turn/completed", + payload=SimpleNamespace( + turn=SimpleNamespace( + id=self.id, + status=SimpleNamespace(value="completed"), + duration_ms=1, + error=None, + ), + ), + ) + finally: + self.thread.stream_closed = True + + +class _Thread: + def __init__(self, codex, thread_id): + self.codex, self.id = codex, thread_id + self.inputs = [] + self.options = [] + self._busy = False + self.interrupts = 0 + self.stream_closed = False + + @property + def busy(self): + return self._busy + + @busy.setter + def busy(self, value): + self._busy = value + self.codex.notify_status(self.id, "active" if value else "idle") + + async def turn(self, inputs, **kwargs): + assert not self.busy, "overlapping turns on one thread" + self.inputs.append(inputs) + self.options.append(kwargs) + self.busy = True + return _Turn(self, inputs) + + async def read(self): + return SimpleNamespace( + thread=SimpleNamespace( + status=SimpleNamespace(root=SimpleNamespace(type="active" if self.busy else "idle")), + ), + ) + + async def compact(self): + self.busy = True + if self.codex.compact_hook: + await self.codex.compact_hook(self) + else: + self.busy = False + + +class _Codex: + instances = [] + + def __init__(self, config): + self.config = config + self.threads = {} + self.loaded = set() + self.unsubscribed = [] + self.closed = 0 + self.logins = [] + self.opened = [] + self.run_hook = None + self.compact_hook = None + self.notifications = asyncio.Queue() + self._client = SimpleNamespace(request=self.request, next_notification=self.next_notification) + self.instances.append(self) + + async def account(self): + await asyncio.sleep(0) + return SimpleNamespace(account=object()) + + async def login_api_key(self, key): + self.logins.append(key) + + async def close(self): + from openai_codex import TransportClosedError + + self.closed += 1 + self.notifications.put_nowait(TransportClosedError("closed")) + + def notify_status(self, thread_id, status): + self.notifications.put_nowait( + SimpleNamespace( + method="thread/status/changed", + payload=SimpleNamespace(thread_id=thread_id, status=SimpleNamespace(root=SimpleNamespace(type=status))), + ), + ) + + async def next_notification(self): + notification = await self.notifications.get() + if isinstance(notification, Exception): + raise notification + return notification + + async def thread_start(self, **options): + thread = _Thread(self, f"thread-{len(self.threads) + 1}") + self.threads[thread.id] = thread + self.loaded.add(thread.id) + self.opened.append(("start", thread.id, options)) + return thread + + async def thread_resume(self, thread_id, **options): + self.loaded.add(thread_id) + self.opened.append(("resume", thread_id, options)) + return self.threads[thread_id] + + async def thread_fork(self, thread_id, **options): + forked = await self.thread_start(**options) + self.opened[-1] = ("fork", thread_id, options) + return forked + + async def request(self, method, params, **_kwargs): + if method == "thread/loaded/list": + return SimpleNamespace(data=list(self.loaded), next_cursor=None) + assert method == "thread/unsubscribe" + self.unsubscribed.append(params["threadId"]) + return SimpleNamespace(status="unsubscribed") def _wrapper(tmp_path, **kwargs): + context = ApplicationContext(workspace_dir=str(tmp_path)) job = _Job() - config = SimpleNamespace( - workspace_dir=str(tmp_path), - mem_session_dir="mem_session", - environment={}, - components={ComponentEnum.AS_LLM: {}}, - model_dump=lambda **_kwargs: { - "workspace_dir": str(tmp_path), - "enable_logo": False, - "log_to_console": False, - "log_to_file": False, - "jobs": {}, - "components": {}, - }, + context.jobs[job.name] = job + server = CodexAppServer( + name="default", + app_context=context, + auth_mode="oauth", + codex_home=kwargs.pop("codex_home", None), ) - context = SimpleNamespace(app_config=config, components={}, jobs={job.name: job}) + context.components = {ComponentEnum.CODEX_APP_SERVER: {"default": server}} return CodexAgentWrapper(app_context=context, **kwargs), job -def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - - config = wrapper._mcp_server_config( # pylint: disable=protected-access - {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, - ) - - assert config["command"] - assert config["enabled_tools"] == ["search"] - assert config["args"].count("--job") == 1 - assert config["args"][config["args"].index("--job") + 1] == "search" - assert "reme.components.agent_wrapper.codex_mcp_server" in config["args"] - assert config["args"][config["args"].index("--config") + 1] == str(tmp_path / "custom.yaml") - assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" - - -def test_mcp_config_serializes_injected_job_kwargs(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - - config = wrapper._mcp_server_config( # pylint: disable=protected-access - { - "job_tools": ["search"], - "tool_context_id": "ctx-1", - "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]}, - }, - ) - - raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] - assert json.loads(raw) == {"_allowed_paths": ["daily/2025-06-01/note.md"]} - - plain = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access - assert "--injected-job-kwargs" not in plain["args"] - - -def test_prepare_config_merges_injected_job_kwargs_with_tool_context(): - prepared = _prepare_config( - {"jobs": {"selected": {"backend": "base"}}}, - ["selected"], - "ctx-1", - {"_allowed_paths": ["daily/2025-06-01/note.md"]}, - ) - - assert prepared["service"]["injected_job_kwargs"] == { - "_allowed_paths": ["daily/2025-06-01/note.md"], - "tool_context_id": "ctx-1", - } - - -def test_thread_config_preserves_other_mcp_servers(tmp_path): - wrapper, _job = _wrapper(tmp_path) - config = wrapper._thread_config( # pylint: disable=protected-access - { - "job_tools": ["search"], - "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}}, - }, - ) - - assert "docs" in config["mcp_servers"] - assert len(config["mcp_servers"]) == 2 - assert next(name for name in config["mcp_servers"] if name != "docs").startswith("reme_jobs_") - - -@pytest.mark.asyncio -async def test_thread_config_injects_proxy_only_into_codex_shell_policy(tmp_path): - wrapper, _job = _wrapper(tmp_path) - proxy = FixedHttpOutboundProxy(url="http://127.0.0.1:18080") - await proxy.start() - wrapper.app_context.components = {ComponentEnum.OUTBOUND_PROXY: {"default": proxy}} - await wrapper.start() - - config = wrapper._thread_config( # pylint: disable=protected-access - { - "config": { - "shell_environment_policy": { - "inherit": "core", - "set": { - "CUSTOM": "preserved", - "HTTP_PROXY": "http://user-proxy.example:8080", - }, - }, - }, - }, - ) - environment = config["shell_environment_policy"]["set"] - - assert config["shell_environment_policy"]["inherit"] == "core" - assert environment["CUSTOM"] == "preserved" - for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): - assert environment[key] == proxy.http_url - - auth = wrapper._resolve_auth_config("oauth") # pylint: disable=protected-access - client_config = wrapper._build_client_config(auth) # pylint: disable=protected-access - assert "HTTP_PROXY" not in client_config.env - assert "HTTPS_PROXY" not in client_config.env - - await wrapper.close() - await proxy.close() - - -def test_mcp_config_rejects_background_jobs(tmp_path): - wrapper, _job = _wrapper(tmp_path) - wrapper.app_context.jobs["watch"] = BackgroundJob(name="watch", app_context=wrapper.app_context) - - with pytest.raises(TypeError, match="non-stream request jobs"): - wrapper._mcp_server_config({"job_tools": ["watch"]}) - - -def test_prepare_config_reuses_selected_stdio_mcp_service(): - prepared = _prepare_config( - { - "service": {"backend": "http"}, - "jobs": { - "selected": {"backend": "base", "enable_serve": False}, - "helper": {"backend": "base"}, - "watch": {"backend": "background"}, - }, - }, - ["selected"], - "ctx-1", - ) - - assert prepared["service"] == { - "backend": "mcp", - "transport": "stdio", - "jobs": ["selected"], - "tool_error_on_failure": True, - "injected_job_kwargs": {"tool_context_id": "ctx-1"}, - } - assert set(prepared["jobs"]) == {"selected", "helper"} - assert prepared["jobs"]["selected"]["enable_serve"] is True - - -def test_prepare_config_rejects_missing_or_background_selected_jobs(): - with pytest.raises(KeyError, match="missing"): - _prepare_config({"jobs": {}}, ["missing"]) - with pytest.raises(KeyError, match="watch"): - _prepare_config({"jobs": {"watch": {"backend": "background"}}}, ["watch"]) - - -@pytest.mark.asyncio -async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - config_path = tmp_path / "bridge.json" - config_path.write_text( - json.dumps( - { - "service": {"backend": "mcp"}, - "workspace_dir": str(tmp_path / "workspace"), - "jobs": { - "empty": { - "backend": "base", - "description": "Return an empty response", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - }, - ), - encoding="utf-8", - ) - transport = StdioTransport( - command=sys.executable, - args=[ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(config_path), - "--workspace", - str(tmp_path / "workspace"), - "--job", - "empty", - ], - cwd=str(Path(__file__).resolve().parents[2]), - keep_alive=False, - ) - - async with Client(transport, timeout=None, init_timeout=STDIO_INIT_TIMEOUT) as client: - tools = await asyncio.wait_for(client.list_tools(), timeout=10) - - assert [tool.name for tool in tools] == ["empty"] - - -@pytest.mark.asyncio -async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): - config_path = tmp_path / "bridge.json" - config_path.write_text( - json.dumps( - { - "service": {"backend": "mcp"}, - "workspace_dir": str(tmp_path / "workspace"), - "jobs": { - "empty": { - "backend": "base", - "description": "Empty", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - }, - ), - encoding="utf-8", - ) - proc = await asyncio.create_subprocess_exec( - sys.executable, - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(config_path), - "--workspace", - str(tmp_path / "workspace"), - "--job", - "empty", - cwd=str(Path(__file__).resolve().parents[2]), - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - request = { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2025-06-18", - "capabilities": {}, - "clientInfo": {"name": "test", "version": "1"}, - }, - } +@pytest_asyncio.fixture +async def system(tmp_path, monkeypatch): + _Codex.instances = [] + monkeypatch.setattr("openai_codex.AsyncCodex", _Codex) + first, job = _wrapper(tmp_path, name="first") + context = first.app_context + second = CodexAgentWrapper(name="second", app_context=context) + server = context.components[ComponentEnum.CODEX_APP_SERVER]["default"] + monkeypatch.setattr(server, "_mcp_url", AsyncMock(return_value="http://127.0.0.1/mcp")) try: - assert proc.stdin is not None and proc.stdout is not None and proc.stderr is not None - proc.stdin.write((json.dumps(request) + "\n").encode()) - await proc.stdin.drain() - first_line = await asyncio.wait_for(proc.stdout.readline(), timeout=STDIO_INIT_TIMEOUT) - message = json.loads(first_line) - assert message["jsonrpc"] == "2.0" - assert message["id"] == 1 + yield first, second, server, job finally: - if proc.returncode is None: - proc.terminate() + await first.close() + await second.close() + await server.close() + + +@pytest.mark.asyncio +async def test_wrappers_share_one_lazy_client_and_keep_independent_sessions(system): + first, second, server, _ = system + await first.start() + assert not _Codex.instances + results = await asyncio.gather(first.reply("first"), second.reply("second")) + assert len(_Codex.instances) == 1 + assert results[0]["session_id"] != results[1]["session_id"] + assert first.app_server is second.app_server is server + client = _Codex.instances[0] + await first.close() + assert client.closed == 0 + assert (await second.reply("still running"))["result"] == "still running" + await server.close() + assert client.closed == 1 + + +@pytest.mark.asyncio +async def test_different_threads_run_concurrently_and_same_thread_serializes_across_wrappers(system): + first, second, server, _ = system + thread_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + entered, release = asyncio.Event(), asyncio.Event() + observed = [] + + async def run_hook(thread, inputs): + observed.append((thread.id, inputs)) + if inputs == "hold": + entered.set() + await release.wait() + + client.run_hook = run_hook + async with asyncio.TaskGroup() as group: + held = group.create_task(first.reply("hold", session_id=thread_id)) + await asyncio.wait_for(entered.wait(), 2) + queued = group.create_task(second.reply("same", session_id=thread_id)) try: - await asyncio.wait_for(proc.wait(), timeout=10) - except asyncio.TimeoutError: - proc.kill() - await proc.wait() - stdout = first_line + await proc.stdout.read() - stderr = (await proc.stderr.read()).decode() - assert b"Loading config" not in stdout - assert b"INFO" not in stdout - assert b"WARNING" not in stdout - assert b"2026-" not in stdout - assert "Failed to parse JSONRPC message" not in stderr - assert "Invalid JSON" not in stderr + independent = await asyncio.wait_for(second.reply("other"), 2) + assert independent["session_id"] != thread_id + assert not queued.done() + assert (thread_id, "same") not in observed + finally: + release.set() + assert held.result()["session_id"] == queued.result()["session_id"] == thread_id + assert len(server._threads) == 2 -def test_sdk_responds_to_interactive_approval_server_requests(): - from openai_codex.client import CodexClient +@pytest.mark.asyncio +@pytest.mark.parametrize("approval_mode", [None, "auto_review"]) +async def test_resume_fork_and_sdk_options_remain_explicit(system, approval_mode): + from openai_codex import ApprovalMode, Sandbox + from openai_codex.types import Personality, ReasoningEffort - client = CodexClient() - assert client._default_approval_handler( - "item/commandExecution/requestApproval", - {}, - ) == { # pylint: disable=protected-access - "decision": "accept", - } - assert client._default_approval_handler( - "item/fileChange/requestApproval", - {}, - ) == { # pylint: disable=protected-access - "decision": "accept", + first, _, _, _ = system + original = await first.reply( + "one", + model="custom", + system_prompt="instructions", + personality="friendly", + approval_mode=approval_mode, + ) + thread_id = original["session_id"] + resumed = await first.reply("two", resume=thread_id, effort="high") + forked = await first.reply("three", session_id=thread_id, fork_session=True) + assert resumed["session_id"] == thread_id + assert forked["session_id"] != thread_id + client = _Codex.instances[0] + options = client.opened[0][2] + assert options["model"] == "custom" + assert options["developer_instructions"] == "instructions" + assert options["approval_mode"] == (ApprovalMode.auto_review if approval_mode else ApprovalMode.deny_all) + assert options["sandbox"] == Sandbox.full_access + assert options["personality"] == Personality.friendly + assert client.threads[thread_id].options[-1]["effort"] == ReasoningEffort.high + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changes", + [ + {"tool_context_id": "other"}, + {"job_tools": []}, + {"injected_job_kwargs": {"_allowed_paths": ["other.md"]}}, + ], +) +async def test_loaded_thread_constraints_cannot_be_changed_by_another_wrapper(system, changes): + first, second, server, _ = system + binding = { + "job_tools": ["search"], + "tool_context_id": "one", + "injected_job_kwargs": {"_allowed_paths": ["one.md"]}, } + original = await first.reply("one", **binding) + thread_id = original["session_id"] + token = server._threads[thread_id].token + with pytest.raises(ValueError, match="cannot change"): + await second.reply("two", session_id=thread_id, **{**binding, **changes}) + assert token in server.jobs_mcp.scopes + await second.reply("three", session_id=thread_id, **binding) + + +@pytest.mark.asyncio +async def test_unloaded_thread_rebinds_tools_and_revokes_old_token(system): + first, _, server, _ = system + original = await first.reply("one", job_tools=["search"], tool_context_id="one") + thread_id = original["session_id"] + token = server._threads[thread_id].token + _Codex.instances[0].loaded.clear() + await first.reply("two", session_id=thread_id, job_tools=["search"], tool_context_id="two") + assert token not in server.jobs_mcp.scopes + assert server._threads[thread_id].token != token + + +@pytest.mark.asyncio +async def test_cancelling_reply_interrupts_only_its_turn_and_keeps_shared_server(system): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + entered = asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await asyncio.Event().wait() + + client = _Codex.instances[0] + client.run_hook = hold + task = asyncio.create_task(first.reply("cancel", session_id=thread_id)) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + client.run_hook = None + assert client.threads[thread_id].interrupts == 1 + assert client.closed == 0 + assert (await second.reply("okay"))["result"] == "okay" + assert not client.threads[thread_id].busy + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["turn", "compact"]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_cancelled_dispatch_settles_before_resuming_the_thread(system, monkeypatch, operation, fail): + first, second, server, _ = system + thread_id = (await first.reply("initial"))["session_id"] + thread = _Codex.instances[0].threads[thread_id] + original = getattr(thread, operation) + entered, release = asyncio.Event(), asyncio.Event() + + async def dispatch(*args, **kwargs): + entered.set() + await release.wait() + if fail: + raise ValueError("dispatch rejected") + return await original(*args, **kwargs) + + monkeypatch.setattr(thread, operation, dispatch) + request = ( + first.compact_session(thread_id) if operation == "compact" else first.reply("cancel", session_id=thread_id) + ) + task = asyncio.create_task(request) + queued = None + try: + await asyncio.wait_for(entered.wait(), 2) + task.cancel("original cancellation") + await asyncio.sleep(0) + task.cancel("cancel again") + await asyncio.sleep(0) + queued = asyncio.create_task(second.reply("later", session_id=thread_id)) + await asyncio.sleep(0.01) + assert not queued.done() + monkeypatch.setattr(thread, operation, original) + release.set() + with pytest.raises(asyncio.CancelledError, match="original cancellation"): + await asyncio.wait_for(task, 2) + assert (await asyncio.wait_for(queued, 2))["result"] == "later" + assert thread.interrupts == int(operation == "turn" and not fail) + await asyncio.wait_for(server.close(), 2) + finally: + release.set() + task.cancel() + if queued is not None: + queued.cancel() + await asyncio.gather(task, *([queued] if queued is not None else []), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_stream_close_interrupts_and_routes_chunks_per_session(system): + first, second, _, _ = system + stream = first.reply_stream("early") + start = await anext(stream) + thread_id = start.session_id + await stream.aclose() + thread = _Codex.instances[0].threads[thread_id] + assert thread.interrupts == 1 + assert thread.stream_closed + chunks = [chunk async for chunk in second.reply_stream("complete")] + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.REPLY_END + assert all(chunk.session_id != thread_id for chunk in chunks) + + +@pytest.mark.asyncio +async def test_compaction_holds_thread_lock_until_it_finishes(system): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + entered, release = asyncio.Event(), asyncio.Event() + + async def compact_hook(thread): + entered.set() + + async def finish(): + await release.wait() + thread.busy = False + + asyncio.create_task(finish()) + + _Codex.instances[0].compact_hook = compact_hook + async with asyncio.TaskGroup() as group: + group.create_task(first.compact_session(thread_id)) + await asyncio.wait_for(entered.wait(), 2) + queued = group.create_task(second.reply("later", session_id=thread_id)) + try: + await asyncio.sleep(0.02) + assert not queued.done() + finally: + release.set() + assert queued.result()["session_id"] == thread_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_compaction_acknowledgement_does_not_release_the_thread_before_it_starts(system, monkeypatch, cancelled): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + thread = client.threads[thread_id] + acknowledged = asyncio.Event() + + async def compact(): + acknowledged.set() + + monkeypatch.setattr(thread, "compact", compact) + async with asyncio.TaskGroup() as group: + compacting = group.create_task(first.compact_session(thread_id)) + await asyncio.wait_for(acknowledged.wait(), 2) + if cancelled: + compacting.cancel() + queued = group.create_task(second.reply("later", session_id=thread_id)) + try: + await asyncio.sleep(0.01) + # Cancellation may end the caller during cleanup; its barrier must still block the next turn. + if not cancelled: + assert not compacting.done() + assert not queued.done() + client.notify_status(thread_id, "idle") + await asyncio.sleep(0) + if not cancelled: + assert not compacting.done() + assert not queued.done() + finally: + client.notify_status(thread_id, "active") + client.notify_status(thread_id, "idle") + if cancelled: + assert compacting.cancelled() + assert queued.result()["session_id"] == thread_id + + +@pytest.mark.asyncio +async def test_close_waits_for_active_request_without_closing_other_wrappers(system): + first, second, server, _ = system + await first.reply("initial") + entered, release = asyncio.Event(), asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await release.wait() + + client = _Codex.instances[0] + client.run_hook = hold + async with asyncio.TaskGroup() as group: + group.create_task(first.reply("held")) + await asyncio.wait_for(entered.wait(), 2) + closing = group.create_task(first.close()) + try: + await asyncio.sleep(0) + assert not closing.done() + assert client.closed == 0 + finally: + release.set() + client.run_hook = None + await second.reply("still works") + assert server.is_started + + +class _StructuredModel(BaseModel): + ok: bool + + +@pytest.mark.asyncio +async def test_structured_reply_and_native_inputs(system): + from openai_codex import LocalImageInput, TextInput + + first, _, _, _ = system + result = await first.reply('{"ok": true}', output_schema=_StructuredModel) + assert result["structured_output"] == {"ok": True} + inputs = [TextInput("describe"), LocalImageInput("image.png")] + await first.reply(inputs) + assert _Codex.instances[0].threads["thread-2"].inputs[0] is inputs + with pytest.raises(NotImplementedError): + await anext(first.reply_stream("no", output_schema=_StructuredModel)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kwargs", [{"resume": "a", "session_id": "b"}, {"fork_session": True}]) +async def test_invalid_session_selection_fails_without_starting_client(system, kwargs): + first, _, _, _ = system + with pytest.raises(ValueError): + await first.reply("invalid", **kwargs) + assert not _Codex.instances + + +@pytest.mark.asyncio +@pytest.mark.parametrize("job_class", [BackgroundJob, StreamJob]) +async def test_unsupported_job_types_are_rejected(system, job_class): + first, _, _, _ = system + first.app_context.jobs["unsupported"] = job_class(name="unsupported") + with pytest.raises(TypeError, match="non-stream"): + await first.reply("invalid", job_tools=["unsupported"]) + + +@pytest.mark.parametrize("name", sorted(CodexAgentWrapper._SERVER_OPTIONS)) +def test_process_options_must_be_configured_on_server(name): + with pytest.raises(TypeError, match="codex_app_server"): + CodexAgentWrapper(**{name: "invalid"}) + + +def test_server_config_authentication_environment_and_home(tmp_path): + context = ApplicationContext(workspace_dir=str(tmp_path), environment={"CUSTOM": "preserved"}) + server = CodexAppServer( + name="one", + app_context=context, + auth_mode="api_key", + api_key="key", + base_url="https://proxy.test/v1", + ) + config = server._client_config() + assert config.env["CUSTOM"] == "preserved" + assert config.env["CODEX_HOME"] == str(tmp_path / ".codex") + assert 'forced_login_method="api"' in config.config_overrides + assert 'model_provider="reme"' in config.config_overrides + assert 'model_providers.reme.base_url="https://proxy.test/v1"' in config.config_overrides + assert "model_providers.reme.requires_openai_auth=true" in config.config_overrides + assert "model_providers.reme.supports_websockets=false" in config.config_overrides + oauth = CodexAppServer(app_context=context, auth_mode="oauth", api_key="ignored", base_url="ignored") + assert CodexAppServer(codex_home="~/.codex").session_path == Path.home() / ".codex" + assert CodexAppServer(app_context=context, codex_home="other-home").session_path == tmp_path / "other-home" + assert oauth._client_config().config_overrides == ('forced_login_method="chatgpt"',) + with pytest.raises(ValueError, match="requires a non-empty"): + CodexAppServer(app_context=context, auth_mode="api_key")._client_config() + + +def test_custom_provider_websocket_support_and_explicit_overrides_are_preserved(tmp_path): + server = CodexAppServer( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + auth_mode="api_key", + api_key="key", + base_url="https://proxy.test/v1", + supports_websockets=True, + config_overrides=['model_provider="custom"'], + ) + overrides = server._client_config().config_overrides + assert "model_providers.reme.supports_websockets=true" in overrides + assert overrides.index('model_provider="custom"') > overrides.index('model_provider="reme"') + + +def test_server_materializes_model_catalog_inside_its_private_home(tmp_path): + context = ApplicationContext(workspace_dir=str(tmp_path)) + catalog = {"models": [{"slug": "main", "auto_review_model_override": "reviewer"}, {"slug": "reviewer"}]} + server = CodexAppServer(name="catalog", app_context=context, model_catalog=catalog) + config = server._client_config() + path = server.session_path / "reme-model-catalog.json" + assert json.loads(path.read_text(encoding="utf-8")) == catalog + assert f"model_catalog_json={json.dumps(str(path))}" in config.config_overrides + assert path.is_relative_to(tmp_path) + with pytest.raises(ValueError, match="not both"): + CodexAppServer( + app_context=context, + model_catalog=catalog, + config_overrides=['model_catalog_json="external.json"'], + )._client_config() + + +@pytest.mark.asyncio +async def test_client_initialization_failure_is_closed_and_retryable(system, monkeypatch): + first, _, server, _ = system + monkeypatch.setattr(_Codex, "account", AsyncMock(side_effect=RuntimeError("login failed"))) + with pytest.raises(RuntimeError, match="login failed"): + await first.reply("fails") + assert _Codex.instances[0].closed == 1 + assert server._codex is None + assert server._active == 0 + + +def test_default_and_demo_configuration_use_named_shared_servers(monkeypatch): + monkeypatch.delenv("CODEX_HOME", raising=False) + monkeypatch.delenv("CODEX_MODEL_NAME", raising=False) + monkeypatch.setenv("LLM_API_KEY", "demo-api-key") + monkeypatch.setenv("LLM_BASE_URL", "https://demo.invalid/v1") + for name in ("default", "demo"): + config = resolve_app_config(config=name, log_config=False) + wrappers = config["components"]["agent_wrapper"] + assert wrappers["codex"]["app_server"] == "default" + assert "auth_mode" not in wrappers["codex"] + servers = config["components"]["codex_app_server"] + assert servers["default"]["auth_mode"] == "api_key" + assert servers["default"]["codex_home"] == ".codex" + assert servers["default"]["supports_websockets"] is False + assert wrappers["codex"]["model"] == "qwen3.8-max" + assert wrappers["codex"]["approval_mode"] == "deny_all" + assert wrappers["codex"]["effort"] == "xhigh" + models = servers["default"]["model_catalog"]["models"] + assert {model["slug"] for model in models} == {"qwen3.8-max", "qwen3.8-flash"} + assert all(model["auto_review_model_override"] == "qwen3.8-flash" for model in models) + assert all(model["default_reasoning_level"] == "xhigh" for model in models) + assert all(model["multi_agent_reasoning_effort"] == "xhigh" for model in models) + assert all([level["effort"] for level in model["supported_reasoning_levels"]] == ["xhigh"] for model in models) + assert set(servers["default"]["config_overrides"]) == { + 'model_reasoning_effort="xhigh"', + 'review_model="qwen3.8-max"', + 'agents.default_subagent_model="qwen3.8-max"', + 'agents.default_subagent_reasoning_effort="xhigh"', + 'memories.extract_model="qwen3.8-flash"', + 'memories.consolidation_model="qwen3.8-max"', + } + if name == "demo": + assert servers["default"]["api_key"] == "demo-api-key" + assert servers["default"]["base_url"] == "https://demo.invalid/v1" + else: + assert servers["oauth"]["codex_home"] == "~/.codex" + + +@pytest.mark.asyncio +async def test_scoped_mcp_calls_original_jobs_and_isolates_catalog_arguments_and_failures(): + gateway = ScopedMCP() + job = _Job() + first = gateway.register([job], {"tool_context_id": "one", "_allowed_paths": ["one.md"]}) + second = gateway.register([job], {"tool_context_id": "two", "_allowed_paths": ["two.md"]}) + empty = gateway.register([], {}) + try: + url = await gateway.listen() + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {first}"})) as one: + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {second}"})) as two: + tools = await one.list_tools() + assert [tool.name for tool in tools] == ["search"] + assert "_allowed_paths" not in tools[0].input_schema["properties"] + await one.call_tool("search", {"query": "a"}) + await two.call_tool("search", {"query": "b"}) + await one.call_tool("search", {"query": "c"}) + with pytest.raises(ToolError, match="injected"): + await one.call_tool("search", {"query": "bad", "_allowed_paths": ["escape"]}) + assert [call["tool_context_id"] for call in job.calls] == ["one", "two", "one"] + assert job.calls[1]["_allowed_paths"] == ["two.md"] + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {empty}"})) as client: + assert not await client.list_tools() + with pytest.raises(ToolError): + await client.call_tool("search", {"query": "not allowed"}) + async with httpx.AsyncClient() as client: + assert (await client.post(url, json={})).status_code == 403 + assert (await client.post(url, headers={"Authorization": "Bearer unknown"}, json={})).status_code == 403 + finally: + await gateway.close() + assert gateway._task is None + assert not gateway.scopes + + +@pytest.mark.asyncio +async def test_mounted_mcp_transport_lifespan_can_restart(): + gateway = ScopedMCP() + app = gateway.app + for _ in range(2): + async with app.lifespan(app): + assert gateway.app is app + + +@pytest.mark.asyncio +async def test_scoped_mcp_can_restart_its_loopback_listener(): + gateway = ScopedMCP() + job = _Job() + try: + for _ in range(2): + token = gateway.register([job], {}) + url = await gateway.listen() + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {token}"})) as client: + await client.call_tool("search", {"query": "same application"}) + await gateway.close() + assert len(job.calls) == 2 + finally: + await gateway.close() + + +@pytest.mark.asyncio +async def test_http_service_mounts_scoped_tools_without_another_listener(tmp_path, monkeypatch): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False}, + components={ + "codex_app_server": {"default": {"backend": "sdk", "auth_mode": "oauth"}}, + "agent_wrapper": {"one": {"backend": "codex"}, "two": {"backend": "codex"}}, + }, + ) + service = app.context.service + service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + writer = SimpleNamespace(close=lambda: None, wait_closed=AsyncMock()) + monkeypatch.setattr(asyncio, "open_connection", AsyncMock(return_value=(None, writer))) + job = _Job() + token = server.jobs_mcp.register([job], {"tool_context_id": "parent"}) + + def factory(**kwargs): + return httpx.AsyncClient(transport=httpx.ASGITransport(app=service.service), **kwargs) + + async with service.service.router.lifespan_context(service.service): + url = await server._mcp_url() + async with Client( + StreamableHttpTransport( + url, + headers={"Authorization": f"Bearer {token}"}, + httpx_client_factory=factory, + ), + ) as client: + await client.call_tool("search", {"query": "same application"}) + assert job.calls == [{"query": "same application", "tool_context_id": "parent"}] + assert server.jobs_mcp._server is None + assert app.context.components["agent_wrapper"]["one"].app_server is server + assert app.context.components["agent_wrapper"]["two"].app_server is server + assert not server.is_started + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tls_options", [{"ssl_certfile": "server.pem"}, {"ssl_keyfile": "server.key"}]) +async def test_tls_http_service_uses_managed_loopback_for_scoped_tools(tmp_path, monkeypatch, tls_options): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False, **tls_options}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + service = app.context.service + service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + probe = AsyncMock(side_effect=AssertionError("must not probe the TLS listener")) + monkeypatch.setattr(asyncio, "open_connection", probe) + job = _Job() + token = server.jobs_mcp.register([job], {"tool_context_id": "parent"}) + listener_task = None + async with service.service.router.lifespan_context(service.service): + url = await server._mcp_url() + assert url.startswith("http://127.0.0.1:") + assert await server._mcp_url() == url + listener_task = server.jobs_mcp._task + assert listener_task is not None + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {token}"})) as client: + await client.call_tool("search", {"query": "TLS host"}) + assert job.calls == [{"query": "TLS host", "tool_context_id": "parent"}] + probe.assert_not_awaited() + assert listener_task.done() + assert server.jobs_mcp._task is None + assert not server.jobs_mcp.scopes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("option", ["cwd", "app_server"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_call_time_wrapper_options_are_rejected_instead_of_silently_ignored(system, option, stream): + first, _, _, _ = system + with pytest.raises(TypeError, match="wrapper options"): + if stream: + await anext(first.reply_stream("invalid", **{option: "ignored"})) + else: + await first.reply("invalid", **{option: "ignored"}) + assert not _Codex.instances + + +@pytest.mark.asyncio +async def test_mounted_tools_wait_for_the_main_listener_before_returning_a_url(tmp_path): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + sock.setblocking(False) + port = sock.getsockname()[1] + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "port": port, "web_enabled": False}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + app.context.service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + lookup = asyncio.create_task(server._mcp_url()) + listener = None + + async def accept(_reader, writer): + writer.close() + await writer.wait_closed() + + try: + await asyncio.sleep(0.02) + assert not lookup.done() + listener = await asyncio.start_server(accept, sock=sock) + url = await asyncio.wait_for(lookup, 2) + assert url == f"http://127.0.0.1:{port}{server.jobs_mcp.path}" + assert server.jobs_mcp._task is None + finally: + lookup.cancel() + await asyncio.gather(lookup, return_exceptions=True) + if listener is not None: + listener.close() + await listener.wait_closed() + sock.close() + + +@pytest.mark.asyncio +async def test_replaced_runtime_uses_loopback_when_its_gateway_is_not_mounted(tmp_path, monkeypatch): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + app.context.service.build_service(app) + replacement = CodexAppServer(app_context=app.context) + app.context.service.service.mount("/assets", app=ScopedMCP()) + listen = AsyncMock(return_value="http://127.0.0.1:1234/local") + monkeypatch.setattr(replacement.jobs_mcp, "listen", listen) + assert await replacement._mcp_url() == "http://127.0.0.1:1234/local" + listen.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_cleanup_preserves_the_original_request_error(system, monkeypatch): + first, _, server, _ = system + await first.reply("initial") + client = _Codex.instances[0] + + async def fail(_thread, _inputs): + raise ValueError("original failure") + + client.run_hook = fail + monkeypatch.setattr(server, "wait_idle", AsyncMock(side_effect=RuntimeError("cleanup failure"))) + with pytest.raises(ValueError, match="original failure"): + await first.reply("failed") + + +@pytest.mark.asyncio +async def test_cancelled_waiter_keeps_the_running_threads_tool_binding(system): + first, second, server, _ = system + binding = {"job_tools": ["search"], "tool_context_id": "one"} + thread_id = (await first.reply("initial", **binding))["session_id"] + token = server._threads[thread_id].token + entered, release = asyncio.Event(), asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await release.wait() + + client = _Codex.instances[0] + client.run_hook = hold + async with asyncio.TaskGroup() as group: + group.create_task(first.reply("held", session_id=thread_id, **binding)) + await asyncio.wait_for(entered.wait(), 2) + waiting = group.create_task(second.reply("queued", session_id=thread_id, **binding)) + try: + await asyncio.sleep(0) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + assert token in server.jobs_mcp.scopes + finally: + release.set() + client.run_hook = None + await second.reply("still authorized", session_id=thread_id, **binding) + assert "queued" not in client.threads[thread_id].inputs + + +@pytest.mark.asyncio +async def test_disconnected_turn_is_not_replayed_and_next_request_reconnects(system): + from openai_codex import TransportClosedError + + first, _, server, _ = system + await first.reply("initial") + old_client = _Codex.instances[0] + + async def disconnect(_thread, _inputs): + raise TransportClosedError("disconnected") + + old_client.run_hook = disconnect + with pytest.raises(TransportClosedError): + await first.reply("do not replay") + assert server._failed + assert old_client.threads["thread-2"].inputs == ["do not replay"] + assert (await first.reply("new request"))["result"] == "new request" + assert len(_Codex.instances) == 2 + assert old_client.closed == 1 + assert _Codex.instances[1].threads["thread-1"].inputs == ["new request"] + + +@pytest.mark.asyncio +async def test_initial_transport_failure_does_not_poison_client_retry(system, monkeypatch): + from openai_codex import TransportClosedError + + first, _, server, _ = system + monkeypatch.setattr( + _Codex, + "account", + AsyncMock( + side_effect=[ + TransportClosedError("initialize failed"), + SimpleNamespace(account=object()), + ], + ), + ) + with pytest.raises(TransportClosedError): + await first.reply("failed initialization") + assert not server._failed + assert (await first.reply("retry"))["result"] == "retry" + assert _Codex.instances[0].closed == 1 + + +@pytest.mark.asyncio +async def test_thread_pruning_waits_for_complete_pagination_and_preserves_tokens_on_error(system, monkeypatch): + first, _, server, _ = system + one = (await first.reply("one", job_tools=["search"]))["session_id"] + two = (await first.reply("two", job_tools=["search"]))["session_id"] + token = server._threads[two].token + client = _Codex.instances[0] + pages = [SimpleNamespace(data=[one], next_cursor="next"), RuntimeError("incomplete listing")] + request = AsyncMock(side_effect=pages) + monkeypatch.setattr(client._client, "request", request) + await server._prune_threads(client) + assert token in server.jobs_mcp.scopes + request.side_effect = [ + SimpleNamespace(data=[one], next_cursor="next"), + SimpleNamespace(data=[], next_cursor=None), + ] + await server._prune_threads(client) + assert one in server._threads + assert two not in server._threads + assert token not in server.jobs_mcp.scopes + + +@pytest.mark.asyncio +async def test_stale_unload_snapshot_does_not_revoke_a_resumed_threads_tools(system, monkeypatch): + first, second, server, _ = system + binding = {"job_tools": ["search"], "tool_context_id": "one"} + thread_id = (await first.reply("initial", **binding))["session_id"] + state = server._threads[thread_id] + token = state.token + client = _Codex.instances[0] + original_request = client.request + resume_snapshot, stale_snapshot = asyncio.Event(), asyncio.Event() + resume_query, stale_query = asyncio.Event(), asyncio.Event() + queries = 0 + + async def request(method, params, **kwargs): + nonlocal queries + if method == "thread/loaded/list": + queries += 1 + if queries == 1: + resume_snapshot.set() + await resume_query.wait() + return SimpleNamespace(data=[thread_id], next_cursor=None) + if queries == 2: + stale_snapshot.set() + await stale_query.wait() + return SimpleNamespace(data=[], next_cursor=None) + return await original_request(method, params, **kwargs) + + monkeypatch.setattr(client._client, "request", request) + async with asyncio.TaskGroup() as group: + resumed = group.create_task(second.reply("resume", session_id=thread_id, **binding)) + await asyncio.wait_for(resume_snapshot.wait(), 2) + client.loaded.clear() + pruning = group.create_task(server._prune_threads(client)) + await asyncio.wait_for(stale_snapshot.wait(), 2) + try: + resume_query.set() + await asyncio.wait_for(asyncio.shield(resumed), 2) + finally: + resume_query.set() + stale_query.set() + await pruning + assert server._threads.get(thread_id) is state + assert token in server.jobs_mcp.scopes + + +@pytest.mark.asyncio +async def test_tool_arguments_are_snapshotted_before_the_caller_mutates_them(system): + first, _, server, job = system + allowed = ["one.md"] + result = await first.reply("one", job_tools=["search"], injected_job_kwargs={"_allowed_paths": allowed}) + allowed[:] = ["other.md"] + token = server._threads[result["session_id"]].token + await server.jobs_mcp.scopes[token][0].fn(query="read") + assert job.calls[0]["_allowed_paths"] == ["one.md"] def test_event_to_chunks_maps_content_usage_and_completion(): @@ -397,89 +1109,6 @@ def test_event_to_chunks_preserves_new_turn_scoped_notifications(): assert chunk.metadata == {"codex_method": "turn/diff/updated"} -@dataclass -class _TurnResult: - final_response: str - status: str = "completed" - - -@pytest.mark.asyncio -async def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - - class FakeThread: - id = "thread-1" - - async def run(self, inputs, **kwargs): - assert inputs == "answer" - assert kwargs["output_schema"] == {"type": "object"} - return _TurnResult(final_response=json.dumps({"ok": True})) - - class FakeCodex: - def __init__(self, _config): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _tb): - return None - - async def close(self): - return None - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - result = await wrapper.reply("answer", output_schema={"type": "object"}) - await wrapper.close() - - assert result["session_id"] == "thread-1" - assert result["structured_output"] == {"ok": True} - - -@pytest.mark.asyncio -async def test_reply_accepts_latest_sdk_run_input(tmp_path, monkeypatch): - from openai_codex import LocalImageInput, TextInput - - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - inputs = [TextInput("describe this image"), LocalImageInput("image.png")] - observed = {} - - class FakeThread: - id = "thread-1" - - async def run(self, run_input, **_kwargs): - observed["input"] = run_input - return _TurnResult(final_response="done") - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - result = await wrapper.reply(inputs) - await wrapper.close() - - assert observed["input"] is inputs - assert result["last_message"] == "done" - - def test_codex_skills_add_all_without_deleting_existing_content(tmp_path): wrapper, _job = _wrapper(tmp_path) for name in ("reme_memory", "qwenpaw_memory"): @@ -560,603 +1189,6 @@ def test_codex_skills_do_not_modify_codex_home(tmp_path): assert marker.read_text(encoding="utf-8") == "keep" -def test_effective_mcp_config_snapshot_is_private_and_removed_on_close(tmp_path): - wrapper, _job = _wrapper(tmp_path) - config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access - snapshot = Path(config["args"][config["args"].index("--config") + 1]) - assert snapshot.exists() - assert stat.S_IMODE(snapshot.stat().st_mode) == 0o600 - assert json.loads(snapshot.read_text(encoding="utf-8"))["workspace_dir"] == str(tmp_path) - - async def close_started_wrapper(): - await wrapper.start() - await wrapper.close() - - asyncio.run(close_started_wrapper()) - assert not snapshot.exists() - - -def test_effective_mcp_config_snapshot_needs_no_fchmod(tmp_path, monkeypatch): - """A platform without ``os.fchmod`` (Windows) still gets a usable snapshot.""" - monkeypatch.delattr(os, "fchmod", raising=False) - - wrapper, _job = _wrapper(tmp_path) - config = wrapper._mcp_server_config({"job_tools": ["search"]}) - snapshot = Path(config["args"][config["args"].index("--config") + 1]) - - assert snapshot.exists() - assert json.loads(snapshot.read_text(encoding="utf-8"))["workspace_dir"] == str(tmp_path) - - -@pytest.mark.asyncio -async def test_effective_snapshot_exposes_parent_only_custom_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - app_config = ApplicationConfig( - workspace_dir=str(tmp_path), - enable_logo=False, - log_to_console=False, - log_to_file=False, - service={"backend": "mcp"}, - jobs={ - "only_custom": { - "backend": "base", - "description": "Parent-only inline job", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - "referenced_helper": { - "backend": "base", - "description": "A normal job that custom jobs may reference.", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - ) - job = _Job("only_custom") - context = SimpleNamespace(app_config=app_config, components={}, jobs={"only_custom": job}) - wrapper = CodexAgentWrapper(app_context=context) - server_config = wrapper._mcp_server_config({"job_tools": ["only_custom"]}) # pylint: disable=protected-access - transport = StdioTransport( - command=server_config["command"], - args=server_config["args"], - cwd=server_config["cwd"], - keep_alive=False, - ) - - await wrapper.start() - try: - async with Client(transport, timeout=None, init_timeout=STDIO_INIT_TIMEOUT) as client: - tools = await asyncio.wait_for(client.list_tools(), timeout=10) - snapshot = Path(server_config["args"][server_config["args"].index("--config") + 1]) - assert snapshot.exists() - assert set(json.loads(snapshot.read_text(encoding="utf-8"))["jobs"]) == { - "only_custom", - "referenced_helper", - } - finally: - await wrapper.close() - - assert [tool.name for tool in tools] == ["only_custom"] - assert not snapshot.exists() - - -@pytest.mark.asyncio -async def test_open_thread_defaults_to_full_access(tmp_path): - from openai_codex import ApprovalMode, Sandbox - - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeCodex: - async def thread_start(self, **kwargs): - observed.update(kwargs) - return SimpleNamespace(id="thread-1") - - await wrapper._open_thread(FakeCodex(), {}) # pylint: disable=protected-access - - assert observed["approval_mode"] == ApprovalMode.auto_review - assert observed["sandbox"] == Sandbox.full_access - - -@pytest.mark.asyncio -async def test_compact_session_uses_native_thread_operation(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeThread: - id = "thread-1" - - async def compact(self): - observed["compacted"] = True - - async def start(): - observed["started"] = True - - async def resume(session_id): - observed["resume"] = session_id - return FakeThread() - - async def get_codex(): - return SimpleNamespace(thread_resume=resume) - - monkeypatch.setattr(wrapper, "start", start) - monkeypatch.setattr(wrapper, "_get_codex", get_codex) - - await wrapper.compact_session("thread-1") - - assert observed["started"] is True - assert observed["resume"] == "thread-1" - assert observed["compacted"] is True - - -@pytest.mark.asyncio -async def test_open_thread_forwards_latest_sdk_options(tmp_path): - from openai_codex.types import Personality, ThreadSource, ThreadStartSource - - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeCodex: - async def thread_start(self, **kwargs): - observed["start"] = kwargs - return SimpleNamespace(id="thread-1") - - async def thread_resume(self, _thread_id, **kwargs): - observed["resume"] = kwargs - return SimpleNamespace(id="thread-1") - - async def thread_fork(self, _thread_id, **kwargs): - observed["fork"] = kwargs - return SimpleNamespace(id="thread-2") - - codex = FakeCodex() - await wrapper._open_thread( # pylint: disable=protected-access - codex, - { - "personality": "friendly", - "service_name": "reme", - "session_start_source": "startup", - "thread_source": "user", - }, - ) - await wrapper._open_thread( # pylint: disable=protected-access - codex, - {"session_id": "thread-1", "personality": "pragmatic"}, - ) - await wrapper._open_thread( # pylint: disable=protected-access - codex, - {"session_id": "thread-1", "fork_session": True, "thread_source": "subagent"}, - ) - - assert observed["start"]["personality"] == Personality.friendly - assert observed["start"]["service_name"] == "reme" - assert observed["start"]["session_start_source"] == ThreadStartSource.startup - assert observed["start"]["thread_source"] == ThreadSource.user - assert observed["resume"]["personality"] == Personality.pragmatic - assert observed["fork"]["thread_source"] == ThreadSource.subagent - - -@pytest.mark.asyncio -async def test_resume_reuses_tool_context_and_rejects_context_change(tmp_path): - wrapper, _job = _wrapper(tmp_path) - - class FakeCodex: - async def thread_start(self, **_kwargs): - return SimpleNamespace(id="thread-1") - - async def thread_resume(self, _thread_id, **_kwargs): - return SimpleNamespace(id="thread-1") - - await wrapper._open_thread(FakeCodex(), {"tool_context_id": "ctx-a"}) # pylint: disable=protected-access - await wrapper._open_thread( # pylint: disable=protected-access - FakeCodex(), - {"resume": "thread-1", "tool_context_id": "ctx-a"}, - ) - with pytest.raises(ValueError, match="cannot change"): - await wrapper._open_thread( # pylint: disable=protected-access - FakeCodex(), - {"resume": "thread-1", "tool_context_id": "ctx-b"}, - ) - - -class _StructuredModel(BaseModel): - ok: bool - - -@pytest.mark.parametrize("schema", [_StructuredModel(ok=True), str]) -def test_output_schema_rejects_instances_and_arbitrary_classes(tmp_path, schema): - wrapper, _job = _wrapper(tmp_path) - with pytest.raises(TypeError, match="JSON schema dict or BaseModel class"): - wrapper._merged_kwargs({"output_schema": schema}) # pylint: disable=protected-access - - -def test_output_schema_normalizes_model_class_and_preserves_dict(tmp_path): - wrapper, _job = _wrapper(tmp_path) - schema = {"type": "object", "properties": {"ok": {"type": "boolean"}}} - - assert ( - wrapper._merged_kwargs({"output_schema": _StructuredModel})["output_schema"] # pylint: disable=protected-access - == _StructuredModel.model_json_schema() - ) - assert ( - wrapper._merged_kwargs({"output_schema": schema})["output_schema"] is schema - ) # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - clients = [] - observed_schemas = [] - close_count = 0 - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **kwargs): - observed_schemas.append(kwargs["output_schema"]) - return _TurnResult(final_response=json.dumps({"ok": True})) - - class FakeCodex: - def __init__(self, _config): - clients.append(self) - - async def __aenter__(self): - return self - - async def close(self): - nonlocal close_count - close_count += 1 - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - result = await wrapper.reply("first", output_schema=_StructuredModel) - await wrapper.close() - - assert result["structured_output"] == {"ok": True} - assert observed_schemas == [_StructuredModel.model_json_schema()] - assert len(clients) == 1 - assert close_count == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("schema", [{}, _StructuredModel]) -async def test_reply_stream_rejects_output_schema(tmp_path, schema): - wrapper, _job = _wrapper(tmp_path) - - with pytest.raises(NotImplementedError, match="Structured output is not supported"): - await anext(wrapper.reply_stream("answer", output_schema=schema)) - - -@pytest.mark.asyncio -async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - stream_closed = False - interrupt_count = 0 - - class FakeTurn: - id = "turn-1" - - async def stream(self): - nonlocal stream_closed - try: - yield SimpleNamespace( - method="turn/started", - payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)), - ) - await asyncio.Event().wait() - finally: - stream_closed = True - - async def interrupt(self): - nonlocal interrupt_count - interrupt_count += 1 - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - stream = wrapper.reply_stream("answer") - first = await anext(stream) - assert first.chunk_type == ChunkEnum.REPLY_START - await stream.aclose() - await wrapper.close() - - assert interrupt_count == 1 - assert stream_closed - - -@pytest.mark.asyncio -async def test_reply_stream_does_not_record_token_usage(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - recorded_usages = [] - usage = TokenUsageBreakdown( - cachedInputTokens=0, - inputTokens=3, - outputTokens=5, - reasoningOutputTokens=0, - totalTokens=8, - ) - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="thread/tokenUsage/updated", - payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)), - ) - yield SimpleNamespace( - method="turn/completed", - payload=SimpleNamespace( - turn=SimpleNamespace( - id=self.id, - status=SimpleNamespace(value="completed"), - duration_ms=1, - error=None, - ), - ), - ) - - async def interrupt(self): - raise AssertionError("Completed turns must not be interrupted") - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append) - - chunks = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - assert any(chunk.chunk_type == ChunkEnum.USAGE for chunk in chunks) - assert not recorded_usages - - -@pytest.mark.asyncio -async def test_close_waits_for_active_turn(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - turn_started = asyncio.Event() - release_turn = asyncio.Event() - client_closed = asyncio.Event() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - turn_started.set() - await release_turn.wait() - return _TurnResult(final_response="done") - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - client_closed.set() - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - reply_task = asyncio.create_task(wrapper.reply("answer")) - await turn_started.wait() - close_task = asyncio.create_task(wrapper.close()) - await asyncio.sleep(0) - - assert not client_closed.is_set() - assert not close_task.done() - - release_turn.set() - result = await reply_task - await close_task - - assert result["last_message"] == "done" - assert client_closed.is_set() - - -@pytest.mark.asyncio -async def test_component_start_keeps_optional_client_lazy(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="") - - def fail_if_constructed(_config): - raise AssertionError("Codex client should be lazy") - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", fail_if_constructed) - - await wrapper.start() - await wrapper.close() - - assert wrapper._codex is None - - -@pytest.mark.asyncio -async def test_client_config_is_fixed_for_component_lifetime(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="one") - clients = [] - - class FakeCodex: - def __init__(self, _config): - clients.append(self) - - async def login_api_key(self, _api_key): - return None - - async def close(self): - return None - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - assert await wrapper._get_codex() is clients[0] # pylint: disable=protected-access - with pytest.raises(TypeError, match="configured on the wrapper: api_key"): - await wrapper.reply("answer", api_key="two") - await wrapper.close() - - assert len(clients) == 1 - - -@pytest.mark.parametrize( - ("name", "value"), - [ - ("auth_mode", "oauth"), - ("base_url", "https://example.test/v1"), - ("codex_bin", "/tmp/codex"), - ("codex_home", "/tmp/codex-home"), - ("config_overrides", ['model="test"']), - ("cwd", "/tmp"), - ("experimental_api", False), - ("launch_args_override", ["codex", "app-server"]), - ], -) -@pytest.mark.asyncio -async def test_reply_rejects_call_time_client_options(tmp_path, name, value): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - - with pytest.raises(TypeError, match=f"configured on the wrapper: {name}"): - await wrapper.reply("answer", **{name: value}) - - assert wrapper._codex is None - - -def test_constructor_rejects_launch_args_override(tmp_path): - with pytest.raises(TypeError, match="configure codex_bin instead"): - _wrapper(tmp_path, launch_args_override=["codex", "app-server"]) - - -def test_oauth_mode_ignores_api_credentials_and_forces_chatgpt(tmp_path, monkeypatch): - wrapper, _job = _wrapper( - tmp_path, - auth_mode="oauth", - api_key="explicit-key", - base_url="https://explicit.example.test/v1", - ) - wrapper.app_context.app_config.environment = {"TOOL_ENV": "preserved"} - monkeypatch.setenv("CODEX_API_KEY", "ambient-key") - monkeypatch.setenv("CODEX_BASE_URL", "https://ambient.example.test/v1") - - auth = wrapper._resolve_auth_config( # pylint: disable=protected-access - "oauth", - "explicit-key", - "https://explicit.example.test/v1", - ) - config = wrapper._build_client_config(auth) # pylint: disable=protected-access - - assert auth.mode == "oauth" - assert auth.api_key == "" - assert auth.base_url == "" - assert config.env["TOOL_ENV"] == "preserved" - assert "CODEX_HOME" not in wrapper.app_context.app_config.environment - assert 'forced_login_method="chatgpt"' in config.config_overrides - assert not any(value.startswith("openai_base_url=") for value in config.config_overrides) - - -@pytest.mark.asyncio -async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch): - wrapper, _job = _wrapper( - tmp_path, - auth_mode="api_key", - api_key="explicit-key", - base_url="https://proxy.example.test/v1", - ) - observed = {} - - class FakeCodex: - def __init__(self, config): - observed["config"] = config - - async def login_api_key(self, api_key): - observed["api_key"] = api_key - - async def close(self): - observed["closed"] = True - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - await wrapper._get_codex() # pylint: disable=protected-access - await wrapper.close() - - config = observed["config"] - assert observed["api_key"] == "explicit-key" - assert 'openai_base_url="https://proxy.example.test/v1"' in config.config_overrides - assert 'forced_login_method="api"' in config.config_overrides - assert observed["closed"] is True - - -def test_auth_selection_uses_explicit_wrapper_options(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - wrapper.app_context.app_config.components[ComponentEnum.AS_LLM]["default"] = SimpleNamespace( - credential={"api_key": "default-key", "base_url": "https://default.example.test/v1"}, - ) - for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY"): - monkeypatch.setenv(name, "ambient-key") - for name in ("CODEX_BASE_URL", "OPENAI_BASE_URL", "LLM_BASE_URL"): - monkeypatch.setenv(name, "https://ambient.example.test/v1") - - auth = wrapper._resolve_auth_config( # pylint: disable=protected-access - "api_key", - "configured-key", - ) - - assert auth.mode == "api_key" - assert auth.api_key == "configured-key" - assert auth.base_url == "" - - with pytest.raises(ValueError, match="requires a non-empty API key"): - wrapper._resolve_auth_config( # pylint: disable=protected-access - "api_key", - ) - - @pytest.mark.parametrize("review_status", ["approved", "denied"]) def test_event_to_chunks_maps_approval_started_and_completed(review_status): action = {"type": "futureApprovalAction", "value": "preserved"} @@ -1192,464 +1224,3 @@ def test_event_to_chunks_maps_approval_started_and_completed(review_status): assert completed_chunk.metadata["status"] == "completed" assert completed_chunk.metadata["review"]["status"] == review_status assert completed_chunk.metadata["decision_source"] == {"type": "guardian"} - - -def test_codex_home_expands_user_directory(tmp_path): - wrapper, _job = _wrapper(tmp_path, codex_home="~/.codex") - assert wrapper.session_path == Path.home() / ".codex" - - -def test_named_default_mcp_config_remains_supported(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="default") - assert wrapper._mcp_config_source({}) == "default" # pylint: disable=protected-access - - -def test_default_config_provides_codex_oauth_wrapper(monkeypatch): - monkeypatch.delenv("CODEX_HOME", raising=False) - config = resolve_app_config(log_config=False) - oauth = config["components"]["agent_wrapper"]["codex_oauth"] - codex = config["components"]["agent_wrapper"]["codex"] - assert oauth["backend"] == "codex" - assert oauth["auth_mode"] == "oauth" - assert oauth["codex_home"] == "~/.codex" - assert oauth["sandbox"] == "full-access" - assert "api_key" not in oauth - assert codex["auth_mode"] == "api_key" - assert codex["sandbox"] == "full-access" - - -class _RecordingClient: - def __init__(self, fail: bool = False): - self.requests = [] - self.fail = fail - self.loaded_threads = ["thread-1"] - - async def request(self, method, params, *, response_model): - self.requests.append({"method": method, "params": params, "response_model": response_model}) - if self.fail: - raise RuntimeError("unsubscribe failed") - if method == "thread/loaded/list": - return SimpleNamespace(data=self.loaded_threads, next_cursor=None) - return SimpleNamespace(status="unsubscribed") - - -def _assert_unsubscribed(client, thread_id="thread-1"): - from openai_codex.generated.v2_all import ThreadUnsubscribeResponse - - assert client.requests == [ - { - "method": "thread/unsubscribe", - "params": {"threadId": thread_id}, - "response_model": ThreadUnsubscribeResponse, - }, - ] - - -def _fake_codex(client, thread): - class FakeCodex: - def __init__(self, _config): - self._client = client - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return thread - - async def thread_resume(self, _thread_id, **_kwargs): - return thread - - return FakeCodex - - -@pytest.mark.asyncio -async def test_reply_releases_thread_after_turn(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - await wrapper.close() - - assert result["last_message"] == "done" - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_releases_thread_when_turn_fails(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - raise ValueError("turn failed") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(ValueError, match="turn failed"): - await wrapper.reply("answer") - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_ignores_unsubscribe_failure(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient(fail=True) - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - - assert result["last_message"] == "done" - assert len(client.requests) == 1 - # The thread is still loaded, so its tool-context guard must survive. - assert "thread-1" in wrapper._thread_tool_contexts # pylint: disable=protected-access - - await wrapper.close() - - -@pytest.mark.asyncio -async def test_failed_release_keeps_tool_context_guard_on_resume(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient(fail=True) - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer", tool_context_id="ctx-a") - assert result["session_id"] == "thread-1" - assert len(client.requests) == 1 - - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - await wrapper.close() - - -@pytest.mark.asyncio -async def test_release_threads_disabled_keeps_threads_loaded(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth", release_threads=False) - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - - assert result["last_message"] == "done" - assert not client.requests - assert "thread-1" in wrapper._thread_tool_contexts # pylint: disable=protected-access - - await wrapper.close() - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_on_completion(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="turn/completed", - payload=SimpleNamespace( - turn=SimpleNamespace( - id="turn-1", - status=SimpleNamespace(value="completed"), - duration_ms=1, - error=None, - ), - ), - ) - - async def interrupt(self): - raise AssertionError("Completed turns must not be interrupted") - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - chunks = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - assert chunks[-1].chunk_type == ChunkEnum.REPLY_END - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_on_early_close(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - interrupt_count = 0 - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="turn/started", - payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)), - ) - await asyncio.Event().wait() - - async def interrupt(self): - nonlocal interrupt_count - interrupt_count += 1 - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - stream = wrapper.reply_stream("answer") - first = await anext(stream) - assert first.chunk_type == ChunkEnum.REPLY_START - await stream.aclose() - await wrapper.close() - - assert interrupt_count == 1 - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_compact_session_releases_thread(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def compact(self): - return None - - async def start(): - return None - - async def resume(_thread_id): - return FakeThread() - - async def get_codex(): - return SimpleNamespace(_client=client, thread_resume=resume) - - monkeypatch.setattr(wrapper, "start", start) - monkeypatch.setattr(wrapper, "_get_codex", get_codex) - - await wrapper.compact_session("thread-1") - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_when_turn_start_fails(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - raise RuntimeError("turn/start failed") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(RuntimeError, match="turn/start failed"): - _ = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_when_turn_start_is_cancelled(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - raise asyncio.CancelledError("cancelled during turn/start") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(asyncio.CancelledError): - _ = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("original", "changed"), - [ - ({"tool_context_id": "ctx-a"}, {"tool_context_id": "ctx-b"}), - ({"injected_job_kwargs": {"tenant": "a"}}, {"injected_job_kwargs": {"tenant": "b"}}), - ], -) -async def test_cancelled_reply_keeps_loaded_thread_constraints(tmp_path, monkeypatch, original, changed): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - entered = asyncio.Event() - - class FakeThread: - id = "thread-1" - calls = 0 - - async def run(self, _inputs, **_kwargs): - self.calls += 1 - if self.calls > 1: - return _TurnResult(final_response="old MCP constraints") - entered.set() - await asyncio.Event().wait() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - task = asyncio.create_task(wrapper.reply("answer", **original)) - try: - await asyncio.wait_for(entered.wait(), timeout=5) - finally: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - try: - _assert_unsubscribed(client) - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", **changed) - finally: - await wrapper.close() - - -@pytest.mark.asyncio -async def test_unloaded_thread_can_resume_with_new_constraints(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - try: - await wrapper.reply("answer", tool_context_id="ctx-a") - assert "thread-1" in wrapper._thread_tool_contexts - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-a") - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - # Unsubscribe schedules delayed unload; only a later loaded-list query - # establishes that a resume will create a fresh runtime and MCP config. - client.loaded_threads.clear() - result = await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - assert result["last_message"] == "done" - assert json.loads(wrapper._thread_tool_contexts["thread-1"])["tool_context_id"] == "ctx-b" - finally: - await wrapper.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("fail_second_page", [False, True]) -async def test_prune_thread_constraints_reads_all_pages_before_removing_entries(tmp_path, fail_second_page): - from openai_codex.generated.v2_all import ThreadLoadedListResponse - - wrapper, _job = _wrapper(tmp_path) - wrapper._thread_tool_contexts = {"first": "a", "last": "b", "unloaded": "c"} - cursors = [] - - class Client: - async def request(self, method, params, *, response_model): - assert method == "thread/loaded/list" - assert response_model is ThreadLoadedListResponse - cursors.append(params["cursor"]) - if params["cursor"] is None: - return ThreadLoadedListResponse(data=["first"], nextCursor="page-2") - if fail_second_page: - raise RuntimeError("loaded-list unavailable") - return ThreadLoadedListResponse(data=["last"]) - - await wrapper._prune_thread_tool_contexts(SimpleNamespace(_client=Client())) - - assert cursors == [None, "page-2"] - expected = {"first": "a", "last": "b"} - if fail_second_page: - expected["unloaded"] = "c" - assert wrapper._thread_tool_contexts == expected diff --git a/tests/unit/test_common_steps.py b/tests/unit/test_common_steps.py index 02c8a870..9738a36f 100644 --- a/tests/unit/test_common_steps.py +++ b/tests/unit/test_common_steps.py @@ -3,6 +3,7 @@ # pylint: disable=protected-access from contextlib import chdir as _temp_chdir +import asyncio import tempfile import warnings @@ -11,11 +12,13 @@ import pytest from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext from reme.components.file_store import LocalFileStore -from reme.schema import FileFrontMatter, FileLink, FileNode, TraverseGraph +from reme.enumeration import ChunkEnum +from reme.schema import FileFrontMatter, FileLink, FileNode, StreamChunk, TraverseGraph from reme.steps.common.add import AddStep from reme.steps.common.health_check import _file_graph_status from reme.steps.common.llm_demo import LLMDemoStep from reme.steps.common.python_execute import PythonExecuteStep +from reme.steps.common.stream_llm_demo import StreamLLMDemoStep from reme.steps.index import traverse as traverse_mod from reme.utils.local_backend import LocalProcessBackend @@ -147,11 +150,17 @@ class _FakeAgentWrapper(BaseAgentWrapper): def __init__(self): super().__init__() self.last_kwargs = None + self.stream_chunks = [] async def reply(self, inputs, **kwargs) -> dict: self.last_kwargs = kwargs return {"result": "ok"} + async def reply_stream(self, inputs, **kwargs): + self.last_kwargs = kwargs + for chunk in self.stream_chunks: + yield chunk + @pytest.mark.asyncio async def test_llm_demo_always_registers_add_tool(): @@ -165,6 +174,27 @@ async def test_llm_demo_always_registers_add_tool(): assert "job_tools" not in resp.metadata +@pytest.mark.asyncio +async def test_stream_llm_demo_preserves_backend_chunk_metadata(): + """Forward retry flags, tool identities, usage, and completion status without rebuilding chunks.""" + wrapper = _FakeAgentWrapper() + wrapper.stream_chunks = [ + StreamChunk(chunk_type=ChunkEnum.ERROR, chunk="retrying", metadata={"will_retry": True}), + StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, tool_call_id="call-1", tool_call_name="add", chunk={"a": 1}), + StreamChunk(chunk_type=ChunkEnum.CONTENT, session_id="session-1", block_id="text-1", chunk="hello"), + StreamChunk(chunk_type=ChunkEnum.USAGE, input_tokens=3, output_tokens=1), + StreamChunk(chunk_type=ChunkEnum.REPLY_END, metadata={"status": "completed"}), + ] + queue = asyncio.Queue() + response = await StreamLLMDemoStep()(query="hello", stream_queue=queue, agent_wrapper=wrapper) + forwarded = [queue.get_nowait() for _ in wrapper.stream_chunks] + assert all(actual is expected for actual, expected in zip(forwarded, wrapper.stream_chunks)) + assert queue.empty() + assert response.success + assert response.answer == "hello" + assert response.metadata["session_id"] == "session-1" + + def test_file_graph_health_reports_neo4j_cached_counts(): """Neo4j file graph health should not be reported as an empty local graph.""" diff --git a/tests/unit/test_http_client.py b/tests/unit/test_http_client.py index bd506cb0..7b477f46 100644 --- a/tests/unit/test_http_client.py +++ b/tests/unit/test_http_client.py @@ -2,8 +2,14 @@ # pylint: disable=protected-access +import json + +import httpx +import pytest + from reme.components.client.http_client import HttpClient from reme.constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT +from reme.enumeration import ChunkEnum def test_default_client_uses_loopback_address(): @@ -34,3 +40,30 @@ def test_format_for_display_shows_metadata_when_requested(): text = '{"answer":"0.4.0.7","success":true,"metadata":{"version":"0.4.0.7"}}' assert client._format_for_display(text) == '0.4.0.7\n✅ {"version": "0.4.0.7"}' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("will_retry", [True, False, None, "false"]) +async def test_stream_retry_notifications_continue_but_terminal_errors_raise(will_retry): + """A temporary Codex reconnect must not prevent CLI callers from receiving the answer.""" + events = [ + {"chunk_type": "error", "chunk": "upstream error", "metadata": {"will_retry": will_retry}}, + {"chunk_type": "content", "chunk": "recovered answer"}, + {"chunk_type": "done", "done": True}, + ] + body = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + + def handle(_request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + client = HttpClient() + client.client = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + async with client: + if will_retry is True: + chunks = [chunk async for chunk in client.stream_chunks("stream_llm_demo")] + assert [chunk.chunk_type for chunk in chunks] == [ChunkEnum.ERROR, ChunkEnum.CONTENT] + assert chunks[0].metadata == {"will_retry": True} + assert chunks[-1].chunk == "recovered answer" + else: + with pytest.raises(RuntimeError, match="upstream error"): + _ = [chunk async for chunk in client.stream_chunks("stream_llm_demo")]