From d323356e6f11c0faf4bf30fae7a2edc30e18fbb9 Mon Sep 17 00:00:00 2001 From: cangen185-sudo Date: Wed, 4 Nov 2026 15:25:11 +0800 Subject: [PATCH 1/5] fix(file-io): refuse edits after incomplete bounded reads (#663) * fix(file-io): refuse edits after incomplete bounded reads * style(file-io): keep complete-read signature formatter-stable * docs(file-io): clarify bounded-read safety guarantee --------- Co-authored-by: cangen185-sudo <234582933+cangen185-sudo@users.noreply.github.com> --- reme/steps/file_io/_file_io.py | 17 ++++- reme/steps/file_io/edit.py | 3 +- tests/unit/test_edit_read_limits.py | 102 ++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_edit_read_limits.py diff --git a/reme/steps/file_io/_file_io.py b/reme/steps/file_io/_file_io.py index 22e945b0..33ff1876 100644 --- a/reme/steps/file_io/_file_io.py +++ b/reme/steps/file_io/_file_io.py @@ -106,15 +106,28 @@ def _decode_known_file(data: bytes, file_extension: str) -> tuple[str, str]: # --------------------------------------------------------------------------- -async def read_file_safe(file_path, max_bytes: int = MAX_FILE_READ_BYTES) -> tuple[str, str]: +async def read_file_safe( + file_path, + max_bytes: int = MAX_FILE_READ_BYTES, + *, + require_complete: bool = False, +) -> tuple[str, str]: """Read file in byte mode and decode using extension-aware strategy. - Returns ``(text, encoding)``. + Returns ``(text, encoding)``. With ``require_complete``, reject files above + the byte limit or reads that leave unread bytes at the EOF check. This only + prevents returning a truncated bounded read for write-back; it does not + support oversized-file editing, provide a stable snapshot against concurrent + changes, or make subsequent writes atomic. """ stat = await aiofiles.os.stat(str(file_path)) + if require_complete and stat.st_size > max_bytes: + raise ValueError(f"file size {stat.st_size} exceeds read limit of {max_bytes} bytes") read_size = min(stat.st_size, max_bytes) async with aiofiles.open(str(file_path), "rb") as f: data = await f.read(read_size) + if require_complete and await f.read(1): + raise ValueError("file changed during read; refusing to return incomplete contents") return _decode_known_file(data, Path(file_path).suffix) diff --git a/reme/steps/file_io/edit.py b/reme/steps/file_io/edit.py index 31a0fd76..62ec1a64 100644 --- a/reme/steps/file_io/edit.py +++ b/reme/steps/file_io/edit.py @@ -9,6 +9,7 @@ from ._base import _FileIOStep from ._file_io import get_path_lock, read_file_safe, write_file_safe from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ...components import R +from ...constants import MAX_FILE_READ_BYTES def _replace_normalized_newlines(body: str, old: str, new: str) -> tuple[str, int]: @@ -100,7 +101,7 @@ class EditStep(_FileIOStep): return None try: - raw_text, encoding = await read_file_safe(target) + raw_text, encoding = await read_file_safe(target, max_bytes=MAX_FILE_READ_BYTES, require_complete=True) except Exception as e: # pylint: disable=broad-except self._fail(f"read failed: {e}", path=str(target)) return None diff --git a/tests/unit/test_edit_read_limits.py b/tests/unit/test_edit_read_limits.py new file mode 100644 index 00000000..6543b3d4 --- /dev/null +++ b/tests/unit/test_edit_read_limits.py @@ -0,0 +1,102 @@ +"""Edits must not write back a prefix produced by a bounded file read.""" + +import asyncio +from functools import partial +from types import SimpleNamespace + +import pytest + +from reme.components import ApplicationContext +from reme.steps.file_io import _file_io, edit as edit_module + + +def _small_edit_limit(monkeypatch): + """Exercise the real bounded reader with small files, without allocating 200 MiB.""" + monkeypatch.setattr(edit_module, "MAX_FILE_READ_BYTES", 64, raising=False) + monkeypatch.setattr(edit_module, "read_file_safe", partial(_file_io.read_file_safe, max_bytes=64)) + + +@pytest.mark.parametrize("suffix", [".md", ".txt"]) +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16"]) +def test_edit_rejects_truncated_read_without_writing(tmp_path, monkeypatch, suffix, encoding): + """An occurrence in the prefix must not cause the unread tail to be discarded.""" + _small_edit_limit(monkeypatch) + path = tmp_path / ("note" + suffix) + source = ("target" + "x" * 80 + "TAIL").encode(encoding) + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert path.read_bytes() == source + assert step.context.response.success is False + assert "limit" in step.context.response.answer + + asyncio.run(run()) + + +@pytest.mark.parametrize("suffix", [".md", ".txt"]) +@pytest.mark.parametrize("encoding,filler", [("utf-8", 58), ("utf-16", 25)]) +def test_edit_at_read_limit_preserves_encoding(tmp_path, monkeypatch, suffix, encoding, filler): + """The exact byte limit remains editable, with the existing write rules.""" + _small_edit_limit(monkeypatch) + text = "target" + "x" * filler + source = text.encode(encoding) + assert len(source) == 64 + path = tmp_path / ("note" + suffix) + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert step.context.response.success is True + expected = text.replace("target", "edited") + ("\n" if suffix == ".md" else "") + assert path.read_bytes() == expected.encode(encoding) + + asyncio.run(run()) + + +def test_edit_rejects_growth_after_stat(tmp_path, monkeypatch): + """A stale small stat cannot turn a larger actual read into a destructive edit.""" + _small_edit_limit(monkeypatch) + path = tmp_path / "note.txt" + source = b"target" + b"x" * 80 + b"TAIL" + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def stale_stat(_path): + return SimpleNamespace(st_size=8) + + monkeypatch.setattr(_file_io.aiofiles.os, "stat", stale_stat) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert step.context.response.success is False + assert "incomplete contents" in step.context.response.answer + assert path.read_bytes() == source + + asyncio.run(run()) + + +@pytest.mark.parametrize("size", [0, 63, 64, 65]) +def test_read_helper_complete_mode_and_default_truncation(tmp_path, size): + """Ordinary limited reads keep their behavior; complete reads reject only overflow.""" + path = tmp_path / "note.txt" + source = b"x" * size + path.write_bytes(source) + + async def run(): + text, _ = await _file_io.read_file_safe(path, max_bytes=64) + assert text == "x" * min(size, 64) + if size > 64: + with pytest.raises(ValueError, match="limit"): + await _file_io.read_file_safe(path, max_bytes=64, require_complete=True) + else: + text, _ = await _file_io.read_file_safe(path, max_bytes=64, require_complete=True) + assert text == "x" * size + assert path.read_bytes() == source + + asyncio.run(run()) From fe812c916febbeffe85bc64df44a95ed192309b2 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:51:14 +0800 Subject: [PATCH 2/5] feat: add DashScope web search job (#710) --- docs/en/configuration.md | 15 ++ docs/zh/configuration.md | 15 ++ example.env | 3 + pyproject.toml | 1 + reme/config/default.yaml | 20 +++ reme/steps/common/__init__.py | 2 + reme/steps/common/dashscope_search.py | 71 +++++++++ tests/unit/test_dashscope_search_step.py | 182 +++++++++++++++++++++++ 8 files changed, 309 insertions(+) create mode 100644 reme/steps/common/dashscope_search.py create mode 100644 tests/unit/test_dashscope_search_step.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index fedf7ac0..8b4a2e78 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -113,6 +113,21 @@ Built-in registrations include `openai`, `anthropic`, `dashscope`, `deepseek`, ` File operations, BM25 search, wikilink traversal, and `proactive_read` do not require an LLM. Evolution workflows such as `auto_memory`, `auto_resource`, `auto_dream`, and proactive refresh do. +## Web search + +With `reme-ai[core]` installed, `dashscope_search` calls `qwen3.8-flash` through the DashScope SDK, independently of LLM settings. +Set the key and start the service in one terminal: + +```bash +export DASHSCOPE_API_KEY=sk-xxx +reme start config=default +``` + +From another terminal, call `reme dashscope_search query="Latest Qwen models"`. +A missing key does not prevent service startup; calls to this Job return a failure. +The result's `search_info.search_results` preserves provider source fields; `content` is separate model-generated text. +Page bodies and snippets are not guaranteed. See `reme/config/default.yaml` for the model, search strategy, and timeout. + ## Decision Model Configuration `as_decision_model` is a provider-independent interface for probabilistic decisions. The current `agentscope` adapter diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 6cec4614..7053a74c 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -117,6 +117,21 @@ components: 基础文件操作、BM25 检索、wikilink 遍历和 `proactive_read` 不需要 LLM。`auto_memory`、`auto_resource`、 `auto_dream` 和 proactive refresh 等演化流程需要可用 LLM。 +## 联网搜索 + +安装 `reme-ai[core]` 后，默认 `dashscope_search` Job 使用 DashScope SDK 调用 `qwen3.8-flash`，独立于上述 LLM 配置。 +在一个终端设置 Key 并启动服务： + +```bash +export DASHSCOPE_API_KEY=sk-xxx +reme start config=default +``` + +在另一个终端调用 `reme dashscope_search query="Qwen 最新模型"`。 +未配置 Key 不影响服务启动，但此 Job 调用会返回失败。 +结果的 `search_info.search_results` 保留供应商返回的来源字段，`content` 是单独的模型生成正文； +网页正文和摘要不保证提供。模型、搜索策略和超时配置见 `reme/config/default.yaml`。 + ## 决策模型配置 `as_decision_model` 是不依赖供应商的概率决策模型接口；当前的 `agentscope` 适配器通过 AgentScope `2.0.11` 调用 TypeSafe API。 diff --git a/example.env b/example.env index b8e9e3d9..fa78dd5a 100644 --- a/example.env +++ b/example.env @@ -3,6 +3,9 @@ EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 LLM_API_KEY=sk-xxx LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +# Optional dashscope_search job (qwen3.8-flash web search). +# DASHSCOPE_API_KEY=sk-xxx + # Optional decision models: install reme-ai[core,decision]. # Setting these variables alone does not enable an as_decision_model component. # DECISION_MODEL_API_KEY=your-typesafe-api-key diff --git a/pyproject.toml b/pyproject.toml index e508f170..b00a962f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ web = [ ] core = [ "reme-ai[as]", + "dashscope>=1.26.0", "pillow>=13.0.0", "claude-agent-sdk>=0.2.174", "openai-codex>=0.163.0", diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 32233ce4..d2d481ba 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -502,6 +502,26 @@ jobs: expand_links: true max_links_per_direction: 10 + dashscope_search: + backend: base + description: >- + Search the web with qwen3.8-flash; requires DASHSCOPE_API_KEY. + Returns raw search_info sources and separate model-generated content; page bodies and snippets are not guaranteed. + parameters: + type: object + properties: + query: + type: string + description: "web search query" + required: + - query + steps: + - backend: dashscope_search_step + model: qwen3.8-flash + api_key: ${DASHSCOPE_API_KEY:-} + search_strategy: turbo + timeout: 60 + node_search: backend: base description: "Digest node recall — given a candidate abstraction's name+description, surface existing digest nodes similar enough to either dedup against or link to as related." diff --git a/reme/steps/common/__init__.py b/reme/steps/common/__init__.py index 18d2c7b0..86240fe4 100644 --- a/reme/steps/common/__init__.py +++ b/reme/steps/common/__init__.py @@ -3,6 +3,7 @@ from .add import AddStep from .app_config import AppConfigStep from .chat import ChatStep +from .dashscope_search import DashScopeSearchStep from .decision_demo import DecisionDemoStep from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep @@ -19,6 +20,7 @@ __all__ = [ "AddStep", "AppConfigStep", "ChatStep", + "DashScopeSearchStep", "DecisionDemoStep", "DemoEchoStep1", "DemoEchoStep2", diff --git a/reme/steps/common/dashscope_search.py b/reme/steps/common/dashscope_search.py new file mode 100644 index 00000000..230aac83 --- /dev/null +++ b/reme/steps/common/dashscope_search.py @@ -0,0 +1,71 @@ +"""Web search through DashScope, preserving provider search results.""" + +import asyncio +import os +from contextlib import aclosing +from http import HTTPStatus + +from ..base_step import BaseStep +from ...components import R + + +@R.register("dashscope_search_step") +class DashScopeSearchStep(BaseStep): + """Return raw search sources and separate model-generated content.""" + + async def execute(self): + assert self.context is not None + response = self.context.response + query = self.context.get("query") + query = "" if query is None else str(query).strip() + api_key = (self.kwargs.get("api_key") or os.getenv("DASHSCOPE_API_KEY", "")).strip() + if not query or not api_key: + response.success = False + response.answer = "query is required" if not query else "DASHSCOPE_API_KEY is required" + return response + + from aiohttp import ClientSession + from dashscope import AioMultiModalConversation + + model = self.kwargs.get("model", "qwen3.8-flash") + search_info, sources, content, usage = {}, {}, [], {} + request_id = "" + async with asyncio.timeout(self.kwargs.get("timeout", 60)), ClientSession(trust_env=True) as session: + stream = await AioMultiModalConversation.call( + api_key=api_key, + session=session, + model=model, + messages=[{"role": "user", "content": [{"text": query}]}], + enable_search=True, + enable_thinking=False, + stream=True, + incremental_output=True, + search_options={ + "forced_search": True, + "enable_source": True, + "search_strategy": self.kwargs.get("search_strategy", "turbo"), + }, + ) + async with aclosing(stream): + async for chunk in stream: + if chunk.status_code != HTTPStatus.OK: + raise RuntimeError(f"DashScope {chunk.code}: {chunk.message} (request_id={chunk.request_id})") + request_id = chunk.request_id or request_id + if any((chunk.usage or {}).values()): + usage = dict(chunk.usage) + output = chunk.output or {} + info = output.get("search_info") or {} + search_info.update({key: value for key, value in info.items() if key != "search_results"}) + for item in info.get("search_results") or []: + sources.setdefault(item["url"], item) + for choice in output.get("choices") or []: + text = (choice.get("message") or {}).get("content") or [] + if isinstance(text, str): + content.append(text) + else: + content.extend(part["text"] for part in text if "text" in part) + + search_info["search_results"] = list(sources.values()) + response.answer = {"search_info": search_info, "content": "".join(content)} + response.metadata.update(model=model, request_id=request_id, usage=usage) + return response diff --git a/tests/unit/test_dashscope_search_step.py b/tests/unit/test_dashscope_search_step.py new file mode 100644 index 00000000..04bc0c20 --- /dev/null +++ b/tests/unit/test_dashscope_search_step.py @@ -0,0 +1,182 @@ +"""DashScope web search contracts, with no provider requests.""" + +import asyncio +from unittest.mock import AsyncMock + +import pytest +from dashscope import AioMultiModalConversation +from dashscope.api_entities.dashscope_response import DashScopeAPIResponse, MultiModalConversationResponse + +from reme.application import Application +from reme.components.job import BaseJob +from reme.config.config_parser import _load_config +from reme.steps.common.dashscope_search import DashScopeSearchStep + + +@pytest.fixture(name="provider") +def mock_provider(monkeypatch): + """Supply SDK frames and record stream cleanup.""" + frames, closed = [], [] + + async def stream(): + try: + for frame in frames: + yield frame + finally: + closed.append(True) + + call = AsyncMock(side_effect=lambda **_kwargs: stream()) + monkeypatch.setattr(AioMultiModalConversation, "call", call) + monkeypatch.setenv("DASHSCOPE_API_KEY", "test-key") + return frames, closed, call + + +def chunk(output=None, **kwargs): + """Build an SDK response, including source-only frames.""" + return MultiModalConversationResponse.from_api_response( + DashScopeAPIResponse( + status_code=200, + request_id="request-1", + output=output or {}, + **kwargs, + ), + ) + + +@pytest.mark.asyncio +async def test_default_job_starts_without_key_and_rejects_call(monkeypatch, tmp_path): + """Default config loads and the registered job starts without credentials.""" + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + call = AsyncMock(side_effect=AssertionError("must not contact provider")) + monkeypatch.setattr(AioMultiModalConversation, "call", call) + config = _load_config("default.yaml")["jobs"]["dashscope_search"] + app = Application( + workspace_dir=str(tmp_path), + service={"backend": "cli"}, + jobs={"dashscope_search": config}, + enable_logo=False, + log_to_file=False, + ) + await app.start() + try: + response = await app.run_job("dashscope_search", query="latest Qwen models") + assert not response.success + assert response.answer == "DASHSCOPE_API_KEY is required" + assert config["parameters"]["required"] == ["query"] + call.assert_not_awaited() + finally: + await app.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("query", [None, "", " "]) +async def test_empty_query_does_not_call_provider(provider, query): + """Reject missing or blank queries before contacting DashScope.""" + response = await DashScopeSearchStep()(query=query) + assert not response.success + assert response.answer == "query is required" + provider[2].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stream_preserves_sources_and_separates_generated_content(provider): + """Source-only, repeated, empty, and later source frames retain original fields and order.""" + frames, closed, call = provider + first = {"index": 1, "title": "One", "url": "https://example.com/1", "unknown": {"raw": True}} + second = {"index": 2, "title": "Two", "url": "https://example.com/2", "snippet": "provider snippet"} + frames.extend( + [ + chunk({"search_info": {"search_results": [first], "provider": "web"}}), + chunk( + { + "search_info": {"search_results": [dict(first, index=9), second]}, + "choices": [{"message": {"role": "assistant", "content": [{"text": "Generated "}]}}], + }, + ), + chunk( + { + "search_info": {"search_results": []}, + "choices": [{"message": {"role": "assistant", "content": "content"}}], + }, + usage={"input_tokens": 12, "output_tokens": 2, "plugins": {"search": {"count": 1}}}, + ), + chunk(), + ], + ) + response = await DashScopeSearchStep(api_key="configured-key")(query=0) + + assert response.success + assert response.answer == { + "search_info": {"search_results": [first, second], "provider": "web"}, + "content": "Generated content", + } + assert response.metadata == { + "model": "qwen3.8-flash", + "request_id": "request-1", + "usage": {"input_tokens": 12, "output_tokens": 2, "characters": 0, "plugins": {"search": {"count": 1}}}, + } + assert closed == [True] + session = call.await_args.kwargs["session"] + assert session.closed + call.assert_awaited_once_with( + api_key="configured-key", + session=session, + model="qwen3.8-flash", + messages=[{"role": "user", "content": [{"text": "0"}]}], + enable_search=True, + enable_thinking=False, + stream=True, + incremental_output=True, + search_options={"forced_search": True, "enable_source": True, "search_strategy": "turbo"}, + ) + + +@pytest.mark.asyncio +async def test_empty_results_remain_empty(provider): + """An ordinary response with no search hits does not invent sources.""" + provider[0].append(chunk({"search_info": {"search_results": []}})) + response = await DashScopeSearchStep()(query="no hits") + assert response.success + assert response.answer == {"search_info": {"search_results": []}, "content": ""} + + +@pytest.mark.asyncio +async def test_provider_error_after_sources_fails_job_and_closes_stream(provider): + """A failed stream cannot be published as a successful partial search.""" + frames, closed, _ = provider + frames.extend( + [ + chunk({"search_info": {"search_results": [{"url": "https://example.com"}]}}), + MultiModalConversationResponse( + status_code=429, + code="Throttling", + message="Try later", + request_id="failed-request", + ), + ], + ) + job = BaseJob() + job.step_specs = [(DashScopeSearchStep, {})] + response = await job(query="query") + assert not response.success + assert "Throttling: Try later" in response.answer + assert "failed-request" in response.answer + assert closed == [True] + + +@pytest.mark.asyncio +async def test_timeout_closes_stream(monkeypatch): + """Bound the entire stream, including waiting for another frame.""" + closed = [] + + async def stream(): + try: + yield chunk() + await asyncio.Event().wait() + finally: + closed.append(True) + + monkeypatch.setattr(AioMultiModalConversation, "call", AsyncMock(return_value=stream())) + with pytest.raises(TimeoutError): + await DashScopeSearchStep(api_key="test-key", timeout=0.01)(query="query") + assert closed == [True] 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 3/5] 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")] From bc91a0c7d217c8e6c02b35cc1db17d2824211ba7 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:51:30 +0800 Subject: [PATCH 4/5] feat(studio): add Codex session history and proactive messages (#711) * feat(studio): enable native Codex chat in proactive workspace * feat(studio): add Codex session history and proactive messages * feat(resource): add dated resource retention * feat(jobs): track in-memory execution and step status Unify cumulative invocation counts with optional execution records, expose job_status, and track nested step failures. Remove redundant proactive timing and share manual/cron step definitions. * fix(codex): anchor turn usage before context compaction * fix: recover interrupted sessions and job failure status --- docs/en/auto_resource.md | 36 ++ docs/en/configuration.md | 6 + docs/en/memory_as_file.md | 2 + docs/en/operations.md | 31 + docs/en/proactive.md | 5 +- docs/zh/auto_resource.md | 32 + docs/zh/configuration.md | 6 + docs/zh/memory_as_file.md | 4 +- docs/zh/operations.md | 31 + docs/zh/proactive.md | 5 +- .../agent_wrapper/base_agent_wrapper.py | 21 +- .../agent_wrapper/codex_agent_wrapper.py | 114 +++- .../agent_wrapper/codex_session_history.py | 188 ++++++ reme/components/application_context.py | 114 +++- reme/components/codex_app_server.py | 175 ++++++ .../file_catalog/base_file_catalog.py | 8 + .../file_catalog/local_file_catalog.py | 70 +++ reme/components/job/background_job.py | 33 +- reme/components/job/base_job.py | 70 ++- reme/components/job/cron_job.py | 12 +- reme/components/job/stream_job.py | 20 +- reme/components/runtime_context.py | 11 +- reme/config/default.yaml | 94 ++- reme/config/proactive.yaml | 227 ++++--- reme/schema/__init__.py | 7 + reme/schema/agent_session.py | 36 ++ reme/schema/application_config.py | 5 + reme/schema/job_execution.py | 21 + reme/schema/proactive.py | 1 - reme/steps/base_step.py | 40 +- reme/steps/common/__init__.py | 6 + reme/steps/common/chat.py | 27 +- reme/steps/common/job_status.py | 19 + reme/steps/common/session.py | 53 ++ reme/steps/evolve/proactive/agenda.py | 3 - reme/steps/evolve/proactive/extract.py | 11 +- reme/steps/evolve/proactive/plan.py | 3 - reme/steps/evolve/proactive/topics.py | 3 - reme/steps/index/__init__.py | 2 + reme/steps/index/cleanup_catalog.py | 22 + reme/utils/evaluation_interface.py | 4 +- reme_studio/README.md | 26 + reme_studio/README_ZH.md | 23 + reme_studio/app/api.ts | 53 +- reme_studio/app/chat-stream.ts | 127 +++- reme_studio/app/globals.css | 19 + reme_studio/app/i18n.ts | 9 + reme_studio/app/session-sync.ts | 123 ++++ reme_studio/app/store.ts | 40 +- reme_studio/app/types.ts | 25 + reme_studio/app/workspace-chat.tsx | 93 ++- reme_studio/app/workspace-persistence.ts | 2 +- reme_studio/app/workspace.tsx | 6 +- reme_studio/demo/api.ts | 4 + reme_studio/tests/chat-stream.test.mjs | 477 +++++++++++++++ reme_studio/tests/session-sync.test.mjs | 211 +++++++ .../tests/workspace-persistence.test.mjs | 14 +- tests/integration/test_codex_token_usage.py | 372 ++++++++++++ tests/unit/test_agent_session_api.py | 392 ++++++++++++ tests/unit/test_benchmark_docker.py | 7 +- tests/unit/test_codex_agent_wrapper.py | 319 +++++++++- tests/unit/test_config_parser.py | 17 + tests/unit/test_evaluation_interface.py | 23 +- tests/unit/test_job.py | 30 +- tests/unit/test_job_status.py | 559 ++++++++++++++++++ tests/unit/test_proactive_refresh.py | 1 + tests/unit/test_resource_cleanup.py | 280 +++++++++ tests/unit/test_workspace_web_steps.py | 45 +- 68 files changed, 4527 insertions(+), 348 deletions(-) create mode 100644 reme/components/agent_wrapper/codex_session_history.py create mode 100644 reme/schema/agent_session.py create mode 100644 reme/schema/job_execution.py create mode 100644 reme/steps/common/job_status.py create mode 100644 reme/steps/common/session.py create mode 100644 reme/steps/index/cleanup_catalog.py create mode 100644 reme_studio/app/session-sync.ts create mode 100644 reme_studio/tests/session-sync.test.mjs create mode 100644 tests/integration/test_codex_token_usage.py create mode 100644 tests/unit/test_agent_session_api.py create mode 100644 tests/unit/test_job_status.py create mode 100644 tests/unit/test_resource_cleanup.py diff --git a/docs/en/auto_resource.md b/docs/en/auto_resource.md index a3c0e5e0..06fe442a 100644 --- a/docs/en/auto_resource.md +++ b/docs/en/auto_resource.md @@ -59,6 +59,42 @@ whose class-level matcher accepts it. `AutoImageResourceStep` handles image suff final fallback. A new modality can therefore add a registered processor, its prompt, and one `dispatch_steps` entry without changing the router. +## Resource Retention and TTL + +The default configuration does not automatically delete source resources. `config=proactive` explicitly enables this +retention convention: + +- Put temporary resources under `resource/YYYY-MM-DD/`, with any nested paths or file extensions. +- Keep permanent material elsewhere in the workspace, for example under `archive/`. +- Loose root files, non-date directories, invalid dates, and date-directory symlinks are preserved for compatibility. + +`resource_watch_loop` recursively monitors every extension and only updates `file_catalog: resource`; it does not call +an LLM or create daily cards. Every hour, `resource_cleanup_cron` runs `cleanup_catalog_step`. The local file catalog +deletes entire expired directories, including uncataloged files and empty directories, then removes their catalog entries +and persists the result. Source directories determine eligibility; an incomplete catalog does not prevent cleanup, and +cached entries never authorize deletion elsewhere. The cleanup root must stay inside the workspace and contain no symlinks. +Cleanup requires descriptor-relative safe directory deletion (Linux and macOS); unsupported platforms fail without deleting files. + +The default `ttl_days` is 7. Expiry uses calendar dates in the application `timezone`, independent of file mtimes. +On 2026-11-05, directories dated 2026-10-29 or earlier expire; 2026-10-30 and later remain. +Cron first runs at the next hour after startup. Cleanup does not run while the service is stopped. + +Start the service in one terminal and choose the retention period: + +```bash +reme start config=proactive jobs.resource_cleanup_cron.ttl_days=7 +``` + +`ttl_days` must be a positive integer. Set `jobs.resource_cleanup_cron.enabled=false` to disable cleanup. Other +configurations without this Job retain their existing behavior. Failed deletions retain the directory's catalog entries +and fail that execution; later scheduled runs check again. See +[`LocalFileCatalog.cleanup_expired`](https://github.com/agentscope-ai/ReMe/blob/main/reme/components/file_catalog/local_file_catalog.py). +The existing `file_catalog.delete()` still removes catalog entries only. + +Cleanup directly affects only dated resource directories and their catalog entries; it does not rewrite other directories +or references. If another configuration also enables Auto Resource, its existing deletion handlers still remove associated +daily cards. + ## Image Resources Text and image resources share the same agent-wrapper and note-writing tools. Image inputs add a native AgentScope diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 8b4a2e78..292fe6c4 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -74,6 +74,7 @@ Keep secrets in `.env` or the process environment, never in committed configurat | `plugins` | `[]` | Installed plugins enabled for this Application | | `service` | HTTP | Service configuration | | `jobs` | default Jobs | Job configurations by name | +| `job_history_limit` | `100` | Completed executions retained per Job in application memory; running records are always retained | | `components` | defaults | Components grouped by type and name | `session_dir` must remain workspace-relative. @@ -92,6 +93,11 @@ Restart the service to apply the override. The separate `auto_dream` API remains `enable_serve` independently controls service exposure: `enabled=true, enable_serve=false` keeps a Job available for local calls. Background and cron jobs are never service-exposed. +`jobs..track_execution` defaults to `true` for background/cron Jobs and `false` for base/stream Jobs. +Set it explicitly to override either default, for example `reme start jobs.chat.track_execution=true`. +Execution records are memory-only and disappear when the Application is recreated. `job_history_limit=0` retains only +running executions. See [Job execution status](./operations.md#job-execution-status) for queries and step paths. + ## LLM The default LLM uses an OpenAI-compatible interface: diff --git a/docs/en/memory_as_file.md b/docs/en/memory_as_file.md index 0393f30b..c375c5b6 100644 --- a/docs/en/memory_as_file.md +++ b/docs/en/memory_as_file.md @@ -47,6 +47,8 @@ system state -> metadata/ Each layer solves a different problem. `session/` and `resource/` preserve source records. Files under `resource/` remain unchanged at their original path. +Explicitly enabled [resource TTL](./auto_resource.md#resource-retention-and-ttl), such as `config=proactive`, deletes +expired date directories; keep permanent material outside `resource/`. Standard Auto Memory records retain conversation messages while intentionally omitting tool-result and base64 data blocks; this keeps recalled output and binary payloads from masquerading as user-provided evidence. Generated Agent runtime state instead lives under `mem_session/`. diff --git a/docs/en/operations.md b/docs/en/operations.md index f5e62823..ab1ba8f5 100644 --- a/docs/en/operations.md +++ b/docs/en/operations.md @@ -26,6 +26,37 @@ reme app_config - `status` estimates stateful component memory and process RSS; - `app_config` returns the effective configuration with secrets redacted. +## Job execution status + +Query the running service from another terminal: + +```bash +reme job_status +reme job_status name=dream_cron limit=10 +reme job_status status=running +``` + +Results map Job names to execution lists, newest first, with at most `limit` records per Job (default 20). +Each record includes an application-local `id`, UTC `started_at`/`finished_at`, `status`, a bounded `error` summary, +`step_path`, `step_started_at`, and `failed_step`. Status is `running`, `succeeded`, `failed`, `stopped` (cooperative +background shutdown), or `cancelled`. A failed response also records a failure even when no exception was raised. +Terminal ERROR chunks also mark a run as failed; notifications with `will_retry=true` remain recoverable. +Completed records clear the active step path; failures retain the innermost failed step path. + +Background Jobs record each body invocation, including supervisor restarts. Cron Jobs record each trigger; waiting +for a schedule does not create an execution. A tracked Job that has not run returns an empty list. Watcher Jobs remain +`running` while waiting for changes; a path such as `["watch_changes_step", "update_index_step"]` shows an active handler +and returns to `["watch_changes_step"]` afterward. Step paths describe invocation boundaries, not progress inside a Step. +Executor-backed Jobs finish their records when the worker body exits, even if the outer Future was cancelled earlier. + +Background/cron tracking is on by default; base/stream tracking is opt-in through `jobs..track_execution`. +The `job_status` query itself is untracked by default. Records contain no request payloads or full responses, live only +in application memory, and retain the latest 100 completed executions per Job by default. Configure +[`job_history_limit`](./configuration.md#application-fields) to change that limit; active executions are always retained. + +The Python evaluation helper `check_job_count(name, app.context)` reports cumulative starts, including failed and +cancelled calls. Its count is independent of execution tracking, query filters, and history retention. + ## Logs and common symptoms `log_to_console` and `log_to_file` control logging. For startup failures, inspect the first exception rather than later client connection errors. diff --git a/docs/en/proactive.md b/docs/en/proactive.md index ab9bcaf9..565c9db0 100644 --- a/docs/en/proactive.md +++ b/docs/en/proactive.md @@ -49,7 +49,7 @@ proactive_refresh_cron: steps: *proactive_refresh_steps ``` -The anchor above is only a compact illustration; `default.yaml` spells out both step lists explicitly. The read job is: +`default.yaml` shares this step list through the same YAML anchor. The read job is: ```yaml proactive_read: @@ -86,6 +86,9 @@ Parameters: ### Refresh cost, files, and opt-out +`metadata.proactive.duration_ms` has been removed. While the service is running, use +`reme job_status name=proactive_refresh_cron` from another terminal to inspect the whole Job's start and finish times. + When no daily Markdown note changed, refresh exits before calling an LLM and does not create a new exposure file. With changed material, extraction normally makes one LLM call and may retry once after an unusable reply. If push candidates remain, planning makes one additional call; agenda generation makes one more when there are multiple candidates. A diff --git a/docs/zh/auto_resource.md b/docs/zh/auto_resource.md index a6abf843..49db7d4b 100644 --- a/docs/zh/auto_resource.md +++ b/docs/zh/auto_resource.md @@ -53,6 +53,38 @@ workspace/ `AutoImageResourceStep` 声明图像后缀匹配规则，`AutoTextResourceStep` 作为最后的 fallback。后续新增模态时， 只需注册新的 processor、提供独立 prompt 并在 `dispatch_steps` 中增加一项，无需修改 router。 +## 资源保留与 TTL + +默认配置不自动删除原始资源。`config=proactive` 显式启用以下保留协议： + +- 临时资源放在 `resource/YYYY-MM-DD/` 下，可包含任意层级和后缀的文件。 +- 永久资料放在 workspace 中 `resource/` 以外的目录，例如 `archive/`。 +- `resource/` 根目录文件、非日期目录、无效日期和日期目录符号链接不参与清理；这些是兼容保留项。 + +`resource_watch_loop` 递归监控全部后缀，只更新 `file_catalog: resource`，不会调用 LLM 或自动生成 daily 卡片。 +`resource_cleanup_cron` 每小时调用 `cleanup_catalog_step`，由 local file catalog 删除过期目录的整个内容（包括尚未登记的文件和空目录）， +再删除对应 catalog 条目并持久化。文件系统是清理依据，不要求 catalog 完整，也不会根据缓存条目删除其他路径。 +清理根目录必须位于 workspace 内，且路径中不能含符号链接。 +清理要求平台支持基于目录描述符的安全删除（Linux、macOS）；不支持的平台会报错并保留文件。 + +默认 `ttl_days: 7`。按应用 `timezone` 中的日历日期计算，目录日期距今天达到 7 天即过期，与文件 mtime 无关。 +例如 2026-11-05 清理时，`resource/2026-10-29/` 及更早日期目录会被删除，10-30 及之后保留。 +Cron 在服务启动后的下一个整点开始执行，停止服务期间不会清理。 + +在一个终端启动服务并设置保留天数： + +```bash +reme start config=proactive jobs.resource_cleanup_cron.ttl_days=7 +``` + +`ttl_days` 必须是正整数。通过 `jobs.resource_cleanup_cron.enabled=false` 关闭清理；未配置该任务的其他配置不受影响。 +删除失败会保留该目录的 catalog 条目并标记本次任务失败，后续调度重新检查。 +实现见 [`LocalFileCatalog.cleanup_expired`](https://github.com/agentscope-ai/ReMe/blob/main/reme/components/file_catalog/local_file_catalog.py)。 +现有 `file_catalog.delete()` 仍然只删除 catalog 条目，不删除源文件。 + +该流程只直接清理 resource 日期目录和对应 catalog，不改写其他目录或引用。 +若在其他配置中同时启用 Auto Resource，删除事件仍会按其规则清理关联 daily 卡片。 + ## 图像资源 文本和图像共用 agent-wrapper 与笔记写入工具；图片输入只是在理解提示词旁加入原生 AgentScope 图像块。 diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 7053a74c..1bdffbfd 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -79,6 +79,7 @@ base_url: ${LLM_BASE_URL:-https://example.com/v1} | `plugins` | `[]` | 为当前 Application 启用的已安装插件 | | `service` | HTTP | 服务端配置 | | `jobs` | 默认 Job | Job 名到 Job 配置的映射 | +| `job_history_limit` | `100` | 每个 Job 在应用内存中保留的已结束执行记录数；运行中的记录始终保留 | | `components` | 默认组件 | 按类型和名称组织的组件配置 | `session_dir` 必须保持 workspace-relative。其他 workspace 子目录也应使用清晰、稳定的相对名称。 @@ -96,6 +97,11 @@ reme start jobs.dream_cron.enabled=false `enable_serve` 独立控制接口暴露：`enabled=true, enable_serve=false` 允许本地调用。 Background 和 cron Job 始终不暴露服务接口。 +`jobs..track_execution` 对 background/cron Job 默认开启，对 base/stream Job 默认关闭。 +可以显式覆盖，例如 `reme start jobs.chat.track_execution=true` 为前台聊天开启记录。 +执行记录仅保存在内存中，重新创建 Application 后清空。`job_history_limit=0` 仅保留运行中的记录。 +查询方式与 Step 路径见 [Job 执行状态](./operations.md#job-执行状态)。 + ## LLM 配置 默认 LLM 使用 OpenAI-compatible 接口： diff --git a/docs/zh/memory_as_file.md b/docs/zh/memory_as_file.md index 2788c199..2f56c0a3 100644 --- a/docs/zh/memory_as_file.md +++ b/docs/zh/memory_as_file.md @@ -42,7 +42,9 @@ system state -> metadata/ 这四层解决的是不同问题。 -`session/` 和 `resource/` 保存来源记录。`resource/` 文件保持原路径和原内容；标准 Auto Memory 保留对话消息，但会有意排除 +`session/` 和 `resource/` 保存来源记录。`resource/` 文件保持原路径和原内容；显式启用的 +[资源 TTL](./auto_resource.md#资源保留与-ttl)（如 `config=proactive`）会删除过期日期目录，永久资料应放在 `resource/` 之外。 +标准 Auto Memory 保留对话消息，但会有意排除 tool-result 和 base64 data block，防止召回结果和二进制负载被误当成用户证据。Agent 运行时生成状态则放在 `mem_session/`。 `daily/` 是浅加工层。它把当天发生的对话和资源整理成更适合阅读的 daily note：什么事情发生了、有哪些结论、留下了哪些后续任务、对应原文在哪里。 diff --git a/docs/zh/operations.md b/docs/zh/operations.md index 27571bf9..248f5506 100644 --- a/docs/zh/operations.md +++ b/docs/zh/operations.md @@ -25,6 +25,37 @@ reme app_config - `status`：状态组件的内存估算和进程 RSS； - `app_config`：隐藏密钥后的实际生效配置。 +## Job 执行状态 + +在另一个终端查询已运行的服务： + +```bash +reme job_status +reme job_status name=dream_cron limit=10 +reme job_status status=running +``` + +结果为 Job 名到执行记录列表的映射，最新执行在前，每个 Job 最多返回 `limit` 条，默认 20 条。 +每条记录包含应用内递增的 `id`、UTC `started_at`/`finished_at`、`status`、限长的错误摘要 `error`、 +当前 `step_path`、`step_started_at` 和 `failed_step`。状态为 `running`、`succeeded`、`failed`、 +`stopped`（后台收到停止信号后正常退出）或 `cancelled`。即使没有抛出异常，失败的 Response 也会记录为失败。 +终止执行的 ERROR chunk 也会标记为失败；`will_retry=true` 的通知仍允许恢复。 +执行结束后清空当前 Step 路径，失败时保留最内层失败路径。 + +后台 Job 每次运行及监护重启各记一条；Cron 每次触发记一条，等待调度不算执行。 +开启记录但尚未执行的 Job 返回空列表。Watcher 等待事件时仍为 `running`；处理变更时可显示 +`["watch_changes_step", "update_index_step"]`，处理完成后恢复为 `["watch_changes_step"]`。 +Step 路径表示执行边界，不代表 Step 内部的处理进度。线程池 Job 的记录以执行体真正退出为准， +外层 Future 被取消时，工作线程可能仍在运行。 + +Background/cron 默认开启记录，base/stream 通过 `jobs..track_execution` 按需开启。 +`job_status` 查询自身默认不记录。记录不保存请求参数或完整响应，仅保存在应用内存中， +默认每个 Job 保留最近 100 条已结束记录。通过 [`job_history_limit`](./configuration.md#application-字段) +修改保留数量，运行中的记录始终保留。 + +Python 评测接口 `check_job_count(name, app.context)` 返回累计启动次数，包含失败和取消的调用， +不受执行记录开关、查询筛选和历史保留数量影响。 + ## 日志 `log_to_console` 和 `log_to_file` 控制日志目标。排查启动失败时先查看第一条异常，而不是后续客户端连接错误。 diff --git a/docs/zh/proactive.md b/docs/zh/proactive.md index c97262be..23acb433 100644 --- a/docs/zh/proactive.md +++ b/docs/zh/proactive.md @@ -46,7 +46,7 @@ proactive_refresh_cron: steps: *proactive_refresh_steps ``` -上面的 anchor 只是为了紧凑展示；`default.yaml` 中显式写出了两组 steps。读取 job 为： +`default.yaml` 也通过这个 YAML anchor 共享同一组 steps。读取 job 为： ```yaml proactive_read: @@ -83,6 +83,9 @@ proactive_read: ### Refresh 成本、文件与关闭方式 +`metadata.proactive.duration_ms` 已移除。服务运行期间，在另一个终端使用 +`reme job_status name=proactive_refresh_cron` 查询整条 Job 的开始和结束时间。 + 没有 daily Markdown 发生变化时，refresh 会在调用 LLM 前结束，也不会生成新的曝光文件。有变化素材时，extract 通常调用一次 LLM；回复不可用时最多重试一次。存在 push candidates 时，plan 会再调用一次；候选数大于一个时， agenda 会再调用一次。因此默认链路单轮最多调用四次 LLM。 diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 69f1ace6..e0bdebfd 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -11,7 +11,7 @@ from pydantic import BaseModel from ..base_component import BaseComponent from ..outbound_proxy import BaseOutboundProxy from ...enumeration import ChunkEnum, ComponentEnum -from ...schema import StreamChunk, TokenUsage +from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage from ...utils import global_counter_add_many if TYPE_CHECKING: @@ -254,5 +254,24 @@ class BaseAgentWrapper(BaseComponent): """Request compaction of one persisted agent session.""" raise NotImplementedError(f"{type(self).__name__} does not support session compaction") + async def read_session(self, session_id: str) -> SessionSnapshot: + """Read display history without waiting for the session to become idle.""" + raise NotImplementedError(f"{type(self).__name__} does not support session reading") + + async def get_session_status(self, session_id: str) -> SessionStatus: + """Read activity without resuming or waiting for the session.""" + raise NotImplementedError(f"{type(self).__name__} does not support session status") + + async def append_session( + self, + session_id: str, + content: str, + *, + message_id: str, + timeout: float = 30, + ) -> SessionMessage: + """Append an idempotent proactive assistant message after the active turn finishes.""" + raise NotImplementedError(f"{type(self).__name__} does not support session appending") + async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream agent events as unified StreamChunk objects.""" diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index d7d5afb5..cb29a3fb 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator, AsyncIterator from contextlib import asynccontextmanager +from dataclasses import dataclass from functools import partial import json import os @@ -15,13 +16,30 @@ 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 +from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage if TYPE_CHECKING: from openai_codex import AsyncThread, RunInput from openai_codex.types import Notification +@dataclass +class _TurnUsage: + """Invocation-local usage derived from monotonic thread counters.""" + + baseline: TokenUsage + usage: TokenUsage | None = None + + def update(self, total: TokenUsage) -> None: + """Keep reported spending separate from synthetic context estimates.""" + current = TokenUsage( + input_tokens=max(0, total.input_tokens - self.baseline.input_tokens), + output_tokens=max(0, total.output_tokens - self.baseline.output_tokens), + ) + if current.total_tokens or self.usage is not None: + self.usage = current + + @R.register("codex") class CodexAgentWrapper(BaseAgentWrapper): """Configure agents while sharing transport, authentication and session locks.""" @@ -65,10 +83,6 @@ class CodexAgentWrapper(BaseAgentWrapper): 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: - return TokenUsage.from_provider(usage) - 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) @@ -179,28 +193,49 @@ class CodexAgentWrapper(BaseAgentWrapper): return to_jsonable_python(value, by_alias=True) + async def _turn_stream(self, turn: Any, usage: _TurnUsage) -> AsyncIterator[Notification]: + """Observe only this turn's reported spending, including compaction calls.""" + stream = turn.stream() + try: + async for event in stream: + if event.method == "thread/tokenUsage/updated": + if event.payload.turn_id != turn.id: + continue + total = self.app_server.update_token_usage( + turn.thread_id, + TokenUsage.from_provider(event.payload.token_usage.total), + ) + usage.update(total) + yield event + finally: + await stream.aclose() + async def reply(self, inputs: RunInput, **kwargs) -> dict: """Run an isolated or explicitly resumed conversation on the shared server.""" + from openai_codex._run import _collect_async_turn_result + kwargs = self._reply_kwargs(kwargs) async with self._session(kwargs) as thread: + usage = _TurnUsage(await self.app_server.token_usage(thread.id)) turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) + stream = self._turn_stream(turn, usage) try: - result = await turn.run() + result = await _collect_async_turn_result(stream, turn_id=turn.id) except BaseException: await self._interrupt(turn) raise + finally: + if usage.usage is not None: + self._record_token_usage(usage.usage) + await stream.aclose() final_response = result.final_response or "" response = { "session_id": thread.id, "last_message": final_response, "result": final_response, "turn": self._serialize(result), - "usage": None, + "usage": usage.usage.model_dump() if usage.usage is not None else None, } - if (raw_usage := getattr(result, "usage", None)) is not None: - usage = self._codex_usage(raw_usage.last) - response["usage"] = usage.model_dump() - self._record_token_usage(usage) if kwargs.get("output_schema") is not None: try: response["structured_output"] = json.loads(final_response) @@ -213,6 +248,28 @@ class CodexAgentWrapper(BaseAgentWrapper): async with self._session({"session_id": session_id}, compact=True) as thread: await self.app_server.compact(thread) + async def read_session(self, session_id: str) -> SessionSnapshot: + """Read history without acquiring the generation lease.""" + await self.start() + return await self.app_server.read_session(session_id) + + async def get_session_status(self, session_id: str) -> SessionStatus: + """Report native thread activity, including approval/input waits.""" + await self.start() + return await self.app_server.get_session_status(session_id) + + async def append_session( + self, + session_id: str, + content: str, + *, + message_id: str, + timeout: float = 30, + ) -> SessionMessage: + """Append a proactive message through the shared session owner.""" + await self.start() + return await self.app_server.append_session(session_id, content, message_id=message_id, timeout=timeout) + 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 @@ -231,8 +288,14 @@ class CodexAgentWrapper(BaseAgentWrapper): @classmethod # pylint: disable=too-many-return-statements - def _event_to_chunks(cls, event: Notification, session_id: str) -> list[StreamChunk]: - """Convert one Codex app-server notification to unified stream chunks.""" + def _event_to_chunks( + cls, + event: Notification, + session_id: str, + *, + usage: TokenUsage | None = None, + ) -> list[StreamChunk]: + """Map notifications; usage replaces the cumulative counts for its turn block.""" method, payload = event.method, event.payload make_chunk = partial(cls._chunk, session_id=session_id) if method == "turn/started": @@ -299,15 +362,16 @@ class CodexAgentWrapper(BaseAgentWrapper): ), ] if method == "thread/tokenUsage/updated": - usage = payload.token_usage.last - normalized = cls._codex_usage(usage) + if usage is None: + return [] return [ make_chunk( ChunkEnum.USAGE, - chunk=normalized.model_dump(), - input_tokens=normalized.input_tokens, - output_tokens=normalized.output_tokens, - metadata={"usage": normalized.model_dump()}, + block_id=f"usage:{payload.turn_id}", + chunk=usage.model_dump(), + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + metadata={"usage": usage.model_dump()}, ), ] if method == "error": @@ -348,17 +412,25 @@ class CodexAgentWrapper(BaseAgentWrapper): async def reply_stream(self, inputs: RunInput, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream notifications routed by the SDK to this conversation's turn.""" kwargs = self._reply_kwargs(kwargs, stream=True) + emitted_usage = None async with self._session(kwargs) as thread: + usage = _TurnUsage(await self.app_server.token_usage(thread.id)) turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) - stream = turn.stream() + stream = self._turn_stream(turn, usage) 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): + if event.method == "thread/tokenUsage/updated": + if usage.usage == emitted_usage: + continue + emitted_usage = usage.usage + for chunk in self._event_to_chunks(event, thread.id, usage=usage.usage): yield chunk finally: + if usage.usage is not None: + self._record_token_usage(usage.usage) if not completed: await self._interrupt(turn) await stream.aclose() diff --git a/reme/components/agent_wrapper/codex_session_history.py b/reme/components/agent_wrapper/codex_session_history.py new file mode 100644 index 00000000..7ae33dfb --- /dev/null +++ b/reme/components/agent_wrapper/codex_session_history.py @@ -0,0 +1,188 @@ +"""Codex display-history projection and file-native proactive message records.""" + +import datetime +import json +import os +from pathlib import Path +import re +from uuid import uuid4 + +from ...enumeration import ChunkEnum +from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage + +_IDENTIFIER = re.compile(r"[A-Za-z0-9_-]{1,128}\Z") + + +def validate_identifier(value: str) -> str: + """Keep session and caller-provided message identifiers out of path syntax.""" + if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value): + raise ValueError("Session/message IDs must contain 1-128 letters, digits, underscores or hyphens") + return value + + +def session_status(thread) -> SessionStatus: + """Normalize native status without mistaking unloaded/error sessions for idle.""" + status = thread.status.root + return SessionStatus( + session_id=thread.id, + status={"notLoaded": "not_loaded", "systemError": "error"}.get(status.type, status.type), + active_flags=[getattr(flag, "value", flag) for flag in getattr(status, "active_flags", [])], + ) + + +def read_records(path: Path) -> list[dict]: + """Read the application-owned provenance file; corruption fails explicitly.""" + if not path.exists(): + return [] + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def write_records(path: Path, records: list[dict]) -> None: + """Atomically publish records before/after an injection; never rewrite a Codex rollout.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{uuid4().hex}.tmp") + try: + with temporary.open("w", encoding="utf-8") as stream: + for record in records: + stream.write(json.dumps(record, ensure_ascii=False) + "\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def injected_ids(thread, home: Path) -> set[str]: + """Recover accepted injections from the native rollout, including an interrupted publication. + + Codex 0.163 persists injected response items but omits them from its public + display-history RPCs. Only our reserved IDs are read from this internal format. + """ + raw_path = getattr(thread, "path", None) + if not raw_path: + raise ValueError("Appending requires a persisted Codex session") + path = Path(raw_path).resolve() + if not path.is_relative_to(home.resolve()): + raise ValueError("Codex session path is outside its configured home") + if not path.exists(): + return set() + found = set() + lines = path.read_text(encoding="utf-8").splitlines(keepends=True) + for index, line in enumerate(lines): + try: + entry = json.loads(line) + except json.JSONDecodeError: + if index == len(lines) - 1 and not line.endswith("\n"): + break # An unrelated native append may still be writing its final line. + raise + if entry.get("type") != "response_item": + continue + payload = entry.get("payload") + if not isinstance(payload, dict) or payload.get("type") != "message" or payload.get("role") != "assistant": + continue + item_id = payload.get("id", "") + if isinstance(item_id, str) and item_id.startswith("reme-proactive-"): + found.add(item_id) + return found + + +def read_token_usage(thread, home: Path) -> TokenUsage: + """Recover reported cumulative spending before starting a restored native session's next turn.""" + raw_path = getattr(thread, "path", None) + if not raw_path: + return TokenUsage() + path = Path(raw_path).resolve() + if not path.is_relative_to(home.resolve()): + raise ValueError("Codex session path is outside its configured home") + usage = TokenUsage() + with path.open(encoding="utf-8") as stream: + for line in stream: + try: + record = json.loads(line) + except json.JSONDecodeError: + if not line.endswith("\n"): + break + raise + payload = record.get("payload", {}) + if record.get("type") == "event_msg" and payload.get("type") == "token_count": + if info := payload.get("info"): + usage = TokenUsage.from_provider(info["total_token_usage"]) + return usage + + +def new_record(message_id: str, content: str, thread) -> dict: + """Anchor proactive messages after the last completed native turn.""" + return { + "message": SessionMessage( + id=f"reme-proactive-{message_id}", + role="assistant", + content=content, + source="proactive", + created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(), + turn_id=thread.turns[-1].id if thread.turns else None, + ).model_dump(mode="json"), + "injected": False, + } + + +def session_snapshot(thread, records: list[dict], recovered_ids: set[str]) -> SessionSnapshot: + """Combine native turn groups and durable proactive messages in insertion order.""" + proactive: dict[str | None, list[SessionMessage]] = {} + for record in records: + message = SessionMessage.model_validate(record["message"]) + if record["injected"] or message.id in recovered_ids: + proactive.setdefault(message.turn_id, []).append(message) + messages = list(proactive.pop(None, [])) + for turn in thread.turns: + users: list[str] = [] + chunks: list[StreamChunk] = [] + for wrapped in turn.items: + item = wrapped.root + if item.type == "userMessage": + users.extend(part.root.text for part in item.content if part.root.type == "text") + elif item.type == "agentMessage": + chunks.append(StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=item.text, block_id=item.id)) + elif item.type in {"reasoning", "plan"}: + text = item.text if item.type == "plan" else "\n".join(item.summary or item.content or []) + if text: + chunks.append(StreamChunk(chunk_type=ChunkEnum.THINK, chunk=text, block_id=item.id)) + elif item.type in { + "commandExecution", + "fileChange", + "mcpToolCall", + "dynamicToolCall", + "collabAgentToolCall", + "webSearch", + }: + chunks.append( + StreamChunk( + chunk_type=ChunkEnum.TOOL_RESULT, + chunk=item.model_dump(mode="json", by_alias=True), + tool_call_id=item.id, + tool_call_name=getattr(item, "tool", None) or item.type, + metadata={"state": getattr(getattr(item, "status", None), "value", "completed")}, + ), + ) + if users: + messages.append( + SessionMessage(id=f"{turn.id}:user", role="user", content="\n".join(users), turn_id=turn.id), + ) + if getattr(turn, "error", None): + chunks.append(StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=turn.error.message)) + if chunks: + messages.append( + SessionMessage( + id=f"{turn.id}:assistant", + role="assistant", + turn_id=turn.id, + chunks=chunks, + content="".join(str(chunk.chunk) for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT), + completed=getattr(getattr(turn, "status", None), "value", "completed") != "inProgress", + ), + ) + messages.extend(proactive.pop(turn.id, [])) + # A native history migration/compaction can omit an old anchor. Retain its + # user-owned proactive records rather than silently dropping messages. + for remaining in proactive.values(): + messages.extend(remaining) + return SessionSnapshot(session_id=thread.id, state=session_status(thread), messages=messages) diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 0e0c393f..e5059c44 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -1,9 +1,11 @@ """Application context: shared state container for components, jobs, and service.""" from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING, Any +from datetime import datetime, timezone +from threading import Lock +from typing import TYPE_CHECKING, Any, get_args -from ..schema import ApplicationConfig +from ..schema import ApplicationConfig, JobExecution, JobExecutionStatus from .component_registry import ComponentRegistry, create_application_registry if TYPE_CHECKING: @@ -13,11 +15,10 @@ if TYPE_CHECKING: class ApplicationContext: - """Passive state container holding parsed config and wired components. + """Shared state container holding parsed config and wired components. - The Application class performs the actual wiring (registry lookups and - component instantiation); this class only stores the results so that - components, jobs, and the service can find each other at runtime. + Application performs the wiring; this class stores the results and provides + synchronized access to in-memory job execution snapshots. """ def __init__(self, *, registry: ComponentRegistry | None = None, **kwargs): @@ -30,9 +31,110 @@ class ApplicationContext: self.components: dict[str, dict[str, "BaseComponent"]] = {} self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None + self.job_history: dict[str, list[JobExecution]] = {} + self._job_counts: dict[str, int] = {} + self._job_lock = Lock() + self._job_execution_id = 0 # Application-lifetime shared state. Values remain available across Job and Step # invocations while this Application is running, so Jobs may keep cross-call state here. # This is in-memory state, not durable storage; use workspace files or a store when state # must survive an Application restart. self.metadata: dict[str, Any] = {} + + def start_job_execution(self, name: str, *, track_execution: bool = True) -> JobExecution | None: + """Count every invocation and optionally append a running execution.""" + with self._job_lock: + self._job_counts[name] = self._job_counts.get(name, 0) + 1 + if not track_execution: + return None + self._job_execution_id += 1 + execution = JobExecution(id=self._job_execution_id, started_at=datetime.now(timezone.utc)) + self.job_history.setdefault(name, []).append(execution) + return execution + + def get_job_count(self, name: str) -> int: + """Return cumulative starts, independently of tracking and history retention.""" + with self._job_lock: + if name not in self.jobs: + raise KeyError(f"Job '{name}' not found") + return self._job_counts.get(name, 0) + + def finish_job_execution( + self, + name: str, + execution: JobExecution, + status: JobExecutionStatus, + error: str | None = None, + ) -> None: + """Finish an execution and prune only completed records, oldest first.""" + with self._job_lock: + execution.status = status + execution.finished_at = datetime.now(timezone.utc) + execution.error = error[:2000] if error else None + if status != "failed": + execution.failed_step = None + execution.step_path = [] + execution.step_started_at = None + history = self.job_history[name] + completed = [record for record in history if record.status != "running"] + limit = self.app_config.job_history_limit + keep = {record.id for record in completed[-limit:]} if limit else set() + self.job_history[name] = [record for record in history if record.status == "running" or record.id in keep] + + def enter_job_step(self, execution: JobExecution, name: str) -> tuple[list[str], datetime | None] | None: + """Push a step and return the parent path and its original start time.""" + with self._job_lock: + if execution.status != "running": + return None + previous = execution.step_path, execution.step_started_at + execution.step_path = [*execution.step_path, name] + execution.step_started_at = datetime.now(timezone.utc) + return previous + + def leave_job_step( + self, + execution: JobExecution, + previous: tuple[list[str], datetime | None], + failed: bool, + response_failed: bool | None = None, + ) -> None: + """Restore the parent; None preserves an exception path, False clears a recovered failure.""" + with self._job_lock: + if execution.status != "running": + return + if failed or (response_failed and execution.failed_step is None): + execution.failed_step = list(execution.step_path) + elif response_failed is False: + execution.failed_step = None + execution.step_path, execution.step_started_at = previous + + def get_job_status( + self, + name: str | None = None, + status: JobExecutionStatus | None = None, + limit: int = 20, + ) -> dict[str, list[dict]]: + """Return detached JSON snapshots, newest first, including untriggered tracked jobs.""" + if not isinstance(limit, int) or isinstance(limit, bool) or limit < 1: + raise ValueError("limit must be a positive integer") + if status is not None and status not in get_args(JobExecutionStatus): + raise ValueError(f"Invalid execution status: {status}") + with self._job_lock: + if name is not None: + if name not in self.jobs and name not in self.job_history: + raise KeyError(f"Job '{name}' not found") + names = [name] + else: + names = sorted( + self.job_history.keys() | {key for key, job in self.jobs.items() if job.track_execution}, + ) + result = {} + for key in names: + records = [ + record + for record in reversed(self.job_history.get(key, [])) + if status is None or record.status == status + ] + result[key] = [record.model_dump(mode="json") for record in records[:limit]] + return result diff --git a/reme/components/codex_app_server.py b/reme/components/codex_app_server.py index da465865..91fa7532 100644 --- a/reme/components/codex_app_server.py +++ b/reme/components/codex_app_server.py @@ -7,6 +7,8 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field import json +import hashlib +import math from pathlib import Path import sys from typing import Any, TYPE_CHECKING @@ -17,6 +19,8 @@ from .service.scoped_mcp import ScopedMCP from ..constants import normalize_connect_host from ..enumeration import ComponentEnum from ..utils.async_utils import finish_cleanup +from ..utils.async_utils import complete_in_thread +from ..schema import SessionMessage, SessionSnapshot, SessionStatus, TokenUsage if TYPE_CHECKING: from openai_codex import AsyncCodex, AsyncThread, CodexConfig @@ -31,6 +35,7 @@ class _ThreadState: token: str | None = None completion: asyncio.Future[None] | None = None operation_started: bool = False + token_usage: TokenUsage | None = None @R.register("sdk") @@ -85,6 +90,126 @@ class CodexAppServer(BaseComponent): return path if path.is_absolute() else self.workspace_path / path return self.workspace_path / ".codex" + def _message_path(self, session_id: str) -> Path: + from .agent_wrapper.codex_session_history import validate_identifier + + validate_identifier(session_id) + session_dir = self.app_context.app_config.session_dir if self.app_context else "session" + namespace = hashlib.sha256(str(self.session_path.resolve()).encode()).hexdigest()[:16] + path = (self.workspace_path / session_dir / "proactive" / namespace / f"{session_id}.jsonl").resolve() + if not path.is_relative_to(self.workspace_path.resolve()): + raise ValueError("Session messages must stay inside the workspace") + return path + + @asynccontextmanager + async def _session_client(self) -> AsyncIterator[AsyncCodex]: + """Keep status/history RPCs under lifecycle ownership without taking the turn lock.""" + await self.start() + self._active += 1 + self._idle.clear() + try: + yield await self._get_client() + finally: + self._active -= 1 + if not self._active: + self._idle.set() + + async def get_session_status(self, session_id: str) -> SessionStatus: + """Inspect a thread without resuming it or waiting for an active turn.""" + from .agent_wrapper.codex_session_history import session_status + + self._message_path(session_id) + async with self._session_client() as codex: + thread = (await codex._client.thread_read(session_id)).thread # pylint: disable=protected-access + return session_status(thread) + + async def read_session(self, session_id: str) -> SessionSnapshot: + """Project public native history plus application-authored proactive messages.""" + from .agent_wrapper import codex_session_history as history + + path = self._message_path(session_id) + async with self._session_client() as codex: + response = await codex._client.thread_read( # pylint: disable=protected-access + session_id, + include_turns=True, + ) + thread = response.thread + records = await asyncio.to_thread(history.read_records, path) + recovered = ( + await asyncio.to_thread(history.injected_ids, thread, self.session_path) + if any(not record["injected"] for record in records) + else set() + ) + return history.session_snapshot(thread, records, recovered) + + async def append_session( + self, + session_id: str, + content: str, + *, + message_id: str, + timeout: float = 30, + ) -> SessionMessage: + """Wait for the thread lease, then inject and publish one retry-safe assistant message.""" + from openai_codex.generated.v2_all import ThreadInjectItemsResponse + from .agent_wrapper import codex_session_history as history + + path = self._message_path(session_id) + history.validate_identifier(message_id) + if not isinstance(content, str) or not content.strip() or len(content) > 100_000: + raise ValueError("content must contain 1-100000 characters") + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("timeout must be positive and finite") + + async with asyncio.timeout(timeout): + async with self.thread(session_id) as loaded: + thread = (await loaded.read(include_turns=True)).thread + if history.session_status(thread).status != "idle": + raise RuntimeError("Session must be idle before appending") + accepted = await asyncio.to_thread(history.injected_ids, thread, self.session_path) + records = await asyncio.to_thread(history.read_records, path) + native_id = f"reme-proactive-{message_id}" + record = next((entry for entry in records if entry["message"]["id"] == native_id), None) + if record is not None: + if record["message"]["content"] != content: + raise ValueError("message_id already belongs to different content") + if record["injected"] or native_id in accepted: + if not record["injected"]: + record["injected"] = True + await complete_in_thread(history.write_records, path, records) + return SessionMessage.model_validate(record["message"]) + else: + record = history.new_record(message_id, content, thread) + records.append(record) + # An unaccepted attempt may be retried after later user turns. + record["message"]["turn_id"] = thread.turns[-1].id if thread.turns else None + + async def inject_and_publish(): + # A crash between the RPC and publication leaves a pending record. + # Its reserved native ID is recovered from the rollout on retry. + await complete_in_thread(history.write_records, path, records) + codex = await self._get_client() + await codex._client.request( # pylint: disable=protected-access + "thread/inject_items", + { + "threadId": session_id, + "items": [ + { + "type": "message", + "id": native_id, + "role": "assistant", + "content": [{"type": "output_text", "text": content}], + }, + ], + }, + response_model=ThreadInjectItemsResponse, + ) + record["injected"] = True + await complete_in_thread(history.write_records, path, records) + + await finish_cleanup(inject_and_publish()) + return SessionMessage.model_validate(record["message"]) + def _client_config(self) -> CodexConfig: from openai_codex import CodexConfig @@ -264,6 +389,27 @@ class CodexAppServer(BaseComponent): state.completion = None raise + async def token_usage(self, thread_id: str) -> TokenUsage: + """Read the pre-turn baseline under the conversation lease, recovering it after restart.""" + from .agent_wrapper.codex_session_history import read_token_usage + + state = self._threads[thread_id] + if state.token_usage is None: + codex = await self._get_client() + view = (await codex._client.thread_read(thread_id)).thread # pylint: disable=protected-access + state.token_usage = await asyncio.to_thread(read_token_usage, view, self.session_path) + return state.token_usage + + def update_token_usage(self, thread_id: str, total: TokenUsage) -> TokenUsage: + """Keep cumulative provider counters monotonic across wrappers and duplicate notifications.""" + state = self._threads[thread_id] + previous = state.token_usage or TokenUsage() + state.token_usage = TokenUsage( + input_tokens=max(previous.input_tokens, total.input_tokens), + output_tokens=max(previous.output_tokens, total.output_tokens), + ) + return state.token_usage + async def turn(self, thread: AsyncThread, inputs: Any, **kwargs): """Start a normal SDK turn with completion tracking owned by this server.""" turn = None @@ -293,8 +439,30 @@ class CodexAppServer(BaseComponent): async with self._operation(thread): await thread.compact() + self._threads[thread.id].token_usage = None await finish_cleanup(dispatch()) + async def _fork_messages(self, source_id: str, thread: AsyncThread) -> None: + """Copy inherited proactive provenance while the source conversation is locked.""" + from .agent_wrapper import codex_session_history as history + + records = await asyncio.to_thread(history.read_records, self._message_path(source_id)) + if not records: + return + accepted = set() + if any(not record["injected"] for record in records): + codex = await self._get_client() + source = (await codex._client.thread_read(source_id)).thread # pylint: disable=protected-access + accepted = await asyncio.to_thread(history.injected_ids, source, self.session_path) + # Native forks inherit context but do not preserve injected item IDs. + inherited = [ + {**record, "injected": True} + for record in records + if record["injected"] or record["message"]["id"] in accepted + ] + if inherited: + await complete_in_thread(history.write_records, self._message_path(thread.id), inherited) + @asynccontextmanager async def thread( self, @@ -360,13 +528,16 @@ class CodexAppServer(BaseComponent): source.revision += 1 try: async with source.lock: + state.token_usage = await self.token_usage(thread_id) thread = await codex.thread_fork(thread_id, **options) + await finish_cleanup(self._fork_messages(thread_id, thread)) finally: source.users -= 1 elif thread_id: thread = await codex.thread_resume(thread_id, **options) else: thread = await codex.thread_start(**options) + state.token_usage = TokenUsage() self._threads[thread.id] = state if thread_id and not fork: await self.wait_idle(thread) @@ -375,6 +546,10 @@ class CodexAppServer(BaseComponent): finally: # Interrupt acknowledgement alone does not mean that tool execution stopped. error = sys.exception() + if error is not None: + # A paused/closed consumer can miss usage emitted by native completion. + # Recover from the persisted rollout before the next turn starts. + state.token_usage = None try: async with asyncio.timeout(30 if error is not None else None): await self.wait_idle(thread) diff --git a/reme/components/file_catalog/base_file_catalog.py b/reme/components/file_catalog/base_file_catalog.py index 96c36315..54b1a794 100644 --- a/reme/components/file_catalog/base_file_catalog.py +++ b/reme/components/file_catalog/base_file_catalog.py @@ -38,3 +38,11 @@ class BaseFileCatalog(BaseComponent): @abstractmethod async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]: """Return nodes by paths; None = all; missing paths are skipped.""" + + async def cleanup_expired(self, directory: str, ttl_days: int) -> dict: + """Delete expired YYYY-MM-DD directories and their catalog entries. + + Retention uses directory dates in the application timezone, not file mtimes. + Backends must explicitly implement filesystem cleanup; ``delete`` remains metadata-only. + """ + raise NotImplementedError(f"{type(self).__name__} does not support directory cleanup") diff --git a/reme/components/file_catalog/local_file_catalog.py b/reme/components/file_catalog/local_file_catalog.py index 1cfe1199..29af2cee 100644 --- a/reme/components/file_catalog/local_file_catalog.py +++ b/reme/components/file_catalog/local_file_catalog.py @@ -1,10 +1,18 @@ """Local file catalog backend: in-memory dict persisted as compressed JSONL.""" import asyncio +import os +import shutil +import stat +from contextlib import ExitStack +from datetime import date, datetime +from pathlib import Path +from zoneinfo import ZoneInfo from .base_file_catalog import BaseFileCatalog from ..component_registry import R from ...schema import FileNode +from ...utils.async_utils import complete_in_thread from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -49,6 +57,68 @@ class LocalFileCatalog(BaseFileCatalog): return list(self._nodes.values()) return [self._nodes[p] for p in paths if p in self._nodes] + async def cleanup_expired(self, directory: str, ttl_days: int) -> dict: + """Remove dated directories at least ``ttl_days`` old, then persist the catalog.""" + if not isinstance(ttl_days, int) or isinstance(ttl_days, bool) or ttl_days < 1: + raise ValueError("ttl_days must be a positive integer") + tz_name = self.app_context.app_config.timezone if self.app_context else None + today = datetime.now(ZoneInfo(tz_name) if tz_name else None).date() + async with self._io_lock: + # Keep filesystem deletion, catalog mutation and persistence together even on cancellation. + return await complete_in_thread(self._cleanup_expired, directory, today, ttl_days) + + def _cleanup_expired(self, directory: str, today: date, ttl_days: int) -> dict: + """Inspect source directories, never using cached nodes to authorize deletion.""" + workspace = self.workspace_path.resolve() + relative_dir = Path(directory) + root = workspace / relative_dir + if relative_dir.anchor or str(directory).startswith("~") or ".." in relative_dir.parts or root == workspace: + raise ValueError("cleanup directory must stay inside the workspace") + if root.resolve() != root or not root.is_relative_to(workspace): + raise ValueError("cleanup directory must stay inside the workspace and must not contain symlinks") + deleted_dirs: list[str] = [] + errors: dict[str, str] = {} + if root.exists(): + if not shutil.rmtree.avoids_symlink_attacks: + raise RuntimeError("resource cleanup requires symlink-resistant directory deletion") + with ExitStack() as stack: + # Open each path component without following links; keep deletion anchored if the root moves. + flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + root_fd = os.open(workspace, flags) + stack.callback(os.close, root_fd) + for part in relative_dir.parts: + root_fd = os.open(part, flags, dir_fd=root_fd) + stack.callback(os.close, root_fd) + for name in sorted(os.listdir(root_fd)): + try: + day = date.fromisoformat(name) + except ValueError: + continue + if day.isoformat() != name or (today - day).days < ttl_days: + continue + relative = (relative_dir / name).as_posix() + try: + if not stat.S_ISDIR(os.stat(name, dir_fd=root_fd, follow_symlinks=False).st_mode): + continue + shutil.rmtree(name, dir_fd=root_fd) + except OSError as exc: + errors[relative] = str(exc) + self.logger.warning(f"Failed to clean up {relative}: {exc}") + else: + deleted_dirs.append(relative) + + prefixes = tuple(path + "/" for path in deleted_dirs) + deleted_nodes = [path for path in self._nodes if path in deleted_dirs or path.startswith(prefixes)] + for path in deleted_nodes: + del self._nodes[path] + # Also retry a prior failed write when the source directories have already disappeared. + write_jsonl_zst(self._catalog_file, (n.model_dump_json() for n in self._nodes.values()), self.encoding) + if deleted_dirs: + self.logger.info( + f"Cleaned up {len(deleted_dirs)} dated directories and {len(deleted_nodes)} catalog entries", + ) + return {"deleted_dirs": deleted_dirs, "deleted_nodes": len(deleted_nodes), "errors": errors} + async def _read_jsonl(self) -> None: for line in read_jsonl_zst(self._catalog_file, self.encoding): if stripped := line.strip(): diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index ec8406e6..aa2fe5d3 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -5,12 +5,18 @@ import contextlib import random import threading import time +from contextvars import ContextVar from .base_job import BaseJob from ..component_registry import R from ..runtime_context import RuntimeContext from ...schema import Response +_supervised_job: ContextVar["BackgroundJob | None"] = ContextVar( + "reme_supervised_job", + default=None, +) + @R.register("background") class BackgroundJob(BaseJob): @@ -29,6 +35,8 @@ class BackgroundJob(BaseJob): eventually crashes restarts quickly rather than at the capped delay. """ + _track_execution_default = True + def __init__( self, supervisor: bool = True, @@ -108,7 +116,7 @@ class BackgroundJob(BaseJob): while not self._stop_event.is_set(): started_at = time.monotonic() try: - await self() + await self._run_once() return except Exception as e: if not self.supervisor: @@ -121,12 +129,25 @@ class BackgroundJob(BaseJob): attempt += 1 await self._wait_or_stop(delay) + async def _run_once(self) -> None: + """Track the actual body, including custom __call__ implementations and restarts.""" + context = RuntimeContext(stop_event=self._stop_event, **self.kwargs) + token = _supervised_job.set(self) + try: + with self._execution(context): + response = await self() + if isinstance(response, Response): + context.response = response + finally: + _supervised_job.reset(token) + async def __call__(self, **kwargs) -> Response: """Default body: run steps in order; errors propagate to supervisor.""" - self.check_enabled() - self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(stop_event=self._stop_event, **merged) - for step in self._build_steps(): - await step(context) - return context.response + if _supervised_job.get() is self: + # Consume once so a nested call of this same job gets its own record. + _supervised_job.set(None) + return await self._run_steps(context) + with self._execution(context): + return await self._run_steps(context) diff --git a/reme/components/job/base_job.py b/reme/components/job/base_job.py index eb8314bc..e23efbd2 100644 --- a/reme/components/job/base_job.py +++ b/reme/components/job/base_job.py @@ -1,13 +1,15 @@ """Base job component for sequential step execution.""" +import asyncio +from contextlib import contextmanager from typing import TYPE_CHECKING +from ..application_context import ApplicationContext from ..base_component import BaseComponent from ..component_registry import R -from ..runtime_context import RuntimeContext +from ..runtime_context import RuntimeContext, _current_job_execution from ...enumeration import ComponentEnum from ...schema import ComponentConfig, Response -from ...utils import global_counter_inc if TYPE_CHECKING: from ...steps import BaseStep @@ -37,6 +39,7 @@ class BaseJob(BaseComponent): """Job that executes steps sequentially and returns a Response.""" component_type = ComponentEnum.JOB + _track_execution_default = False def __init__( self, @@ -45,6 +48,7 @@ class BaseJob(BaseComponent): steps: list[ComponentConfig | dict] | None = None, enable_serve: bool = True, enabled: bool = True, + track_execution: bool | None = None, **kwargs, ): super().__init__(**kwargs) @@ -53,6 +57,7 @@ class BaseJob(BaseComponent): self.step_configs = steps or [] self.enabled = enabled self.enable_serve = enable_serve + self.track_execution = self._track_execution_default if track_execution is None else track_execution self.step_specs: list[tuple[type["BaseStep"], dict]] = [] async def _start(self) -> None: @@ -79,28 +84,59 @@ class BaseJob(BaseComponent): # dict(params) copies kwargs so steps cannot mutate the shared spec. return [step_cls(**dict(params)) for step_cls, params in self.step_specs] - def _record_call(self) -> None: - """Increment this job's application-lifetime call counter.""" - metadata = getattr(self.app_context, "metadata", None) - if isinstance(metadata, dict): - global_counter_inc(metadata, ["__job_counter", self.name]) - def check_enabled(self) -> None: """Reject execution of a disabled job.""" if not self.enabled: raise ValueError(f"Job '{self.name}' is disabled") + @contextmanager + def _execution(self, context: RuntimeContext): + """Count one call and optionally track it, preserving the caller's error handling.""" + self.check_enabled() + owner = self.app_context + execution = ( + owner.start_job_execution(self.name, track_execution=self.track_execution) + if isinstance(owner, ApplicationContext) + else None + ) + context.job_execution = execution + token = _current_job_execution.set(execution) + try: + yield + except BaseException as exc: + if execution is not None: + status = "cancelled" if isinstance(exc, asyncio.CancelledError) else "failed" + owner.finish_job_execution(self.name, execution, status, _describe_exception(exc)) + raise + else: + if execution is not None: + if not context.response.success: + status, error = "failed", str(context.response.answer) + elif context.stream_error is not None: + status, error = "failed", context.stream_error + elif context.stop_event is not None and context.stop_event.is_set(): + status, error = "stopped", None + else: + status, error = "succeeded", None + owner.finish_job_execution(self.name, execution, status, error) + finally: + _current_job_execution.reset(token) + + async def _run_steps(self, context: RuntimeContext) -> Response: + """Execute fresh steps against a single invocation context.""" + for step in self._build_steps(): + await step(context) + return context.response + async def __call__(self, **kwargs) -> Response: """Run all steps in order, capturing any failure into the response.""" - self.check_enabled() - self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(**merged) - try: - for step in self._build_steps(): - await step(context) - except Exception as e: - self.logger.exception(f"Failed to execute job: {e}") - context.response.success = False - context.response.answer = _describe_exception(e) + with self._execution(context): + try: + await self._run_steps(context) + except Exception as e: + self.logger.exception(f"Failed to execute job: {e}") + context.response.success = False + context.response.answer = _describe_exception(e) return context.response diff --git a/reme/components/job/cron_job.py b/reme/components/job/cron_job.py index 71948597..9b88cdeb 100644 --- a/reme/components/job/cron_job.py +++ b/reme/components/job/cron_job.py @@ -35,10 +35,13 @@ class CronJob(BackgroundJob): return max(0.0, (nxt - now).total_seconds()) async def _execute_steps(self) -> Response: - context = RuntimeContext(**self.kwargs) - for step in self._build_steps(): - await step(context) - return context.response + context = RuntimeContext(stop_event=self._stop_event, **self.kwargs) + with self._execution(context): + return await self._run_steps(context) + + async def _run_once(self) -> None: + # The scheduler waits between executions; only individual triggers are recorded. + await self() async def __call__(self, **kwargs) -> Response: self.check_enabled() @@ -48,7 +51,6 @@ class CronJob(BackgroundJob): if self._stop_event.is_set(): break try: - self._record_call() await self._execute_steps() except Exception as exc: self.logger.exception(f"Cron job '{self.name}' failed: {exc}") diff --git a/reme/components/job/stream_job.py b/reme/components/job/stream_job.py index b947a12a..5f437c7e 100644 --- a/reme/components/job/stream_job.py +++ b/reme/components/job/stream_job.py @@ -1,6 +1,6 @@ """Streaming job for real-time output delivery.""" -from .base_job import BaseJob +from .base_job import BaseJob, _describe_exception from ..component_registry import R from ..runtime_context import RuntimeContext from ...enumeration import ChunkEnum @@ -12,14 +12,14 @@ class StreamJob(BaseJob): async def __call__(self, **kwargs) -> None: """Run steps; emit failures as ERROR chunks, then a terminal DONE marker.""" - self.check_enabled() - self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(**merged) - try: - for step in self._build_steps(): - await step(context) - except Exception as e: - await context.add_stream_string(str(e), ChunkEnum.ERROR) - # Always emit DONE so consumers can detach even after an error. - await context.add_stream_done() + with self._execution(context): + try: + await self._run_steps(context) + except Exception as e: + context.response.success = False + context.response.answer = _describe_exception(e) + await context.add_stream_string(str(e), ChunkEnum.ERROR) + # Always emit DONE so consumers can detach even after an error. + await context.add_stream_done() diff --git a/reme/components/runtime_context.py b/reme/components/runtime_context.py index 355b94ce..10ad66f1 100644 --- a/reme/components/runtime_context.py +++ b/reme/components/runtime_context.py @@ -1,9 +1,12 @@ """Per-request runtime context shared across steps and jobs.""" import asyncio +from contextvars import ContextVar from ..enumeration import ChunkEnum -from ..schema import Response, StreamChunk +from ..schema import JobExecution, Response, StreamChunk + +_current_job_execution: ContextVar[JobExecution | None] = ContextVar("reme_job_execution", default=None) class RuntimeContext: @@ -24,6 +27,10 @@ class RuntimeContext: self.stream_queue: asyncio.Queue | None = stream_queue self.stop_event: asyncio.Event | None = stop_event self.data: dict = kwargs + self.job_execution: JobExecution | None = _current_job_execution.get() + self.step_error: Exception | None = None + # Providers can emit terminal errors without raising or changing the Response. + self.stream_error: str | None = None # ----- Data dict access ---------------------------------------------- @@ -66,6 +73,8 @@ class RuntimeContext: """Put a chunk on the stream queue; raises when streaming is disabled.""" if self.stream_queue is None: raise RuntimeError("Stream queue not initialized") + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True: + self.stream_error = str(chunk.chunk) await self.stream_queue.put(chunk) async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 677adf8c..985a7805 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -72,7 +72,7 @@ jobs: dream_cron: backend: cron cron: "0 23 * * *" - steps: + steps: &dream_steps - backend: dream_extract_step file_catalog: dream scan_days: 2 @@ -100,7 +100,7 @@ jobs: type: string description: "YYYY-MM-DD to refresh; defaults to today in the application's timezone" default: "" - steps: + steps: &proactive_refresh_steps - backend: proactive_extract_step file_catalog: proactive scan_days: 2 @@ -122,24 +122,7 @@ jobs: proactive_refresh_cron: backend: cron cron: "0 18 * * *" # once a day at 18:00 - steps: - - backend: proactive_extract_step - file_catalog: proactive - scan_days: 2 - carry_forward_days: 14 - max_carry_forward_topics: 20 - llm_timeout_seconds: 300 - max_chars_per_file: 60000 - max_total_chars: 300000 - - backend: proactive_topics_step - known_threshold: 0.85 - known_threshold_calibrated_for: text-embedding-v4@1024 - min_push_confidence: 0.5 - max_topics: 10 - - backend: proactive_plan_step - - backend: proactive_agenda_step - - backend: proactive_finish_step - file_catalog: proactive + steps: *proactive_refresh_steps auto_dream: backend: base @@ -167,16 +150,7 @@ jobs: type: integer description: "estimated bytes of changed file content per extraction call, taken from file size; it excludes prompt overhead and a single oversized file still gets its own call, so it is a batching budget rather than a hard request limit" default: 40000 - steps: - - backend: dream_extract_step - file_catalog: dream - scan_days: 2 - max_units: 5 - max_chars_per_batch: 40000 - - backend: dream_integrate_step - - backend: dream_finish_step - file_catalog: dream - - backend: auto_tag_step + steps: *dream_steps auto_memory: backend: base @@ -297,6 +271,45 @@ jobs: steps: - backend: version_step + session_read: + backend: base + description: "Read session history, including proactive assistant messages." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + required: [session_id] + steps: + - backend: session_read_step + agent_wrapper: default + + session_status: + backend: base + description: "Inspect session activity without waiting for generation." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + required: [session_id] + steps: + - backend: session_status_step + agent_wrapper: default + + session_append: + backend: base + description: "Append a proactive assistant message once the session is idle; retry with the same message_id." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + message_id: {type: string, pattern: "^[A-Za-z0-9_-]{1,128}$", description: "stable retry/idempotency identifier"} + content: {type: string, minLength: 1, maxLength: 100000} + wait_timeout: {type: number, exclusiveMinimum: 0, default: 30} + required: [session_id, message_id, content] + steps: + - backend: session_append_step + agent_wrapper: default + app_config: backend: base description: "return the effective application config with secrets redacted" @@ -336,6 +349,27 @@ jobs: steps: - backend: health_check_step + job_status: + backend: base + description: "Query in-memory job executions and current step paths, newest first per job." + parameters: + type: object + properties: + name: + type: string + description: "Optional job name; defaults to all tracked jobs" + status: + type: string + enum: [running, succeeded, failed, stopped, cancelled] + description: "Optional execution status filter" + limit: + type: integer + minimum: 1 + default: 20 + description: "Maximum executions returned per job" + steps: + - backend: job_status_step + status: backend: base description: "report memory estimates for stateful data components and process RSS" diff --git a/reme/config/proactive.yaml b/reme/config/proactive.yaml index 7f189229..878126ec 100644 --- a/reme/config/proactive.yaml +++ b/reme/config/proactive.yaml @@ -1,13 +1,15 @@ -# Minimal Studio workspace with one Codex chat agent using API-key authentication. +# Studio workspace with MCP and a full-access 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 +# Temporary assets: resource/YYYY-MM-DD/...; keep permanent assets outside resource/. +# The resource catalog tracks all suffixes; dated directories expire after 7 calendar days. service: backend: http web_enabled: true - mcp_enabled: false + mcp_enabled: true jobs: # Keep Markdown search and graph data in sync with the workspace files. @@ -26,9 +28,40 @@ jobs: - backend: update_index_step persist: false + resource_watch_loop: + backend: background + watch_dirs: [resource_dir] + watch_suffixes: [] # Catalog every file type, without interpreting resources into daily cards. + steps: + - backend: init_changes_step + monitor_type: file_catalog + monitor_name: resource + dispatch_steps: + - backend: update_catalog_step + file_catalog: resource + - backend: watch_changes_step + monitor_type: file_catalog + monitor_name: resource + dispatch_steps: + - backend: update_catalog_step + file_catalog: resource + + resource_cleanup_cron: + backend: cron + cron: "0 * * * *" # Hourly in the application timezone; no startup cleanup. + ttl_days: 7 + steps: + - backend: cleanup_catalog_step + file_catalog: resource + chat: backend: stream - description: "Stream a read-only agent conversation over the ReMe workspace." + description: "Stream a full-access Codex conversation with the prome assistant." + system_prompt: | + 你是 prome（proactive ReMe），主人的工作与生活管家。 + 始终称呼用户为“主人”，使用主人的语言自然交流，回答清晰、简洁、务实。 + 临时资源放在 resource/YYYY-MM-DD/ 下，按资源清理任务的 ttl_days 到期后自动删除（默认 7 天），修改文件不会续期。 + 永久资料放在 workspace 中 resource/ 以外的目录；不要把需要长期保存的原件放入临时日期目录。 parameters: type: object properties: @@ -39,6 +72,46 @@ jobs: steps: - backend: chat_step agent_wrapper: default + job_tools: [] + + session_read: + backend: base + description: "Read session history, including proactive assistant messages." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + required: [session_id] + steps: + - backend: session_read_step + agent_wrapper: default + + session_status: + backend: base + description: "Inspect session activity without waiting for generation." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + required: [session_id] + steps: + - backend: session_status_step + agent_wrapper: default + + session_append: + backend: base + description: "Append a proactive assistant message once the session is idle; retry with the same message_id." + parameters: + type: object + properties: + session_id: {type: string, description: "agent session identifier"} + message_id: {type: string, pattern: "^[A-Za-z0-9_-]{1,128}$", description: "stable retry/idempotency identifier"} + content: {type: string, minLength: 1, maxLength: 100000} + wait_timeout: {type: number, exclusiveMinimum: 0, default: 30} + required: [session_id, message_id, content] + steps: + - backend: session_append_step + agent_wrapper: default app_config: backend: base @@ -119,86 +192,6 @@ jobs: 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: @@ -207,29 +200,77 @@ components: api_key: ${LLM_API_KEY:-} base_url: ${LLM_BASE_URL:-} codex_home: .codex + supports_websockets: false + experimental_api: true + config_overrides: + # Preserve Qwen3.8 Max's actual input limit; put safety headroom in the compaction threshold. + # At 95% usable context, 600000 leaves 334435 tokens below Codex's hard cap. + # Budget 128 KiB of new text at one token/byte, not Codex's byte/4 estimate of 32768. + # Planning allowances: 131072 output + 131072 new text + 32768 estimation/compaction overhead. + # This soft margin does not enforce request limits or guarantee room for arbitrary text/images. + - model_context_window=983616 + - model_auto_compact_token_limit=600000 + - 'model_auto_compact_token_limit_scope="total"' + # Zero disables extra compaction after a completed reply. + - model_post_turn_compact_threshold_percent=0 + - features.async_compaction=false + # 2500 approximate tokens matches the fallback's 10000-byte tool-output limit. + - tool_output_token_limit=2500 + # Leave compact_prompt / experimental_compact_prompt_file unset to use the built-in prompt. + + # Keep review and spawned agents on the configured Qwen endpoint. + - 'model_reasoning_effort="xhigh"' + - 'review_model="${CODEX_MODEL_NAME:-qwen3.8-max}"' + - 'agents.default_subagent_model="${CODEX_MODEL_NAME:-qwen3.8-max}"' + - 'agents.default_subagent_reasoning_effort="xhigh"' + + # Long-term memory generation and injection are separate from context compaction. + - features.memories=true + - memories.generate_memories=true + - memories.use_memories=true + - 'memories.extract_model="${CODEX_MODEL_NAME:-qwen3.8-max}"' + - 'memories.consolidation_model="${CODEX_MODEL_NAME:-qwen3.8-max}"' + + # These capabilities describe Qwen3.8 Max; update them when switching model families. + # Official limits: https://help.aliyun.com/en/model-studio/qwen3-8-max model_catalog: models: - - slug: qwen3.8-max + - slug: ${CODEX_MODEL_NAME:-qwen3.8-max} display_name: Qwen3.8 Max - description: Studio chat + description: Work and life assistant with text and image input default_reasoning_level: xhigh supported_reasoning_levels: + - effort: none + description: Disable thinking + - effort: low + description: Light reasoning + - effort: medium + description: Balanced reasoning - effort: xhigh description: Maximum reasoning context_window: 983616 + max_context_window: 983616 + effective_context_window_percent: 95 input_modalities: [text, image] - supports_parallel_tool_calls: false - supports_reasoning_summaries: false + # The provider returns reasoning content without a reasoning.summary request parameter. + supports_reasoning_summary_parameter: false support_verbosity: false - shell_type: default + shell_type: unified_exec visibility: list supported_in_api: true priority: 0 - base_instructions: "" experimental_supported_tools: [] truncation_policy: mode: bytes limit: 10000 + # Custom catalogs do not inherit Codex's built-in base instruction template. + model_messages: + instructions_template: | + You are prome, a capable assistant working in the user's local ReMe workspace. + Follow developer instructions and applicable AGENTS.md files. + Use available tools to verify workspace facts; never invent file contents or completed actions. + Preserve user-owned memory files and unrelated changes when carrying out tasks. + Complete requested work, report relevant validation limits, and reply clearly in the user's language. agent_wrapper: default: @@ -237,8 +278,8 @@ components: app_server: default model: ${CODEX_MODEL_NAME:-qwen3.8-max} effort: xhigh - approval_mode: deny_all - sandbox: read-only + approval_mode: deny_all # SDK maps this to approval_policy=never; no approval requests. + sandbox: full-access tokenizer: default: @@ -253,6 +294,10 @@ components: default: backend: local + file_catalog: + resource: + backend: local + file_chunker: markdown: backend: markdown diff --git a/reme/schema/__init__.py b/reme/schema/__init__.py index 187ee50b..0699aee5 100644 --- a/reme/schema/__init__.py +++ b/reme/schema/__init__.py @@ -1,6 +1,7 @@ """Schema""" from .application_config import ApplicationConfig, ComponentConfig, JobConfig +from .agent_session import SessionMessage, SessionSnapshot, SessionStatus from .compaction import CallDecision, CompactionOptions, CompactionRequest, CompactionResult from .dream import ( DreamExtractOutput, @@ -18,6 +19,7 @@ from .emb_node import EmbNode from .file_chunk import FileChunk from .file_front_matter import FileFrontMatter from .graph_snapshot import GraphSnapshot, GraphSnapshotEdge, GraphSnapshotNode +from .job_execution import JobExecution, JobExecutionStatus from .file_link import FileLink from .file_node import FileNode from .request import Request @@ -46,12 +48,17 @@ __all__ = [ "GraphSnapshotNode", "IntegrateOutcome", "JobConfig", + "JobExecution", + "JobExecutionStatus", "ProactiveResult", "ProactiveState", "ProactiveStateFile", "ProactiveTopic", "Request", "Response", + "SessionMessage", + "SessionSnapshot", + "SessionStatus", "StreamChunk", "TokenUsage", "TraverseGraph", diff --git a/reme/schema/agent_session.py b/reme/schema/agent_session.py new file mode 100644 index 00000000..a518b757 --- /dev/null +++ b/reme/schema/agent_session.py @@ -0,0 +1,36 @@ +"""Backend-neutral session history and activity contracts.""" + +from typing import Literal + +from pydantic import BaseModel, Field + +from .stream_chunk import StreamChunk + + +class SessionStatus(BaseModel): + """Activity observed by the session owner, independently of an HTTP stream.""" + + session_id: str + status: Literal["active", "idle", "not_loaded", "error"] + active_flags: list[str] = Field(default_factory=list) + + +class SessionMessage(BaseModel): + """One displayed message; proactive provenance survives session reloads.""" + + id: str + role: Literal["user", "assistant"] + content: str + source: Literal["chat", "proactive"] = "chat" + created_at: str | None = None + turn_id: str | None = None + completed: bool = True + chunks: list[StreamChunk] = Field(default_factory=list) + + +class SessionSnapshot(BaseModel): + """Ordered display history, including messages absent from native UI history.""" + + session_id: str + state: SessionStatus + messages: list[SessionMessage] = Field(default_factory=list) diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index dc12737e..cb6a55f9 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -25,6 +25,10 @@ class JobConfig(ComponentConfig): steps: list[ComponentConfig] = Field(default_factory=list, description="Ordered step configs") enable_serve: bool = Field(default=True, description="Whether to expose this job through the service layer") enabled: bool = Field(default=True, description="Whether to start, serve, and allow execution of this job") + track_execution: bool | None = Field( + default=None, + description="Record execution and step status in memory; defaults on for background/cron, off otherwise", + ) class ApplicationConfig(BaseModel): @@ -57,6 +61,7 @@ class ApplicationConfig(BaseModel): service: ComponentConfig = Field(default_factory=ComponentConfig, description="Service endpoint config") jobs: dict[str, JobConfig] = Field(default_factory=dict, description="Job definitions keyed by job name") thread_pool_max_workers: int = Field(default=0, description="Max worker threads; 0 to disable") + job_history_limit: int = Field(default=100, ge=0, description="Completed executions retained per job in memory") components: dict[str, dict[str, ComponentConfig]] = Field( default_factory=dict, description="Component registry keyed by type then name", diff --git a/reme/schema/job_execution.py b/reme/schema/job_execution.py new file mode 100644 index 00000000..70e0ac66 --- /dev/null +++ b/reme/schema/job_execution.py @@ -0,0 +1,21 @@ +"""Application-lifetime job execution snapshots.""" + +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel, Field + +JobExecutionStatus = Literal["running", "succeeded", "failed", "stopped", "cancelled"] + + +class JobExecution(BaseModel): + """One execution; step paths describe the currently active nested call.""" + + id: int + started_at: datetime + finished_at: datetime | None = None + status: JobExecutionStatus = "running" + error: str | None = None + step_path: list[str] = Field(default_factory=list) + step_started_at: datetime | None = None + failed_step: list[str] | None = None diff --git a/reme/schema/proactive.py b/reme/schema/proactive.py index 6b74c2e0..35d523eb 100644 --- a/reme/schema/proactive.py +++ b/reme/schema/proactive.py @@ -102,7 +102,6 @@ class ProactiveState(BaseModel): interests_path: str = "" interests_written: bool = False checkpoint_paths: list[str] = Field(default_factory=list) - duration_ms: int = 0 class ProactiveStateFile(BaseModel): diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index b2042237..4357103a 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -149,12 +149,40 @@ class BaseStep(ComponentMixin, ABC): del self.__dict__[key] self.context = RuntimeContext.from_context(context, **kwargs) assert self.context is not None - if self.input_mapping: - self.context.apply_mapping(self.input_mapping) - result = await self.execute() - if self.output_mapping: - self.context.apply_mapping(self.output_mapping) - return result + execution = self.context.job_execution + previous = None + failed = False + response_failed = None + if execution is not None and self.app_context is not None: + previous = self.app_context.enter_job_step(execution, self.backend or self.name) + self.context.step_error = None + try: + if self.input_mapping: + self.context.apply_mapping(self.input_mapping) + result = await self.execute() + if self.output_mapping: + self.context.apply_mapping(self.output_mapping) + response_failed = not self.context.response.success or self.context.stream_error is not None + return result + except Exception as exc: + # Preserve the innermost path while an error is propagated or wrapped. + # A later independent failure replaces an error handled inside a parent. + if previous is not None: + previous_error = self.context.step_error + cause = exc.__cause__ or (None if exc.__suppress_context__ else exc.__context__) + failed = previous_error is None or (previous_error is not exc and previous_error is not cause) + self.context.step_error = exc + raise + finally: + if previous is not None: + self.app_context.leave_job_step( + execution, + previous, + failed, + response_failed, + ) + if not previous[0]: + self.context.step_error = None def prompt_format(self, prompt_name: str, **kwargs) -> str: """Format a named prompt template with the given kwargs.""" diff --git a/reme/steps/common/__init__.py b/reme/steps/common/__init__.py index 86240fe4..111b58a8 100644 --- a/reme/steps/common/__init__.py +++ b/reme/steps/common/__init__.py @@ -8,12 +8,14 @@ from .decision_demo import DecisionDemoStep from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep from .help import HelpStep +from .job_status import JobStatusStep from .llm_demo import LLMDemoStep from .python_execute import PythonExecuteStep from .shell import ShellStep from .stream_demo import StreamDemoStep1, StreamDemoStep2 from .stream_llm_demo import StreamLLMDemoStep from .status import StatusStep +from .session import SessionAppendStep, SessionReadStep, SessionStatusStep from .version import VersionStep __all__ = [ @@ -26,6 +28,7 @@ __all__ = [ "DemoEchoStep2", "HealthCheckStep", "HelpStep", + "JobStatusStep", "LLMDemoStep", "PythonExecuteStep", "ShellStep", @@ -33,5 +36,8 @@ __all__ = [ "StreamDemoStep2", "StreamLLMDemoStep", "StatusStep", + "SessionAppendStep", + "SessionReadStep", + "SessionStatusStep", "VersionStep", ] diff --git a/reme/steps/common/chat.py b/reme/steps/common/chat.py index 17feda80..c9a34f68 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -1,8 +1,5 @@ """Workspace-aware streaming chat step for the ReMe web interface.""" -import datetime -import zoneinfo - from ..base_step import BaseStep from ...components import R from ...enumeration import ChunkEnum @@ -10,7 +7,7 @@ from ...enumeration import ChunkEnum @R.register("chat_step") class ChatStep(BaseStep): - """Stream a read-only agent conversation over the current workspace.""" + """Stream a workspace conversation with configurable, read-only-by-default job tools.""" DEFAULT_SYSTEM_PROMPT = """You are ReMe Agent, an assistant for a local-first memory workspace. Use the available ReMe tools when workspace facts are needed. Cite workspace-relative file paths @@ -18,24 +15,6 @@ when referring to notes. Never invent file contents, and do not claim to have ch this chat intentionally provides read-only tools. Reply in the user's language.""" READ_ONLY_TOOLS = ["search", "list", "read", "read_image", "frontmatter_read", "stat", "traverse"] - def _system_prompt(self) -> str: - """Append request-time environment facts to the configured prompt.""" - timezone = self.app_context.app_config.timezone if self.app_context is not None else None - try: - current = datetime.datetime.now(zoneinfo.ZoneInfo(timezone)) if timezone else datetime.datetime.now() - except (zoneinfo.ZoneInfoNotFoundError, ValueError): - current = datetime.datetime.now() - - assert self.context is not None - base_prompt = str(self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT).rstrip() - return ( - f"{base_prompt}\n\n" - "\n" - f"Current date: {current.date().isoformat()}\n" - f"Current working directory: {self.agent_wrapper.cwd.resolve(strict=False)}\n" - "" - ) - async def execute(self): assert self.context is not None query = str(self.context.get("query") or "").strip() @@ -47,8 +26,8 @@ this chat intentionally provides read-only tools. Reply in the user's language." raise RuntimeError("chat_step requires an agent_wrapper") wrapper_kwargs = { - "system_prompt": self._system_prompt(), - "job_tools": self.READ_ONLY_TOOLS, + "system_prompt": str(self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT), + "job_tools": self.kwargs.get("job_tools", self.READ_ONLY_TOOLS), "builtin_tools": [], } if session_id := self.context.get("session_id"): diff --git a/reme/steps/common/job_status.py b/reme/steps/common/job_status.py new file mode 100644 index 00000000..0ff677d8 --- /dev/null +++ b/reme/steps/common/job_status.py @@ -0,0 +1,19 @@ +"""Query in-memory execution and step status.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("job_status_step") +class JobStatusStep(BaseStep): + """Return recent executions grouped by job name.""" + + async def execute(self): + assert self.context is not None + assert self.app_context is not None + self.context.response.answer = self.app_context.get_job_status( + name=self.context.get("name"), + status=self.context.get("status"), + limit=self.context.get("limit", 20), + ) + return self.context.response diff --git a/reme/steps/common/session.py b/reme/steps/common/session.py new file mode 100644 index 00000000..8431ff8d --- /dev/null +++ b/reme/steps/common/session.py @@ -0,0 +1,53 @@ +"""Backend-neutral session jobs; unsupported wrappers fail explicitly.""" + +from ..base_step import BaseStep +from ...components import R + + +class _SessionStep(BaseStep): + """Resolve the configured wrapper and return a typed session result as JSON.""" + + operation = "" + + async def execute(self): + assert self.context is not None + if self.agent_wrapper is None: + raise RuntimeError("Session jobs require an agent_wrapper") + session_id = self.context.get("session_id") + if not isinstance(session_id, str) or not session_id: + raise ValueError("session_id is required") + kwargs = {} + if self.operation == "append_session": + kwargs = { + "content": self.context.get("content"), + "message_id": self.context.get("message_id"), + "timeout": float(self.context.get("wait_timeout", 30)), + } + try: + result = await getattr(self.agent_wrapper, self.operation)(session_id, **kwargs) + except NotImplementedError: + self.context.response.metadata["code"] = "session_not_supported" + raise + self.context.response.answer = result.model_dump(mode="json") + return self.context.response + + +@R.register("session_read_step") +class SessionReadStep(_SessionStep): + """Read ordered session history.""" + + operation = "read_session" + + +@R.register("session_status_step") +class SessionStatusStep(_SessionStep): + """Inspect session activity without blocking behind generation.""" + + operation = "get_session_status" + + +@R.register("session_append_step") +class SessionAppendStep(_SessionStep): + """Append a proactive assistant message after the current turn finishes.""" + + operation = "append_session" diff --git a/reme/steps/evolve/proactive/agenda.py b/reme/steps/evolve/proactive/agenda.py index bd15becb..919a4a93 100644 --- a/reme/steps/evolve/proactive/agenda.py +++ b/reme/steps/evolve/proactive/agenda.py @@ -18,7 +18,6 @@ interests.yaml render, extending the v2 shape with ``agenda``/``suppressed``. import asyncio import json -import time from ...base_step import BaseStep from .._evolve import agent_reply_result_text, passthrough_response @@ -68,7 +67,6 @@ class ProactiveAgendaStep(BaseStep): assert self.context is not None if self.context.get(self.skip_key): return passthrough_response(self, self.skip_key) - started = time.monotonic() raw_state = self.context.get("proactive") if not raw_state: self.context.response.success = True @@ -165,7 +163,6 @@ class ProactiveAgendaStep(BaseStep): rel_path = norm_path(interests_path.relative_to(ws).as_posix()) state.interests_path = rel_path state.interests_written = bool(state.interests_written or written) - state.duration_ms = state.duration_ms + int((time.monotonic() - started) * 1000) self._store(state) answer = ( f"Agenda: {len(agenda)} item(s), suppressed={len(suppressed)}, push={state.push}, " diff --git a/reme/steps/evolve/proactive/extract.py b/reme/steps/evolve/proactive/extract.py index 3c8466b7..cfc077b0 100644 --- a/reme/steps/evolve/proactive/extract.py +++ b/reme/steps/evolve/proactive/extract.py @@ -2,7 +2,6 @@ import asyncio import json -import time from ...base_step import BaseStep from .._evolve import agent_reply_result_text, passthrough_response @@ -58,7 +57,6 @@ class ProactiveExtractStep(BaseStep): assert self.context is not None if self.context.get(self.skip_key): return passthrough_response(self, self.skip_key) - started = time.monotonic() day = today(self, str(self.context.get("date", "") or "")) ws = workspace_dir(self) daily = daily_dir(self) @@ -114,7 +112,7 @@ class ProactiveExtractStep(BaseStep): # 4) Zero-consumption early exit (A4 row 1). if not state.changed_paths: state.early_exit = "no_new_evidence" - return self._finish(state, started, "Skipped: no new evidence; 0 LLM calls") + return self._finish(state, "Skipped: no new evidence; 0 LLM calls") # 5) LLM channel (F4.4): structurally one reply per round plus at most # one parse-failure retry, so no persistent budget is needed (v5 R5). @@ -122,7 +120,7 @@ class ProactiveExtractStep(BaseStep): if wrapper is None: state.early_exit = "no_agent_wrapper" self.logger.warning(f"[{self.name}] no agent_wrapper available; skipping round") - return self._finish(state, started, "Skipped: no agent_wrapper configured") + return self._finish(state, "Skipped: no agent_wrapper configured") # 6) One LLM call with sectioned output (A3); retry once on parse failure. # meta is None when this round must NOT checkpoint (timeout, or two @@ -138,7 +136,7 @@ class ProactiveExtractStep(BaseStep): f"Extracted {len(state.follow_ups)} follow_up(s), {len(state.extends)} extend(s), " f"{len(state.updates)} update(s) from {len(state.changed_paths)} changed file(s)" ) - return self._finish(state, started, answer) + return self._finish(state, answer) def _build_messages(self, state: ProactiveState, ws, day: str, daily: str, material_blob: str) -> tuple[str, str]: carry_forward_json = json.dumps( @@ -257,8 +255,7 @@ class ProactiveExtractStep(BaseStep): self.context["proactive"] = data self.context.response.metadata["proactive"] = data - def _finish(self, state: ProactiveState, started: float, answer: str): - state.duration_ms = int((time.monotonic() - started) * 1000) + def _finish(self, state: ProactiveState, answer: str): self._store(state) self.context.response.success = True self.context.response.answer = answer diff --git a/reme/steps/evolve/proactive/plan.py b/reme/steps/evolve/proactive/plan.py index e29b3af3..8cbf4ad1 100644 --- a/reme/steps/evolve/proactive/plan.py +++ b/reme/steps/evolve/proactive/plan.py @@ -18,7 +18,6 @@ the LLM. import asyncio import json -import time from ...base_step import BaseStep from .._evolve import agent_reply_result_text, passthrough_response @@ -128,7 +127,6 @@ class ProactivePlanStep(BaseStep): assert self.context is not None if self.context.get(self.skip_key): return passthrough_response(self, self.skip_key) - started = time.monotonic() raw_state = self.context.get("proactive") if not raw_state: self.context.response.success = True @@ -174,7 +172,6 @@ class ProactivePlanStep(BaseStep): cards.append(card) state.scenario_cards = cards - state.duration_ms = state.duration_ms + int((time.monotonic() - started) * 1000) self._store(state) answer = f"Plan: {len(cards)} scenario card(s), fallback={fallbacks}, llm_calls={state.plan_llm_calls}" self.context.response.success = True diff --git a/reme/steps/evolve/proactive/topics.py b/reme/steps/evolve/proactive/topics.py index 9bf5df75..9d03fe70 100644 --- a/reme/steps/evolve/proactive/topics.py +++ b/reme/steps/evolve/proactive/topics.py @@ -16,7 +16,6 @@ workflow continues on exact comparison (BM25-only deployments). import datetime as dt import re -import time from ...base_step import BaseStep from .._evolve import passthrough_response @@ -78,7 +77,6 @@ class ProactiveTopicsStep(BaseStep): assert self.context is not None if self.context.get(self.skip_key): return passthrough_response(self, self.skip_key) - started = time.monotonic() if not self.context.get("proactive"): self.context.response.success = True self.context.response.answer = "Skipped topics: no proactive extract state in context" @@ -264,7 +262,6 @@ class ProactiveTopicsStep(BaseStep): state.file_skip_reason = file_skip_reason state.interests_path = rel_path state.interests_written = written - state.duration_ms = int((time.monotonic() - started) * 1000) self._store(state) answer = ( f"Topics: {len(topics_out)} rendered, push={push}, " diff --git a/reme/steps/index/__init__.py b/reme/steps/index/__init__.py index 2780b251..a226190d 100644 --- a/reme/steps/index/__init__.py +++ b/reme/steps/index/__init__.py @@ -4,6 +4,7 @@ from ._source_format import normalize_posix_path from .bm25_search import Bm25SearchStep from .clear_paths import ClearPathsStep from .clear_store import ClearStoreStep +from .cleanup_catalog import CleanupCatalogStep from .draft import AddDraftStep, ReadAllDraftStep from .graph_snapshot import GraphSnapshotStep from .log_changes import LogChangesStep @@ -30,6 +31,7 @@ __all__ = [ "ChangeApplyStep", "ClearPathsStep", "ClearStoreStep", + "CleanupCatalogStep", "DEFAULT_LOW_POWER_POLL_MS", "DEFAULT_WATCH_DEBOUNCE_MS", "DEFAULT_WATCH_STEP_MS", diff --git a/reme/steps/index/cleanup_catalog.py b/reme/steps/index/cleanup_catalog.py new file mode 100644 index 00000000..edd4d616 --- /dev/null +++ b/reme/steps/index/cleanup_catalog.py @@ -0,0 +1,22 @@ +"""Explicit retention cleanup for dated resource directories.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("cleanup_catalog_step") +class CleanupCatalogStep(BaseStep): + """Ask a file catalog to remove expired directories under ``resource_dir``.""" + + async def execute(self): + assert self.context is not None + if self.file_catalog is None: + raise RuntimeError("cleanup_catalog_step requires file_catalog") + ttl_days = self.context.get("ttl_days", self.kwargs.get("ttl_days", 7)) + result = await self.file_catalog.cleanup_expired(str(self.config_value("resource_dir")), ttl_days) + self.context.response.metadata.update(result) + self.context.response.success = not result["errors"] + self.context.response.answer = f"Deleted {len(result['deleted_dirs'])} expired resource directories" + if result["errors"]: + self.context.response.answer += f"; failed to delete {len(result['errors'])} directories" + return self.context.response diff --git a/reme/utils/evaluation_interface.py b/reme/utils/evaluation_interface.py index 1f6819d4..edef5c4f 100644 --- a/reme/utils/evaluation_interface.py +++ b/reme/utils/evaluation_interface.py @@ -29,9 +29,7 @@ def check_job_count(job_name: str, app_context: "ApplicationContext") -> int: current Application instance. Unknown job names use the same ``KeyError`` contract as :meth:`Application.run_job`. """ - if job_name not in app_context.jobs: - raise KeyError(f"Job '{job_name}' not found") - return global_counter_get(app_context.metadata, ["__job_counter", job_name]) + return app_context.get_job_count(job_name) class JobCountTracker: diff --git a/reme_studio/README.md b/reme_studio/README.md index 2cbeea61..467240cd 100644 --- a/reme_studio/README.md +++ b/reme_studio/README.md @@ -213,6 +213,32 @@ The built-in chat job is read-only with respect to the workspace. It can search as a substitute for explicitly saving durable knowledge. Chat requires a valid Agent/model configuration; file browsing, editing, graphs, and service inspection do not. +With `config=proactive`, Studio uses a Codex agent with full tool access and supports proactive messages. It synchronizes history when opening a +conversation, after streaming, on focus, and every five seconds while the conversation is visible. Hidden pages and +locally streaming conversations stop history polling. Proactive messages have a distinct background and a “Proactive reminder” label. + +Start the service in one terminal: + +```bash +reme start config=proactive +``` + +From another terminal, use the conversation's session identifier: + +```bash +reme session_status session_id=SESSION_ID +reme session_append session_id=SESSION_ID message_id=reminder-1 content="This information may be useful to you." wait_timeout=30 +reme session_read session_id=SESSION_ID +``` + +`session_append` waits for the current turn's lease and appends an assistant message without starting model generation. +Retry with the same `message_id` and content; reusing an ID with different content fails. Provenance records live under +`{session_dir}/proactive/` in the workspace, filling the injected-message gap in Codex's native display-history APIs. +Writes and ordinary chat must use the same ReMe app-server; session locks are not shared across processes. +Only persistent Codex sessions are supported; other wrappers explicitly raise `NotImplementedError` for these three APIs. +Busy waits time out. Once injection starts, cancellation finishes publication before releasing the session lock. +Recovering interrupted publication reads reserved message IDs from the Codex rollout without rewriting native files. + ### 7. Service status Open **Settings** from the top bar to inspect the backend that Studio is actually using. diff --git a/reme_studio/README_ZH.md b/reme_studio/README_ZH.md index 3351d3a9..c0b07e14 100644 --- a/reme_studio/README_ZH.md +++ b/reme_studio/README_ZH.md @@ -187,6 +187,29 @@ Chat 会在同一个标签工作区中打开，因此对话进行时仍可保留 内置 chat job 对工作区只读：它可以搜索和读取上下文，但不应替代明确保存长期知识的操作。Agent 对话需要有效的 Agent/模型配置；文件浏览、编辑、图谱和服务检查不需要模型。 +使用 `config=proactive` 时，Studio 的 Codex agent 启用完整工具访问权限，并支持主动消息。Studio 会在打开对话、流结束、页面重新获得焦点时同步历史， +并每 5 秒刷新当前可见的对话；隐藏页面和正在流式接收的对话停止补拉。主动消息使用独立背景色和“主动提醒”标签。 + +先在一个终端启动服务： + +```bash +reme start config=proactive +``` + +在另一个终端中，使用对话对应的 session id 调用： + +```bash +reme session_status session_id=SESSION_ID +reme session_append session_id=SESSION_ID message_id=reminder-1 content="主人，这条信息可能与你有关。" wait_timeout=30 +reme session_read session_id=SESSION_ID +``` + +`session_append` 等待同会话的当前轮次结束，只追加 assistant 消息，不启动模型生成。重试必须使用相同的 +`message_id` 和内容；同一 ID 对应不同内容会报错。来源记录保存在工作区的 `{session_dir}/proactive/` 下， +用于补齐 Codex 原生历史查询未返回的注入消息。写入与普通聊天必须经过同一个 ReMe app-server；不支持跨进程共用会话锁。 +仅支持持久化 Codex 会话；其他 wrapper 的这三个接口明确返回 `NotImplementedError`。忙碌等待有超时；一旦开始注入， +取消请求会先完成注入结果的持久化再释放会话锁。恢复中断的发布会只读检查 Codex rollout 中的保留消息 ID，不改写原生文件。 + ### 7. 服务状态 从顶部栏打开 **Settings**，可以检查 Studio 实际连接的后端。 diff --git a/reme_studio/app/api.ts b/reme_studio/app/api.ts index 82a07997..bddb16c0 100644 --- a/reme_studio/app/api.ts +++ b/reme_studio/app/api.ts @@ -5,8 +5,9 @@ import type { ReMeHealth, ReMeResponse, StreamChunk, + SessionSnapshot, } from "./types"; -import { decodeSseEvent } from "./chat-stream"; +import { readChatStream } from "./chat-stream"; import { healthFromResponse } from "./health-status"; import { translate, useLanguageStore, type TranslationKey } from "./i18n"; import { @@ -65,6 +66,29 @@ export async function getReMeVersion(): Promise { return String((await callReMe("version")).answer); } +export async function readChatSession( + sessionId: string, + signal: AbortSignal, +): Promise { + const response = await fetch(`${REME_API_URL}/session_read`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ session_id: sessionId }), + signal, + }); + if (response.status === 404) return undefined; + const payload = (await response.json()) as ReMeResponse; + if (payload.metadata?.code === "session_not_supported") return undefined; + if (!response.ok || !payload.success) { + throw new Error( + typeof payload.answer === "string" + ? payload.answer + : message("requestFailed", response.status), + ); + } + return payload.answer; +} + export async function getReMeStatus(): Promise> { return callReMe("status"); } @@ -135,21 +159,16 @@ export async function streamChat( }); if (!response.ok || !response.body) throw new Error(message("agentUnavailable", response.status)); - - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ""; - while (true) { - const { value, done } = await reader.read(); - buffer += decoder.decode(value, { stream: !done }); - const events = buffer.split(/\r?\n\r?\n/); - buffer = events.pop() || ""; - for (const event of events) { - const chunk = decodeSseEvent(event); - if (chunk) onChunk(chunk); - } - if (done) break; + if ( + response.headers.get("content-type")?.split(";")[0].trim().toLowerCase() !== + "text/event-stream" + ) { + await response.body.cancel(); + throw new Error(message("invalidResponse", response.status)); } - const finalChunk = decodeSseEvent(buffer); - if (finalChunk) onChunk(finalChunk); + await readChatStream( + response.body, + onChunk, + translate(useLanguageStore.getState().language, "chatInterrupted"), + ); } diff --git a/reme_studio/app/chat-stream.ts b/reme_studio/app/chat-stream.ts index 381ad5aa..6b30b5ac 100644 --- a/reme_studio/app/chat-stream.ts +++ b/reme_studio/app/chat-stream.ts @@ -34,6 +34,41 @@ export function decodeSseEvent(raw: string): StreamChunk | null { return JSON.parse(data) as StreamChunk; } +export async function readChatStream( + stream: ReadableStream, + onChunk: (chunk: StreamChunk) => void, + interruptedMessage: string, +): Promise { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + const emit = (event: string): boolean => { + const chunk = decodeSseEvent(event); + if (!chunk) return false; + onChunk(chunk); + return chunk.chunk_type === "done" || chunk.done; + }; + try { + while (true) { + const { value, done } = await reader.read(); + buffer += decoder.decode(value, { stream: !done }); + const events = buffer.split(/\r?\n\r?\n/); + buffer = events.pop() || ""; + for (const event of events) { + if (emit(event)) return; + } + if (done) { + if (emit(buffer)) return; + throw new Error(interruptedMessage); + } + } + } finally { + // Cleanup must not replace the original transport or parsing error. + await reader.cancel().catch(() => {}); + reader.releaseLock(); + } +} + const payloadText = (payload: StreamPayload): string => typeof payload === "string" ? payload : JSON.stringify(payload, null, 2); @@ -49,12 +84,27 @@ const appendError = (blocks: ChatBlock[], text: string): ChatBlock[] => { export function finishChatMessage( message: ChatMessage, error?: string, + { finishTools = true }: { finishTools?: boolean } = {}, ): ChatMessage { + const failed = + Boolean(error) || + (finishTools && message.blocks?.some((block) => block.type === "error")); let blocks = (message.blocks || []).map((block): ChatBlock => { if (block.type === "content" || block.type === "error") return block; + // Claude emits reply_end before executing the tools in that message. + if ( + !finishTools && + block.type === "tool" && + block.status !== "done" && + block.status !== "error" + ) + return block; return { ...block, - status: block.status === "error" ? "error" : "done", + status: + block.status === "error" || (failed && block.status !== "done") + ? "error" + : "done", expanded: false, }; }); @@ -127,7 +177,7 @@ function applyDetail(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { ...current, payloads: [...current.payloads, ...payloads], status: done ? "done" : "streaming", - expanded: true, + expanded: current.expanded, mediaType: chunk.media_type || current.mediaType, inputTokens: chunk.input_tokens ?? current.inputTokens, outputTokens: chunk.output_tokens ?? current.outputTokens, @@ -169,8 +219,32 @@ function applyTool(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { const current = index >= 0 ? (blocks[index] as ToolBlock) : undefined; const hasPayload = payloadText(chunk.chunk).length > 0; const result = chunk.chunk_type === "tool_result"; - const state = String(chunk.metadata?.state || "").toLowerCase(); - const failed = state.includes("error") || state.includes("fail"); + const payload = + typeof chunk.chunk === "object" && !Array.isArray(chunk.chunk) + ? chunk.chunk + : undefined; + const snapshot = + payload?.id === toolId && typeof payload.type === "string" + ? payload + : undefined; + const claudeResult = + result && payload?.tool_use_id === toolId && "content" in payload + ? payload + : undefined; + const state = String( + chunk.metadata?.state || snapshot?.status || "", + ).toLowerCase(); + const failed = + current?.status === "error" || + claudeResult?.is_error === true || + state.includes("error") || + state.includes("fail") || + state === "declined"; + const completed = + failed || + current?.status === "done" || + claudeResult !== undefined || + (state !== "" && state !== "inprogress"); const nextBlock: ToolBlock = { id, type: "tool", @@ -189,12 +263,12 @@ function applyTool(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { : current?.resultPayloads || [], status: failed ? "error" + : completed + ? "done" : result - ? state - ? "done" - : "running" + ? "running" : "calling", - expanded: true, + expanded: completed ? false : current?.expanded ?? true, mediaType: chunk.media_type || current?.mediaType, metadata: { ...current?.metadata, ...chunk.metadata }, }; @@ -228,14 +302,36 @@ export function applyStreamChunk( ], } : message; - return finishChatMessage(completed); + return finishChatMessage(completed, undefined, { finishTools: false }); } if (chunk.chunk_type === "done" || chunk.done) { return finishChatMessage(message); } if (chunk.chunk_type === "reply_start") return message; + if ( + chunk.chunk_type === "data" && + typeof chunk.metadata?.codex_method === "string" && + chunk.metadata.codex_method.startsWith("hook/") + ) + return message; - const blocks = message.blocks || []; + let blocks = message.blocks || []; + if ( + chunk.chunk_type === "think" || + (chunk.chunk_type === "content" && payloadText(chunk.chunk)) || + chunk.chunk_type === "tool_call" + ) { + const thinkingId = + chunk.chunk_type === "think" ? detailId(blocks, chunk) : undefined; + blocks = blocks.map( + (block): ChatBlock => + block.type === "think" && + block.status === "streaming" && + block.id !== thinkingId + ? { ...block, status: "done", expanded: false } + : block, + ); + } if (chunk.chunk_type === "content") return { ...message, blocks: applyContent(blocks, chunk) }; if (chunk.chunk_type === "tool_call" || chunk.chunk_type === "tool_result") { @@ -265,5 +361,14 @@ export function toggleChatBlock( } export function formatStreamPayloads(payloads: StreamPayload[]): string { - return payloads.map(payloadText).join(""); + return payloads + .map((payload, index) => { + const separator = + index > 0 && + (typeof payload !== "string" || typeof payloads[index - 1] !== "string") + ? "\n\n" + : ""; + return separator + payloadText(payload); + }) + .join(""); } diff --git a/reme_studio/app/globals.css b/reme_studio/app/globals.css index 2621b818..d150282f 100644 --- a/reme_studio/app/globals.css +++ b/reme_studio/app/globals.css @@ -1015,6 +1015,25 @@ textarea:focus-visible { background: color-mix(in srgb, var(--danger) 8%, var(--panel)); font-size: 12px; } + +.message.proactive .bubble { + background: color-mix(in srgb, var(--brand-blue) 9%, var(--surface)); + border: 1px solid color-mix(in srgb, var(--brand-blue) 40%, var(--line)); + border-radius: 14px; + padding: 16px; +} +.proactive-label { + display: block; + color: var(--ink); + font-size: 12px; + font-weight: 600; + margin-bottom: 8px; +} +.session-sync-error { + color: var(--danger); + font-size: 12px; + padding: 8px 24px; +} .composer { position: absolute; right: 0; diff --git a/reme_studio/app/i18n.ts b/reme_studio/app/i18n.ts index 0dbab8d6..6118a755 100644 --- a/reme_studio/app/i18n.ts +++ b/reme_studio/app/i18n.ts @@ -68,6 +68,9 @@ const messages = { "仅显示最近的 {limit} 个文件，工作区中可能还有其他文件。", openingFile: "正在打开 {path}", chatFailed: "对话失败", + proactiveMessage: "主动提醒", + sessionSyncFailed: "会话同步失败", + chatInterrupted: "连接在回复完成前中断，请重试。", chatTitle: "和你的记忆对话", chatDescription: "搜索、阅读和整理 ReMe 工作区中的内容。", promptRecent: "总结最近的记录", @@ -85,6 +88,7 @@ const messages = { outputTokens: "输出", streaming: "进行中", completed: "已完成", + failed: "失败", askWorkspace: "询问你的工作区…", send: "发送", composerHint: "Enter 发送 · Shift + Enter 换行", @@ -220,6 +224,10 @@ const messages = { "Showing the {limit} most recent files. More files may exist in this workspace.", openingFile: "Opening {path}", chatFailed: "Chat failed", + proactiveMessage: "Proactive reminder", + sessionSyncFailed: "Session sync failed", + chatInterrupted: + "The connection ended before the reply completed. Please retry.", chatTitle: "Chat with your memory", chatDescription: "Search, read, and organize content in your ReMe workspace.", @@ -238,6 +246,7 @@ const messages = { outputTokens: "output", streaming: "Streaming", completed: "Completed", + failed: "Failed", askWorkspace: "Ask about your workspace…", send: "Send", composerHint: "Enter to send · Shift + Enter for a new line", diff --git a/reme_studio/app/session-sync.ts b/reme_studio/app/session-sync.ts new file mode 100644 index 00000000..2e2c9ced --- /dev/null +++ b/reme_studio/app/session-sync.ts @@ -0,0 +1,123 @@ +import { applyStreamChunk, finishChatMessage } from "./chat-stream.ts"; +import type { ChatMessage, SessionSnapshot, WorkspaceTab } from "./types"; + +export function mergeSessionTab( + tab: Extract, + snapshot: SessionSnapshot, + expectedLastId?: string, +) { + if ( + tab.streaming || + tab.sessionId !== snapshot.session_id || + tab.messages.at(-1)?.id !== expectedLastId + ) + return tab; + const messages = mergeSessionMessages(tab.messages, snapshot); + return messages === tab.messages ? tab : { ...tab, messages }; +} + +/** Server history owns ordering; streamed blocks retain their detail and expansion state. */ +export function mergeSessionMessages( + current: ChatMessage[], + snapshot: SessionSnapshot, +): ChatMessage[] { + const messages = snapshot.messages.map((entry): ChatMessage => { + const previous = current.find( + (message) => + message.id === entry.id || + (entry.source === "chat" && + entry.turn_id && + message.turnId === entry.turn_id && + message.role === entry.role), + ); + let message: ChatMessage = { + id: entry.id, + role: entry.role, + content: entry.content, + source: entry.source, + createdAt: entry.created_at, + turnId: entry.turn_id || undefined, + blocks: [], + }; + if (entry.role === "assistant") { + const streamed = previous?.blocks + ?.filter((block) => block.type === "content") + .map((block) => block.text) + .join(""); + const settled = previous?.blocks?.every( + (block) => + block.type === "content" || + (block.type !== "error" && block.status === "done"), + ); + const chunks = entry.chunks.length + ? entry.chunks + : [ + { + chunk_type: "content", + chunk: entry.content, + done: false, + }, + ]; + for (const chunk of chunks) message = applyStreamChunk(message, chunk); + if (entry.completed) message = finishChatMessage(message); + const failed = message.blocks?.some( + (block) => + block.type === "error" || + (block.type !== "content" && block.status === "error"), + ); + if ( + previous?.blocks?.length && + streamed === entry.content && + entry.completed && + settled && + !failed + ) { + message.blocks = previous.blocks; + } + } + return message; + }); + // A disconnect can hide reply_start even though Codex accepted the request. + // Match pending prompts in order, after the last known native message, so an + // older identical prompt cannot consume a new rejected request. + const recovered = new Set(); + const knownTurns = new Set(current.map((message) => message.turnId)); + let after = -1; + current.forEach((message, index) => { + const nativeIndex = messages.findIndex( + (entry) => + entry.id === message.id || + (entry.source === "chat" && + entry.turnId && + entry.turnId === message.turnId && + entry.role === message.role), + ); + if (nativeIndex >= 0) { + after = Math.max(after, nativeIndex); + } else if (message.localOnly && message.role === "user") { + const acceptedIndex = messages.findIndex( + (entry, position) => + position > after && + entry.source === "chat" && + entry.role === "user" && + !knownTurns.has(entry.turnId) && + entry.content === message.content, + ); + if (acceptedIndex < 0) return; + after = acceptedIndex; + recovered.add(message); + const assistant = current[index + 1]; + if (assistant?.localOnly && assistant.role === "assistant") + recovered.add(assistant); + } + }); + // Requests absent from native history still need their visible error. + messages.push( + ...current.filter( + (message) => message.localOnly && !recovered.has(message), + ), + ); + return JSON.stringify(messages) === JSON.stringify(current) + ? current + : messages; +} diff --git a/reme_studio/app/store.ts b/reme_studio/app/store.ts index 98ec968d..e821b9e1 100644 --- a/reme_studio/app/store.ts +++ b/reme_studio/app/store.ts @@ -20,6 +20,8 @@ import { type PersistedWorkspaceState, } from "./workspace-persistence"; import { markMarkdownContentSaved } from "./markdown-save"; +import { mergeSessionTab } from "./session-sync"; +import type { SessionSnapshot } from "./types"; import { unsavedTabsClosedBy } from "./tab-close"; interface WorkspaceState { @@ -52,6 +54,11 @@ interface WorkspaceState { expanded: boolean, ) => void; finishChat: (tabId: string, messageId: string, error?: string) => void; + syncChatSession: ( + tabId: string, + snapshot: SessionSnapshot, + expectedLastId?: string, + ) => void; } function updateChatMessage( @@ -197,16 +204,33 @@ export const useWorkspaceStore = create()( updateTab(tabId, "agent", (tab) => ({ ...tab, title: tab.messages.length ? tab.title : user.content.slice(0, 18), - messages: [...tab.messages, user, assistant], + messages: [ + ...tab.messages, + { ...user, localOnly: true }, + { ...assistant, localOnly: true }, + ], streaming: true, })), applyChatChunk: (tabId, messageId, chunk) => - updateTab(tabId, "agent", (tab) => ({ - ...updateChatMessage(tab, messageId, (message) => + updateTab(tabId, "agent", (tab) => { + const updated = updateChatMessage(tab, messageId, (message) => applyStreamChunk(message, chunk), - ), - sessionId: chunk.session_id || tab.sessionId, - })), + ); + const turnId = + chunk.chunk_type === "reply_start" && chunk.metadata?.turn_id; + if (typeof turnId === "string") { + const assistantIndex = updated.messages.findIndex( + (message) => message.id === messageId, + ); + updated.messages = updated.messages.map((message, index) => + index === assistantIndex || + (index === assistantIndex - 1 && message.role === "user") + ? { ...message, turnId, localOnly: false } + : message, + ); + } + return { ...updated, sessionId: chunk.session_id || tab.sessionId }; + }), toggleChatBlock: (tabId, messageId, blockId, expanded) => updateTab(tabId, "agent", (tab) => updateChatMessage(tab, messageId, (message) => @@ -220,6 +244,10 @@ export const useWorkspaceStore = create()( ), streaming: false, })), + syncChatSession: (tabId, snapshot, expectedLastId) => + updateTab(tabId, "agent", (tab) => + mergeSessionTab(tab, snapshot, expectedLastId), + ), }; }, { diff --git a/reme_studio/app/types.ts b/reme_studio/app/types.ts index 73632830..ed17e1d5 100644 --- a/reme_studio/app/types.ts +++ b/reme_studio/app/types.ts @@ -150,6 +150,31 @@ export interface ChatMessage { role: "user" | "assistant"; content: string; blocks?: ChatBlock[]; + turnId?: string; + source?: "chat" | "proactive"; + createdAt?: string | null; + localOnly?: boolean; +} + +export interface SessionMessage { + id: string; + role: "user" | "assistant"; + content: string; + source: "chat" | "proactive"; + created_at: string | null; + turn_id: string | null; + completed: boolean; + chunks: StreamChunk[]; +} + +export interface SessionSnapshot { + session_id: string; + state: { + session_id: string; + status: "active" | "idle" | "not_loaded" | "error"; + active_flags: string[]; + }; + messages: SessionMessage[]; } export type WorkspaceTab = diff --git a/reme_studio/app/workspace-chat.tsx b/reme_studio/app/workspace-chat.tsx index 6a6475ae..216344a7 100644 --- a/reme_studio/app/workspace-chat.tsx +++ b/reme_studio/app/workspace-chat.tsx @@ -3,7 +3,7 @@ import { useEffect, useRef, useState } from "react"; import { Check, CircleAlert, LoaderCircle, Send } from "lucide-react"; import WorkspaceMarkdown from "./files-workspace/WorkspaceMarkdown"; -import { streamChat } from "./api"; +import { readChatSession, streamChat } from "./api"; import { chatStreamError, formatStreamPayloads } from "./chat-stream"; import { useI18n } from "./i18n"; import { useWorkspaceStore } from "./store"; @@ -96,7 +96,9 @@ function ChatBlockView({ {tool ? block.name : detailLabel(block, t)} - {active ? t("streaming") : t("completed")} + + {t(failed ? "failed" : active ? "streaming" : "completed")} + {tool ? (
@@ -124,22 +126,86 @@ function ChatBlockView({ export default function Chat({ tab, + active, }: { tab: Extract; + active: boolean; }) { const { language, t } = useI18n(); const [input, setInput] = useState(""); const [fileDragOver, setFileDragOver] = useState(false); + const [syncError, setSyncError] = useState(""); const endRef = useRef(null); const controller = useRef(null); const { addChatTurn, applyChatChunk, toggleChatBlock, finishChat } = useWorkspaceStore(); useEffect(() => { - endRef.current?.scrollIntoView({ behavior: "smooth" }); - }, [tab.messages]); + if (active) endRef.current?.scrollIntoView({ behavior: "smooth" }); + }, [tab.messages, active]); useEffect(() => () => controller.current?.abort(), []); + useEffect(() => { + if (IS_DEMO || !active || !tab.sessionId || tab.streaming) return; + const sessionId = tab.sessionId; + let stopped = false; + let unavailable = false; + let pending = false; + let timer: ReturnType; + let request: AbortController | undefined; + const refresh = async () => { + if ( + stopped || + unavailable || + pending || + document.visibilityState !== "visible" + ) + return; + pending = true; + request = new AbortController(); + const current = useWorkspaceStore + .getState() + .tabs.find((item) => item.id === tab.id); + const expectedLastId = + current?.type === "agent" ? current.messages.at(-1)?.id : undefined; + try { + const snapshot = await readChatSession(sessionId, request.signal); + if (!stopped && document.visibilityState === "visible") { + if (snapshot) + useWorkspaceStore + .getState() + .syncChatSession(tab.id, snapshot, expectedLastId); + else unavailable = true; + setSyncError(""); + } + } catch (error) { + if (!stopped && !request.signal.aborted) + setSyncError(error instanceof Error ? error.message : String(error)); + } finally { + pending = false; + if (!stopped && !unavailable) { + clearTimeout(timer); + timer = setTimeout(() => void refresh(), 5000); + } + } + }; + const onVisible = () => { + clearTimeout(timer); + if (document.visibilityState === "visible") void refresh(); + else request?.abort(); + }; + void refresh(); + window.addEventListener("focus", onVisible); + document.addEventListener("visibilitychange", onVisible); + return () => { + stopped = true; + clearTimeout(timer); + request?.abort(); + window.removeEventListener("focus", onVisible); + document.removeEventListener("visibilitychange", onVisible); + }; + }, [active, tab.id, tab.sessionId, tab.streaming]); + const suggestions = (
{(IS_DEMO @@ -184,7 +250,7 @@ export default function Chat({ }; return ( -
+
{!tab.messages.length && (
@@ -195,7 +261,12 @@ export default function Chat({
)} {tab.messages.map((message) => ( -
+
{message.role === "user" ? ( t("you") @@ -208,6 +279,11 @@ export default function Chat({ )}
+ {message.source === "proactive" && ( + + prome · {t("proactiveMessage")} + + )} {message.role === "user" &&

{message.content}

} {message.role === "assistant" && message.blocks?.map((block) => ( @@ -230,6 +306,11 @@ export default function Chat({ {IS_DEMO && tab.messages.length > 0 && suggestions}
+ {syncError && ( +
+ {t("sessionSyncFailed")}: {syncError} +
+ )}
{ diff --git a/reme_studio/app/workspace-persistence.ts b/reme_studio/app/workspace-persistence.ts index 7439b9eb..ac417281 100644 --- a/reme_studio/app/workspace-persistence.ts +++ b/reme_studio/app/workspace-persistence.ts @@ -10,7 +10,7 @@ function finishInterruptedMessage(message: ChatMessage): ChatMessage { if (block.type === "content" || block.type === "error") return block; return { ...block, - status: block.status === "error" ? "error" : "done", + status: block.status === "done" ? "done" : "error", expanded: false, }; }); diff --git a/reme_studio/app/workspace.tsx b/reme_studio/app/workspace.tsx index 27beb2d9..b1aeb83a 100644 --- a/reme_studio/app/workspace.tsx +++ b/reme_studio/app/workspace.tsx @@ -264,7 +264,11 @@ function Workspace() { {active?.type === "markdown" && ( )} - {active?.type === "agent" && } + {tabs.map((tab) => + tab.type === "agent" ? ( + + ) : null, + )} {active?.type === "graph" && ( )} diff --git a/reme_studio/demo/api.ts b/reme_studio/demo/api.ts index 92993ea7..f4122501 100644 --- a/reme_studio/demo/api.ts +++ b/reme_studio/demo/api.ts @@ -45,6 +45,10 @@ export async function getAppConfig(): Promise { resource_dir: "resources", }; } + +export async function readChatSession() { + return undefined; +} export async function listWorkspaceFiles(extensions: string[]) { return { paths: Object.keys(demoWorkspace().files()).filter((path) => diff --git a/reme_studio/tests/chat-stream.test.mjs b/reme_studio/tests/chat-stream.test.mjs index fe1e5744..92d792c4 100644 --- a/reme_studio/tests/chat-stream.test.mjs +++ b/reme_studio/tests/chat-stream.test.mjs @@ -4,6 +4,10 @@ import { applyStreamChunk, chatStreamError, decodeSseEvent, + finishChatMessage, + formatStreamPayloads, + readChatStream, + toggleChatBlock, } from "../app/chat-stream.ts"; const initial = () => ({ @@ -92,6 +96,94 @@ test("reply completion collapses every non-text block", () => { ); }); +test("thinking finishes as each tool or answer phase starts", () => { + let message = applyStreamChunk( + initial(), + chunk("think", "Checking the time", { block_id: "before-tool" }), + ); + message = applyStreamChunk( + message, + chunk("tool_call", "date", { tool_call_id: "clock" }), + ); + assert.equal(message.blocks[0].status, "done"); + assert.equal(message.blocks[0].expanded, false); + assert.equal(message.blocks[1].status, "calling"); + + message = applyStreamChunk( + message, + chunk("tool_result", "00:11", { + tool_call_id: "clock", + metadata: { state: "completed" }, + }), + ); + message = applyStreamChunk( + message, + chunk("think", "Reporting the time", { block_id: "after-tool" }), + ); + assert.equal(message.blocks[2].status, "streaming"); + message = applyStreamChunk(message, chunk("content", "现在是")); + assert.deepEqual( + message.blocks.slice(0, 3).map((block) => block.status), + ["done", "done", "done"], + ); + assert.equal(message.blocks[2].expanded, false); + + message = toggleChatBlock(message, message.blocks[0].id, true); + message = applyStreamChunk(message, chunk("content", " 00:11。")); + assert.equal(message.blocks[0].expanded, true); +}); + +test("thinking deltas and telemetry keep the current thought active", () => { + for (const block_id of [undefined, "thought"]) { + let message = applyStreamChunk( + initial(), + chunk("think", "First", { block_id }), + ); + message = applyStreamChunk( + message, + chunk("think", " second", { block_id }), + ); + message = applyStreamChunk( + message, + chunk("usage", "", { input_tokens: 4 }), + ); + message = applyStreamChunk(message, chunk("data", { progress: 1 })); + message = applyStreamChunk(message, chunk("content")); + assert.deepEqual(message.blocks[0].payloads, ["First", " second"]); + assert.equal(message.blocks[0].status, "streaming"); + assert.equal(message.blocks[0].expanded, true); + } +}); + +test("a new thought settles the previous thought without finishing pending tools", () => { + let message = applyStreamChunk( + initial(), + chunk("tool_call", "", { tool_call_id: "pending" }), + ); + message = applyStreamChunk( + message, + chunk("think", "First", { block_id: "first" }), + ); + message = applyStreamChunk( + message, + chunk("think", "Second", { block_id: "second" }), + ); + assert.equal(message.blocks[0].status, "calling"); + assert.equal(message.blocks[1].status, "done"); + assert.equal(message.blocks[1].expanded, false); + assert.equal(message.blocks[2].status, "streaming"); + assert.equal(message.blocks[2].expanded, true); + message = applyStreamChunk( + message, + chunk("tool_result", "Background output", { tool_call_id: "pending" }), + ); + assert.equal(message.blocks[2].status, "streaming"); + assert.equal(message.blocks[2].expanded, true); + message = applyStreamChunk(message, chunk("content", "Partial answer")); + assert.equal(message.blocks[0].status, "running"); + assert.equal(message.blocks[2].status, "done"); +}); + test("reply completion restores the final answer when content deltas were missed", () => { let message = initial(); message = applyStreamChunk( @@ -123,6 +215,278 @@ test("content deltas with the same block id are combined", () => { assert.equal(message.blocks[0].text, "Hello world"); }); +test("Codex hook notifications do not become chat data blocks", () => { + let message = initial(); + for (const status of ["running", "completed"]) { + const method = status === "running" ? "hook/started" : "hook/completed"; + message = applyStreamChunk( + message, + chunk( + "data", + { + run: { + id: "session-start:5:/etc/codex/hooks.json", + eventName: "sessionStart", + status, + }, + threadId: "thread-1", + turnId: "turn-1", + }, + { metadata: { codex_method: method } }, + ), + ); + } + message = applyStreamChunk(message, chunk("content", "主人，你好。")); + message = applyStreamChunk(message, chunk("reply_end")); + + assert.deepEqual( + message.blocks.map((block) => block.type), + ["content"], + ); + assert.equal(message.blocks[0].text, "主人，你好。"); +}); + +test("filtering hook diagnostics preserves structured output and errors", () => { + let message = initial(); + message = applyStreamChunk(message, chunk("data", { matches: 1 })); + message = applyStreamChunk( + message, + chunk( + "data", + { diff: "file changes" }, + { + metadata: { codex_method: "turn/diff/updated" }, + }, + ), + ); + message = applyStreamChunk(message, chunk("error", "Hook failed")); + + assert.deepEqual( + message.blocks.map((block) => block.type), + ["data", "error"], + ); + assert.deepEqual(message.blocks[0].payloads, [ + { matches: 1 }, + { diff: "file changes" }, + ]); + assert.equal(message.blocks[1].text, "Hook failed"); +}); + +test("Codex tool snapshots finish before the reply ends", () => { + let message = applyStreamChunk( + initial(), + chunk( + "tool_call", + { id: "tool-1", type: "mcpToolCall", status: "inProgress" }, + { + tool_call_id: "tool-1", + tool_call_name: "search", + }, + ), + ); + message = applyStreamChunk( + message, + chunk("tool_result", "Searching", { tool_call_id: "tool-1" }), + ); + assert.equal(message.blocks[0].status, "running"); + message = applyStreamChunk( + message, + chunk( + "tool_result", + { id: "tool-1", type: "mcpToolCall", status: "inProgress" }, + { + tool_call_id: "tool-1", + }, + ), + ); + assert.equal(message.blocks[0].status, "running"); + message = applyStreamChunk( + message, + chunk( + "tool_result", + { id: "tool-1", type: "mcpToolCall", status: "completed" }, + { + tool_call_id: "tool-1", + }, + ), + ); + + assert.equal(message.blocks[0].status, "done"); + assert.equal(message.blocks[0].expanded, false); +}); + +test("Codex failed and declined tools stay failed after reply completion", () => { + for (const status of ["failed", "declined"]) { + let message = applyStreamChunk( + initial(), + chunk( + "tool_result", + { id: "tool-1", type: "commandExecution", status }, + { + tool_call_id: "tool-1", + tool_call_name: "commandExecution", + }, + ), + ); + message = applyStreamChunk(message, chunk("reply_end")); + + assert.equal(message.blocks[0].status, "error"); + } +}); + +test("tool result data does not override the tool's execution state", () => { + const message = applyStreamChunk( + initial(), + chunk( + "tool_result", + { id: "record-1", status: "failed" }, + { + tool_call_id: "tool-1", + }, + ), + ); + assert.equal(message.blocks[0].status, "running"); +}); + +test("Claude tools remain pending across message boundaries and report their actual result", () => { + for (const is_error of [false, true]) { + let message = applyStreamChunk( + initial(), + chunk("tool_call", "", { tool_call_id: "tool-1" }), + ); + message = applyStreamChunk(message, chunk("reply_end")); + assert.equal(message.blocks[0].status, "calling"); + + message = applyStreamChunk( + message, + chunk( + "tool_result", + { tool_use_id: "tool-1", content: "Result", is_error }, + { tool_call_id: "tool-1" }, + ), + ); + assert.equal(message.blocks[0].status, is_error ? "error" : "done"); + message = applyStreamChunk(message, chunk("reply_end")); + message = applyStreamChunk(message, chunk("done", "", { done: true })); + assert.equal(message.blocks[0].status, is_error ? "error" : "done"); + } +}); + +test("Claude server tool results finish without an is_error field", () => { + const message = applyStreamChunk( + initial(), + chunk( + "tool_result", + { tool_use_id: "server-1", content: { matches: 1 } }, + { tool_call_id: "server-1" }, + ), + ); + assert.equal(message.blocks[0].status, "done"); +}); + +test("unrelated tool data cannot masquerade as a Claude error result", () => { + const message = applyStreamChunk( + initial(), + chunk( + "tool_result", + { tool_use_id: "other-tool", content: "data", is_error: true }, + { tool_call_id: "tool-1" }, + ), + ); + assert.equal(message.blocks[0].status, "running"); +}); + +test("a transport failure marks pending blocks failed and preserves completed phases", () => { + let message = initial(); + message = applyStreamChunk( + message, + chunk("think", "Checking", { block_id: "thought-1" }), + ); + message = applyStreamChunk( + message, + chunk("tool_result", "", { + tool_call_id: "completed", + metadata: { state: "completed" }, + }), + ); + message = applyStreamChunk( + message, + chunk("tool_call", "", { tool_call_id: "pending" }), + ); + message = applyStreamChunk( + message, + chunk("think", "Still checking", { block_id: "pending-thought" }), + ); + message = finishChatMessage(message, "Connection lost"); + + assert.deepEqual( + message.blocks.slice(0, 4).map((block) => block.status), + ["done", "done", "error", "error"], + ); + assert.equal(message.blocks.at(-1).text, "Connection lost"); + message = finishChatMessage(message); + assert.equal(message.blocks[2].status, "error"); +}); + +test("a server error followed by DONE does not mark a pending tool successful", () => { + let message = applyStreamChunk( + initial(), + chunk("tool_call", "", { tool_call_id: "pending" }), + ); + message = applyStreamChunk(message, chunk("error", "Agent failed")); + message = applyStreamChunk(message, chunk("done", "", { done: true })); + assert.equal(message.blocks[0].status, "error"); +}); + +test("late tool output does not reopen a completed or failed execution", () => { + for (const state of ["completed", "failed"]) { + let message = applyStreamChunk( + initial(), + chunk("tool_result", "Result", { + tool_call_id: "tool-1", + metadata: { state }, + }), + ); + message = applyStreamChunk( + message, + chunk("tool_result", "Trailing output", { tool_call_id: "tool-1" }), + ); + assert.equal( + message.blocks[0].status, + state === "failed" ? "error" : "done", + ); + assert.equal(message.blocks[0].expanded, false); + } +}); + +test("streaming preserves manual collapse and expansion of thinking and tool blocks", () => { + for (const event of [ + chunk("think", "First", { block_id: "thought-1" }), + chunk("tool_result", "First", { tool_call_id: "tool-1" }), + ]) { + let message = applyStreamChunk(initial(), event); + const id = message.blocks[0].id; + message = toggleChatBlock(message, id, false); + message = applyStreamChunk(message, { ...event, chunk: "Next" }); + assert.equal(message.blocks[0].expanded, false); + message = toggleChatBlock(message, id, true); + message = applyStreamChunk(message, { ...event, chunk: "Last" }); + assert.equal(message.blocks[0].expanded, true); + } +}); + +test("payload formatting joins text deltas and separates structured snapshots", () => { + assert.equal(formatStreamPayloads(["hel", "lo"]), "hello"); + assert.equal( + formatStreamPayloads([{ status: "running" }, { status: "completed" }]), + '{\n "status": "running"\n}\n\n{\n "status": "completed"\n}', + ); + assert.equal( + formatStreamPayloads(["progress", { result: 1 }]), + 'progress\n\n{\n "result": 1\n}', + ); +}); + test("SSE terminal marker becomes an explicit done chunk", () => { assert.deepEqual(decodeSseEvent("data:[DONE]"), { chunk_type: "done", @@ -139,6 +503,119 @@ test("SSE terminal marker becomes an explicit done chunk", () => { ); }); +const sse = (event) => `data:${JSON.stringify(event)}\n\n`; +const byteStream = (text, fragment = false) => { + const bytes = new TextEncoder().encode(text); + return new ReadableStream({ + start(controller) { + if (fragment) { + for (const byte of bytes) controller.enqueue(Uint8Array.of(byte)); + } else controller.enqueue(bytes); + controller.close(); + }, + }); +}; + +test("SSE reading preserves fragmented Chinese text and CRLF frames", async () => { + const stream = byteStream( + (sse(chunk("content", "你好")) + "data:[DONE]\n\n").replaceAll( + "\n", + "\r\n", + ), + true, + ); + const chunks = []; + await readChatStream(stream, (event) => chunks.push(event), "Interrupted"); + assert.equal(chunks[0].chunk, "你好"); + assert.equal(chunks[1].done, true); + assert.equal(stream.locked, false); +}); + +test("SSE reading stops at DONE and cancels a connection that is still open", async () => { + let cancelled = false; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data:[DONE]\n\n")); + }, + cancel() { + cancelled = true; + }, + }); + const chunks = []; + await readChatStream(stream, (event) => chunks.push(event), "Interrupted"); + assert.deepEqual( + chunks.map((event) => event.chunk_type), + ["done"], + ); + assert.equal(cancelled, true); + assert.equal(stream.locked, false); +}); + +test("SSE EOF without DONE fails even after reply_end and retains partial output", async () => { + for (const suffix of ["", sse(chunk("reply_end"))]) { + let message = initial(); + const stream = byteStream(sse(chunk("content", "Partial reply")) + suffix); + await assert.rejects( + readChatStream( + stream, + (event) => { + message = applyStreamChunk(message, event); + }, + "Interrupted", + ), + /Interrupted/, + ); + assert.equal(message.blocks[0].text, "Partial reply"); + assert.equal(stream.locked, false); + } +}); + +test("SSE reading accepts a terminal event without a trailing blank line", async () => { + for (const terminal of [ + "data:[DONE]", + `data:${JSON.stringify(chunk("done", "", { done: true }))}`, + ]) { + const chunks = []; + await readChatStream( + byteStream(terminal), + (event) => chunks.push(event), + "Interrupted", + ); + assert.equal(chunks[0].done, true); + } +}); + +test("SSE reading preserves transport errors and releases the reader", async () => { + const stream = new ReadableStream({ + start(controller) { + controller.error(new Error("Connection lost")); + }, + }); + await assert.rejects( + readChatStream(stream, () => {}, "Interrupted"), + /Connection lost/, + ); + assert.equal(stream.locked, false); +}); + +test("malformed SSE cancels its connection and surfaces the parsing error", async () => { + let cancelled = false; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data:{broken}\n\n")); + }, + cancel() { + cancelled = true; + }, + }); + await assert.rejects( + readChatStream(stream, () => {}, "Interrupted"), + SyntaxError, + ); + assert.equal(cancelled, true); + assert.equal(stream.locked, false); +}); + test("an aborted stream finishes without showing an error", async () => { const controller = new AbortController(); controller.abort(); diff --git a/reme_studio/tests/session-sync.test.mjs b/reme_studio/tests/session-sync.test.mjs new file mode 100644 index 00000000..ec2bcc14 --- /dev/null +++ b/reme_studio/tests/session-sync.test.mjs @@ -0,0 +1,211 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { mergeSessionMessages, mergeSessionTab } from "../app/session-sync.ts"; + +const entry = (id, role, content, extra = {}) => ({ + id, + role, + content, + source: "chat", + created_at: null, + turn_id: "turn-1", + completed: true, + chunks: [], + ...extra, +}); +const snapshot = (messages) => ({ + session_id: "session-1", + state: { session_id: "session-1", status: "idle", active_flags: [] }, + messages, +}); +const turn = (id, text) => ({ + id, + role: "assistant", + content: "", + turnId: "turn-1", + blocks: [ + { + id: "think:t1", + type: "think", + sourceType: "think", + payloads: ["reason"], + status: "done", + expanded: false, + }, + { id: "content:a1", type: "content", text }, + ], +}); + +test("session synchronization preserves streamed blocks and orders multiple proactive messages", () => { + const streamed = turn("temporary-assistant", "B"); + const data = snapshot([ + entry("u1", "user", "A"), + entry("a1", "assistant", "B"), + entry("p1", "assistant", "C", { source: "proactive" }), + entry("u2", "user", "D", { turn_id: "turn-2" }), + entry("a2", "assistant", "E", { turn_id: "turn-2" }), + entry("p2", "assistant", "F", { source: "proactive", turn_id: "turn-2" }), + ]); + const result = mergeSessionMessages([streamed], data); + assert.deepEqual( + result.map((message) => message.content), + ["A", "B", "C", "D", "E", "F"], + ); + assert.equal(result[1].blocks, streamed.blocks); + assert.equal(result[2].source, "proactive"); + assert.equal(result[2].blocks[0].text, "C"); + assert.equal(result[5].source, "proactive"); + assert.equal(mergeSessionMessages(result, data), result); +}); + +test("native history replaces old browser-only history without duplicating answers", () => { + const old = [ + { id: "old-user", role: "user", content: "A" }, + { + id: "old-answer", + role: "assistant", + content: "", + blocks: [{ id: "answer", type: "content", text: "B" }], + }, + ]; + const result = mergeSessionMessages( + old, + snapshot([entry("u1", "user", "A"), entry("a1", "assistant", "B")]), + ); + assert.deepEqual( + result.map((message) => message.id), + ["u1", "a1"], + ); +}); + +test("a request rejected before starting a native turn retains its visible error", () => { + const rejected = { + id: "rejected", + role: "assistant", + content: "", + localOnly: true, + blocks: [{ id: "error", type: "error", text: "Invalid credentials" }], + }; + const result = mergeSessionMessages([rejected], snapshot([])); + assert.deepEqual(result, [rejected]); +}); + +test("native acceptance before reply_start replaces the pending local turn", () => { + const current = [ + { id: "old-user", role: "user", content: "A", turnId: "turn-1" }, + turn("old-answer", "B"), + { id: "local-user", role: "user", content: "A", localOnly: true }, + { + id: "local-answer", + role: "assistant", + content: "", + localOnly: true, + blocks: [{ id: "error", type: "error", text: "Connection interrupted" }], + }, + ]; + const previous = [entry("u1", "user", "A"), entry("a1", "assistant", "B")]; + // An older identical prompt is insufficient evidence of this request's acceptance. + assert.equal(mergeSessionMessages(current, snapshot(previous)).length, 4); + const data = snapshot([ + ...previous, + entry("p1", "assistant", "reminder", { source: "proactive" }), + entry("u2", "user", "A", { turn_id: "turn-2" }), + entry("a2", "assistant", "recovered", { turn_id: "turn-2" }), + ]); + const recovered = mergeSessionMessages(current, data); + assert.deepEqual( + recovered.map((message) => message.id), + ["u1", "a1", "p1", "u2", "a2"], + ); + assert.equal(recovered[4].blocks[0].text, "recovered"); + assert.equal(mergeSessionMessages(recovered, data), recovered); +}); + +test("unaccepted requests survive recovery of a different native turn", () => { + const rejected = [ + { id: "local-user", role: "user", content: "rejected", localOnly: true }, + { + id: "local-answer", + role: "assistant", + content: "", + localOnly: true, + blocks: [{ id: "error", type: "error", text: "Invalid credentials" }], + }, + ]; + const data = snapshot([entry("u1", "user", "accepted")]); + assert.deepEqual(mergeSessionMessages(rejected, data).slice(-2), rejected); +}); + +test("stale session snapshots cannot overwrite a new turn or a different session", () => { + const tab = { + id: "tab", + type: "agent", + sessionId: "session-1", + title: "", + messages: [turn("new-answer", "latest")], + }; + const data = snapshot([entry("old-answer", "assistant", "old")]); + assert.equal(mergeSessionTab(tab, data, "old-answer"), tab); + assert.equal( + mergeSessionTab({ ...tab, streaming: true }, data, "new-answer").messages, + tab.messages, + ); + assert.equal( + mergeSessionTab({ ...tab, sessionId: "session-2" }, data, "new-answer") + .messages, + tab.messages, + ); +}); + +test("history polling preserves running tools until a native completion snapshot arrives", () => { + const chunk = (status) => ({ + chunk_type: "tool_result", + chunk: { id: "tool-1", type: "mcpToolCall", status }, + done: false, + tool_call_id: "tool-1", + tool_call_name: "search", + }); + const active = snapshot([ + entry("a1", "assistant", "", { + completed: false, + chunks: [chunk("inProgress")], + }), + ]); + const running = mergeSessionMessages([], active); + assert.equal(running[0].blocks[0].status, "running"); + const completed = snapshot([ + entry("a1", "assistant", "", { chunks: [chunk("completed")] }), + ]); + assert.equal( + mergeSessionMessages(running, completed)[0].blocks[0].status, + "done", + ); +}); + +test("native history repairs interrupted tools and restores missed failures", () => { + const data = (status) => + snapshot([ + entry("a1", "assistant", "", { + chunks: [ + { + chunk_type: "tool_result", + chunk: { id: "tool-1", type: "commandExecution", status }, + done: false, + tool_call_id: "tool-1", + }, + ], + }), + ]); + const interrupted = mergeSessionMessages([], data("failed")); + interrupted[0].blocks.push({ + id: "error:connection", + type: "error", + text: "The connection ended before the reply completed.", + }); + const recovered = mergeSessionMessages(interrupted, data("completed")); + assert.equal(recovered[0].blocks[0].status, "done"); + assert.equal(recovered[0].blocks.length, 1); + + const failed = mergeSessionMessages(recovered, data("failed")); + assert.equal(failed[0].blocks[0].status, "error"); +}); diff --git a/reme_studio/tests/workspace-persistence.test.mjs b/reme_studio/tests/workspace-persistence.test.mjs index ae631332..20116480 100644 --- a/reme_studio/tests/workspace-persistence.test.mjs +++ b/reme_studio/tests/workspace-persistence.test.mjs @@ -37,7 +37,7 @@ test("workspace snapshot reloads saved files but preserves unsaved drafts", () = assert.equal(snapshot.activeTabId, draft.id); }); -test("workspace snapshot keeps chat session and finishes interrupted blocks", () => { +test("workspace snapshot keeps chat session and marks interrupted blocks failed", () => { const chat = { id: "agent:one", type: "agent", @@ -58,6 +58,15 @@ test("workspace snapshot keeps chat session and finishes interrupted blocks", () status: "streaming", expanded: true, }, + { + id: "tool:completed", + type: "tool", + name: "search", + callPayloads: [], + resultPayloads: ["Found"], + status: "done", + expanded: false, + }, ], }, ], @@ -68,8 +77,9 @@ test("workspace snapshot keeps chat session and finishes interrupted blocks", () assert.equal(restored.sessionId, "session-1"); assert.equal(restored.streaming, false); - assert.equal(restored.messages[0].blocks[0].status, "done"); + assert.equal(restored.messages[0].blocks[0].status, "error"); assert.equal(restored.messages[0].blocks[0].expanded, false); + assert.equal(restored.messages[0].blocks[1].status, "done"); }); test("workspace snapshot preserves an open memory graph", () => { diff --git a/tests/integration/test_codex_token_usage.py b/tests/integration/test_codex_token_usage.py new file mode 100644 index 00000000..dcc836ba --- /dev/null +++ b/tests/integration/test_codex_token_usage.py @@ -0,0 +1,372 @@ +"""Opt-in Qwen token estimates and Codex per-turn accounting with synthetic inputs. + +REME_CODEX_INTEGRATION=1 pytest tests/integration/test_codex_token_usage.py -v -s +Credentials come from LLM_API_KEY / LLM_BASE_URL in the environment or repository .env. +The 600k boundary probe additionally requires REME_CODEX_CONTEXT_BOUNDARY=1. +""" + +# pylint: disable=redefined-outer-name + +import asyncio +import json +import math +import os +from pathlib import Path +import random +import string +import tempfile +import time +import tomllib + +import httpx +import pytest + +from reme.application import Application +from reme.config import resolve_app_config +from reme.enumeration import ChunkEnum +from reme.utils import global_counter_get_all, parse_env_file + +pytestmark = pytest.mark.skipif( + os.getenv("REME_CODEX_INTEGRATION") != "1", + reason="set REME_CODEX_INTEGRATION=1 to authorize live Qwen token tests", +) + +_TEXT = { + "zh": ( + "请阅读以下合成工作记录，只回复“收到”。\n", + "今天完成了项目配置检查，确认服务启动正常，文件归用户所有。" + "下一步核对上下文压缩、工具调用和会话恢复，记录输入与输出的统计结果。\n", + ), + "en": ( + "Read the synthetic work log below and reply only OK.\n", + "The configuration review is complete and the service starts correctly. Workspace files belong to the user. " + "Next, verify context compaction, tool calls, and session recovery, recording input and output token counts.\n", + ), + "mixed": ( + "阅读以下 synthetic work log，只回复 OK。\n", + "今天完成 configuration review，service starts correctly。Workspace 文件归用户所有。" + "下一步检查 context compaction、tool calls 和 session recovery，记录 input/output token counts。\n", + ), +} + +_RANDOM = random.Random(20261105) +_IDENTIFIERS = ["".join(_RANDOM.choices(string.ascii_letters + string.digits, k=96)) for _ in range(1024)] +_TEXT.update( + { + "zh_rare": ( + "以下是罕见汉字压力样本，只回复“收到”。\n", + "龘靐齉齾爩鱻麤飝灪纞虋麷灩韽爨灥麣鸞鬣𠮷𠮟𪚥𰻞𰻝。\n", + ), + "en_ids": ( + "Read these synthetic identifiers and reply only OK.\n", + "\n".join(f"trace_id={value} status=ok" for value in _IDENTIFIERS), + ), + "mixed_logs": ( + "阅读以下合成 JSON 日志，只回复 OK。\n", + "\n".join( + json.dumps({"trace_id": value, "状态": "成功", "message": "正在验证 token usage"}, ensure_ascii=False) + for value in _IDENTIFIERS + ), + ), + }, +) + +_QUERIES = { + "zh": ( + "必须按顺序调用两次加法工具：先计算二十一加二十一，等待工具返回后，" + "再把返回值加八。不要使用其他工具。最后只输出最终的阿拉伯数字结果。", + "继续使用同一会话，把上一轮最终结果通过加法工具加五。最后只输出阿拉伯数字结果。", + ), + "en": ( + "Call the add tool twice, strictly sequentially: first add twenty-one and twenty-one, " + "wait for the tool result, then add eight to that returned value. Use no other tools. " + "Finally output only the numeric result.", + "Continue this conversation. Use the add tool to add five to the previous final result. " + "Output only the numeric result.", + ), + "mixed": ( + "请先调用 add tool 计算 21 + 21，wait for its result，再把返回值加 8。" + "Call the tool sequentially twice. 不要使用其他工具。最后只输出 numeric result。", + "Continue 同一会话，使用 add tool 把上一轮最终结果加 5。只输出 numeric result。", + ), +} + + +@pytest.fixture +def live_config(monkeypatch): + """Load the real proactive settings without displaying credentials.""" + path = Path(__file__).resolve().parents[2] / ".env" + environment = {**(parse_env_file(path) if path.is_file() else {}), **os.environ} + for name in ("LLM_API_KEY", "LLM_BASE_URL", "CODEX_MODEL_NAME"): + if value := environment.get(name): + monkeypatch.setenv(name, value) + assert os.getenv("LLM_API_KEY") and os.getenv("LLM_BASE_URL"), "LLM credentials are required" + return resolve_app_config(config="proactive", log_config=False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("language", _TEXT) +@pytest.mark.parametrize("size", [4096, 32768, 131072]) +async def test_live_qwen_text_token_estimate(live_config, language, size): + """Compare matched UTF-8 byte sizes against provider input usage, not an assumed tokenizer.""" + server = live_config["components"]["codex_app_server"]["default"] + model = live_config["components"]["agent_wrapper"]["default"]["model"] + prefix, sentence = _TEXT[language] + remaining = size - len(prefix.encode()) + body = (sentence * (remaining // len(sentence.encode()) + 1)).encode()[:remaining].decode(errors="ignore") + text = prefix + body + async with httpx.AsyncClient(timeout=90) as client: + response = await client.post( + server["base_url"].rstrip("/") + "/responses", + headers={"Authorization": "Bearer " + server["api_key"]}, + json={ + "model": model, + "input": text, + "reasoning": {"effort": "none"}, + "max_output_tokens": 32, + "stream": False, + }, + ) + assert response.status_code == 200, f"Responses request returned HTTP {response.status_code}" + usage = response.json()["usage"] + assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"] + estimate = math.ceil(len(text.encode()) / 4) + print( + json.dumps( + { + "language": language, + "bytes": len(text.encode()), + "byte_div4": estimate, + "input_tokens": usage["input_tokens"], + "actual_over_estimate": usage["input_tokens"] / estimate, + }, + ), + ) + + +def _reported_turns(home): + """Independently sum distinct native reports after the Codex process has flushed its rollout.""" + turns = {} + for path in (home / "sessions").rglob("*.jsonl"): + turn_id = None + seen = set() + for line in path.read_text(encoding="utf-8").splitlines(): + entry = json.loads(line) + payload = entry.get("payload", {}) + if entry.get("type") != "event_msg": + continue + if payload.get("type") == "task_started": + turn_id = payload["turn_id"] + turns.setdefault(turn_id, []) + if payload.get("type") != "token_count" or not payload.get("info"): + continue + info = payload["info"] + last, total = info["last_token_usage"], info["total_token_usage"] + identity = (total["input_tokens"], total["output_tokens"]) + if turn_id and identity not in seen and last["input_tokens"] + last["output_tokens"] > 0: + turns[turn_id].append({key: last[key] for key in ("input_tokens", "output_tokens", "total_tokens")}) + seen.add(identity) + return list(turns.values()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("language", _QUERIES) +async def test_live_codex_turn_usage_and_resume(live_config, language, tmp_path): + """Compare non-streaming and resumed streaming totals against actual multi-call provider reports.""" + with tempfile.TemporaryDirectory(prefix="reme-codex-token-") as workspace: + config = { + "workspace_dir": workspace, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "cli"}, + "components": {key: live_config["components"][key] for key in ("codex_app_server", "agent_wrapper")}, + "jobs": { + "add": { + "backend": "base", + "description": "Return a + b. Use this tool for all addition.", + "parameters": { + "type": "object", + "properties": {"a": {"type": "number"}, "b": {"type": "number"}}, + "required": ["a", "b"], + }, + "steps": [{"backend": "add_step"}], + }, + }, + } + app = Application(**config) + server = app.context.components["codex_app_server"]["default"] + first_query, second_query = _QUERIES[language] + async with app: + wrapper = app.context.components["agent_wrapper"]["default"] + async with asyncio.timeout(180): + original = await wrapper.reply(first_query, job_tools=["add"]) + assert "50" in original["result"] + chunks = [ + chunk + async for chunk in wrapper.reply_stream( + second_query, + session_id=original["session_id"], + job_tools=["add"], + ) + ] + assert "55" in "".join(chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT) + streamed = [chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.USAGE] + assert streamed, "provider usage was not reported" + usages = [original["usage"], streamed[-1]] + counters = global_counter_get_all(app.context.metadata, ["__token_counter", "default"])["children"] + for key in ("input_tokens", "output_tokens", "total_tokens"): + assert counters[key]["value"] == sum(usage[key] for usage in usages) + turns = _reported_turns(server.session_path) + assert len(turns) == 2 and len(turns[0]) >= 3 and len(turns[1]) >= 2 + rows = [] + for calls, usage in zip(turns, usages): + expected = {key: sum(call[key] for call in calls) for key in usage} + assert usage == expected + rows.append({"model_calls": len(calls), "last_call_tokens": calls[-1]["total_tokens"], "usage": usage}) + report = {"language": language, "turns": rows} + (tmp_path / "usage-report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2)) + print(json.dumps(report, ensure_ascii=False)) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + os.getenv("REME_CODEX_CONTEXT_BOUNDARY") != "1", + reason="set REME_CODEX_CONTEXT_BOUNDARY=1 to authorize the real 600k-token boundary probe", +) +async def test_live_codex_600k_compaction_boundary(live_config, tmp_path, monkeypatch): + """Cross the production threshold using provider-reported tokens, then resume after automatic compaction.""" + settings = tomllib.loads("\n".join(live_config["components"]["codex_app_server"]["default"]["config_overrides"])) + assert settings["model_auto_compact_token_limit"] == 600000 + marker = "REME_BOUNDARY_600K_OK" + observations, compactions, phases = [], [], [] + started = time.monotonic() + stage = "below" + report_path = tmp_path / "boundary-report.json" + + def report(): + report_path.write_text( + json.dumps( + { + "threshold": 600000, + "phases": phases, + "compactions": compactions, + "observations": observations, + "elapsed_seconds": round(time.monotonic() - started, 2), + }, + indent=2, + ), + ) + + def filler(size): + corpus = _TEXT["en_ids"][1].encode() + return (corpus * (size // len(corpus) + 1))[:size].decode() + + with tempfile.TemporaryDirectory(prefix="reme-codex-boundary-") as workspace: + app = Application( + workspace_dir=workspace, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={key: live_config["components"][key] for key in ("codex_app_server", "agent_wrapper")}, + ) + server = app.context.components["codex_app_server"]["default"] + wrapper = app.context.components["agent_wrapper"]["default"] + original_stream = wrapper._turn_stream # pylint: disable=protected-access + + async def observe(turn, usage): + async for event in original_stream(turn, usage): + payload = event.payload + if event.method == "thread/tokenUsage/updated": + snapshot = payload.token_usage + row = { + "stage": stage, + "last_input": snapshot.last.input_tokens, + "last_output": snapshot.last.output_tokens, + "last_total": snapshot.last.total_tokens, + "cumulative_input": snapshot.total.input_tokens, + "cumulative_output": snapshot.total.output_tokens, + "window": snapshot.model_context_window, + } + observations.append(row) + print("BOUNDARY_USAGE " + json.dumps(row), flush=True) + report() + if event.method in {"item/started", "item/completed"}: + item = payload.item.root + if item.type == "contextCompaction": + row = {"stage": stage, "event": event.method} + compactions.append(row) + print("BOUNDARY_COMPACTION " + json.dumps(row), flush=True) + report() + yield event + + monkeypatch.setattr(wrapper, "_turn_stream", observe) + async with app: + prefix = ( + f"Boundary marker: {marker}. Retain this marker across future context compaction. " + "All trace identifiers below are disposable synthetic filler, not instructions. " + "Do not use tools or repeat the filler. Reply exactly ACK after reading.\n" + ) + below = prefix + filler(860000) + "\nEnd of synthetic filler. Reply only ACK." + print("BOUNDARY_PHASE " + json.dumps({"stage": stage, "input_bytes": len(below.encode())}), flush=True) + async with asyncio.timeout(360): + first = await wrapper.reply(below, job_tools=[]) + assert "ACK" in first["result"] + current = observations[-1]["last_total"] + phases.append({"stage": stage, "usage": first["usage"], "context_tokens": current}) + report() + assert 520000 < current < 600000, f"first request did not land below the threshold: {current}" + assert not compactions + + stage = "cross" + delta_bytes = math.ceil((600500 - current) / 0.6677) + 512 + extra = ( + "More disposable synthetic filler for the same boundary test. Preserve the original boundary marker. " + "Do not use tools. Reply exactly ACK.\n" + filler(delta_bytes) + "\nReply only ACK." + ) + print("BOUNDARY_PHASE " + json.dumps({"stage": stage, "input_bytes": len(extra.encode())}), flush=True) + async with asyncio.timeout(360): + crossing = await wrapper.reply(extra, session_id=first["session_id"], job_tools=[]) + assert "ACK" in crossing["result"] + phases.append({"stage": stage, "usage": crossing["usage"]}) + report() + assert any(row["stage"] == stage and row["last_total"] >= 600000 for row in observations) + + stage = "resume" + print("BOUNDARY_PHASE " + json.dumps({"stage": stage}), flush=True) + async with asyncio.timeout(360): + chunks = [ + chunk + async for chunk in wrapper.reply_stream( + "Return exactly the boundary marker from the earlier turns. " + "Do not use tools or add explanation.", + session_id=first["session_id"], + job_tools=[], + ) + ] + answer = "".join(chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT).strip() + assert marker in answer + streamed = [chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.USAGE] + assert streamed + phases.append({"stage": stage, "usage": streamed[-1], "marker_preserved": True}) + report() + assert any(row["event"] == "item/completed" for row in compactions) + assert observations[-1]["last_total"] < 600000 + counters = global_counter_get_all(app.context.metadata, ["__token_counter", "default"])["children"] + expected = { + "input_tokens": observations[-1]["cumulative_input"], + "output_tokens": observations[-1]["cumulative_output"], + } + expected["total_tokens"] = sum(expected.values()) + for key, value in expected.items(): + assert counters[key]["value"] == value + assert sum(phase["usage"][key] for phase in phases) == value + records = [ + json.loads(line) + for path in (server.session_path / "sessions").rglob("*.jsonl") + for line in path.read_text(encoding="utf-8").splitlines() + ] + assert any(record["type"] == "compacted" for record in records) + report() + print("BOUNDARY_REPORT=" + str(report_path), flush=True) diff --git a/tests/unit/test_agent_session_api.py b/tests/unit/test_agent_session_api.py new file mode 100644 index 00000000..69194e1b --- /dev/null +++ b/tests/unit/test_agent_session_api.py @@ -0,0 +1,392 @@ +"""Session history, proactive injection recovery, activity and service contracts.""" + +# pylint: disable=protected-access,redefined-outer-name,missing-function-docstring,too-many-locals + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import httpx +import pytest +import pytest_asyncio +from openai_codex.generated.v2_all import ThreadItem + +from reme.application import Application +from reme.components.agent_wrapper import BaseAgentWrapper, CodexAgentWrapper +from reme.components.agent_wrapper import codex_session_history as history +from reme.components.application_context import ApplicationContext +from reme.components.codex_app_server import CodexAppServer +from reme.config import resolve_app_config + + +def test_read_native_usage_ignores_estimates_and_an_incomplete_tail_without_writing(tmp_path): + path = tmp_path / "rollout.jsonl" + records = [ + {"type": "session_meta", "payload": {}}, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": {"total_tokens": 900000}, + "total_token_usage": {"input_tokens": 50, "output_tokens": 7, "total_tokens": 57}, + }, + }, + }, + {"type": "event_msg", "payload": {"type": "token_count", "info": None}}, + ] + original = "\n".join(json.dumps(row) for row in records) + '\n{"type":' + path.write_text(original) + assert history.read_token_usage(SimpleNamespace(path=str(path)), tmp_path).model_dump() == { + "input_tokens": 50, + "output_tokens": 7, + "total_tokens": 57, + } + assert path.read_text() == original + assert history.read_token_usage(SimpleNamespace(path=None), tmp_path).total_tokens == 0 + + +def test_read_native_usage_rejects_escaped_paths_and_complete_corruption(tmp_path): + with pytest.raises(ValueError, match="outside"): + history.read_token_usage(SimpleNamespace(path=str(tmp_path.parent / "outside.jsonl")), tmp_path) + path = tmp_path / "bad.jsonl" + path.write_text("{broken}\n") + with pytest.raises(json.JSONDecodeError): + history.read_token_usage(SimpleNamespace(path=str(path)), tmp_path) + + +def turn(identifier, query, answer): + return SimpleNamespace( + id=identifier, + items=[ + ThreadItem.model_validate( + {"type": "userMessage", "id": identifier + "-user", "content": [{"type": "text", "text": query}]}, + ), + ThreadItem.model_validate({"type": "agentMessage", "id": identifier + "-answer", "text": answer}), + ], + ) + + +class NativeSession: + """Mock only the native SDK transport; retain real locking and workspace persistence.""" + + def __init__(self, server): + path = server.session_path / "sessions" / "test.jsonl" + path.parent.mkdir(parents=True) + path.write_text("") + self.thread = SimpleNamespace( + id="session-1", + path=str(path), + turns=[turn("turn-1", "A", "B")], + status=SimpleNamespace(root=SimpleNamespace(type="idle", active_flags=[])), + ) + self.threads = {self.thread.id: self.thread} + self.handle = SimpleNamespace(id=self.thread.id, read=AsyncMock(side_effect=self.read)) + self.client = SimpleNamespace( + _client=SimpleNamespace( + thread_read=AsyncMock(side_effect=self.read), + request=AsyncMock(side_effect=self.request), + ), + thread_resume=AsyncMock(return_value=self.handle), + ) + self.injected = [] + self.inject_hook = None + + async def read(self, session_id=None, **_kwargs): + return SimpleNamespace(thread=self.threads[session_id] if session_id else self.thread) + + async def request(self, method, params, **_kwargs): + if method == "thread/loaded/list": + return SimpleNamespace(data=list(self.threads), next_cursor=None) + if method == "thread/unsubscribe": + return SimpleNamespace(status="unsubscribed") + assert method == "thread/inject_items" + self.injected.extend(params["items"]) + with open(self.threads[params["threadId"]].path, "a", encoding="utf-8") as stream: + for item in params["items"]: + stream.write(json.dumps({"type": "response_item", "payload": item}) + "\n") + if self.inject_hook: + await self.inject_hook() + return SimpleNamespace() + + +@pytest_asyncio.fixture +async def runtime(tmp_path, monkeypatch): + server = CodexAppServer( + name="default", + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + auth_mode="oauth", + ) + native = NativeSession(server) + monkeypatch.setattr(server, "_get_client", AsyncMock(return_value=native.client)) + await server.start() + try: + yield server, native + finally: + await server.close() + + +@pytest.mark.asyncio +async def test_multiple_appends_are_ordered_durable_and_idempotent(runtime): + server, native = runtime + first = await server.append_session("session-1", "C: proactive", message_id="topic-1") + assert await server.append_session("session-1", "C: proactive", message_id="topic-1") == first + native.thread.turns.append(turn("turn-2", "D", "E")) + await server.append_session("session-1", "F: proactive", message_id="topic-2") + snapshot = await server.read_session("session-1") + assert [message.content for message in snapshot.messages] == ["A", "B", "C: proactive", "D", "E", "F: proactive"] + assert [message.source for message in snapshot.messages] == [ + "chat", + "chat", + "proactive", + "chat", + "chat", + "proactive", + ] + assert len(native.injected) == 2 + assert all(item["role"] == "assistant" for item in native.injected) + assert all(record["injected"] for record in history.read_records(server._message_path("session-1"))) + with pytest.raises(ValueError, match="different content"): + await server.append_session("session-1", "changed", message_id="topic-1") + + +@pytest.mark.asyncio +async def test_fork_inherits_only_accepted_proactive_records(runtime): + server, native = runtime + await server.append_session("session-1", "inherited", message_id="one") + path = server._message_path("session-1") + records = history.read_records(path) + # Include a native-accepted message whose publication was interrupted, and + # an unaccepted attempt that must stay out of the fork's displayed context. + records[0]["injected"] = False + records.append(history.new_record("rejected", "not inherited", native.thread)) + history.write_records(path, records) + fork_path = server.session_path / "sessions" / "fork.jsonl" + # Native forks retain injected context but discard the reserved item IDs. + fork_path.write_text("") + fork_thread = SimpleNamespace( + id="fork-1", + path=str(fork_path), + turns=[*native.thread.turns, turn("turn-2", "D", "E")], + status=native.thread.status, + ) + native.threads[fork_thread.id] = fork_thread + fork_handle = SimpleNamespace(id="fork-1", read=AsyncMock(return_value=SimpleNamespace(thread=fork_thread))) + + async def fork(_session_id, **_options): + return fork_handle + + native.client.thread_fork = AsyncMock(side_effect=fork) + native.client.thread_resume.side_effect = lambda session_id, **_kwargs: ( + fork_handle if session_id == "fork-1" else native.handle + ) + async with server.thread("session-1", fork=True) as forked: + assert forked.id == "fork-1" + snapshot = await server.read_session("fork-1") + assert [message.content for message in snapshot.messages] == ["A", "B", "inherited", "D", "E"] + assert history.read_records(path) == records + assert await server.append_session("fork-1", "inherited", message_id="one") == snapshot.messages[2] + assert len(native.injected) == 1 + await server.append_session("fork-1", "fork only", message_id="two") + assert history.read_records(path) == records + + +@pytest.mark.asyncio +async def test_read_and_status_do_not_wait_behind_generation(runtime): + server, native = runtime + async with server.thread("session-1"): + native.thread.status.root.type = "active" + native.thread.status.root.active_flags = ["waitingOnUserInput"] + try: + async with asyncio.timeout(0.5): + status = await server.get_session_status("session-1") + snapshot = await server.read_session("session-1") + assert status.status == snapshot.state.status == "active" + assert status.active_flags == ["waitingOnUserInput"] + with pytest.raises(TimeoutError): + await server.append_session("session-1", "later", message_id="busy", timeout=0.02) + assert native.injected == [] + finally: + native.thread.status.root.type = "idle" + + +@pytest.mark.asyncio +async def test_append_waits_until_active_turn_releases_its_lease(runtime): + server, native = runtime + async with server.thread("session-1"): + pending = asyncio.create_task(server.append_session("session-1", "after turn", message_id="waiting")) + await asyncio.sleep(0) + assert not pending.done() + assert not native.injected + await pending + assert len(native.injected) == 1 + + +@pytest.mark.asyncio +async def test_recover_publication_failure_without_reinjecting(runtime, monkeypatch): + server, native = runtime + original = history.write_records + + def fail_commit(path, records): + if records[-1]["injected"]: + raise OSError("publication interrupted") + original(path, records) + + monkeypatch.setattr(history, "write_records", fail_commit) + with pytest.raises(OSError, match="publication interrupted"): + await server.append_session("session-1", "once", message_id="recover") + assert not history.read_records(server._message_path("session-1"))[0]["injected"] + assert (await server.read_session("session-1")).messages[-1].content == "once" + monkeypatch.setattr(history, "write_records", original) + await server.append_session("session-1", "once", message_id="recover") + assert len(native.injected) == 1 + + +@pytest.mark.asyncio +async def test_unaccepted_retry_uses_the_current_turn_anchor(runtime): + server, native = runtime + original = native.request + + async def reject(method, params, **kwargs): + if method == "thread/inject_items": + raise ValueError("rejected before acceptance") + return await original(method, params, **kwargs) + + native.client._client.request.side_effect = reject + with pytest.raises(ValueError, match="before acceptance"): + await server.append_session("session-1", "later", message_id="retry") + native.thread.turns.append(turn("turn-2", "D", "E")) + native.client._client.request.side_effect = original + await server.append_session("session-1", "later", message_id="retry") + assert [message.content for message in (await server.read_session("session-1")).messages] == [ + "A", + "B", + "D", + "E", + "later", + ] + + +@pytest.mark.asyncio +async def test_cancelled_append_finishes_publication_before_releasing_lease(runtime): + server, native = runtime + entered, release = asyncio.Event(), asyncio.Event() + + async def hold(): + entered.set() + await release.wait() + + native.inject_hook = hold + pending = asyncio.create_task(server.append_session("session-1", "once", message_id="cancel")) + await entered.wait() + pending.cancel() + await asyncio.sleep(0) + assert not pending.done() + assert server._threads["session-1"].lock.locked() + release.set() + with pytest.raises(asyncio.CancelledError): + await pending + await server.append_session("session-1", "once", message_id="cancel") + assert len(native.injected) == 1 + + +@pytest.mark.asyncio +async def test_session_files_and_native_recovery_are_contained(runtime, tmp_path): + server, native = runtime + for identifier in ("../escape", "/absolute", "a/b", ""): + with pytest.raises(ValueError): + await server.read_session(identifier) + native.thread.path = str(tmp_path.parent / "outside.jsonl") + with pytest.raises(ValueError, match="configured home"): + await server.append_session("session-1", "no", message_id="contained") + assert not native.injected + + +@pytest.mark.asyncio +async def test_non_idle_error_status_is_not_appendable(runtime): + server, native = runtime + native.thread.status.root.type = "systemError" + assert (await server.get_session_status("session-1")).status == "error" + with pytest.raises(RuntimeError, match="idle"): + await server.append_session("session-1", "no", message_id="error") + assert not native.injected + + +@pytest.mark.asyncio +async def test_unloaded_status_is_read_only_and_append_resumes_before_injecting(runtime): + server, native = runtime + native.thread.status.root.type = "notLoaded" + assert (await server.get_session_status("session-1")).status == "not_loaded" + assert (await server.read_session("session-1")).state.status == "not_loaded" + native.client.thread_resume.assert_not_awaited() + + async def resume(*_args, **_kwargs): + native.thread.status.root.type = "idle" + return native.handle + + native.client.thread_resume.side_effect = resume + await server.append_session("session-1", "resumed", message_id="unloaded") + assert native.injected[-1]["content"][0]["text"] == "resumed" + + +class UnsupportedAgent(BaseAgentWrapper): + """Existing wrappers need not implement optional session capabilities.""" + + async def reply(self, inputs, **kwargs): + return {} + + +@pytest.mark.asyncio +async def test_optional_base_capabilities_fail_explicitly(): + agent = UnsupportedAgent() + for operation, kwargs in ( + (agent.read_session, {}), + (agent.get_session_status, {}), + (agent.append_session, {"content": "text", "message_id": "one"}), + ): + with pytest.raises(NotImplementedError, match="UnsupportedAgent"): + await operation("session-1", **kwargs) + + +@pytest.mark.asyncio +async def test_http_session_jobs_and_unsupported_error_code(tmp_path, monkeypatch): + config = resolve_app_config(config="proactive", log_config=False) + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_file=False, + log_to_console=False, + service={"backend": "http", "web_enabled": False, "mcp_enabled": False}, + components={ + "codex_app_server": {"default": {"backend": "sdk", "auth_mode": "oauth"}}, + "agent_wrapper": {"default": {"backend": "codex"}}, + }, + jobs={name: config["jobs"][name] for name in ("session_read", "session_status", "session_append")}, + ) + server = app.context.components["codex_app_server"]["default"] + native = NativeSession(server) + monkeypatch.setattr(server, "_get_client", AsyncMock(return_value=native.client)) + service = app.context.service + service.build_service(app) + service.add_jobs(app) + async with app: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=service.service), + base_url="http://test", + ) as client: + for action in ("session_status", "session_read"): + response = (await client.post("/" + action, json={"session_id": "session-1"})).json() + assert response["success"] + response = ( + await client.post( + "/session_append", + json={"session_id": "session-1", "content": "hello", "message_id": "http", "wait_timeout": 1}, + ) + ).json() + assert response["success"] + assert response["answer"]["source"] == "proactive" + monkeypatch.setattr(CodexAgentWrapper, "read_session", BaseAgentWrapper.read_session) + response = (await client.post("/session_read", json={"session_id": "session-1"})).json() + assert not response["success"] + assert response["metadata"]["code"] == "session_not_supported" diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 3ec8499a..55834089 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -493,17 +493,18 @@ async def test_worker_keeps_startup_error_or_cancellation_when_close_fails(tmp_p def test_worker_response_preserves_failed_job_when_close_fails(tmp_path, monkeypatch, mode, job_fails): import reme.reme as reme_module from types import SimpleNamespace - from reme.utils.counter import global_counter_inc + from reme.components import ApplicationContext class FakeReMe: def __init__(self, **_kwargs): - self.context = SimpleNamespace(components={}, jobs={"build": object()}, metadata={}) + self.context = ApplicationContext(workspace_dir=str(tmp_path)) + self.context.jobs["build"] = object() async def start(self): pass async def run_job(self, job, **_kwargs): - global_counter_inc(self.context.metadata, ["__job_counter", job]) + self.context.start_job_execution(job, track_execution=False) if job_fails == "result": return SimpleNamespace(success=False, answer="job failed", metadata={"diagnostic": "original"}) if job_fails: diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index f4f86937..3dd9ac22 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -3,7 +3,6 @@ # pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,redefined-outer-name import asyncio -from dataclasses import dataclass import json from pathlib import Path import socket @@ -16,7 +15,15 @@ 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 openai_codex.generated.v2_all import ( + AgentMessageThreadItem, + ItemCompletedNotification, + ThreadTokenUsage, + ThreadTokenUsageUpdatedNotification, + TokenUsageBreakdown, + Turn, + TurnCompletedNotification, +) from pydantic import BaseModel from reme.application import Application @@ -27,7 +34,8 @@ 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 Response +from reme.schema import Response, TokenUsage +from reme.utils import global_counter_get_all class _Job: @@ -49,26 +57,13 @@ class _Job: 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.thread_id = thread.id 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 @@ -76,17 +71,33 @@ class _Turn: async def stream(self): try: yield SimpleNamespace(method="turn/started", payload=SimpleNamespace(turn=SimpleNamespace(id=self.id))) + for event in self.thread.codex.usage_events: + yield event + if self.thread.codex.run_hook: + await self.thread.codex.run_hook(self.thread, self.inputs) yield SimpleNamespace( method="item/agentMessage/delta", payload=SimpleNamespace(item_id="item", delta=self.inputs), ) + item = AgentMessageThreadItem(type="agentMessage", id="item", text=str(self.inputs)) + yield SimpleNamespace( + method="item/completed", + payload=ItemCompletedNotification( + threadId=self.thread.id, + turnId=self.id, + item=item, + completedAtMs=0, + ), + ) self.thread.busy = False yield SimpleNamespace( method="turn/completed", - payload=SimpleNamespace( - turn=SimpleNamespace( + payload=TurnCompletedNotification( + threadId=self.thread.id, + turn=Turn( id=self.id, - status=SimpleNamespace(value="completed"), + items=[item], + status="completed", duration_ms=1, error=None, ), @@ -125,6 +136,7 @@ class _Thread: return SimpleNamespace( thread=SimpleNamespace( status=SimpleNamespace(root=SimpleNamespace(type="active" if self.busy else "idle")), + path=getattr(self, "path", None), ), ) @@ -148,9 +160,14 @@ class _Codex: self.logins = [] self.opened = [] self.run_hook = None + self.usage_events = [] self.compact_hook = None self.notifications = asyncio.Queue() - self._client = SimpleNamespace(request=self.request, next_notification=self.next_notification) + self._client = SimpleNamespace( + request=self.request, + next_notification=self.next_notification, + thread_read=self.thread_read, + ) self.instances.append(self) async def account(self): @@ -192,6 +209,9 @@ class _Codex: self.opened.append(("resume", thread_id, options)) return self.threads[thread_id] + async def thread_read(self, thread_id): + return await self.threads[thread_id].read() + async def thread_fork(self, thread_id, **options): forked = await self.thread_start(**options) self.opened[-1] = ("fork", thread_id, options) @@ -236,6 +256,237 @@ async def system(tmp_path, monkeypatch): await server.close() +def _usage_event(last, total, *, turn_id="turn-2", estimate=None): + def breakdown(tokens, total_tokens=None): + return TokenUsageBreakdown( + inputTokens=tokens[0], + outputTokens=tokens[1], + totalTokens=sum(tokens) if total_tokens is None else total_tokens, + cachedInputTokens=0, + reasoningOutputTokens=0, + ) + + return SimpleNamespace( + method="thread/tokenUsage/updated", + payload=ThreadTokenUsageUpdatedNotification( + threadId="thread-1", + turnId=turn_id, + tokenUsage=ThreadTokenUsage( + last=breakdown(last, estimate), + total=breakdown(total), + modelContextWindow=934435, + ), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("baseline,initial_estimate", [((0, 0), False), ((50000, 1000), False), ((50000, 1000), True)]) +async def test_turn_usage_counts_all_calls_without_recounting_restored_or_duplicate_usage( + system, + stream, + baseline, + initial_estimate, +): + first, _, server, _ = system + session_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + server.update_token_usage(session_id, TokenUsage(input_tokens=baseline[0], output_tokens=baseline[1])) + + def total(inputs, outputs): + return baseline[0] + inputs, baseline[1] + outputs + + first_call = _usage_event((10, 2), total(10, 2)) + client.usage_events = [_usage_event((9000, 900), (9000, 900), turn_id="turn-1")] + if initial_estimate: + client.usage_events.append(_usage_event((0, 0), baseline, estimate=800000)) + client.usage_events.extend( + [ + first_call, + first_call, + _usage_event((20, 3), total(30, 5)), + _usage_event((0, 0), total(30, 5), estimate=700000), + first_call, + _usage_event((7, 1), total(37, 6)), + ], + ) + expected = {"input_tokens": 37, "output_tokens": 6, "total_tokens": 43} + if stream: + chunks = [chunk async for chunk in first.reply_stream("multi-call", session_id=session_id)] + updates = [chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.USAGE] + assert updates[-1].chunk == expected + assert [chunk.chunk["total_tokens"] for chunk in updates] == [12, 35, 43] + assert {chunk.block_id for chunk in updates} == {"usage:turn-2"} + else: + response = await first.reply("multi-call", session_id=session_id) + assert response["result"] == "multi-call" + assert response["usage"] == expected + + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert {name: data["value"] for name, data in counter.items()} == expected + + client.usage_events = [_usage_event((2, 1), total(39, 7), turn_id="turn-3")] + response = await first.reply("next turn", session_id=session_id) + assert response["usage"] == {"input_tokens": 2, "output_tokens": 1, "total_tokens": 3} + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert counter["total_tokens"]["value"] == 46 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("estimate_only", [False, True]) +async def test_turn_usage_without_provider_reports_remains_unavailable(system, stream, estimate_only): + first, _, _, _ = system + session_id = (await first.reply("initial"))["session_id"] + if estimate_only: + first.app_server.update_token_usage(session_id, TokenUsage(input_tokens=50000, output_tokens=1000)) + _Codex.instances[0].usage_events = [_usage_event((0, 0), (50000, 1000), estimate=700000)] + if stream: + chunks = [chunk async for chunk in first.reply_stream("no usage", session_id=session_id)] + assert not any(chunk.chunk_type == ChunkEnum.USAGE for chunk in chunks) + else: + assert (await first.reply("no usage", session_id=session_id))["usage"] is None + assert global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name]) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_failed_turn_records_reported_spending_once(system, stream): + first, _, _, _ = system + session_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + first.app_server.update_token_usage(session_id, TokenUsage(input_tokens=5000, output_tokens=500)) + event = _usage_event((10, 2), (5010, 502)) + client.usage_events = [event, event] + + async def fail(*_args): + raise ValueError("model failed after a completed tool round") + + client.run_hook = fail + with pytest.raises(ValueError, match="model failed"): + if stream: + _ = [chunk async for chunk in first.reply_stream("fails", session_id=session_id)] + else: + await first.reply("fails", session_id=session_id) + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert counter["input_tokens"]["value"] == 10 + assert counter["output_tokens"]["value"] == 2 + assert counter["total_tokens"]["value"] == 12 + + +@pytest.mark.asyncio +async def test_closing_stream_keeps_reported_usage_and_interrupts_the_turn(system): + first, _, _, _ = system + session_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + first.app_server.update_token_usage(session_id, TokenUsage(input_tokens=5000, output_tokens=500)) + client.usage_events = [_usage_event((10, 2), (5010, 502))] + stream = first.reply_stream("interrupted", session_id=session_id) + assert (await anext(stream)).chunk_type == ChunkEnum.REPLY_START + assert (await anext(stream)).chunk_type == ChunkEnum.USAGE + await stream.aclose() + assert client.threads[session_id].interrupts == 1 + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert counter["total_tokens"]["value"] == 12 + + +@pytest.mark.asyncio +async def test_cancelling_reply_keeps_already_reported_usage(system): + first, _, _, _ = system + session_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + first.app_server.update_token_usage(session_id, TokenUsage(input_tokens=5000, output_tokens=500)) + client.usage_events = [_usage_event((10, 2), (5010, 502))] + started = asyncio.Event() + + async def hold(*_args): + started.set() + await asyncio.Event().wait() + + client.run_hook = hold + task = asyncio.create_task(first.reply("cancel", session_id=session_id)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert counter["total_tokens"]["value"] == 12 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_compaction_replaying_previous_usage_under_current_turn_does_not_recount_it(system, stream): + first, _, server, _ = system + session_id = (await first.reply("initial"))["session_id"] + server.update_token_usage(session_id, TokenUsage(input_tokens=1189426, output_tokens=220)) + _Codex.instances[0].usage_events = [ + _usage_event((610679, 61), (1189426, 220)), + _usage_event((606779, 471), (1796205, 691)), + _usage_event((0, 0), (1796205, 691), estimate=23616), + _usage_event((61232, 23), (1857437, 714)), + ] + expected = {"input_tokens": 668011, "output_tokens": 494, "total_tokens": 668505} + if stream: + chunks = [chunk async for chunk in first.reply_stream("compaction", session_id=session_id)] + usages = [chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.USAGE] + assert [usage["total_tokens"] for usage in usages] == [607250, 668505] + assert usages[-1] == expected + else: + assert (await first.reply("compaction", session_id=session_id))["usage"] == expected + counter = global_counter_get_all(first.app_context.metadata, ["__token_counter", first.name])["children"] + assert {key: value["value"] for key, value in counter.items()} == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("interrupted", [False, True]) +async def test_usage_baseline_recovers_after_lost_state_or_missed_stream_usage(system, interrupted): + first, second, server, _ = system + session_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + if interrupted: + stream = first.reply_stream("interrupted", session_id=session_id) + assert (await anext(stream)).chunk_type == ChunkEnum.REPLY_START + assert (await anext(stream)).chunk_type == ChunkEnum.CONTENT + # Native completion can overtake a paused consumer and persist unseen usage. + path = server.session_path / "rollout.jsonl" + usage = {"total_token_usage": {"input_tokens": 5000, "output_tokens": 500}} + record = {"type": "event_msg", "payload": {"type": "token_count", "info": usage}} + path.write_text(json.dumps(record) + "\n") + client.threads[session_id].path = str(path) + if interrupted: + client.threads[session_id].busy = False + await stream.aclose() + else: + server._threads.pop(session_id) + turn_id = "turn-3" if interrupted else "turn-2" + client.usage_events = [ + _usage_event((1000, 100), (5000, 500), turn_id=turn_id), + _usage_event((10, 2), (5010, 502), turn_id=turn_id), + ] + assert (await second.reply("restored", session_id=session_id))["usage"] == { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + } + + +@pytest.mark.asyncio +async def test_fresh_thread_does_not_require_its_rollout_to_be_materialized(system, monkeypatch): + first, _, server, _ = system + original_start = _Codex.thread_start + pending = server.session_path / "not-materialized.jsonl" + + async def start(client, **options): + thread = await original_start(client, **options) + thread.path = str(pending) + return thread + + monkeypatch.setattr(_Codex, "thread_start", start) + assert (await first.reply("first turn"))["usage"] is None + assert not pending.exists() + + @pytest.mark.asyncio async def test_wrappers_share_one_lazy_client_and_keep_independent_sessions(system): first, second, server, _ = system @@ -306,6 +557,7 @@ async def test_resume_fork_and_sdk_options_remain_explicit(system, approval_mode options = client.opened[0][2] assert options["model"] == "custom" assert options["developer_instructions"] == "instructions" + assert options["base_instructions"] is None 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 @@ -379,7 +631,9 @@ async def test_cancelling_reply_interrupts_only_its_turn_and_keeps_shared_server 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] + client = _Codex.instances[0] + thread = client.threads[thread_id] + server.update_token_usage(thread_id, TokenUsage(input_tokens=100, output_tokens=10)) original = getattr(thread, operation) entered, release = asyncio.Event(), asyncio.Event() @@ -388,6 +642,14 @@ async def test_cancelled_dispatch_settles_before_resuming_the_thread(system, mon await release.wait() if fail: raise ValueError("dispatch rejected") + if operation == "compact": + path = server.session_path / "rollout.jsonl" + path.write_text( + '{"type":"event_msg","payload":{"type":"token_count","info":' + '{"total_token_usage":{"input_tokens":500,"output_tokens":50}}}}\n', + ) + thread.path = str(path) + client.usage_events = [_usage_event((10, 2), (510, 52))] return await original(*args, **kwargs) monkeypatch.setattr(thread, operation, dispatch) @@ -409,7 +671,10 @@ async def test_cancelled_dispatch_settles_before_resuming_the_thread(system, mon 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" + result = await asyncio.wait_for(queued, 2) + assert result["result"] == "later" + if operation == "compact" and not fail: + assert result["usage"] == {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12} assert thread.interrupts == int(operation == "turn" and not fail) await asyncio.wait_for(server.close(), 2) finally: @@ -1063,7 +1328,7 @@ def test_event_to_chunks_maps_content_usage_and_completion(): ) usage_event = SimpleNamespace( method="thread/tokenUsage/updated", - payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)), + payload=SimpleNamespace(turn_id="turn-1"), ) completed_event = SimpleNamespace( method="turn/completed", @@ -1073,7 +1338,11 @@ def test_event_to_chunks_maps_content_usage_and_completion(): ) content = CodexAgentWrapper._event_to_chunks(content_event, "thread-1") # pylint: disable=protected-access - usage_chunks = CodexAgentWrapper._event_to_chunks(usage_event, "thread-1") # pylint: disable=protected-access + usage_chunks = CodexAgentWrapper._event_to_chunks( # pylint: disable=protected-access + usage_event, + "thread-1", + usage=TokenUsage.from_provider(usage), + ) completed = CodexAgentWrapper._event_to_chunks(completed_event, "thread-1") # pylint: disable=protected-access assert content[0].chunk_type == ChunkEnum.CONTENT diff --git a/tests/unit/test_config_parser.py b/tests/unit/test_config_parser.py index 35b1c174..25246c2a 100644 --- a/tests/unit/test_config_parser.py +++ b/tests/unit/test_config_parser.py @@ -137,6 +137,23 @@ def test_default_config_registers_manual_and_scheduled_proactive_refresh(): assert manual["steps"] == scheduled["steps"] +@pytest.mark.parametrize( + "manual,scheduled", + [("auto_dream", "dream_cron"), ("proactive_refresh", "proactive_refresh_cron")], +) +def test_shared_default_steps_keep_job_overrides_isolated(manual, scheduled): + """YAML aliases share definitions without coupling loaded job data or CLI overrides.""" + jobs = _load_config("default.yaml")["jobs"] + original = jobs[scheduled]["steps"][0]["scan_days"] + jobs[manual]["steps"][0]["scan_days"] = original + 1 + assert jobs[scheduled]["steps"][0]["scan_days"] == original + + replacement = [{"backend": "add_step"}] + resolved = resolve_app_config(log_config=False, jobs={manual: {"steps": replacement}})["jobs"] + assert resolved[manual]["steps"] == replacement + assert resolved[scheduled]["steps"] == jobs[scheduled]["steps"] + + def test_default_config_keeps_frontmatter_chunk_metadata_opt_in(): """Markdown frontmatter-to-chunk metadata is disabled by default for compatibility.""" cfg = _load_config("default.yaml") diff --git a/tests/unit/test_evaluation_interface.py b/tests/unit/test_evaluation_interface.py index fd58b2a9..55c6fdb7 100644 --- a/tests/unit/test_evaluation_interface.py +++ b/tests/unit/test_evaluation_interface.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest from reme.components.job import BaseJob, StreamJob +from reme.components import ApplicationContext from reme.utils import global_counter_add from reme.utils.evaluation_interface import ( check_agent_token_count, @@ -16,11 +17,11 @@ from reme.utils.evaluation_interface import ( ) -def test_check_job_count_reads_registered_base_job_count(): - """check_job_count returns the number of completed BaseJob invocations.""" +def test_check_job_count_reads_registered_base_job_count(tmp_path): + """check_job_count returns the number of started BaseJob invocations.""" async def run(): - app_context = SimpleNamespace(metadata={}, jobs={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = BaseJob(name="search", app_context=app_context) app_context.jobs[job.name] = job @@ -32,14 +33,14 @@ def test_check_job_count_reads_registered_base_job_count(): asyncio.run(run()) -def test_check_job_count_reads_custom_job_by_name(): +def test_check_job_count_reads_custom_job_by_name(tmp_path): """check_job_count reads subclassed job counters without depending on inheritance.""" async def run(): class ProjectStreamJob(StreamJob): """Project-specific StreamJob subclass used to exercise MRO lookup.""" - app_context = SimpleNamespace(metadata={}, jobs={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = ProjectStreamJob(name="chat", app_context=app_context) app_context.jobs[job.name] = job @@ -50,19 +51,19 @@ def test_check_job_count_reads_custom_job_by_name(): asyncio.run(run()) -def test_check_job_count_rejects_unknown_job_name(): +def test_check_job_count_rejects_unknown_job_name(tmp_path): """Unknown job names raise KeyError, matching Application.run_job.""" - app_context = SimpleNamespace(metadata={}, jobs={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) with pytest.raises(KeyError, match="Job 'missing' not found"): check_job_count("missing", app_context) -def test_track_job_counts_returns_calls_made_inside_context(): +def test_track_job_counts_returns_calls_made_inside_context(tmp_path): """The context manager reports only the calls made in its body.""" async def run(): - app_context = SimpleNamespace(metadata={}, jobs={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) search = BaseJob(name="search", app_context=app_context) app_context.jobs[search.name] = search @@ -76,11 +77,11 @@ def test_track_job_counts_returns_calls_made_inside_context(): asyncio.run(run()) -def test_track_job_counts_updates_results_when_body_raises(): +def test_track_job_counts_updates_results_when_body_raises(tmp_path): """Calls made before an exception are still included in the delta.""" async def run(): - app_context = SimpleNamespace(metadata={}, jobs={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) search = BaseJob(name="search", app_context=app_context) app_context.jobs[search.name] = search counts = {} diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 191a27d9..b3982ffc 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -10,6 +10,7 @@ import pytest from reme.enumeration import ComponentEnum from reme.application import Application +from reme.components import ApplicationContext from reme.components.base_component import BaseComponent from reme.components.component_registry import ComponentRegistry from reme.components.job.background_job import BackgroundJob @@ -18,7 +19,6 @@ from reme.components.job.cron_job import CronJob from reme.components.job.stream_job import StreamJob from reme.components.job import cron_job as cron_job_module from reme.schema import ComponentConfig -from reme.utils import global_counter_get # -- helpers ------------------------------------------------------------------ @@ -156,50 +156,54 @@ def test_stream_job_merges_config_kwargs_into_context(): # -- Job call counters ------------------------------------------------------- -def test_base_job_records_calls_by_name(): +def test_base_job_records_calls_by_name(tmp_path): async def run(): - app_context = SimpleNamespace(metadata={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = BaseJob(name="search", app_context=app_context) + app_context.jobs[job.name] = job await job() await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "search"]) == 2 + assert app_context.get_job_count("search") == 2 asyncio.run(run()) -def test_stream_job_subclass_records_calls_by_job_name(): +def test_stream_job_subclass_records_calls_by_job_name(tmp_path): async def run(): class ProjectStreamJob(StreamJob): pass - app_context = SimpleNamespace(metadata={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = ProjectStreamJob(name="chat", app_context=app_context) + app_context.jobs[job.name] = job await job(stream_queue=asyncio.Queue()) - assert global_counter_get(app_context.metadata, ["__job_counter", "chat"]) == 1 + assert app_context.get_job_count("chat") == 1 asyncio.run(run()) -def test_background_job_records_calls_by_name(): +def test_background_job_records_calls_by_name(tmp_path): async def run(): - app_context = SimpleNamespace(metadata={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = BackgroundJob(name="watch", app_context=app_context) + app_context.jobs[job.name] = job await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "watch"]) == 1 + assert app_context.get_job_count("watch") == 1 asyncio.run(run()) -def test_cron_job_records_each_triggered_execution(): +def test_cron_job_records_each_triggered_execution(tmp_path): async def run(): - app_context = SimpleNamespace(metadata={}) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) job = CronJob(name="nightly", cron="* * * * *", app_context=app_context) + app_context.jobs[job.name] = job job._stop_event = asyncio.Event() waits = 0 @@ -215,7 +219,7 @@ def test_cron_job_records_each_triggered_execution(): await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "nightly"]) == 1 + assert app_context.get_job_count("nightly") == 1 asyncio.run(run()) diff --git a/tests/unit/test_job_status.py b/tests/unit/test_job_status.py new file mode 100644 index 00000000..4ee06ed3 --- /dev/null +++ b/tests/unit/test_job_status.py @@ -0,0 +1,559 @@ +"""Execution status across job entry points, nested steps, retries, and shutdown.""" + +# pylint: disable=protected-access + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import httpx +import pytest +import yaml +from pydantic import ValidationError + +from reme.application import Application +from reme.components import ApplicationContext +from reme.components.agent_wrapper import BaseAgentWrapper +from reme.components.job import BackgroundJob, BaseJob, CronJob, StreamJob +from reme.components.runtime_context import RuntimeContext +from reme.components.service import HttpService +from reme.enumeration import ChunkEnum +from reme.schema import ApplicationConfig, JobConfig, Response, StreamChunk +from reme.steps.base_step import BaseStep +from reme.utils.evaluation_interface import check_job_count + + +class _CallbackStep(BaseStep): + """Run a test callback through the actual Step invocation boundary.""" + + async def execute(self): + return await self.kwargs["callback"](self) + + +@pytest.fixture(name="context") +def _context(tmp_path): + """Use real application state without starting services or workspace components.""" + owner = ApplicationContext(workspace_dir=str(tmp_path)) + owner.registry.register(_CallbackStep, "callback_step") + return owner + + +def _job(context, cls=BaseJob, name="test", **kwargs): + """Register a job and let normal startup resolve any configured steps.""" + job = cls(name=name, app_context=context, **kwargs) + context.jobs[name] = job + return job + + +@pytest.mark.parametrize("cls,default", [(BaseJob, False), (StreamJob, False), (BackgroundJob, True), (CronJob, True)]) +@pytest.mark.parametrize("option", [None, True, False]) +def test_tracking_defaults_and_explicit_overrides(context, cls, default, option): + """Typed defaults select by backend, and explicit switches never leak to job inputs.""" + config = JobConfig(track_execution=option) + kwargs = {"cron": "* * * * *"} if cls is CronJob else {} + job = _job(context, cls, **kwargs, **config.model_dump(exclude={"backend"})) + assert job.track_execution is (default if option is None else option) + assert "track_execution" not in job.kwargs + + +@pytest.mark.asyncio +async def test_foreground_tracking_is_opt_in_and_application_scoped(context): + """Untracked calls leave no history, while tracked calls remain application-local.""" + job = _job(context) + await job() + assert context.get_job_status() == {} + job.track_execution = True + app = object.__new__(Application) + app.context = context + assert (await app.run_job("test")).success + record = context.get_job_status()["test"][0] + assert record["status"] == "succeeded" + assert record["finished_at"] >= record["started_at"] + assert not ApplicationContext().get_job_status() + assert check_job_count("test", context) == 2 + + +@pytest.mark.asyncio +async def test_nested_steps_report_live_path_and_restore_parent(context): + """A watcher exposes its handler while active and then returns to its waiting step.""" + observations = [] + + async def child(step): + observations.append(context.get_job_status()["test"][0]) + assert step.context.data == {"request": "value"} + + async def parent(step): + before = context.get_job_status()["test"][0] + await step.dispatch_steps([{"backend": "callback_step", "callback": child}]) + after = context.get_job_status()["test"][0] + assert before["step_path"] == after["step_path"] == ["callback_step"] + assert before["step_started_at"] == after["step_started_at"] + + job = _job(context, track_execution=True, steps=[{"backend": "callback_step", "callback": parent}]) + await job.start() + await job(request="value") + assert observations[0]["step_path"] == ["callback_step", "callback_step"] + assert observations[0]["status"] == "running" + record = context.get_job_status()["test"][0] + assert record["step_path"] == [] + assert record["step_started_at"] is None + assert record["failed_step"] is None + # Snapshots remain detached after the underlying execution has ended. + assert observations[0]["finished_at"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raise_error", [True, False]) +async def test_failure_retains_innermost_step_for_exceptions_and_failed_responses(context, raise_error): + """Both error channels identify the child, without replacing it with the parent path.""" + + async def child(step): + if raise_error: + raise ValueError("failed child") + step.context.response.success = False + step.context.response.answer = "failed child" + + async def parent(step): + await step.dispatch_steps([{"backend": "callback_step", "callback": child}]) + + job = _job(context, track_execution=True, steps=[{"backend": "callback_step", "callback": parent}]) + await job.start() + response = await job() + assert not response.success + record = context.get_job_status()["test"][0] + assert record["status"] == "failed" + assert record["failed_step"] == ["callback_step", "callback_step"] + assert "failed child" in record["error"] + assert record["step_path"] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["recovered", "parent_failed", "sibling_failed", "wrapped_child_failed"]) +async def test_recovered_step_errors_do_not_hide_the_actual_terminal_failure(context, outcome): + """Caught child errors cannot leave a stale failure path on success or a later failure.""" + + async def fail(_step): + raise ValueError("child failed") + + context.registry.register(_CallbackStep, "first_step") + context.registry.register(_CallbackStep, "second_step") + + async def parent(step): + try: + await step.dispatch_steps([{"backend": "first_step", "callback": fail}]) + except ValueError as exc: + if outcome == "wrapped_child_failed": + raise RuntimeError("parent wrapped the child error") from exc + if outcome == "parent_failed": + raise RuntimeError("parent failed independently") + if outcome == "sibling_failed": + await step.dispatch_steps([{"backend": "second_step", "callback": fail}]) + + job = _job(context, track_execution=True, steps=[{"backend": "callback_step", "callback": parent}]) + await job.start() + response = await job() + record = context.get_job_status()["test"][0] + expected = { + "recovered": None, + "parent_failed": ["callback_step"], + "sibling_failed": ["callback_step", "second_step"], + "wrapped_child_failed": ["callback_step", "first_step"], + } + assert response.success is (outcome == "recovered") + assert record["failed_step"] == expected[outcome] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", [False, True]) +async def test_recovered_response_failure_does_not_own_a_later_failure(context, nested): + """A successful recovery clears the obsolete path before the final failure.""" + for name in ("recovered_failure", "recovery", "terminal_failure"): + context.registry.register(_CallbackStep, name) + + async def fail(step): + step.context.response.success = False + step.context.response.answer = step.backend + + async def recover(step): + step.context.response.success = True + + steps = [ + {"backend": "recovered_failure", "callback": fail}, + {"backend": "recovery", "callback": recover}, + {"backend": "terminal_failure", "callback": fail}, + ] + + async def parent(step): + await step.dispatch_steps(steps) + + job = _job( + context, + track_execution=True, + steps=[{"backend": "callback_step", "callback": parent}] if nested else steps, + ) + await job.start() + assert not (await job()).success + record = context.get_job_status()["test"][0] + assert record["error"] == "terminal_failure" + assert record["failed_step"] == (["callback_step"] if nested else []) + ["terminal_failure"] + + +@pytest.mark.asyncio +async def test_stream_failure_records_status_and_preserves_error_done_protocol(context): + """ERROR plus DONE remains the client contract even when the exception is handled.""" + + async def fail(_step): + raise RuntimeError("stream failed") + + job = _job(context, StreamJob, track_execution=True, steps=[{"backend": "callback_step", "callback": fail}]) + await job.start() + app = object.__new__(Application) + app.context = context + chunks = [chunk async for chunk in app.run_stream_job("test")] + assert [chunk.chunk_type for chunk in chunks] == [ChunkEnum.ERROR, ChunkEnum.DONE] + assert chunks[-1].done + record = context.get_job_status()["test"][0] + assert record["status"] == "failed" + assert record["error"] == "RuntimeError: stream failed" + assert record["failed_step"] == ["callback_step"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("will_retry", [True, False]) +async def test_chat_error_chunks_match_execution_status_without_changing_stream(context, will_retry): + """A terminal ERROR is a failed run even when ChatStep returns normally.""" + + class ErrorAgent(BaseAgentWrapper): + """Model terminal and recoverable provider notifications without an external service.""" + + async def reply(self, inputs, **kwargs): + raise AssertionError("stream only") + + async def reply_stream(self, inputs, **kwargs): + yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk="provider failed", metadata={"will_retry": will_retry}) + yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk="partial or recovered") + yield StreamChunk(chunk_type=ChunkEnum.REPLY_END) + + context.components["agent_wrapper"] = {"default": ErrorAgent()} + job = _job(context, StreamJob, track_execution=True, steps=[{"backend": "chat_step"}]) + await job.start() + queue = asyncio.Queue() + await job(query="hello", stream_queue=queue) + chunks = [queue.get_nowait() for _ in range(queue.qsize())] + assert [chunk.chunk_type for chunk in chunks] == [ + ChunkEnum.ERROR, + ChunkEnum.CONTENT, + ChunkEnum.REPLY_END, + ChunkEnum.DONE, + ] + record = context.get_job_status()["test"][0] + assert record["status"] == ("succeeded" if will_retry else "failed") + assert record["error"] == (None if will_retry else "provider failed") + assert record["failed_step"] == (None if will_retry else ["chat_step"]) + + +@pytest.mark.asyncio +async def test_late_child_step_exit_cannot_reopen_a_completed_record(context): + """Cleanup from a child task cannot restore a step path after the Job has ended.""" + entered, release = asyncio.Event(), asyncio.Event() + tasks = [] + + async def child(_step): + entered.set() + await release.wait() + raise RuntimeError("child finished after its job") + + async def parent(step): + tasks.append( + asyncio.create_task( + _CallbackStep(app_context=context, backend="child_step", callback=child)(step.context), + ), + ) + await entered.wait() + + job = _job(context, track_execution=True, steps=[{"backend": "callback_step", "callback": parent}]) + await job.start() + try: + assert (await job()).success + completed = context.get_job_status()["test"][0] + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + assert context.get_job_status()["test"][0] == completed + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_same_name_concurrent_runs_and_cancellation_are_isolated(context): + """Cancelling one invocation leaves its sibling's live step and result intact.""" + + async def block(step): + step.context["entered"].set() + await step.context["release"].wait() + + job = _job(context, track_execution=True, steps=[{"backend": "callback_step", "callback": block}]) + await job.start() + entered = [asyncio.Event(), asyncio.Event()] + release = asyncio.Event() + tasks = [asyncio.create_task(job(entered=event, release=release)) for event in entered] + try: + await asyncio.wait_for(asyncio.gather(*(event.wait() for event in entered)), 2) + records = context.get_job_status()["test"] + assert len({record["id"] for record in records}) == 2 + assert all(record["step_path"] == ["callback_step"] for record in records) + tasks[0].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[0] + records = context.get_job_status()["test"] + assert [record["status"] for record in records] == ["running", "cancelled"] + release.set() + await tasks[1] + assert [record["status"] for record in context.get_job_status()["test"]] == ["succeeded", "cancelled"] + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("child_tracked", [True, False]) +async def test_nested_job_has_its_own_tracking_switch(context, child_tracked): + """An untracked child cannot overwrite its tracked caller's current step.""" + + async def child(step): + assert (step.context.job_execution is not None) is child_tracked + assert context.get_job_status("parent")["parent"][0]["step_path"] == ["callback_step"] + + child_job = _job( + context, + name="child", + track_execution=child_tracked, + steps=[{"backend": "callback_step", "callback": child}], + ) + await child_job.start() + + async def parent(step): + await step.run_job("child") + + parent_job = _job( + context, + name="parent", + track_execution=True, + steps=[{"backend": "callback_step", "callback": parent}], + ) + await parent_job.start() + assert (await parent_job()).success + assert bool(context.get_job_status("child")["child"]) is child_tracked + + +@pytest.mark.asyncio +async def test_background_restarts_have_separate_records_and_graceful_stop(context): + """Supervisor attempts are recorded once each, and cooperative shutdown is stopped.""" + attempts = 0 + entered = asyncio.Event() + + async def body(step): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("restart me") + entered.set() + await step.context.stop_event.wait() + + job = _job(context, BackgroundJob, backoff_base=0, steps=[{"backend": "callback_step", "callback": body}]) + await job.start() + try: + await asyncio.wait_for(entered.wait(), 2) + assert [record["status"] for record in context.get_job_status()["test"]] == ["running", "failed"] + finally: + await job.close() + assert [record["status"] for record in context.get_job_status()["test"]] == ["stopped", "failed"] + assert check_job_count("test", context) == 2 + + +@pytest.mark.asyncio +async def test_custom_background_body_is_tracked_by_supervisor(context): + """Overriding __call__ still gets lifecycle tracking and a step-aware runtime context.""" + + async def body(step): + assert step.context.job_execution is not None + assert context.get_job_status()["test"][0]["step_path"] == ["custom_step"] + step.context.response.success = False + step.context.response.answer = "custom failure" + + class CustomJob(BackgroundJob): + """Simulate an extension providing its own job body.""" + + async def __call__(self, **kwargs): + runtime = RuntimeContext() + await _CallbackStep(app_context=context, backend="custom_step", callback=body)(runtime) + return runtime.response + + job = _job(context, CustomJob) + await job.start() + try: + await asyncio.wait_for(job._task, 2) + record = context.get_job_status()["test"][0] + assert record["status"] == "failed" + assert record["error"] == "custom failure" + assert record["failed_step"] == ["custom_step"] + assert check_job_count("test", context) == 1 + finally: + await job.close() + + +@pytest.mark.asyncio +async def test_custom_background_super_call_preserves_runtime_overrides(context): + """The supervisor's wrapper must not replace arguments supplied by an overridden body.""" + + async def body(step): + assert step.context["value"] == "from-body" + + class CustomJob(BackgroundJob): + """An extension can still specialize inputs before invoking the configured steps.""" + + async def __call__(self, **kwargs): + return await super().__call__(value="from-body") + + job = _job(context, CustomJob, value="from-config", steps=[{"backend": "callback_step", "callback": body}]) + await job.start() + try: + await asyncio.wait_for(job._task, 2) + records = context.get_job_status()["test"] + assert len(records) == 1 + assert records[0]["status"] == "succeeded" + assert check_job_count("test", context) == 1 + finally: + await job.close() + + +@pytest.mark.asyncio +async def test_cron_records_triggers_not_scheduler_and_continues_after_failure(context): + """Waiting creates no run, while each trigger gets a terminal result and fresh steps.""" + triggers = 0 + waits = 0 + + async def body(_step): + nonlocal triggers + triggers += 1 + if triggers == 1: + raise ValueError("first trigger failed") + + job = _job(context, CronJob, cron="* * * * *", steps=[{"backend": "callback_step", "callback": body}]) + + async def wait(_delay): + nonlocal waits + waits += 1 + if waits == 1: + assert context.get_job_status() == {"test": []} + elif waits == 3: + job._stop_event.set() + + job._wait_or_stop = wait + job._next_fire_delay = lambda: 0 + await job.start() + try: + await asyncio.wait_for(job._task, 2) + records = context.get_job_status()["test"] + assert [record["status"] for record in records] == ["succeeded", "failed"] + assert records[1]["failed_step"] == ["callback_step"] + assert check_job_count("test", context) == 2 + finally: + await job.close() + + +@pytest.mark.asyncio +async def test_worker_future_cancellation_does_not_claim_thread_body_has_stopped(context): + """A timed-out executor Future cannot publish completion before its worker actually exits.""" + entered, release = threading.Event(), threading.Event() + + class ThreadJob(BackgroundJob): + """Keep the worker alive after the outer Future has been cancelled.""" + + async def __call__(self, **kwargs): + entered.set() + if not release.wait(timeout=3): + raise TimeoutError("worker was not released") + return Response() + + context.thread_pool = ThreadPoolExecutor(max_workers=1) + job = _job(context, ThreadJob, use_thread_pool=True, close_timeout=0.01) + try: + await job.start() + assert await asyncio.to_thread(entered.wait, 2) + await job.close() + assert context.get_job_status()["test"][0]["status"] == "running" + finally: + release.set() + context.thread_pool.shutdown(wait=True) + assert context.get_job_status()["test"][0]["status"] == "stopped" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", [0, 2]) +async def test_retention_preserves_running_records_and_snapshots_are_detached(context, limit): + """Only completed records count towards the cap, and clients cannot mutate stored paths.""" + context.app_config.job_history_limit = limit + active = context.start_job_execution("test") + context.enter_job_step(active, "waiting") + job = _job(context, track_execution=True) + for _ in range(4): + await job() + records = context.get_job_status(limit=20)["test"] + assert len(records) == limit + 1 + assert records[-1]["status"] == "running" + records[-1]["step_path"].append("tampered") + assert context.get_job_status(status="running")["test"][0]["step_path"] == ["waiting"] + context.finish_job_execution("test", active, "cancelled") + assert len(context.get_job_status()["test"]) == limit + assert check_job_count("test", context) == 5 + + +@pytest.mark.parametrize("tracked", [True, False]) +def test_concurrent_starts_count_every_call_independently_of_retained_history(context, tracked): + """Worker threads share one atomic count even with optional and bounded execution records.""" + context.app_config.job_history_limit = 2 + job = _job(context, track_execution=tracked) + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(lambda _: asyncio.run(job()), range(32))) + assert check_job_count("test", context) == 32 + assert len(context.get_job_status("test")["test"]) == (2 if tracked else 0) + + +@pytest.mark.asyncio +async def test_default_query_job_works_through_http_and_step_internal_calls(context): + """The default config/registration produce a structured foreground query without self-history.""" + config_path = Path(__file__).parents[2] / "reme" / "config" / "default.yaml" + config = yaml.safe_load(config_path.read_text(encoding="utf-8"))["jobs"]["job_status"] + query = _job(context, name="job_status", **config) + await query.start() + work = _job(context, name="work", track_execution=True) + await work() + app = object.__new__(Application) + app.context = context + service = HttpService(web_enabled=False, mcp_enabled=False) + service.build_service(app) + service.add_jobs(app) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=service.service), base_url="http://test") as client: + response = await client.post("/job_status", json={"name": "work", "status": "succeeded", "limit": 1}) + assert response.status_code == 200 + assert response.json()["answer"]["work"][0]["status"] == "succeeded" + caller = _CallbackStep(app_context=context, callback=None) + assert (await caller.run_job("job_status")).success + assert "job_status" not in context.get_job_status() + + +@pytest.mark.parametrize("kwargs", [{"limit": 0}, {"limit": True}, {"limit": "1"}, {"status": "unknown"}]) +def test_query_rejects_invalid_filters(context, kwargs): + """Direct callers receive the same validation as the service's documented schema.""" + with pytest.raises(ValueError): + context.get_job_status(**kwargs) + with pytest.raises(KeyError, match="not found"): + context.get_job_status(name="missing") + + +def test_retention_limit_is_validated_by_application_config(): + """A negative retention limit cannot silently slice histories incorrectly.""" + with pytest.raises(ValidationError): + ApplicationConfig(job_history_limit=-1) diff --git a/tests/unit/test_proactive_refresh.py b/tests/unit/test_proactive_refresh.py index 08d3f2e4..72139007 100644 --- a/tests/unit/test_proactive_refresh.py +++ b/tests/unit/test_proactive_refresh.py @@ -624,6 +624,7 @@ async def test_refresh_chain_e2e(tmp_path): resp_finish = await finish(context) assert resp_extract.success and resp_topics.success and resp_finish.success + assert "duration_ms" not in resp_finish.metadata["proactive"] state = context.get("proactive") assert state["llm_calls"] == 1 assert state["push"] is True diff --git a/tests/unit/test_resource_cleanup.py b/tests/unit/test_resource_cleanup.py new file mode 100644 index 00000000..5168e6cf --- /dev/null +++ b/tests/unit/test_resource_cleanup.py @@ -0,0 +1,280 @@ +"""Dated resource retention, catalog persistence, and proactive configuration.""" + +import asyncio +import shutil +import threading +from datetime import datetime, timezone +from functools import wraps +from pathlib import Path + +import pytest + +from reme.components.application_context import ApplicationContext +from reme.components.file_catalog import LocalFileCatalog +from reme.components.runtime_context import RuntimeContext +from reme.config.config_parser import resolve_app_config +from reme.schema import ApplicationConfig, FileNode +from reme.steps.index.cleanup_catalog import CleanupCatalogStep +from reme.steps.index.init_changes import InitChangesStep + + +@pytest.fixture(autouse=True) +def fixed_clock(monkeypatch): + """The same instant falls on different calendar days in Shanghai and Los Angeles.""" + + class FixedDatetime(datetime): + """Return a fixed instant while honoring the requested timezone.""" + + @classmethod + def now(cls, tz=None): + return cls(2026, 11, 5, 0, 30, tzinfo=timezone.utc).astimezone(tz) + + monkeypatch.setattr("reme.components.file_catalog.local_file_catalog.datetime", FixedDatetime) + + +def _catalog(workspace, **kwargs): + app = ApplicationContext(workspace_dir=str(workspace), **kwargs) + catalog = LocalFileCatalog(name="resource", app_context=app) + app.components["file_catalog"] = {"resource": catalog} + return catalog + + +def _write(workspace, path): + target = workspace / path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"resource") + return FileNode(path=path, st_mtime=target.stat().st_mtime) + + +@pytest.mark.asyncio +async def test_retention_uses_directory_dates_and_persists_only_successful_deletions(tmp_path): + """At day seven remove the entire dated tree, including uncataloged files and empty folders.""" + catalog = _catalog(tmp_path) + expired = "resource/2026-10-29/sub/report.pdf" + retained = [ + "resource/2026-10-30/recent.txt", + "resource/2026-11-06/future.png", + "resource/2026-02-30/invalid.txt", + "resource/20261001/noncanonical.txt", + "resource/archive/2026-10-01/permanent.txt", + "resource/loose.txt", + "archive/2026-10-01/permanent.txt", + "resource-extra/2026-10-01/other.txt", + ] + await catalog.upsert([_write(tmp_path, path) for path in [expired, *retained]]) + _write(tmp_path, "resource/2026-10-29/untracked.bin") + (tmp_path / "resource/2026-10-01").mkdir() + _write(tmp_path, "resource/2026-10-02") # A date-named file is not a directory. + + result = await catalog.cleanup_expired("resource", 7) + + assert result == { + "deleted_dirs": ["resource/2026-10-01", "resource/2026-10-29"], + "deleted_nodes": 1, + "errors": {}, + } + assert not (tmp_path / "resource/2026-10-29").exists() + assert all((tmp_path / path).exists() for path in retained) + assert (tmp_path / "resource/2026-10-02").is_file() + reloaded = _catalog(tmp_path) + await reloaded.load() + assert {node.path for node in await reloaded.get_nodes()} == set(retained) + assert (await catalog.cleanup_expired("resource", 7))["deleted_dirs"] == [] + + +@pytest.mark.asyncio +async def test_cleanup_uses_configured_timezone_and_resource_directory(tmp_path): + """The configured resource root and local date determine eligibility.""" + catalog = _catalog(tmp_path, timezone="America/Los_Angeles", resource_dir="assets/tmp") + node = _write(tmp_path, "assets/tmp/2026-10-29/source.bin") + await catalog.upsert([node]) + step = CleanupCatalogStep(app_context=catalog.app_context, file_catalog="resource") + response = await step(RuntimeContext(ttl_days=7)) + assert response.success is True + assert response.metadata["deleted_dirs"] == [] + response = await step(RuntimeContext(ttl_days=6)) + assert response.metadata["deleted_dirs"] == ["assets/tmp/2026-10-29"] + assert not (tmp_path / node.path).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl_days", [0, -1, True, 1.5, "7", None]) +async def test_cleanup_rejects_invalid_ttl_before_deleting(tmp_path, ttl_days): + """Invalid retention settings cannot authorize deletions.""" + catalog = _catalog(tmp_path) + node = _write(tmp_path, "resource/2026-10-01/source.txt") + with pytest.raises(ValueError, match="ttl_days"): + await catalog.cleanup_expired("resource", ttl_days) + assert (tmp_path / node.path).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("directory", ["", ".", "..", "resource/../archive", "~/resource", "/tmp"]) +async def test_cleanup_rejects_unsafe_roots(tmp_path, directory): + """Cleanup accepts only a non-root workspace-relative directory.""" + catalog = _catalog(tmp_path) + with pytest.raises(ValueError, match="inside the workspace"): + await catalog.cleanup_expired(directory, 7) + + +@pytest.mark.asyncio +async def test_cleanup_never_follows_directory_symlinks(tmp_path): + """Dated symlinks are skipped, inner links are unlinked, and linked roots are refused.""" + catalog = _catalog(tmp_path) + permanent = _write(tmp_path, "archive/2026-10-01/keep.txt") + root = tmp_path / "resource" + root.mkdir() + (root / "2026-10-01").symlink_to(tmp_path / "archive/2026-10-01", target_is_directory=True) + (root / "2026-10-02").mkdir() + (root / "2026-10-02/link").symlink_to(tmp_path / "archive", target_is_directory=True) + (tmp_path / "linked").symlink_to(root, target_is_directory=True) + + with pytest.raises(ValueError, match="symlinks"): + await catalog.cleanup_expired("linked", 7) + result = await catalog.cleanup_expired("resource", 7) + assert result["deleted_dirs"] == ["resource/2026-10-02"] + assert (root / "2026-10-01").is_symlink() + assert (tmp_path / permanent.path).is_file() + + +@pytest.mark.asyncio +async def test_cleanup_reports_failures_without_pruning_failed_catalog_entries(tmp_path, monkeypatch): + """A failed directory remains retryable while successful deletes are persisted.""" + catalog = _catalog(tmp_path) + nodes = [_write(tmp_path, f"resource/{day}/source.txt") for day in ["2026-10-01", "2026-10-02"]] + await catalog.upsert(nodes) + original = shutil.rmtree + + @wraps(original) + def remove(path, **kwargs): + if Path(path).name == "2026-10-01": + raise PermissionError("blocked") + original(path, **kwargs) + + monkeypatch.setattr("reme.components.file_catalog.local_file_catalog.shutil.rmtree", remove) + step = CleanupCatalogStep(app_context=catalog.app_context, file_catalog="resource", ttl_days=7) + response = await step(RuntimeContext()) + assert response.success is False + assert response.metadata["errors"] == {"resource/2026-10-01": "blocked"} + assert response.metadata["deleted_dirs"] == ["resource/2026-10-02"] + reloaded = _catalog(tmp_path) + await reloaded.load() + assert [node.path for node in await reloaded.get_nodes()] == [nodes[0].path] + + +@pytest.mark.asyncio +async def test_cancellation_finishes_deletion_and_catalog_persistence(tmp_path, monkeypatch): + """Cancellation cannot leave a deletion worker publishing outside the catalog lock.""" + catalog = _catalog(tmp_path) + await catalog.upsert([_write(tmp_path, "resource/2026-10-01/source.txt")]) + entered, release = threading.Event(), threading.Event() + original = shutil.rmtree + + @wraps(original) + def remove(path, **kwargs): + entered.set() + assert release.wait(5) + original(path, **kwargs) + + monkeypatch.setattr("reme.components.file_catalog.local_file_catalog.shutil.rmtree", remove) + task = asyncio.create_task(catalog.cleanup_expired("resource", 7)) + try: + async with asyncio.timeout(5): + while not entered.is_set(): + await asyncio.sleep(0.01) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + reloaded = _catalog(tmp_path) + await reloaded.load() + assert await reloaded.get_nodes() == [] + assert not (tmp_path / "resource/2026-10-01").exists() + + +@pytest.mark.asyncio +async def test_cleanup_cannot_follow_a_root_swapped_after_validation(tmp_path, monkeypatch): + """Replacing resource/ with a link during deletion cannot delete permanent material.""" + catalog = _catalog(tmp_path) + await catalog.upsert([_write(tmp_path, "resource/2026-10-01/temporary.txt")]) + permanent = _write(tmp_path, "archive/2026-10-01/permanent.txt") + root = tmp_path / "resource" + original = shutil.rmtree + + @wraps(original) + def remove(path, **kwargs): + root.rename(tmp_path / "moved-resource") + root.symlink_to(tmp_path / "archive", target_is_directory=True) + original(path, **kwargs) + + monkeypatch.setattr("reme.components.file_catalog.local_file_catalog.shutil.rmtree", remove) + await catalog.cleanup_expired("resource", 7) + assert (tmp_path / permanent.path).is_file() + + +@pytest.mark.asyncio +async def test_cleanup_retries_catalog_persistence_after_a_write_failure(tmp_path, monkeypatch): + """Retrying after a failed dump persists removals even though source directories are already gone.""" + catalog = _catalog(tmp_path) + await catalog.upsert([_write(tmp_path, "resource/2026-10-01/source.txt")]) + await catalog.dump() + import reme.components.file_catalog.local_file_catalog as module + + original = module.write_jsonl_zst + + def fail(*_args, **_kwargs): + raise OSError("disk full") + + monkeypatch.setattr(module, "write_jsonl_zst", fail) + with pytest.raises(OSError, match="disk full"): + await catalog.cleanup_expired("resource", 7) + monkeypatch.setattr(module, "write_jsonl_zst", original) + await catalog.cleanup_expired("resource", 7) + reloaded = _catalog(tmp_path) + await reloaded.load() + assert await reloaded.get_nodes() == [] + + +@pytest.mark.asyncio +async def test_cleanup_refuses_platforms_without_safe_directory_deletion(tmp_path, monkeypatch): + """Unsupported platforms fail before deleting source files.""" + catalog = _catalog(tmp_path) + node = _write(tmp_path, "resource/2026-10-01/source.txt") + monkeypatch.setattr(shutil.rmtree, "avoids_symlink_attacks", False) + with pytest.raises(RuntimeError, match="symlink-resistant"): + await catalog.cleanup_expired("resource", 7) + assert (tmp_path / node.path).is_file() + + +@pytest.mark.asyncio +async def test_proactive_config_catalogs_every_suffix_and_cleans_up_through_registered_step(tmp_path): + """Exercise the configured init -> catalog -> cleanup path without starting any model providers.""" + config = ApplicationConfig(**resolve_app_config(config="proactive", log_config=False, workspace_dir=str(tmp_path))) + app = ApplicationContext(**config.model_dump()) + catalog = LocalFileCatalog(name="resource", app_context=app) + app.components["file_catalog"] = {"resource": catalog} + paths = ["resource/2026-10-01/source.pdf", "resource/2026-11-05/source.unknown"] + for path in paths: + _write(tmp_path, path) + watcher = config.jobs["resource_watch_loop"] + kwargs = watcher.model_extra + assert kwargs["watch_suffixes"] == [] + for producer in watcher.steps: + assert producer.model_extra["monitor_name"] == "resource" + assert producer.model_extra["dispatch_steps"] == [ + {"backend": "update_catalog_step", "file_catalog": "resource"}, + ] + await InitChangesStep(app_context=app, **watcher.steps[0].model_dump())(RuntimeContext(**kwargs)) + assert {node.path for node in await catalog.get_nodes()} == set(paths) + cleanup = config.jobs["resource_cleanup_cron"] + assert cleanup.backend == "cron" + assert cleanup.model_extra["cron"] == "0 * * * *" + cleanup_config = cleanup.steps[0].model_dump() + step_cls = app.registry.get("step", cleanup_config["backend"]) + response = await step_cls(app_context=app, **cleanup_config)(RuntimeContext(**cleanup.model_extra)) + assert response.success is True + assert response.metadata["deleted_dirs"] == ["resource/2026-10-01"] + assert [node.path for node in await catalog.get_nodes()] == [paths[1]] diff --git a/tests/unit/test_workspace_web_steps.py b/tests/unit/test_workspace_web_steps.py index 39254470..febf3768 100644 --- a/tests/unit/test_workspace_web_steps.py +++ b/tests/unit/test_workspace_web_steps.py @@ -3,13 +3,14 @@ import asyncio import os from datetime import datetime -from zoneinfo import ZoneInfo import pytest from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext +from reme.components.job import StreamJob from reme.components.runtime_context import RuntimeContext +from reme.config.config_parser import _load_config from reme.enumeration import ChunkEnum from reme.schema import StreamChunk from reme.steps.common.chat import ChatStep @@ -136,25 +137,47 @@ async def test_chat_step_streams_rich_chunks_and_resumes_session(tmp_path): "traverse", ] assert agent.reply_kwargs["builtin_tools"] == [] - system_prompt = agent.reply_kwargs["system_prompt"] - assert f"Current date: {datetime.now(ZoneInfo('Asia/Shanghai')).date().isoformat()}" in system_prompt - assert f"Current working directory: {tmp_path.resolve()}" in system_prompt + assert agent.reply_kwargs["system_prompt"] == ChatStep.DEFAULT_SYSTEM_PROMPT @pytest.mark.asyncio -async def test_chat_step_appends_environment_context_to_system_prompt_override(tmp_path): - """A caller prompt override retains request-time date and cwd context.""" +async def test_chat_step_passes_system_prompt_override_verbatim(tmp_path): + """A caller prompt reaches the wrapper without extra environment text.""" app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") agent = _StreamingAgent(app_context=app_context) await ChatStep(agent_wrapper=agent, app_context=app_context)( query="hello", - system_prompt="Custom prompt.", + system_prompt="Custom prompt.\n", stream_queue=asyncio.Queue(), ) - system_prompt = agent.reply_kwargs["system_prompt"] - assert system_prompt.startswith("Custom prompt.\n\n") - assert "Current date:" in system_prompt - assert f"Current working directory: {tmp_path.resolve()}" in system_prompt + assert agent.reply_kwargs["system_prompt"] == "Custom prompt.\n" + + +@pytest.mark.asyncio +async def test_proactive_chat_uses_native_tools_and_preserves_configured_prompt(tmp_path): + """The proactive stream opts out of ReMe tools and adds only its configured prompt.""" + + config = _load_config("proactive") + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + agent = _StreamingAgent(app_context=app_context) + app_context.components["agent_wrapper"] = {"default": agent} + chat = StreamJob(name="chat", app_context=app_context, **config["jobs"]["chat"]) + queue = asyncio.Queue() + + await chat.start() + try: + await chat(query="Help plan my day", stream_queue=queue) + assert agent.reply_kwargs["job_tools"] == [] + assert agent.reply_kwargs["system_prompt"] == config["jobs"]["chat"]["system_prompt"] + chunks = [await queue.get() for _ in range(4)] + assert [chunk.chunk_type for chunk in chunks] == [ + ChunkEnum.TOOL_CALL, + ChunkEnum.CONTENT, + ChunkEnum.REPLY_END, + ChunkEnum.DONE, + ] + finally: + await chat.close() From 36c23674ad8c599dae20b0687af40acd8ecb990e Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 5 Nov 2026 13:10:43 +0800 Subject: [PATCH 5/5] fix: preserve QwenPaw dependency compatibility (#712) * fix: preserve QwenPaw dependency compatibility * ci: use declared dependencies in wheel smoke test --- .github/workflows/_build-python-packages.yml | 1 + pyproject.toml | 4 ++-- scripts/sync_qwenpaw_agentscope.py | 21 +++++++---------- tests/unit/test_sync_qwenpaw_agentscope.py | 24 +++++++++++++++++--- 4 files changed, 32 insertions(+), 18 deletions(-) diff --git a/.github/workflows/_build-python-packages.yml b/.github/workflows/_build-python-packages.yml index b4bfcf85..74ba547f 100644 --- a/.github/workflows/_build-python-packages.yml +++ b/.github/workflows/_build-python-packages.yml @@ -64,6 +64,7 @@ jobs: python -m zipfile -l "${REME_WHEEL}" | (! grep 'reme_studio/') python -m venv "${RUNNER_TEMP}/reme-package-smoke" "${RUNNER_TEMP}/reme-package-smoke/bin/python" -m pip install "${REME_WHEEL}[as]" + "${RUNNER_TEMP}/reme-package-smoke/bin/python" -m pip check cd "${RUNNER_TEMP}" "${RUNNER_TEMP}/reme-package-smoke/bin/python" -c "import reme" diff --git a/pyproject.toml b/pyproject.toml index b00a962f..4fb54cf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ classifiers = [ ] dependencies = [ - "aiofiles>=25.1.0", + "aiofiles>=24.1.0", "croniter>=6.2.4", "fastapi>=0.162.0", "fastmcp>=4.1.4,<5", @@ -31,7 +31,7 @@ dependencies = [ "mistletoe>=1.6.0", "numpy>=2.4.6; python_version < '3.12'", "numpy>=2.5.4; python_version >= '3.12'", - "openai>=3.40.0", + "openai>=2.33.0", "psutil>=7.2.2", "pydantic>=2.14.1", "python-frontmatter>=1.3.0", diff --git a/scripts/sync_qwenpaw_agentscope.py b/scripts/sync_qwenpaw_agentscope.py index 71603219..f7b22655 100644 --- a/scripts/sync_qwenpaw_agentscope.py +++ b/scripts/sync_qwenpaw_agentscope.py @@ -8,7 +8,7 @@ from pathlib import Path from packaging.requirements import Requirement from packaging.version import Version -_PIN = re.compile(r'"agentscope\[model-ollama\]==[^"\n]+"') +_PIN = re.compile(r'"agentscope\[(model-ollama|classifier-jev)\]==[^"\n]+"') def _exact_agentscope_version(dependencies: list[str]) -> str: @@ -18,14 +18,9 @@ def _exact_agentscope_version(dependencies: list[str]) -> str: raise ValueError(f"Expected one AgentScope dependency, found {len(matches)}") requirement = matches[0] specifiers = list(requirement.specifier) - if ( - requirement.extras != {"model-ollama"} - or requirement.marker is not None - or len(specifiers) != 1 - or specifiers[0].operator != "==" - ): + if requirement.marker is not None or len(specifiers) != 1 or specifiers[0].operator != "==": raise ValueError( - f"Expected an exact agentscope[model-ollama] pin, found {requirement}", + f"Expected an exact AgentScope pin, found {requirement}", ) version = Version(specifiers[0].version) if version.is_prerelease or version.is_devrelease: @@ -38,16 +33,16 @@ def sync_pin(qwenpaw_text: str, reme_text: str) -> tuple[str, str]: qwenpaw = tomllib.loads(qwenpaw_text) reme = tomllib.loads(reme_text) version = _exact_agentscope_version(qwenpaw["project"]["dependencies"]) - _exact_agentscope_version(reme["project"]["optional-dependencies"]["as"]) + for extra in ("as", "decision"): + _exact_agentscope_version(reme["project"]["optional-dependencies"][extra]) matches = list(_PIN.finditer(reme_text)) - if len(matches) != 1: + if len(matches) != 2 or {match.group(1) for match in matches} != {"model-ollama", "classifier-jev"}: raise ValueError( - f"Expected one replaceable ReMe AgentScope pin, found {len(matches)}", + f"Expected replaceable ReMe as and decision pins, found {len(matches)}", ) - replacement = f'"agentscope[model-ollama]=={version}"' return ( version, - reme_text[: matches[0].start()] + replacement + reme_text[matches[0].end() :], + _PIN.sub(lambda match: f'"agentscope[{match.group(1)}]=={version}"', reme_text), ) diff --git a/tests/unit/test_sync_qwenpaw_agentscope.py b/tests/unit/test_sync_qwenpaw_agentscope.py index 71f13199..b04022e2 100644 --- a/tests/unit/test_sync_qwenpaw_agentscope.py +++ b/tests/unit/test_sync_qwenpaw_agentscope.py @@ -5,16 +5,34 @@ import pytest from scripts.sync_qwenpaw_agentscope import sync_pin QWENPAW = '[project]\ndependencies = ["agentscope[model-ollama]==2.0.8"]\n' -REME = '[project.optional-dependencies]\nas = ["agentscope[model-ollama]==2.0.7.post1"]\n' +REME = ( + "[project.optional-dependencies]\n" + 'as = ["agentscope[model-ollama]==2.0.7.post1"]\n' + 'decision = ["agentscope[classifier-jev]==2.0.7.post1"]\n' +) -def test_sync_pin_updates_only_reme_pin(): - """A stable QwenPaw pin replaces ReMe's exact pin.""" +def test_sync_pin_updates_both_reme_pins(): + """The model and decision extras must resolve to the same stable version.""" version, updated = sync_pin(QWENPAW, REME) assert version == "2.0.8" assert updated == REME.replace("2.0.7.post1", "2.0.8") +def test_sync_pin_accepts_qwenpaw_classifier_extra(): + """Additional host extras do not change ReMe's independently selected extras.""" + upstream = '[project]\ndependencies = ["agentscope[model-ollama,classifier-jev]==2.0.11"]\n' + version, updated = sync_pin(upstream, REME) + assert version == "2.0.11" + assert updated == REME.replace("2.0.7.post1", "2.0.11") + + +def test_sync_pin_preserves_already_aligned_manifest(): + """A repeated sync does not rewrite an aligned manifest.""" + current = REME.replace("2.0.7.post1", "2.0.8") + assert sync_pin(QWENPAW, current) == ("2.0.8", current) + + @pytest.mark.parametrize( "upstream", [