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()