From 2a307f80a8c2c168612c820d58ff9faba9114843 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:04:27 +0800 Subject: [PATCH] refactor(agent)!: unify lossless streaming and collected replies (#713) * refactor(agent)!: unify lossless streaming and collected replies * fix(agent): complete reply consumer migration * fix(agent): migrate Home consumers to canonical replies --- benchmark/memlens/run.py | 7 +- benchmark/mobilemem/run.py | 7 +- benchmark/pibench/bridge_reme.py | 6 +- docs/en/framework.md | 23 + docs/zh/framework.md | 20 + plugins/auto-fin/src/reme_auto_fin/base.py | 7 +- plugins/auto-fin/src/reme_auto_fin/topic.py | 4 +- plugins/auto-fin/tests/test_auto_fin.py | 12 +- .../beam-judge/src/judge_beam/llm_judge.py | 5 +- plugins/beam/src/reme_beam/agentic_answer.py | 43 +- .../daily_paper/src/reme_daily_paper/base.py | 5 +- plugins/daily_paper/tests/test_daily_paper.py | 7 +- plugins/dingtalk/src/reme_dingtalk/wait.py | 13 +- plugins/dingtalk/tests/test_dingtalk.py | 35 +- plugins/home/README.md | 5 +- .../home/src/reme_home/agent/case_study.py | 13 +- .../src/reme_home/agent/structured_report.py | 26 +- .../src/reme_home/agent/workflow_common.py | 11 +- .../runtime_docs/validation-artifact-guide.md | 8 +- plugins/home/tests/test_bounded_wrapper.py | 3 +- plugins/home/tests/test_candidate_gate.py | 21 +- plugins/home/tests/test_case_study.py | 13 +- .../tests/test_construction_equivalence.py | 16 +- plugins/home/tests/test_evaluation.py | 33 +- plugins/home/tests/test_evaluation_history.py | 10 +- .../home/tests/test_evaluation_repetition.py | 2 +- plugins/home/tests/test_home_eval.py | 6 +- plugins/home/tests/test_sampling.py | 14 +- plugins/home/tests/test_structured_report.py | 69 ++- plugins/home/tests/test_workflows.py | 13 + plugins/lme-judge/src/judge_lme/llm_judge.py | 3 +- plugins/lme/src/reme_lme/agentic_answer.py | 43 +- .../src/reme_memlens/agentic_answer.py | 8 +- .../src/reme_mobilemem/agentic_answer.py | 8 +- .../agent_wrapper/as_agent_wrapper.py | 267 +++++----- .../agent_wrapper/base_agent_wrapper.py | 49 +- .../agent_wrapper/cc_agent_wrapper.py | 470 +++++++++--------- .../agent_wrapper/codex_agent_wrapper.py | 232 +++++---- .../agent_wrapper/codex_session_history.py | 93 +++- reme/components/client/http_client.py | 44 +- reme/components/job/base_job.py | 2 +- reme/components/job/stream_job.py | 1 + reme/components/runtime_context.py | 4 + reme/enumeration/chunk_enum.py | 5 +- reme/schema/__init__.py | 4 + reme/schema/agent_reply.py | 61 +++ reme/schema/stream_chunk.py | 66 ++- reme/steps/common/chat.py | 23 +- reme/steps/common/llm_demo.py | 9 +- reme/steps/common/stream_llm_demo.py | 30 +- reme/steps/evolve/_evolve.py | 19 +- reme/steps/evolve/auto_memory.py | 5 +- reme/steps/evolve/auto_tag.py | 3 +- reme/steps/evolve/base_auto_resource.py | 3 +- reme/steps/evolve/dream/integrate.py | 2 +- reme/steps/evolve/proactive/agenda.py | 2 +- reme/steps/evolve/proactive/plan.py | 2 +- reme/utils/common_utils.py | 18 +- reme_studio/app/chat-stream.ts | 271 ++++++---- reme_studio/app/i18n.ts | 4 + reme_studio/app/session-sync.ts | 12 + reme_studio/app/store.ts | 3 +- reme_studio/app/types.ts | 56 ++- reme_studio/app/workspace-chat.tsx | 26 +- reme_studio/tests/chat-stream.test.mjs | 181 ++++++- reme_studio/tests/session-sync.test.mjs | 88 ++++ tests/integration/test_agent_session.py | 42 +- tests/integration/test_codex_agent_wrapper.py | 14 +- tests/unit/auto_resource_test_support.py | 8 +- tests/unit/test_agent_reply.py | 437 ++++++++++++++++ tests/unit/test_agent_reply_consumers.py | 163 ++++++ tests/unit/test_agent_stream.py | 270 ++++++++++ tests/unit/test_auto_dream.py | 20 +- tests/unit/test_auto_memory_cc.py | 6 +- tests/unit/test_auto_memory_direct_images.py | 10 +- .../test_auto_resource_review_regressions.py | 2 +- tests/unit/test_auto_tag.py | 4 +- tests/unit/test_background_steps.py | 2 +- tests/unit/test_base_agent_wrapper.py | 30 +- tests/unit/test_cc_agent_wrapper.py | 36 +- tests/unit/test_codex_agent_wrapper.py | 22 +- tests/unit/test_common_steps.py | 39 +- tests/unit/test_dream_checkpoint_versions.py | 12 +- tests/unit/test_evolve_utils.py | 28 +- tests/unit/test_http_client.py | 96 +++- tests/unit/test_injected_job_kwargs.py | 2 +- tests/unit/test_job_status.py | 17 +- tests/unit/test_lme_judge.py | 19 +- tests/unit/test_memlens_benchmark.py | 40 +- tests/unit/test_mobilemem_benchmark.py | 40 +- tests/unit/test_proactive_refresh.py | 4 +- tests/unit/test_token_usage.py | 116 ++--- tests/unit/test_workspace_web_steps.py | 14 +- 93 files changed, 3026 insertions(+), 1041 deletions(-) create mode 100644 reme/schema/agent_reply.py create mode 100644 tests/unit/test_agent_reply.py create mode 100644 tests/unit/test_agent_reply_consumers.py create mode 100644 tests/unit/test_agent_stream.py diff --git a/benchmark/memlens/run.py b/benchmark/memlens/run.py index 433270fc..41434075 100644 --- a/benchmark/memlens/run.py +++ b/benchmark/memlens/run.py @@ -265,15 +265,14 @@ def collect_usage(target: list | None = None, receipt: Path | None = None): def observe_usage(app) -> None: - """Observe owned clients and distinguish native interrupted replies from answers.""" + """Observe owned clients and reject incomplete agent replies.""" for wrapper in app.context.components.get("agent_wrapper", {}).values(): original_reply = wrapper.reply async def reply(*args, _reply=original_reply, **kwargs): result = await _reply(*args, **kwargs) - last = result.get("last_message") or {} - if last.get("finished_reason") == "interrupted" or last.get("error") is not None: - raise RuntimeError("Native Agent was interrupted; see the application log.") + if result["status"] != "completed": + raise RuntimeError(result["error"] or f"Native Agent reply ended: {result['status']}") return result wrapper.reply = reply diff --git a/benchmark/mobilemem/run.py b/benchmark/mobilemem/run.py index 28962ab4..c406eb6d 100644 --- a/benchmark/mobilemem/run.py +++ b/benchmark/mobilemem/run.py @@ -291,15 +291,14 @@ def observe_usage(app) -> None: def observe_completion(app) -> None: - """Keep native interrupted replies unscored; successful replies are unchanged.""" + """Keep incomplete native replies unscored; completed replies are unchanged.""" for wrapper in app.context.components.get("agent_wrapper", {}).values(): original = wrapper.reply async def reply(inputs, *args, _reply=original, **kwargs): result = await _reply(inputs, *args, **kwargs) - last = result.get("last_message") or {} - if str(last.get("finished_reason")) == "interrupted" or last.get("error") is not None: - raise RuntimeError("Native Agent was interrupted; its state remains saved.") + if result["status"] != "completed": + raise RuntimeError(result["error"] or f"Native Agent reply ended: {result['status']}") return result wrapper.reply = reply diff --git a/benchmark/pibench/bridge_reme.py b/benchmark/pibench/bridge_reme.py index dadd7877..8bc49509 100755 --- a/benchmark/pibench/bridge_reme.py +++ b/benchmark/pibench/bridge_reme.py @@ -787,8 +787,10 @@ class ReMeBridge: result = await self.agent_wrapper.reply(user_message, **reply_kwargs) - reply_text = result.get("result", "") - session_id = result.get("session_id", "") + from reme.steps.evolve._evolve import agent_reply_result_text + + reply_text = agent_reply_result_text(result) + session_id = result["session_id"] if session_id: self.agent_session_id = session_id diff --git a/docs/en/framework.md b/docs/en/framework.md index fb6757fa..f2e95200 100644 --- a/docs/en/framework.md +++ b/docs/en/framework.md @@ -557,6 +557,29 @@ Current default components in `reme/config/default.yaml`: | `keyword_index` | `default` | `bm25` | BM25 keyword index. | | `file_store` | `default` | `local` | Combines file_graph and keyword_index; defaults to `embedding_store: ""`. | +All three agent backends execute through `reply_stream()`. `reply()` collects that same stream into an +[`AgentReply`](../../reme/schema/agent_reply.py) dictionary: `session_id`, final text in `result`, +`structured_output`, invocation-scoped `usage`, `status`, `error`, `raw`, and the complete ordered JSON `events`. +Both modes accept `output_schema`. Only `completed` is successful; `failed`, `interrupted`, and `waiting` preserve +incomplete outcomes. Setup/transport errors and caller cancellation propagate. Both modes record reported token +usage once per invocation; unavailable usage remains null. + +[`StreamChunk`](../../reme/schema/stream_chunk.py) uses `start`, append-only `delta`, and `end` operations. +Non-null start/end payloads replace a block's accumulated content; null leaves it unchanged, and an empty string +clears it. Session, turn, message, block, and tool-call IDs correlate updates. `metadata.native` retains complete +provider events, including unknown fields and diagnostic notifications. `reply_end.chunk` carries the outcome +fields above, excluding `events`; `reply_start`/`reply_end` delimit one invocation, while DONE only closes transport. +`collect_reply(stream)` and `replay_reply(reply)` convert between the async stream and JSON result without losing +payloads, order, native metadata, or terminal state. Collection fails if the terminal outcome is missing or repeated. + +`raw` preserves native results: AgentScope's final message in `raw.message` plus an optional +`raw.structured_response`, Claude's `ResultMessage`, or Codex's `TurnResult`. Persist the full AgentReply for exact +replay; native session history supplies item snapshots rather than the original stream's deltas. +Studio reconciles block snapshots, displays structured output and terminal states, and retains live events through +history refreshes. Raw HTTP `stream_chunks()` includes errors and DONE; CLI/text helpers handle errors separately. +Migration: replace `last_message`/`turn` with `result`/`raw`, and consume the terminal outcome instead of concatenating +all stream payloads. Existing text-only producers use the default `delta` operation. + Codex wrappers reference `app_server: default` or `app_server: oauth`. Multiple wrappers can share one server while keeping their own models, instructions and conversations. Omit `session_id` to start a new thread; pass a returned `session_id` to continue it. Different threads can execute concurrently; operations on one thread are serialized. diff --git a/docs/zh/framework.md b/docs/zh/framework.md index 1910b1b1..f84a13ea 100644 --- a/docs/zh/framework.md +++ b/docs/zh/framework.md @@ -538,6 +538,26 @@ flowchart LR | `keyword_index` | `default` | `bm25` | BM25 关键词索引 | | `file_store` | `default` | `local` | 组合 file_graph、keyword_index；默认 `embedding_store: ""` | +三个 agent 后端统一通过 `reply_stream()` 执行，`reply()` 收集同一条流，返回 +[`AgentReply`](../../reme/schema/agent_reply.py) 字典：`session_id`、最终文本 `result`、`structured_output`、 +本次调用的 `usage`、`status`、`error`、`raw` 和完整、有序的 JSON `events`。两种模式均支持 `output_schema`。 +只有 `completed` 表示成功；`failed`、`interrupted` 和 `waiting` 保留未完成的结果。启动、传输异常和调用方的 +取消继续传播。两种模式均按调用记录一次已报告的 token 用量；无法获取用量时为 null。 + +[`StreamChunk`](../../reme/schema/stream_chunk.py) 使用 `start`、追加内容的 `delta` 和 `end` 操作。 +start/end 的非 null 内容替换对应块的累计内容；null 表示不变，空字符串表示清空。会话、调用、消息、内容块和 +工具调用 ID 用于关联更新。`metadata.native` 保留完整原生事件，包括未知字段和诊断通知。 +`reply_end.chunk` 携带上述结果字段，不含 `events`；`reply_start`/`reply_end` 包含一次 agent 调用内的模型与工具 +循环，DONE 仅关闭传输。`collect_reply(stream)` 与 `replay_reply(reply)` 在异步流和 JSON 结果间无损转换，保留 +payload、顺序、原生元数据和终止状态；缺少或重复的终止结果会明确报错。 + +`raw` 保留原生结果：AgentScope 的最终消息在 `raw.message`，额外结构化结果在 `raw.structured_response`； +Claude 保留 `ResultMessage`，Codex 保留 `TurnResult`。需要原样重放时应保存完整 AgentReply；原生会话历史提供 +item 快照，不含原始流的逐次增量。Studio 合并内容快照，展示结构化输出和终止状态，历史刷新保留已有实时事件。 +HTTP `stream_chunks()` 完整返回错误事件和 DONE，CLI 与文本辅助函数单独处理报错。 +迁移：用 `result`/`raw` 替代 `last_message`/`turn`，消费终止结果，不再拼接所有流式 payload。 +纯文本生产者使用默认 `delta` 操作。 + Codex wrapper 通过 `app_server: default` 或 `app_server: oauth` 引用共享组件。多个 wrapper 可以共享同一个进程， 并分别配置模型、指令和会话。不传 `session_id` 时创建新 thread；传回已有 `session_id` 时续聊。 不同 thread 可以并发，同一 thread 的操作串行。 diff --git a/plugins/auto-fin/src/reme_auto_fin/base.py b/plugins/auto-fin/src/reme_auto_fin/base.py index 6eac42fe..9218cc33 100644 --- a/plugins/auto-fin/src/reme_auto_fin/base.py +++ b/plugins/auto-fin/src/reme_auto_fin/base.py @@ -325,10 +325,9 @@ class AutoFinStep(BaseStep): if tool_context_id: kwargs["tool_context_id"] = tool_context_id result = await self.agent_wrapper.reply(prompt, **kwargs) - if not isinstance(result, dict) or result.get("structured_output") is None: - raise ValueError(f"Auto Fin Agent returned no structured output: {self._preview(result)}") - value = result["structured_output"] - output = value if isinstance(value, model) else model.model_validate(value) + if result["status"] != "completed": + raise RuntimeError(result["error"] or f"Agent reply ended: {result['status']}") + output = model.model_validate(result["structured_output"]) rendered = self._preview(output.model_dump(), AGENT_OUTPUT_LOG_LIMIT) self.logger.info( f"[{self.name}] agent output prompt={prompt_name} schema={model.__name__} " diff --git a/plugins/auto-fin/src/reme_auto_fin/topic.py b/plugins/auto-fin/src/reme_auto_fin/topic.py index 993ebaa3..98c5d711 100644 --- a/plugins/auto-fin/src/reme_auto_fin/topic.py +++ b/plugins/auto-fin/src/reme_auto_fin/topic.py @@ -7,6 +7,8 @@ import json import re from time import perf_counter +from reme.steps.evolve._evolve import agent_reply_result_text + from .base import AGENT_INPUT_LOG_LIMIT, AGENT_OUTPUT_LOG_LIMIT, AutoFinStep @@ -43,7 +45,7 @@ class AutoFinTopicStep(AutoFinStep): started_at = perf_counter() result = await self.agent_wrapper.reply(prompt) try: - ids = self._parse_news_ids(result.get("result") if isinstance(result, dict) else None, topics) + ids = self._parse_news_ids(agent_reply_result_text(result), topics) except ValueError as exc: if attempt: raise ValueError(f"Auto Fin topic Agent returned invalid news IDs: {exc}") from exc diff --git a/plugins/auto-fin/tests/test_auto_fin.py b/plugins/auto-fin/tests/test_auto_fin.py index d84d9f68..28407921 100644 --- a/plugins/auto-fin/tests/test_auto_fin.py +++ b/plugins/auto-fin/tests/test_auto_fin.py @@ -149,7 +149,10 @@ class _TopicAgent(BaseAgentWrapper): async def reply(self, inputs, **kwargs): self.calls.append((str(inputs), kwargs)) - return {"result": f"筛选结果：\n```json\n{json.dumps(self.selected)}\n```\n以上是相关 ID。"} + return { + "status": "completed", + "result": f"筛选结果：\n```json\n{json.dumps(self.selected)}\n```\n以上是相关 ID。", + } @pytest.mark.asyncio @@ -221,7 +224,10 @@ async def test_topic_step_retries_invalid_json_once(tmp_path: Path): async def reply(self, inputs, **kwargs): self.calls.append((str(inputs), kwargs)) - return {"result": ('{"new_ids": "[\\"1\\"]"}' if len(self.calls) == 1 else '```json\n{"黄金": ["1"]}\n```')} + return { + "status": "completed", + "result": ('{"new_ids": "[\\"1\\"]"}' if len(self.calls) == 1 else '```json\n{"黄金": ["1"]}\n```'), + } app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") agent = RetryAgent({"黄金": []}, app_context=app_context) @@ -293,6 +299,7 @@ class _ReportAgent(BaseAgentWrapper): topic = re.search(r"当前主题：(\S+)", prompt) title = f"{topic.group(1)}观察" if topic else "主题新闻观察" return { + "status": "completed", "structured_output": AutoFinReportOutput( title=f"# {title}", description="关注政策变化。", @@ -363,6 +370,7 @@ async def test_research_names_the_note_after_the_topic_not_the_agent_title(tmp_p class _VerboseAgent(BaseAgentWrapper): async def reply(self, _prompt, **_kwargs): return { + "status": "completed", "structured_output": AutoFinReportOutput( title="机器人主题 9-18 研究：" + "量产与政策" * 60, description="关注政策变化。", diff --git a/plugins/beam-judge/src/judge_beam/llm_judge.py b/plugins/beam-judge/src/judge_beam/llm_judge.py index 40e420eb..4c7213b0 100644 --- a/plugins/beam-judge/src/judge_beam/llm_judge.py +++ b/plugins/beam-judge/src/judge_beam/llm_judge.py @@ -25,6 +25,7 @@ import numpy as np from json_repair import repair_json from reme.steps.base_step import BaseStep, Ref +from reme.steps.evolve._evolve import agent_reply_result_text from reme.components.as_embedding import BaseAsEmbedding from reme.enumeration import ComponentEnum @@ -77,7 +78,7 @@ async def _llm_equivalence(agent_wrapper, reference: str, system: str) -> bool: user_prompt = f"First snippet: {reference}\n\nSecond snippet: {system}" result = await agent_wrapper.reply(user_prompt, system_prompt=system_prompt) - raw = (result.get("result") or "").strip().lower() + raw = agent_reply_result_text(result).lower() return "yes" in raw @@ -277,7 +278,7 @@ class BeamRubricJudgeStep(BaseStep): prompt = judge_template.replace("", item).replace("", llm_response) result = await self.agent_wrapper.reply(prompt) - raw = (result.get("result") or "").strip() + raw = agent_reply_result_text(result) try: parsed = _parse_json_response(raw) diff --git a/plugins/beam/src/reme_beam/agentic_answer.py b/plugins/beam/src/reme_beam/agentic_answer.py index e852be3a..ca7d5be2 100644 --- a/plugins/beam/src/reme_beam/agentic_answer.py +++ b/plugins/beam/src/reme_beam/agentic_answer.py @@ -2,7 +2,7 @@ import os -from reme.enumeration import ChunkEnum +from reme.schema import collect_reply from reme.steps.base_step import BaseStep from reme.steps.index._dedup import _ToolContextDedupMixin from reme.utils.counter import global_counter_inc @@ -61,16 +61,24 @@ class BeamAgenticAnswerStep(BaseStep): if injected_job_kwargs := self._injected_job_kwargs(f"{query}(query time: {query_time})"): wrapper_kwargs["injected_job_kwargs"] = injected_job_kwargs - if self.context.stream: - text = await self._stream_reply(query, **wrapper_kwargs) - else: - result = await self.agent_wrapper.reply(query, **wrapper_kwargs) - text = (result.get("result") or "").strip() + try: + if self.context.stream: + result = await collect_reply( + self.agent_wrapper.reply_stream(query, **wrapper_kwargs), + on_chunk=self.context.add_stream_chunk, + ) + else: + result = await self.agent_wrapper.reply(query, **wrapper_kwargs) + finally: + if self.app_context is not None: + self.app_context.metadata.get(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {}).pop(tool_context_id, None) + text = result["result"].strip() + self.context.response.metadata.update(result) self.logger.debug(f"[{self.name}] response: {text!r}") - self.context.response.success = True - self.context.response.answer = text + self.context.response.success = result["status"] == "completed" and self.context.stream_error is None + self.context.response.answer = result["error"] or text self.context.response.metadata.update( { "query": query, @@ -79,23 +87,4 @@ class BeamAgenticAnswerStep(BaseStep): "response": text, }, ) - - if self.app_context is not None: - self.app_context.metadata.get(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {}).pop(tool_context_id, None) return self.context.response - - async def _stream_reply(self, query: str, **wrapper_kwargs) -> str: - """Stream unified chunks to the context stream queue.""" - assert self.context is not None - text_parts: list[str] = [] - - async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - await self.context.add_stream_string(chunk.chunk, chunk.chunk_type) - - if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): - text_parts.append(chunk.chunk) - - if chunk.session_id: - self.context.response.metadata["session_id"] = chunk.session_id - - return "".join(text_parts).strip() diff --git a/plugins/daily_paper/src/reme_daily_paper/base.py b/plugins/daily_paper/src/reme_daily_paper/base.py index b58b8f9d..bf61b2da 100644 --- a/plugins/daily_paper/src/reme_daily_paper/base.py +++ b/plugins/daily_paper/src/reme_daily_paper/base.py @@ -30,8 +30,9 @@ _OutputT = TypeVar("_OutputT", bound=BaseModel) def structured_output(result: dict[str, Any], model: type[_OutputT]) -> _OutputT: """Validate an agent wrapper's structured output.""" - value = result.get("structured_output") - return value if isinstance(value, model) else model.model_validate(value) + if result["status"] != "completed": + raise RuntimeError(result["error"] or f"Agent reply ended: {result['status']}") + return model.model_validate(result["structured_output"]) def strip_frontmatter(body: str) -> str: diff --git a/plugins/daily_paper/tests/test_daily_paper.py b/plugins/daily_paper/tests/test_daily_paper.py index a6d514a9..2c11d81a 100644 --- a/plugins/daily_paper/tests/test_daily_paper.py +++ b/plugins/daily_paper/tests/test_daily_paper.py @@ -82,7 +82,12 @@ class _QueuedAgentWrapper(BaseAgentWrapper): async def reply(self, inputs, **kwargs) -> dict: """Record the request and pop the next structured fixture.""" self.calls.append({"inputs": inputs, "kwargs": kwargs}) - return {"structured_output": self.outputs.pop(0), "result": "ok"} + return self._reply_result( + session_id="session", + result="ok", + raw={}, + structured_output=self.outputs.pop(0), + ) def _paper(arxiv_id: str, *, title: str = "Paper", upvotes: int = 10) -> PaperInfo: diff --git a/plugins/dingtalk/src/reme_dingtalk/wait.py b/plugins/dingtalk/src/reme_dingtalk/wait.py index ce49bc3b..94089642 100644 --- a/plugins/dingtalk/src/reme_dingtalk/wait.py +++ b/plugins/dingtalk/src/reme_dingtalk/wait.py @@ -9,6 +9,7 @@ from urllib.parse import quote_plus from reme.components.agent_wrapper import handle_session_command from reme.steps.base_step import BaseStep +from reme.steps.evolve._evolve import agent_reply_result_text def _session_key(message: Any) -> str: @@ -151,17 +152,11 @@ class DingTalkWaitStep(BaseStep): builtin_tools=self.builtin_tools, job_tools=self.job_tools, ) - if not isinstance(result, dict): - raise TypeError("Agent reply must be a dictionary") - if session_id := result.get("session_id"): + if session_id := result["session_id"]: sessions[key] = session_id - last_message = result.get("last_message") - if isinstance(last_message, dict) and last_message.get("is_error"): - raise RuntimeError("Agent 执行失败") - - answer = result.get("result") - if not isinstance(answer, str) or not (answer := answer.strip()): + answer = agent_reply_result_text(result) + if not answer: raise ValueError("Agent 返回了空回复") response = await asyncio.to_thread(handler.reply_markdown, "ReMe Agent", answer, message) if response is None: diff --git a/plugins/dingtalk/tests/test_dingtalk.py b/plugins/dingtalk/tests/test_dingtalk.py index 282ac25e..ec21a12d 100644 --- a/plugins/dingtalk/tests/test_dingtalk.py +++ b/plugins/dingtalk/tests/test_dingtalk.py @@ -155,7 +155,8 @@ class _AgentWrapper(BaseAgentWrapper): self.reply_calls = [] self.compact_calls = [] self.result_text = "回答" - self.is_error = False + self.status = "completed" + self.error = None async def compact_session(self, session_id): self.compact_calls.append(session_id) @@ -163,11 +164,13 @@ class _AgentWrapper(BaseAgentWrapper): async def reply(self, inputs, **kwargs): self.reply_calls.append((inputs, kwargs)) session_id = kwargs.get("resume") or "session-1" - return { - "session_id": session_id, - "last_message": {"is_error": self.is_error}, - "result": self.result_text, - } + return self._reply_result( + session_id=session_id, + result=self.result_text, + raw={}, + status=self.status, + error=self.error, + ) class _Handler: @@ -403,3 +406,23 @@ async def test_stream_client_reconnects_after_server_request(monkeypatch, tmp_pa step.logger.info.assert_called_once_with( "[DingTalkWaitStep] DingTalk server requested reconnect reason='connection is expired'; reconnecting in 1.0s", ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +@pytest.mark.parametrize("error", [None, "provider failed"]) +async def test_incomplete_reply_is_not_sent_as_a_success(tmp_path, status, error): + """Partial text cannot turn a failed/interrupted/waiting AgentReply into a Markdown reply.""" + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = _AgentWrapper(app_context=app_context) + wrapper.status, wrapper.error = status, error + wrapper.result_text = "partial answer" + step = DingTalkWaitStep(app_context=app_context, agent_wrapper=wrapper) + step.logger = MagicMock() + handler, sessions, message = _Handler(), {}, _message() + key = _session_key(message) + with pytest.raises(RuntimeError, match=error or status): + await step._handle_message(message, key, sessions, handler) + assert not handler.markdown_replies + assert sessions[key] == "session-1" + assert not any("completed DingTalk reply" in call.args[0] for call in step.logger.info.call_args_list) diff --git a/plugins/home/README.md b/plugins/home/README.md index 504cddd1..7923a2b3 100644 --- a/plugins/home/README.md +++ b/plugins/home/README.md @@ -643,9 +643,12 @@ and `LLM_BASE_URL` without changing the evaluation judge's model. Override startup, token accounting and saved session transcripts. The response answer is `{"results": [...]}`; each item contains `case_id`, `success`, and either `report` with its session ID/usage or `error`. Metadata includes `evaluation_path`, the -evaluated commit and total `usage` (null if any case fails or usage is unavailable). +evaluated commit and total `usage` (null if any usage is unavailable). Response success is false if any case fails; completed sibling reports remain available. No case scratch directory or intermediate analysis files are created. +Returned reports include the agent's `status` and `error` in single-case metadata +or each batch result. Only `completed` succeeds; interrupted, waiting and +iteration-exhausted replies retain their report text and reported usage with `success=false`. The host bounds each returned report using `output_token_limit` (default 2,000) with a 10% allowance, rounded down. The agent sees only the configured budget. diff --git a/plugins/home/src/reme_home/agent/case_study.py b/plugins/home/src/reme_home/agent/case_study.py index 119b2652..b3d03855 100644 --- a/plugins/home/src/reme_home/agent/case_study.py +++ b/plugins/home/src/reme_home/agent/case_study.py @@ -35,7 +35,15 @@ CASE_TOOLS = ( "home_compare_memory", ) CODE_TOOLS = ("home_list_code_tree", "home_search_code", "home_read_code_file") -REPORT_METADATA = ("session_id", "usage", "estimated_tokens", "original_estimated_tokens", "truncated") +REPORT_METADATA = ( + "session_id", + "usage", + "status", + "error", + "estimated_tokens", + "original_estimated_tokens", + "truncated", +) _SANDBOX_PROBE = """\ import json import pathlib @@ -341,6 +349,7 @@ class HomeCaseStudyStep(BaseStep): # Preserve the original text response and error handling for singular callers. if "case_id" in self.context: result = await self._investigate(scopes[0], spec.focus, prompts, output_limit) + response.success = result["status"] == "completed" and result["error"] is None response.answer = result["result"] response.metadata.update(validation_path=scopes[0]["validation_path"], case_id=spec.case_ids[0]) response.metadata.update({key: result[key] for key in REPORT_METADATA if key in result}) @@ -353,7 +362,7 @@ class HomeCaseStudyStep(BaseStep): result = await self._investigate(scope, spec.focus, prompts, output_limit) return { "case_id": scope["case_id"], - "success": True, + "success": result["status"] == "completed" and result["error"] is None, "report": result["result"], **{key: result[key] for key in REPORT_METADATA if key in result}, } diff --git a/plugins/home/src/reme_home/agent/structured_report.py b/plugins/home/src/reme_home/agent/structured_report.py index 6ec2bb16..fb88f33a 100644 --- a/plugins/home/src/reme_home/agent/structured_report.py +++ b/plugins/home/src/reme_home/agent/structured_report.py @@ -5,9 +5,12 @@ from __future__ import annotations import json import logging import re +from collections.abc import AsyncGenerator +from contextlib import aclosing from typing import Any, TypeVar from agentscope.message import AssistantMsg, Msg, Usage, UserMsg +from agentscope.types import ReplyFinishedReason from pydantic import BaseModel from .bounded import BoundedReActAgent @@ -24,9 +27,16 @@ class StructuredReportAgent(BoundedReActAgent): super().__init__(*args, **kwargs) self.report_model = report_model - async def reply(self, *args, **kwargs) -> Msg: + async def _reply(self, *args, **kwargs) -> AsyncGenerator[Any, None]: + """Recover final reports through both native reply entry points.""" + async with aclosing(super()._reply(*args, **kwargs)) as stream: + async for item in stream: + if isinstance(item, Msg) and item.finished_reason not in {None, ReplyFinishedReason.INTERRUPTED}: + item = await self._finalize_report(item) + yield item + + async def _finalize_report(self, reply: Msg) -> Msg: """Keep normal wrapper persistence and accounting around report recovery.""" - reply = await super().reply(*args, **kwargs) try: report = self.report_model.model_validate(parse_json_object(reply, report_name=self.name)) except ValueError as exc: @@ -57,7 +67,17 @@ class StructuredReportAgent(BoundedReActAgent): usage = Usage( **{key: getattr(reply.usage, key) + getattr(corrected.usage, key, 0) for key in Usage.model_fields}, ) - reply = AssistantMsg(name=self.name, content=report.model_dump_json(), usage=usage) + reply = reply.model_copy( + update={ + "content": AssistantMsg(name=self.name, content=report.model_dump_json()).content, + "usage": usage, + "finished_reason": ( + ReplyFinishedReason.INTERRUPTED + if corrected.finished_reason == "interrupted" + else reply.finished_reason + ), + }, + ) self.state.context.extend([correction, reply]) return reply return reply.model_copy( diff --git a/plugins/home/src/reme_home/agent/workflow_common.py b/plugins/home/src/reme_home/agent/workflow_common.py index f6820dbf..6324d3b9 100644 --- a/plugins/home/src/reme_home/agent/workflow_common.py +++ b/plugins/home/src/reme_home/agent/workflow_common.py @@ -8,6 +8,7 @@ from uuid import uuid4 from agentscope.tool import BackendBase, DirEntry, Glob, Grep, Read, Toolkit, Write import yaml +from reme.steps.evolve._evolve import agent_reply_result_text from reme.utils.async_utils import complete_in_thread as blocking from ..auto_bash import AutoBash @@ -96,7 +97,7 @@ async def workflow_reply( if not probe.ok(): raise RuntimeError("Home workflow sandbox unavailable: " + probe.stderr.decode(errors="replace")) prompt = yaml.safe_load(Path(__file__).with_name(f"{role}.yaml").read_text(encoding="utf-8")) - return await wrapper.reply( + reply = await wrapper.reply( task, system_prompt=prompt["system_prompt"].format(**scope), toolkit=create_toolkit( @@ -116,6 +117,8 @@ async def workflow_reply( permission_mode="bypass", **(agent_options or {}), ) + agent_reply_result_text(reply) + return reply def prepare_scope(home, role): @@ -204,7 +207,11 @@ async def review_report( permission_mode="bypass", ) report.update({key: reply.get(key) for key in ("session_id", "usage", "result")}) - report.update(report_model.model_validate(parse_json_object(reply["result"], report_name=role)).model_dump()) + report.update( + report_model.model_validate( + parse_json_object(agent_reply_result_text(reply), report_name=role), + ).model_dump(), + ) except Exception as exc: report.update( decision=failure_decision, diff --git a/plugins/home/src/reme_home/runtime_docs/validation-artifact-guide.md b/plugins/home/src/reme_home/runtime_docs/validation-artifact-guide.md index ce1ac525..9087a72f 100644 --- a/plugins/home/src/reme_home/runtime_docs/validation-artifact-guide.md +++ b/plugins/home/src/reme_home/runtime_docs/validation-artifact-guide.md @@ -247,9 +247,11 @@ Call `home_case_study(evaluation_path, case_ids, focus=...)` for up to eight unique cases sharing one focus and evaluation. Cases run sequentially and results preserve input order. Each gets fresh agent state, bounded iterations and advisory time reminders. The answer contains `results`, with a report/session/usage or error -for each case; a failed case does not cancel the rest. Response success is false -if any case fails. Metadata includes `evaluation_path`, evaluated commit and total -usage (null if any call fails or usage is unavailable). `validation_path` remains +for each case; a failed case does not cancel the rest. Report results also include +the agent's `status` and `error`; only `completed` succeeds. Incomplete reports +retain their text and reported usage. Response success is false if any case fails. +Metadata includes `evaluation_path`, evaluated commit and total usage (null if any +usage is unavailable). `validation_path` remains a path alias; legacy `case_id` calls retain a text report. Do not combine aliases. Other inspection tools still use `validation_path`. There is no parallel case dispatcher. Scope and evaluated code revision are host-injected into the structured tools. AutoBash and Python share a mandatory read-only file diff --git a/plugins/home/tests/test_bounded_wrapper.py b/plugins/home/tests/test_bounded_wrapper.py index 2b261947..3bdc1c94 100644 --- a/plugins/home/tests/test_bounded_wrapper.py +++ b/plugins/home/tests/test_bounded_wrapper.py @@ -196,7 +196,8 @@ async def test_multiround_usage_and_tools_continue_past_soft_time_limit(wrapper, } result = await wrapper.reply("Investigate", **options) assert result["usage"] == {"input_tokens": 30, "output_tokens": 15, "total_tokens": 45} - assert result["last_message"]["finished_reason"] == "completed" + assert result["status"] == "completed" + assert result["raw"]["message"]["finished_reason"] == "completed" assert len(executions) == 2 assert executions[-1] > 60 # A tool really executes after the soft deadline. assert len(model.calls) == 3 diff --git a/plugins/home/tests/test_candidate_gate.py b/plugins/home/tests/test_candidate_gate.py index 0dec19d1..584580c2 100644 --- a/plugins/home/tests/test_candidate_gate.py +++ b/plugins/home/tests/test_candidate_gate.py @@ -149,7 +149,10 @@ async def test_nearest_ancestor_and_infrastructure_retry(app, prepared, sandbox, @pytest.mark.asyncio -@pytest.mark.parametrize("output", ["invalid", '{"decision":"reject","reason":"x","extra":true}', "error"]) +@pytest.mark.parametrize( + "output", + ["invalid", '{"decision":"reject","reason":"x","extra":true}', "error", "failed", "interrupted", "waiting"], +) async def test_model_failure_allows_validation_with_saved_diagnostics(app, prepared, sandbox, monkeypatch, output): commit_candidate(prepared) wrapper = set_model(app, "home_candidate_gate", ScriptedModel(final=output)) @@ -159,10 +162,24 @@ async def test_model_failure_allows_validation_with_saved_diagnostics(app, prepa raise RuntimeError("model unavailable") monkeypatch.setattr(wrapper, "reply", broken) + elif output in {"failed", "interrupted", "waiting"}: + monkeypatch.setattr( + wrapper, + "reply", + AsyncMock( + return_value={ + "result": '{"decision":"reject","reason":"Unfinished judgement"}', + "status": output, + "error": None, + }, + ), + ) response = await app.run_job("home_evaluate", max_evaluations=10) assert response.success, response.answer report = response.answer["candidate_gate"] assert report["decision"] == "accept" and report["review_failed"] + if output in {"failed", "interrupted", "waiting"}: + assert output in report["reason"] assert read_json(Path(report["review_path"]))["review_failed"] assert not list((prepared / "tmp").glob("candidate-gate-*")) @@ -216,7 +233,7 @@ async def test_review_change_race_and_cancellation_release_locks(app, prepared, async def paused(*args, **kwargs): entered.set() await release.wait() - return {"result": '{"decision":"accept","reason":"Mechanism change"}'} + return {"result": '{"decision":"accept","reason":"Mechanism change"}', "status": "completed", "error": None} monkeypatch.setattr(wrapper, "reply", paused) task = asyncio.create_task(app.run_job("home_evaluate", max_evaluations=10)) diff --git a/plugins/home/tests/test_case_study.py b/plugins/home/tests/test_case_study.py index 98049f15..6d15e447 100644 --- a/plugins/home/tests/test_case_study.py +++ b/plugins/home/tests/test_case_study.py @@ -225,7 +225,8 @@ async def test_invalid_scope_and_sandbox_failure_precede_llm(app, prepared, case @pytest.mark.asyncio -async def test_iteration_limit_returns_final_report(app, prepared, case_run, sandbox): +@pytest.mark.parametrize("batch", [False, True]) +async def test_iteration_limit_preserves_report_and_failure(app, prepared, case_run, sandbox, batch): run, case_id, _query_id, _commit, _source = case_run wrapper = app.context.components["agent_wrapper"]["home_case_study"] assert wrapper.kwargs["react_config"]["max_iters"] == 60 @@ -233,8 +234,12 @@ async def test_iteration_limit_returns_final_report(app, prepared, case_run, san wrapper.kwargs["react_config"] = {"max_iters": 1} model = RecordingModel() wrapper.as_llm = SimpleNamespace(model=model) - response = await app.run_job("home_case_study", validation_path=str(run), case_id=case_id) - assert response.success and response.answer == "final report" + arguments = {"case_ids": [case_id]} if batch else {"case_id": case_id} + response = await app.run_job("home_case_study", validation_path=str(run), **arguments) + assert not response.success + report = response.answer["results"][0] if batch else {**response.metadata, "report": response.answer} + assert report["report"] == "final report" + assert report["status"] == "failed" and "exceed_max_iters" in report["error"] assert model.calls[-1]["tools"] == [] @@ -572,7 +577,7 @@ async def test_batch_bounds_concurrency_and_preserves_order(app, many_case_run, completed.append(case_id) if case_id == case_ids[1]: raise RuntimeError("one investigator failed") - return {"result": f"report {case_id}"} + return {"result": f"report {case_id}", "status": "completed", "error": None} finally: active.remove(case_id) diff --git a/plugins/home/tests/test_construction_equivalence.py b/plugins/home/tests/test_construction_equivalence.py index 419b08f7..894e9537 100644 --- a/plugins/home/tests/test_construction_equivalence.py +++ b/plugins/home/tests/test_construction_equivalence.py @@ -51,7 +51,13 @@ def real_reviewer(app, prepared, monkeypatch): wrapper = app.context.components["agent_wrapper"]["home_construction_equivalence"] monkeypatch.setattr(wrapper, "reply", MethodType(type(wrapper).reply, wrapper)) gate = app.context.components["agent_wrapper"]["home_candidate_gate"] - monkeypatch.setattr(gate, "reply", AsyncMock(return_value={"result": '{"decision":"accept","reason":"Mechanism"}'})) + monkeypatch.setattr( + gate, + "reply", + AsyncMock( + return_value={"result": '{"decision":"accept","reason":"Mechanism"}', "status": "completed", "error": None}, + ), + ) return wrapper @@ -116,7 +122,7 @@ async def test_large_resources_export_for_both_reviewers(app, prepared, reviewer async def accept(*args, **kwargs): inspect_export("candidate-gate") - return {"result": '{"decision":"accept","reason":"Fixture export check"}'} + return {"result": '{"decision":"accept","reason":"Fixture export check"}', "status": "completed", "error": None} def finish(): inspect_export("construction-equivalence") @@ -375,7 +381,11 @@ async def test_review_checkout_race_and_cancellation_do_not_sample(app, prepared async def paused(*args, **kwargs): entered.set() await release.wait() - return {"result": '{"decision":"construction_unchanged","reason":"Query only"}'} + return { + "result": '{"decision":"construction_unchanged","reason":"Query only"}', + "status": "completed", + "error": None, + } monkeypatch.setattr(reviewer, "reply", paused) task = asyncio.create_task(app.run_job("home_evaluate", sampling=True)) diff --git a/plugins/home/tests/test_evaluation.py b/plugins/home/tests/test_evaluation.py index 1c87688c..94ce5697 100644 --- a/plugins/home/tests/test_evaluation.py +++ b/plugins/home/tests/test_evaluation.py @@ -189,12 +189,22 @@ async def prepared_home(app, inputs, monkeypatch): monkeypatch.setattr(evaluation, "DockerReMeSandboxFactory", FakeFactory) # Run the real judge Job and Step; mock only its model-facing wrapper. wrapper = app.context.components["agent_wrapper"]["judge"] - monkeypatch.setattr(wrapper, "reply", AsyncMock(return_value={"result": '{"score": 0.5}'})) + monkeypatch.setattr( + wrapper, + "reply", + AsyncMock(return_value={"result": '{"score": 0.5}', "status": "completed", "error": None}), + ) reviewer = app.context.components["agent_wrapper"]["home_construction_equivalence"] monkeypatch.setattr( reviewer, "reply", - AsyncMock(return_value={"result": '{"decision":"construction_changed","reason":"Fixture rebuild"}'}), + AsyncMock( + return_value={ + "result": '{"decision":"construction_changed","reason":"Fixture rebuild"}', + "status": "completed", + "error": None, + }, + ), ) return Path(app.config.workspace_dir) / "home" @@ -565,7 +575,16 @@ def test_unsafe_archives_are_rejected_without_partial_extraction(tmp_path, kind) @pytest.mark.parametrize("app", ["lme-judge"], indirect=True) async def test_lme_evaluation_calls_host_judge_job(app, prepared, monkeypatch): wrapper = app.context.components["agent_wrapper"]["judge"] - monkeypatch.setattr(wrapper, "reply", AsyncMock(side_effect=[{"result": "maybe"}, {"result": "yes"}])) + monkeypatch.setattr( + wrapper, + "reply", + AsyncMock( + side_effect=[ + {"result": "maybe", "status": "completed", "error": None}, + {"result": "yes", "status": "completed", "error": None}, + ], + ), + ) response = await app.run_job("home_evaluate") assert response.success, response.answer assert response.answer["mean_query_score"] == 1 @@ -585,7 +604,7 @@ async def test_lme_evaluation_calls_host_judge_job(app, prepared, monkeypatch): @pytest.mark.parametrize("app", ["lme-judge"], indirect=True) async def test_invalid_lme_judge_is_infra_error_after_one_retry(app, prepared): wrapper = app.context.components["agent_wrapper"]["judge"] - wrapper.reply.return_value = {"result": "maybe"} + wrapper.reply.return_value = {"result": "maybe", "status": "completed", "error": None} response = await app.run_job("home_evaluate") assert not response.success assert response.answer["mean_query_score"] is None @@ -597,7 +616,11 @@ async def test_invalid_lme_judge_is_infra_error_after_one_retry(app, prepared): @pytest.mark.asyncio @pytest.mark.parametrize("raw", ["nonsense", '{"score": 2}', '{"score": 0.2}', '{"score": NaN}']) async def test_invalid_beam_judge_is_infra_error(app, prepared, raw): - app.context.components["agent_wrapper"]["judge"].reply.return_value = {"result": raw} + app.context.components["agent_wrapper"]["judge"].reply.return_value = { + "result": raw, + "status": "completed", + "error": None, + } response = await app.run_job("home_evaluate", case_ids=["1"]) assert not response.success assert response.answer["mean_query_score"] is None diff --git a/plugins/home/tests/test_evaluation_history.py b/plugins/home/tests/test_evaluation_history.py index 088039ef..08ca1972 100644 --- a/plugins/home/tests/test_evaluation_history.py +++ b/plugins/home/tests/test_evaluation_history.py @@ -245,7 +245,11 @@ async def test_history_catalog_keeps_successful_low_scores_and_deduplicates(app, manager = history(app) assert isinstance(manager, HomeEvaluationHistory) and manager.is_started assert manager.list_results() == [] - app.context.components["agent_wrapper"]["judge"].reply.return_value = {"result": '{"score": 0}'} + app.context.components["agent_wrapper"]["judge"].reply.return_value = { + "result": '{"score": 0}', + "status": "completed", + "error": None, + } response = await app.run_job("home_evaluate") assert response.success, response.answer results = manager.list_results() @@ -644,7 +648,7 @@ async def run(): app = Application(**json.loads(sys.argv[2])) await app.start() app.context.components["agent_wrapper"]["judge"].reply = AsyncMock( - return_value={"result": '{"score": 0.5}'}) + return_value={"result": '{"score": 0.5}', "status": "completed", "error": None}) app.context.components["home.evaluation_history"]["default"].record = lambda run: os._exit(23) await app.run_job("home_evaluate") asyncio.run(run()) @@ -663,7 +667,7 @@ asyncio.run(run()) await restarted.start() try: judge = restarted.context.components["agent_wrapper"]["judge"] - judge.reply = AsyncMock(return_value={"result": '{"score": 0.5}'}) + judge.reply = AsyncMock(return_value={"result": '{"score": 0.5}', "status": "completed", "error": None}) manager = history(restarted) assert len(manager.list_results()) == 1 assert read_state(prepared)["successful_round"] == 1 diff --git a/plugins/home/tests/test_evaluation_repetition.py b/plugins/home/tests/test_evaluation_repetition.py index 7e6e7c30..0e2bb5fa 100644 --- a/plugins/home/tests/test_evaluation_repetition.py +++ b/plugins/home/tests/test_evaluation_repetition.py @@ -195,7 +195,7 @@ async def test_sampled_recovery_preserves_cases_queries_and_order( monkeypatch.setattr( app.context.components["agent_wrapper"]["judge"], "reply", - AsyncMock(return_value={"result": '{"score": 0.5}'}), + AsyncMock(return_value={"result": '{"score": 0.5}', "status": "completed", "error": None}), ) monkeypatch.setattr("reme_home.sampling.time.time_ns", lambda: 1) monkeypatch.setattr(FakeFactory, "bad_export", True) diff --git a/plugins/home/tests/test_home_eval.py b/plugins/home/tests/test_home_eval.py index 033b106d..ce5f2b09 100644 --- a/plugins/home/tests/test_home_eval.py +++ b/plugins/home/tests/test_home_eval.py @@ -159,7 +159,11 @@ async def test_full_evaluation_records_source_provenance_with_fake_docker( path, _ = write_config(tmp_path, app, spec, commit) monkeypatch.setattr(evaluation, "DockerReMeSandboxFactory", FakeFactory) answer = "yes" if spec["benchmark"]["name"] == "longmemeval" else '{"score": 0.5}' - monkeypatch.setattr(AsAgentWrapper, "reply", AsyncMock(return_value={"result": answer})) + monkeypatch.setattr( + AsAgentWrapper, + "reply", + AsyncMock(return_value={"result": answer, "status": "completed", "error": None}), + ) before = snapshot(home) result = await home_eval.run_evaluation(path) run = Path(result["evaluation"]["result_dir"]) diff --git a/plugins/home/tests/test_sampling.py b/plugins/home/tests/test_sampling.py index 8a5ba91f..9b457291 100644 --- a/plugins/home/tests/test_sampling.py +++ b/plugins/home/tests/test_sampling.py @@ -129,7 +129,11 @@ async def test_baseline_then_adaptive_job_updates_only_selected_cases(app, prepa assert all(record["ewma_loss"] == 0.5 for record in state(prepared)["cases"].values()) evaluation = app.context.components["home.evaluation"]["default"] evaluation.sampling_policy = sampling.SamplingPolicy(max_cases=1, max_queries=1, ewma_alpha=0.5) - app.context.components["agent_wrapper"]["judge"].reply.return_value = {"result": '{"score": 1}'} + app.context.components["agent_wrapper"]["judge"].reply.return_value = { + "result": '{"score": 1}', + "status": "completed", + "error": None, + } before = copy.deepcopy(state(prepared)) commit_candidate(prepared) response = await app.run_job("home_evaluate", sampling=True) @@ -304,7 +308,11 @@ async def test_sampling_cache_rebuild_and_new_round_after_reopen(app, prepared): @pytest.mark.asyncio async def test_fail_fast_started_round_without_summary_is_not_quality_feedback(app, prepared): - app.context.components["agent_wrapper"]["judge"].reply.return_value = {"result": "invalid judge output"} + app.context.components["agent_wrapper"]["judge"].reply.return_value = { + "result": "invalid judge output", + "status": "completed", + "error": None, + } response = await app.run_job("home_evaluate", sampling=True, fail_fast=True) assert not response.success observed = state(prepared) @@ -401,7 +409,7 @@ async def test_query_budget_and_draw_order_reach_actual_execution(app, inputs, m monkeypatch.setattr( app.context.components["agent_wrapper"]["judge"], "reply", - AsyncMock(return_value={"result": '{"score": 0.5}'}), + AsyncMock(return_value={"result": '{"score": 0.5}', "status": "completed", "error": None}), ) monkeypatch.setattr(sampling.time, "time_ns", lambda: 123456) component = app.context.components["home.evaluation"]["default"] diff --git a/plugins/home/tests/test_structured_report.py b/plugins/home/tests/test_structured_report.py index 1bb45653..65e8ffd6 100644 --- a/plugins/home/tests/test_structured_report.py +++ b/plugins/home/tests/test_structured_report.py @@ -4,17 +4,23 @@ import asyncio import json +from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from agentscope.agent import ReActConfig -from agentscope.message import AssistantMsg, UserMsg -from agentscope.model import ChatUsage, StructuredResponse +from agentscope.message import AssistantMsg, Msg, UserMsg +from agentscope.model import ChatUsage, FinishedReason, StructuredResponse +from reme_home.agent import BoundedAgentWrapper, BoundedReActAgent from reme_home.agent.candidate_gate import CandidateGateReport from reme_home.agent.structured_report import StructuredReportAgent, parse_json_object from test_workflows import ScriptedModel +from reme.components.application_context import ApplicationContext +from reme.schema import collect_reply, replay_reply +from reme.utils import global_counter_get + @pytest.mark.parametrize( "text", @@ -87,11 +93,12 @@ def make_agent(model): @pytest.mark.asyncio @pytest.mark.parametrize("missing_usage", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize( "initial", ["Reject: mechanism unchanged.", '{"decision":"reject","reason":"x","extra":true}'], ) -async def test_one_correction_retains_history_disables_tools_and_accounts_usage(missing_usage, initial): +async def test_one_correction_retains_history_disables_tools_and_accounts_usage(missing_usage, streaming, initial): model = ScriptedModel(final=initial) model.extra_body = {"enable_thinking": True, "thinking_budget": 100, "host_option": "retained"} model.generate_structured_output = AsyncMock( @@ -102,7 +109,14 @@ async def test_one_correction_retains_history_disables_tools_and_accounts_usage( ) agent = make_agent(model) agent.state.summary = "Earlier investigation summary." - reply = await agent.reply(UserMsg(name="host", content="Review this candidate.")) + inputs = UserMsg(name="host", content="Review this candidate.") + if streaming: + items = [item async for item in agent.reply_stream(inputs, yield_final_msg=True)] + reply = items[-1] + assert isinstance(reply, Msg) + else: + reply = await agent.reply(inputs) + assert reply.finished_reason == "completed" assert CandidateGateReport.model_validate_json(reply.get_text_content()).decision == "reject" assert len(model.calls) == 1 model.generate_structured_output.assert_awaited_once() @@ -158,3 +172,50 @@ async def test_correction_is_not_retried_and_cancellation_propagates(failure): with pytest.raises(expected): await make_agent(model).reply(UserMsg(name="host", content="Review.")) model.generate_structured_output.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("interrupted", [False, True]) +async def test_wrapper_collects_and_persists_recovered_report(tmp_path, streaming, interrupted): + model = ScriptedModel(final="Reject: mechanism unchanged.") + model.generate_structured_output = AsyncMock( + return_value=StructuredResponse( + content={"decision": "reject", "reason": "Mechanism unchanged"}, + usage=ChatUsage(input_tokens=7, output_tokens=11, time=0), + finished_reason=FinishedReason.INTERRUPTED if interrupted else FinishedReason.COMPLETED, + ), + ) + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = BoundedAgentWrapper(app_context=context, as_llm="", name="review") + wrapper.as_llm = SimpleNamespace(model=model) + options = {"_structured_report_model": CandidateGateReport, "builtin_tools": []} + if streaming: + result = await collect_reply(wrapper.reply_stream("Review this candidate.", **options)) + else: + result = await wrapper.reply("Review this candidate.", **options) + assert result["status"] == ("interrupted" if interrupted else "completed") + assert result["error"] is None + assert CandidateGateReport.model_validate_json(result["result"]).decision == "reject" + assert result["raw"]["message"]["finished_reason"] == result["status"] + assert result["usage"] == {"input_tokens": 9, "output_tokens": 14, "total_tokens": 23} + assert global_counter_get(context.metadata, ["__token_counter", "review", "total_tokens"]) == 23 + assert await collect_reply(replay_reply(result)) == result + saved = (wrapper.session_path / f"{result['session_id']}.jsonl").read_text() + assert "home_report_recovery" in saved and "Mechanism unchanged" in saved + model.generate_structured_output.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reason", [None, "interrupted"]) +async def test_incomplete_native_reply_does_not_request_report_correction(monkeypatch, reason): + reply = AssistantMsg(name="review", content="partial", finished_reason=reason) + + async def incomplete_reply(self, *args, **kwargs): + yield reply + + monkeypatch.setattr(BoundedReActAgent, "_reply", incomplete_reply) + model = ScriptedModel() + model.generate_structured_output = AsyncMock() + assert await make_agent(model).reply() is reply + model.generate_structured_output.assert_not_awaited() diff --git a/plugins/home/tests/test_workflows.py b/plugins/home/tests/test_workflows.py index e26fbec5..6f700cb5 100644 --- a/plugins/home/tests/test_workflows.py +++ b/plugins/home/tests/test_workflows.py @@ -84,6 +84,19 @@ async def baseline(app): return result.answer["result_dir"] +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +async def test_workflow_rejects_incomplete_agent_reports(app, prepared, sandbox, monkeypatch, status): + wrapper = app.context.components["agent_wrapper"]["home_diagnostic"] + monkeypatch.setattr( + wrapper, + "reply", + AsyncMock(return_value={"result": "Unfinished diagnosis", "status": status, "error": None}), + ) + with pytest.raises(RuntimeError, match=status): + await workflow_common.workflow_reply(wrapper, prepared, "diagnostic", "Diagnose this failure.") + + def final_report(home, branch="init", path=None): repo = home / "code/repo" if path is None: diff --git a/plugins/lme-judge/src/judge_lme/llm_judge.py b/plugins/lme-judge/src/judge_lme/llm_judge.py index fb220df4..2dbb99df 100644 --- a/plugins/lme-judge/src/judge_lme/llm_judge.py +++ b/plugins/lme-judge/src/judge_lme/llm_judge.py @@ -3,6 +3,7 @@ import re from reme.steps.base_step import BaseStep +from reme.steps.evolve._evolve import agent_reply_result_text class LmeAnswerJudgeStep(BaseStep): @@ -58,7 +59,7 @@ class LmeAnswerJudgeStep(BaseStep): system_prompt = self.prompt_format(judge_prompt_key) for _ in range(2): # One initial attempt and at most one invalid-output retry. result = await self.agent_wrapper.reply(user_prompt, system_prompt=system_prompt) - raw_answer = (result.get("result") or "").strip() + raw_answer = agent_reply_result_text(result) answer = self._normalize_judgement(raw_answer) if answer in {"yes", "no"}: break diff --git a/plugins/lme/src/reme_lme/agentic_answer.py b/plugins/lme/src/reme_lme/agentic_answer.py index ffa36fd3..b5b19c15 100644 --- a/plugins/lme/src/reme_lme/agentic_answer.py +++ b/plugins/lme/src/reme_lme/agentic_answer.py @@ -2,7 +2,7 @@ import os -from reme.enumeration import ChunkEnum +from reme.schema import collect_reply from reme.steps.base_step import BaseStep from reme.steps.index._dedup import _ToolContextDedupMixin from reme.utils.counter import global_counter_inc @@ -61,16 +61,24 @@ class LmeAgenticAnswerStep(BaseStep): if injected_job_kwargs := self._injected_job_kwargs(f"{query}(query time: {query_time})"): wrapper_kwargs["injected_job_kwargs"] = injected_job_kwargs - if self.context.stream: - text = await self._stream_reply(query, **wrapper_kwargs) - else: - result = await self.agent_wrapper.reply(query, **wrapper_kwargs) - text = (result.get("result") or "").strip() + try: + if self.context.stream: + result = await collect_reply( + self.agent_wrapper.reply_stream(query, **wrapper_kwargs), + on_chunk=self.context.add_stream_chunk, + ) + else: + result = await self.agent_wrapper.reply(query, **wrapper_kwargs) + finally: + if self.app_context is not None: + self.app_context.metadata.get(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {}).pop(tool_context_id, None) + text = result["result"].strip() + self.context.response.metadata.update(result) self.logger.debug(f"[{self.name}] response: {text!r}") - self.context.response.success = True - self.context.response.answer = text + self.context.response.success = result["status"] == "completed" and self.context.stream_error is None + self.context.response.answer = result["error"] or text self.context.response.metadata.update( { "query": query, @@ -79,23 +87,4 @@ class LmeAgenticAnswerStep(BaseStep): "response": text, }, ) - - if self.app_context is not None: - self.app_context.metadata.get(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {}).pop(tool_context_id, None) return self.context.response - - async def _stream_reply(self, query: str, **wrapper_kwargs) -> str: - """Stream unified chunks to the context stream queue.""" - assert self.context is not None - text_parts: list[str] = [] - - async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - await self.context.add_stream_string(chunk.chunk, chunk.chunk_type) - - if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): - text_parts.append(chunk.chunk) - - if chunk.session_id: - self.context.response.metadata["session_id"] = chunk.session_id - - return "".join(text_parts).strip() diff --git a/plugins/memlens/src/reme_memlens/agentic_answer.py b/plugins/memlens/src/reme_memlens/agentic_answer.py index 2115d95d..f02609d1 100644 --- a/plugins/memlens/src/reme_memlens/agentic_answer.py +++ b/plugins/memlens/src/reme_memlens/agentic_answer.py @@ -121,11 +121,11 @@ class MemLensAgenticAnswerStep(BaseStep): resume="", react_config={"max_iters": max_iters}, ) - self.context.response.answer = (result.get("result") or "").strip() + self.context.response.success = result["status"] == "completed" + self.context.response.answer = result["error"] or result["result"].strip() self.context.response.metadata.update( - usage=result.get("usage"), - agentscope_usage=(result.get("last_message") or {}).get("usage"), - agent_session_id=result.get("session_id"), + **result, + agent_session_id=result["session_id"], include_source_sessions=sources, include_images=images, max_iters=max_iters, diff --git a/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py b/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py index 90f19510..8aa6d8a6 100644 --- a/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py +++ b/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py @@ -121,11 +121,11 @@ class MobileMemAgenticAnswerStep(BaseStep): resume="", react_config={"max_iters": max_iters}, ) - self.context.response.answer = (result.get("result") or "").strip() + self.context.response.success = result["status"] == "completed" + self.context.response.answer = result["error"] or result["result"].strip() self.context.response.metadata.update( - usage=result.get("usage"), - agentscope_usage=(result.get("last_message") or {}).get("usage"), - agent_session_id=result.get("session_id"), + **result, + agent_session_id=result["session_id"], include_source_sessions=sources, include_images=images, max_iters=max_iters, diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 38eccbbb..bf6fe58f 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -1,12 +1,12 @@ """AgentScope backend for the unified agent wrapper.""" import asyncio -import json import os import re import subprocess import time from collections.abc import AsyncGenerator +from contextlib import aclosing from pathlib import Path from typing import Any, TYPE_CHECKING from uuid import uuid4 @@ -36,7 +36,7 @@ from agentscope.event import ( ToolResultStartEvent, ToolResultTextDeltaEvent, ) -from agentscope.message import TextBlock, ToolResultState, UserMsg +from agentscope.message import Msg, TextBlock, ToolResultState, UserMsg from agentscope.permission import PermissionBehavior, PermissionContext, PermissionDecision, PermissionMode from agentscope.state import AgentState from agentscope.tool import ( @@ -58,7 +58,8 @@ from .base_agent_wrapper import BaseAgentWrapper from ..as_llm import BaseAsLLM from ..component_registry import R from ...enumeration import ChunkEnum -from ...schema import StreamChunk, TokenUsage +from ...schema import AgentReply, StreamChunk, TokenUsage +from ...schema.stream_chunk import stream_status from ...utils import AsStateHandler if TYPE_CHECKING: @@ -349,47 +350,64 @@ class AsAgentWrapper(BaseAgentWrapper): return agent, inputs - async def reply(self, inputs: Any, **kwargs) -> dict: - kwargs = self._merged_kwargs(kwargs) - agent, inputs = await self._build_agent(inputs, **kwargs) - - await agent.observe(inputs) - last_msg = await agent.reply() + async def _finish_reply(self, agent: Agent, last_msg: Msg, kwargs: dict) -> AgentReply: + """Preserve the native message and optional structured-model response.""" usage = self._agentscope_usage(last_msg.usage) if last_msg.usage is not None else None - await self._dump_state(agent.state) - result = { - "session_id": agent.state.session_id, - "last_message": last_msg.model_dump(), - "result": last_msg.get_text_content(), - } + message = last_msg.model_dump(mode="json") + raw = {"message": message} + reason = message["finished_reason"] + status = "completed" + error = None + if reason is None: + status = "waiting" + elif reason == "interrupted": + status = "interrupted" + elif reason != "completed": + status = "failed" + error = (message.get("error") or {}).get("message") or f"AgentScope reply ended: {reason}" + structured_output = message.get("structured_output") if usage is None: self.logger.error("AgentScope did not return token usage; token accounting is unavailable for this reply.") - output_schema: dict | None = kwargs.get("output_schema") - if output_schema is not None: - assert self.as_llm is not None, "AsAgentWrapper requires a bound as_llm component with a valid model." - model = self.as_llm.model - assert model is not None, "AsAgentWrapper requires a bound as_llm component with a valid model." - res = await model.generate_structured_output( - messages=agent.state.context, - structured_model=output_schema, - ) - result["structured_output"] = res.content - if res.usage is None: - usage = None - self.logger.error( - "AgentScope did not return structured-output token usage; token accounting is unavailable.", + try: + output_schema: dict | None = kwargs.get("output_schema") + if status == "completed" and output_schema is not None: + assert self.as_llm is not None, "AsAgentWrapper requires a bound as_llm component with a valid model." + model = self.as_llm.model + assert model is not None, "AsAgentWrapper requires a bound as_llm component with a valid model." + res = await model.generate_structured_output( + messages=agent.state.context, + structured_model=output_schema, ) - elif usage is not None: - usage = TokenUsage.combine([usage, self._agentscope_usage(res.usage)]) + raw["structured_response"] = self._serialize(res) + structured_output = res.content + if raw["structured_response"]["finished_reason"] == "interrupted": + status = "interrupted" + if res.usage is None: + usage = None + self.logger.error( + "AgentScope did not return structured-output token usage; token accounting is unavailable.", + ) + elif usage is not None: + usage = TokenUsage.combine([usage, self._agentscope_usage(res.usage)]) + finally: + if usage is not None: + self._record_token_usage(usage) - result["usage"] = usage.model_dump() if usage is not None else None - if usage is not None: - self._record_token_usage(usage) - - return result + task = asyncio.current_task() + if task is not None and task.cancelling(): + raise asyncio.CancelledError() + return self._reply_result( + session_id=agent.state.session_id, + result=last_msg.get_text_content(), + structured_output=structured_output, + usage=usage, + status=status, + error=error, + raw=raw, + ) async def compact_session(self, session_id: str) -> None: """Force compression of an AgentScope session.""" @@ -403,103 +421,120 @@ class AsAgentWrapper(BaseAgentWrapper): @classmethod # pylint: disable=too-many-return-statements - def _event_to_chunk(cls, event: Any) -> StreamChunk | None: - """Convert an AgentScope event to a unified StreamChunk. - - Returns ``None`` for events that should be silently skipped - (e.g. ``RequireUserConfirmEvent``). - """ + def _event_to_chunk(cls, event: Any) -> StreamChunk: + """Normalize block operations while retaining native fields and unknown events.""" + native = event.model_dump(mode="json") + operation = ( + "start" if str(event.type).endswith("START") else "end" if str(event.type).endswith("END") else "delta" + ) + metadata = {"native": native} + common = { + "event": operation, + "turn_id": getattr(event, "reply_id", None), + "message_id": getattr(event, "reply_id", None), + "chunk": None, + "metadata": metadata, + } if isinstance(event, ReplyStartEvent): - meta = {"reply_id": event.reply_id, "name": event.name, "role": event.role} - return cls._chunk(ChunkEnum.REPLY_START, session_id=event.session_id, chunk="", metadata=meta) + return cls._chunk(ChunkEnum.REPLY_START, **common, session_id=event.session_id, status="running") if isinstance(event, ReplyEndEvent): + metadata["finished_reason"] = native["finished_reason"] + metadata["error"] = (native.get("error") or {}).get("message") return cls._chunk( ChunkEnum.REPLY_END, + **common, session_id=event.session_id, - chunk="", - metadata={"reply_id": event.reply_id}, + status=stream_status(event.finished_reason) or "waiting", ) - for event_cls, chunk_type, attr in ( - (TextBlockStartEvent, ChunkEnum.CONTENT, None), - (TextBlockDeltaEvent, ChunkEnum.CONTENT, "delta"), - (TextBlockEndEvent, ChunkEnum.CONTENT, None), - (ThinkingBlockStartEvent, ChunkEnum.THINK, None), - (ThinkingBlockDeltaEvent, ChunkEnum.THINK, "delta"), - (ThinkingBlockEndEvent, ChunkEnum.THINK, None), - (DataBlockStartEvent, ChunkEnum.DATA, None), - (DataBlockDeltaEvent, ChunkEnum.DATA, "data"), - (DataBlockEndEvent, ChunkEnum.DATA, None), + for classes, chunk_type in ( + ((TextBlockStartEvent, TextBlockDeltaEvent, TextBlockEndEvent), ChunkEnum.CONTENT), + ((ThinkingBlockStartEvent, ThinkingBlockDeltaEvent, ThinkingBlockEndEvent), ChunkEnum.THINK), + ((DataBlockStartEvent, DataBlockDeltaEvent, DataBlockEndEvent), ChunkEnum.DATA), + ((ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent), ChunkEnum.TOOL_CALL), + ( + (ToolResultStartEvent, ToolResultTextDeltaEvent, ToolResultDataDeltaEvent, ToolResultEndEvent), + ChunkEnum.TOOL_RESULT, + ), ): - if isinstance(event, event_cls): - kwargs = {"block_id": event.block_id, "chunk": getattr(event, attr) if attr else ""} - if isinstance(event, (DataBlockStartEvent, DataBlockDeltaEvent)): - kwargs["media_type"] = event.media_type - return cls._chunk(chunk_type, **kwargs) - - if isinstance(event, ToolCallStartEvent): - payload = {"name": event.tool_call_name, "id": event.tool_call_id} - return cls._chunk( - ChunkEnum.TOOL_CALL, - tool_call_id=event.tool_call_id, - tool_call_name=event.tool_call_name, - chunk=json.dumps(payload), - ) - if isinstance(event, ToolCallDeltaEvent): - return cls._chunk(ChunkEnum.TOOL_CALL, tool_call_id=event.tool_call_id, chunk=event.delta) - if isinstance(event, ToolCallEndEvent): - return cls._chunk(ChunkEnum.TOOL_CALL, tool_call_id=event.tool_call_id, chunk="") - if isinstance(event, ToolResultStartEvent): - return cls._chunk( - ChunkEnum.TOOL_RESULT, - tool_call_id=event.tool_call_id, - tool_call_name=event.tool_call_name, - chunk="", - ) - if isinstance(event, ToolResultTextDeltaEvent): - return cls._chunk(ChunkEnum.TOOL_RESULT, tool_call_id=event.tool_call_id, chunk=event.delta) - if isinstance(event, ToolResultDataDeltaEvent): - return cls._chunk( - ChunkEnum.TOOL_RESULT, - tool_call_id=event.tool_call_id, - chunk=event.data, - media_type=event.media_type, - metadata={"url": event.url} if event.url else {}, - ) - if isinstance(event, ToolResultEndEvent): - return cls._chunk( - ChunkEnum.TOOL_RESULT, - tool_call_id=event.tool_call_id, - chunk="", - metadata={"state": str(event.state)}, - ) + if isinstance(event, classes): + common["chunk"] = getattr(event, "delta", getattr(event, "text", getattr(event, "data", None))) + if getattr(event, "url", None): + common["chunk"] = {"url": event.url} + return cls._chunk( + chunk_type, + **common, + block_id=getattr(event, "block_id", None), + tool_call_id=getattr(event, "tool_call_id", None), + tool_call_name=getattr(event, "tool_call_name", None), + media_type=getattr(event, "media_type", None), + status=stream_status(getattr(event, "state", None)), + ) if isinstance(event, ModelCallStartEvent): - return cls._chunk(ChunkEnum.USAGE, chunk="", metadata={"model_name": getattr(event, "model_name", None)}) + return cls._chunk(ChunkEnum.USAGE, **common) if isinstance(event, ModelCallEndEvent): usage = TokenUsage(input_tokens=event.input_tokens, output_tokens=event.output_tokens) + common["chunk"] = usage.model_dump() + metadata["usage"] = usage.model_dump() return cls._chunk( ChunkEnum.USAGE, - chunk=json.dumps(usage.model_dump()), + **common, input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, - metadata={ - "model_name": getattr(event, "model_name", None), - "usage": usage.model_dump(), - }, ) if isinstance(event, ExceedMaxItersEvent): - return cls._chunk(ChunkEnum.ERROR, chunk="Exceeded max iterations") - return None + common["chunk"] = "Exceeded max iterations" + return cls._chunk(ChunkEnum.ERROR, **common) + common["chunk"] = native + common["event"] = "end" + approval = str(event.type) in {"REQUIRE_USER_CONFIRM", "REQUIRE_EXTERNAL_EXECUTION"} + return cls._chunk( + ChunkEnum.APPROVAL if approval else ChunkEnum.DATA, + **common, + status="waiting" if approval else None, + ) async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: - """Stream agent events as unified StreamChunk objects.""" - kwargs = self._merged_stream_kwargs(kwargs) + """Execute once, streaming native events followed by the complete outcome.""" + kwargs = self._merged_kwargs(kwargs) agent, inputs = await self._build_agent(inputs, **kwargs) - - async for event in agent.reply_stream(inputs): - chunk = self._event_to_chunk(event) - if chunk is not None: - chunk.session_id = chunk.session_id or agent.state.session_id - yield chunk - - await self._dump_state(agent.state) + last_msg, terminal = None, None + model_usages = [] + finishing = False + try: + async with aclosing(agent.reply_stream(inputs, yield_final_msg=True)) as stream: + async for event in stream: + if isinstance(event, Msg): + last_msg = event + continue + if isinstance(event, ModelCallEndEvent): + model_usages.append(self._agentscope_usage(event)) + chunk = self._event_to_chunk(event) + chunk.session_id = chunk.session_id or agent.state.session_id + if isinstance(event, ReplyEndEvent): + terminal = chunk + else: + yield chunk + if last_msg is None: + raise RuntimeError("AgentScope final message not received") + if terminal is None: + if last_msg.finished_reason is not None: + raise RuntimeError("AgentScope reply_end not received") + terminal = self._chunk( + ChunkEnum.REPLY_END, + event="end", + chunk=None, + session_id=agent.state.session_id, + turn_id=last_msg.id, + status="waiting", + ) + finishing = True + result = await self._finish_reply(agent, last_msg, kwargs) + finally: + if not finishing and model_usages: + self._record_token_usage(TokenUsage.combine(model_usages)) + await self._dump_state(agent.state) + terminal.chunk = {key: value for key, value in result.items() if key != "events"} + terminal.status = result["status"] + terminal.metadata["error"] = result["error"] + yield terminal diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index e0bdebfd..f195de38 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -1,6 +1,5 @@ """Base agent wrapper component.""" -from abc import abstractmethod from collections.abc import AsyncGenerator from importlib import metadata from pathlib import Path @@ -11,7 +10,8 @@ from pydantic import BaseModel from ..base_component import BaseComponent from ..outbound_proxy import BaseOutboundProxy from ...enumeration import ChunkEnum, ComponentEnum -from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage +from ...schema import AgentReply, SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage +from ...schema.agent_reply import AgentReplyStatus, collect_reply from ...utils import global_counter_add_many if TYPE_CHECKING: @@ -217,13 +217,6 @@ class BaseAgentWrapper(BaseComponent): merged["output_schema"] = self._normalize_output_schema(merged["output_schema"]) return merged - def _merged_stream_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: - """Merge stream options and reject unsupported structured output.""" - merged = self._merged_kwargs(kwargs) - if merged.get("output_schema") is not None: - raise NotImplementedError("Structured output is not supported by reply_stream()") - return merged - @staticmethod def _chunk(chunk_type: ChunkEnum = ChunkEnum.CONTENT, **kwargs: Any) -> StreamChunk: """Create a StreamChunk with a short backend-friendly call site.""" @@ -246,9 +239,39 @@ class BaseAgentWrapper(BaseComponent): }, ) - @abstractmethod - async def reply(self, inputs: Any, **kwargs) -> dict: - """Send inputs to the agent and return a dict with session_id and last_message.""" + @staticmethod + def _serialize(value: Any) -> Any: + """Keep native model/dataclass fields in a JSON-serializable form.""" + from pydantic_core import to_jsonable_python + + return to_jsonable_python(value, by_alias=True) + + @staticmethod + def _reply_result( + *, + session_id: str, + result: str | None, + raw: dict[str, Any], + status: AgentReplyStatus = "completed", + error: str | None = None, + structured_output: Any = None, + usage: TokenUsage | None = None, + ) -> AgentReply: + """Build the same portable fields for every backend.""" + return AgentReply( + session_id=session_id, + result=result or "", + structured_output=structured_output, + usage=usage.model_dump() if usage is not None else None, + status=status, + error=error, + raw=raw, + events=[], + ) + + async def reply(self, inputs: Any, **kwargs) -> AgentReply: + """Collect the same execution used by streaming; errors and cancellation propagate.""" + return await collect_reply(self.reply_stream(inputs, **kwargs)) async def compact_session(self, session_id: str) -> None: """Request compaction of one persisted agent session.""" @@ -275,3 +298,5 @@ class BaseAgentWrapper(BaseComponent): async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream agent events as unified StreamChunk objects.""" + raise NotImplementedError + yield # pylint: disable=unreachable diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 068175df..250a36a0 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -7,11 +7,12 @@ from contextlib import aclosing from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any, TYPE_CHECKING +from uuid import uuid4 from .base_agent_wrapper import BaseAgentWrapper from ..component_registry import R from ...enumeration import ChunkEnum -from ...schema import StreamChunk, TokenUsage +from ...schema import AgentReply, StreamChunk, TokenUsage if TYPE_CHECKING: from claude_agent_sdk import AssistantMessage, ResultMessage, UserMessage @@ -227,122 +228,83 @@ class CcAgentWrapper(BaseAgentWrapper): raw: dict, session_id: str | None = None, block_states: dict[int, _BlockState] | None = None, - ) -> StreamChunk | None: - """Convert a raw Anthropic streaming event dict to a StreamChunk. - - ``block_states`` maps each content-block index to metadata captured at - ``content_block_start`` for use by later delta and stop events. - - Returns ``None`` for events that should be silently skipped. - """ - event_type = raw.get("type") - - # --- Message-level lifecycle ---------------------------------------- - - if event_type == "message_start": - if block_states is not None: + message_id: str | None = None, + ) -> StreamChunk: + """Normalize model-message and block events, retaining unknown payloads.""" + event_type = raw.get("type", "") + common = {"session_id": session_id, "message_id": message_id, "chunk": None} + if event_type in {"message_start", "message_stop"}: + if event_type == "message_start" and block_states is not None: block_states.clear() - message = raw.get("message", {}) - meta = { - "message_id": message.get("id"), - "model": message.get("model"), - "role": message.get("role"), - } - return cls._chunk(ChunkEnum.REPLY_START, session_id=session_id, chunk="", metadata=meta) - + return cls._chunk( + ChunkEnum.DATA, + **common, + event="start" if event_type == "message_start" else "end", + block_id=f"message:{message_id}", + metadata={"kind": "message", "visibility": "diagnostic", "native": raw}, + ) if event_type == "message_delta": - delta = raw.get("delta", {}) - usage = raw.get("usage", {}) - normalized = cls._claude_usage(usage) + usage = cls._claude_usage(raw.get("usage")) + common["chunk"] = usage.model_dump() return cls._chunk( ChunkEnum.USAGE, - session_id=session_id, - chunk=json.dumps(normalized.model_dump()), - input_tokens=normalized.input_tokens, - output_tokens=normalized.output_tokens, + **common, + event="end", + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, metadata={ - "stop_reason": delta.get("stop_reason"), - "usage": normalized.model_dump(), + "stop_reason": raw.get("delta", {}).get("stop_reason"), + "usage": usage.model_dump(), + "native": raw, }, ) - if event_type == "message_stop": - return cls._chunk(ChunkEnum.REPLY_END, session_id=session_id, chunk="") - - # --- Content-block lifecycle ---------------------------------------- - + index = raw.get("index", 0) + block = raw.get("content_block", {}) if event_type == "content_block_start": - idx, content_block = raw.get("index", 0), raw.get("content_block", {}) - block_type, bid = content_block.get("type", ""), content_block.get("id", "") + state = _BlockState( + block.get("id") or f"{message_id or 'message'}:{index}", + block.get("type", ""), + block.get("name"), + ) if block_states is not None: - block_states[idx] = _BlockState(bid or None, block_type, content_block.get("name")) - - if block_type == "text": - return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=content_block.get("text", "")) - if block_type == "thinking": - return cls._chunk( - ChunkEnum.THINK, - block_id=bid, - chunk=content_block.get("thinking", ""), - ) - if block_type in {"tool_use", "server_tool_use"}: - payload = { - "name": content_block.get("name"), - "id": content_block.get("id"), - } - return cls._chunk( - ChunkEnum.TOOL_CALL, - block_id=bid, - tool_call_id=content_block.get("id"), - tool_call_name=content_block.get("name"), - chunk=json.dumps(payload), - ) - return None - - if event_type == "content_block_delta": + block_states[index] = state + else: + state = block_states.get(index) if block_states is not None else None + if event_type == "content_block_stop" and block_states is not None: + block_states.pop(index, None) + operation = {"content_block_start": "start", "content_block_delta": "delta", "content_block_stop": "end"}.get( + event_type, + ) + if state is not None and operation: + common.update(event=operation, block_id=state.block_id) delta = raw.get("delta", {}) - delta_type = delta.get("type", "") - idx = raw.get("index", 0) - state = block_states.get(idx) if block_states else None - bid = state.block_id if state else None - tool_name = state.tool_name if state else None - - if delta_type == "text_delta": - return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=delta.get("text", "")) - if delta_type == "thinking_delta": - return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk=delta.get("thinking", "")) - if delta_type == "input_json_delta": + if event_type == "content_block_delta" and delta.get("type") not in { + "text_delta", + "thinking_delta", + "input_json_delta", + }: + return cls._chunk(ChunkEnum.DATA, **{**common, "chunk": raw}, metadata={"visibility": "diagnostic"}) + metadata = { + "native": {key: value for key, value in block.items() if key not in {"text", "thinking", "input"}}, + } + if state.block_type == "text": + common["chunk"] = delta.get("text") if operation == "delta" else block.get("text") + return cls._chunk(ChunkEnum.CONTENT, **common, metadata=metadata) + if state.block_type == "thinking": + common["chunk"] = delta.get("thinking") if operation == "delta" else block.get("thinking") + return cls._chunk(ChunkEnum.THINK, **common, metadata=metadata) + if state.block_type in {"tool_use", "server_tool_use"}: + common["chunk"] = delta.get("partial_json") if operation == "delta" else block.get("input") or None return cls._chunk( ChunkEnum.TOOL_CALL, - block_id=bid, - tool_call_id=bid, - tool_call_name=tool_name, - chunk=delta.get("partial_json", ""), + **common, + tool_call_id=state.block_id, + tool_call_name=state.tool_name, + metadata=metadata, ) - return None - - if event_type == "content_block_stop": - idx = raw.get("index", 0) - state = block_states.pop(idx, None) if block_states else None - bid = state.block_id if state else None - block_type = state.block_type if state else None - tool_name = state.tool_name if state else None - - if block_type in {"tool_use", "server_tool_use"}: - return cls._chunk( - ChunkEnum.TOOL_CALL, - block_id=bid, - tool_call_id=bid, - tool_call_name=tool_name, - chunk="", - ) - if block_type == "thinking": - return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk="") - # text or unknown -> CONTENT - return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk="") - - # Ping / other unknown types -> skip - return None + common.update(chunk=raw, event="end") + return cls._chunk(ChunkEnum.DATA, **common, metadata={"visibility": "diagnostic"}) @classmethod def _message_content_to_chunks( @@ -351,48 +313,85 @@ class CcAgentWrapper(BaseAgentWrapper): session_id: str | None = None, visible_tool_call_ids: set[str] | None = None, include_text: bool = False, + message_id: str | None = None, ) -> list[StreamChunk]: - """Convert typed SDK content blocks that are not partial events.""" - from claude_agent_sdk import ServerToolResultBlock, TextBlock, ToolResultBlock + """Use completed SDK blocks as authoritative snapshots, including tool results.""" + from claude_agent_sdk import ( + ServerToolResultBlock, + ServerToolUseBlock, + TextBlock, + ThinkingBlock, + ToolResultBlock, + ToolUseBlock, + ) chunks: list[StreamChunk] = [] content = getattr(msg, "content", None) if isinstance(content, str): - if include_text and content: - chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=content)) + if include_text: + chunks.append( + cls._chunk( + ChunkEnum.CONTENT, + session_id=session_id, + message_id=message_id, + block_id=f"{message_id}:0", + event="end", + chunk=content, + ), + ) return chunks - if content is None: - return chunks - - for block in content: - if include_text and isinstance(block, TextBlock) and block.text: - chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=block.text)) + for index, block in enumerate(content or []): + common = { + "session_id": session_id, + "message_id": message_id, + "event": "end", + "block_id": f"{message_id or 'message'}:{index}", + } + if include_text and isinstance(block, (TextBlock, ThinkingBlock)): + thinking = isinstance(block, ThinkingBlock) + chunks.append( + cls._chunk( + ChunkEnum.THINK if thinking else ChunkEnum.CONTENT, + **common, + chunk=block.thinking if thinking else block.text, + metadata={"native": {"signature": block.signature}} if thinking else {}, + ), + ) + elif include_text and isinstance(block, (ToolUseBlock, ServerToolUseBlock)): + common["block_id"] = block.id + chunks.append( + cls._chunk( + ChunkEnum.TOOL_CALL, + **common, + tool_call_id=block.id, + tool_call_name=block.name, + chunk=block.input, + ), + ) + if visible_tool_call_ids is not None: + visible_tool_call_ids.add(block.id) elif isinstance(block, (ToolResultBlock, ServerToolResultBlock)): - tool_use_id = block.tool_use_id - if visible_tool_call_ids is not None and tool_use_id not in visible_tool_call_ids: + if visible_tool_call_ids is not None and block.tool_use_id not in visible_tool_call_ids: continue - payload: dict[str, Any] = { - "tool_use_id": tool_use_id, - "content": block.content, - } + common["block_id"] = block.tool_use_id + payload = {"tool_use_id": block.tool_use_id, "content": block.content} if isinstance(block, ToolResultBlock): payload["is_error"] = block.is_error chunks.append( cls._chunk( ChunkEnum.TOOL_RESULT, - session_id=session_id, - block_id=tool_use_id, - tool_call_id=tool_use_id, + **common, + tool_call_id=block.tool_use_id, chunk=payload, + status="failed" if payload.get("is_error") else "completed", ), ) - return chunks @classmethod def _result_error_text(cls, msg: "ResultMessage") -> str: """Return the error text used by both the SDK and unified chunks.""" - return "; ".join(msg.errors or []) or str(msg.subtype) + return "; ".join(msg.errors or []) or msg.result or str(msg.subtype) @classmethod def _result_message_to_chunks(cls, msg: "ResultMessage") -> list[StreamChunk]: @@ -418,6 +417,7 @@ class CcAgentWrapper(BaseAgentWrapper): "permission_denials": msg.permission_denials, "deferred_tool_use": (asdict(msg.deferred_tool_use) if msg.deferred_tool_use else None), "api_error_status": msg.api_error_status, + "native": cls._serialize(msg), }, ), ] @@ -437,39 +437,11 @@ class CcAgentWrapper(BaseAgentWrapper): async def compact_session(self, session_id: str) -> None: """Compact a Claude Code session through its native command.""" result = await self.reply("/compact", resume=session_id) - if result["last_message"].get("is_error"): - raise RuntimeError("Claude Code session compaction failed") - - async def reply(self, inputs: Any, **kwargs) -> dict: - from claude_agent_sdk import query, ResultMessage - - kwargs = self._merged_kwargs(kwargs) - opts = self._build_options(inputs, stream=False, **kwargs) - - last_msg = None - async with aclosing(query(prompt=inputs, options=opts)) as stream: - async for msg in stream: - if isinstance(msg, ResultMessage): - last_msg = msg - - if last_msg is None: - raise ValueError("No message received from Claude Code.") - - usage = self._claude_usage(last_msg.usage) if last_msg.usage is not None else None - result = { - "session_id": last_msg.session_id or "", - "last_message": asdict(last_msg), - "result": last_msg.result, - "usage": usage.model_dump() if usage is not None else None, - } - if usage is not None: - self._record_token_usage(usage) - if kwargs.get("output_schema") is not None: - result["structured_output"] = last_msg.structured_output - return result + if result["status"] != "completed": + raise RuntimeError(result["error"] or f"Claude Code session compaction {result['status']}") async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: - """Stream Claude Code events as unified StreamChunk objects.""" + """Stream one invocation, keeping model-message boundaries inside it.""" from claude_agent_sdk import ( AssistantMessage, MirrorErrorMessage, @@ -481,104 +453,146 @@ class CcAgentWrapper(BaseAgentWrapper): UserMessage, ) - kwargs = self._merged_stream_kwargs(kwargs) + kwargs = self._merged_kwargs(kwargs) opts = self._build_options(inputs, stream=True, **kwargs) - - block_states: dict[int, _BlockState] = {} + block_states: dict[str | None, dict[int, _BlockState]] = {} visible_tool_call_ids: set[str] = set() - current_session_id: str | None = None + turn_id = uuid4().hex + partial_message_ids: dict[str | None, str] = {} + current_session_id = None emitted_content = False - emitted_reply_end = False - reply_open = False - expected_trailing_error: str | None = None - + started = False + result: AgentReply | None = None + usage = None + expected_trailing_error = None try: async with aclosing(query(prompt=inputs, options=opts)) as stream: async for msg in stream: + current_session_id = getattr(msg, "session_id", None) or current_session_id + if not started: + started = True + yield self._chunk( + ChunkEnum.REPLY_START, + event="start", + turn_id=turn_id, + session_id=current_session_id, + status="running", + ) if expected_trailing_error is not None and not ( isinstance(msg, SystemMessage) and msg.subtype == "session_state_changed" ): expected_trailing_error = None - + chunks = [] + parent_id = getattr(msg, "parent_tool_use_id", None) if isinstance(msg, StreamEvent): - current_session_id = msg.session_id or current_session_id + if msg.event.get("type") == "message_start": + partial_message_ids[parent_id] = msg.event.get("message", {}).get("id") or msg.uuid chunk = self._raw_event_to_chunk( msg.event, - session_id=msg.session_id, - block_states=block_states, + session_id=current_session_id, + block_states=block_states.setdefault(parent_id, {}), + message_id=partial_message_ids.get(parent_id), ) - if chunk is not None: - chunk.session_id = chunk.session_id or msg.session_id - if chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_id: - visible_tool_call_ids.add(chunk.tool_call_id) - if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: - emitted_content = True - if chunk.chunk_type == ChunkEnum.REPLY_START: - reply_open = True - if chunk.chunk_type == ChunkEnum.REPLY_END: - emitted_reply_end = True - reply_open = False - yield chunk - - elif isinstance(msg, UserMessage): - for chunk in self._message_content_to_chunks(msg, current_session_id, visible_tool_call_ids): - yield chunk - + chunk.metadata["parent_tool_use_id"] = msg.parent_tool_use_id + if chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_id: + visible_tool_call_ids.add(chunk.tool_call_id) + chunks.append(chunk) + elif isinstance(msg, (AssistantMessage, UserMessage)): + message_id = ( + getattr(msg, "message_id", None) + or partial_message_ids.get(parent_id) + or getattr(msg, "uuid", None) + or uuid4().hex + ) + chunks = self._message_content_to_chunks( + msg, + current_session_id, + visible_tool_call_ids, + include_text=isinstance(msg, AssistantMessage), + message_id=message_id, + ) + if isinstance(msg, AssistantMessage): + partial_message_ids.pop(parent_id, None) elif isinstance(msg, ResultMessage): if msg.is_error: expected_trailing_error = ( f"Claude Code returned an error result: {self._result_error_text(msg)}" ) - current_session_id = msg.session_id or current_session_id if not emitted_content and msg.result: - emitted_content = True - yield self._chunk( - ChunkEnum.CONTENT, - session_id=msg.session_id or "", - chunk=msg.result, + chunks.append( + self._chunk( + ChunkEnum.CONTENT, + event="end", + block_id=f"{turn_id}:result", + chunk=msg.result, + ), ) - for chunk in self._result_message_to_chunks(msg): - yield chunk - if reply_open or not emitted_reply_end: - emitted_reply_end = True - reply_open = False - yield self._chunk( - ChunkEnum.REPLY_END, - session_id=current_session_id, - chunk="", - ) - - elif isinstance(msg, AssistantMessage): - current_session_id = msg.session_id or current_session_id - for chunk in self._message_content_to_chunks( - msg, - current_session_id, - visible_tool_call_ids, - include_text=not emitted_content, - ): - if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: - emitted_content = True - yield chunk - + chunks.extend(self._result_message_to_chunks(msg)) + status = "completed" + if msg.terminal_reason in {"aborted_streaming", "aborted_tools"}: + status = "interrupted" + elif msg.is_error or msg.subtype != "success": + status = "failed" + elif msg.deferred_tool_use is not None: + status = "waiting" + usage = self._claude_usage(msg.usage) if msg.usage is not None else None + result = self._reply_result( + session_id=msg.session_id or "", + result=msg.result, + structured_output=msg.structured_output, + usage=usage, + status=status, + error=self._result_error_text(msg) if status == "failed" else None, + raw=self._serialize(msg), + ) elif isinstance(msg, MirrorErrorMessage): - self.logger.warning(f"Claude Code session mirror failed: {msg.error}") - yield self._chunk( - ChunkEnum.DATA, - session_id=current_session_id, - chunk=f"Session mirror failed: {msg.error}", - metadata={ - "event": "session_mirror_error", - "session_key": msg.key, - }, + chunks.append( + self._chunk( + ChunkEnum.DATA, + event="end", + chunk=f"Session mirror failed: {msg.error}", + metadata={"event": "session_mirror_error", "session_key": msg.key}, + ), ) - - elif isinstance(msg, RateLimitEvent) and msg.rate_limit_info.status == "rejected": - yield self._chunk( - ChunkEnum.ERROR, - session_id=msg.session_id, - chunk="Rate limit exceeded", + elif isinstance(msg, RateLimitEvent): + if msg.rate_limit_info.status == "rejected": + chunks.append(self._chunk(ChunkEnum.ERROR, chunk="Rate limit exceeded")) + else: + chunks.append( + self._chunk( + ChunkEnum.DATA, + event="end", + chunk=self._serialize(msg), + metadata={"visibility": "diagnostic"}, + ), ) + if not chunks: + chunks = [self._chunk(ChunkEnum.DATA, chunk=None, metadata={"visibility": "diagnostic"})] + for chunk in chunks: + chunk.metadata.pop("native", None) + chunks[0].metadata["native"] = self._serialize(msg) + for chunk in chunks: + chunk.metadata["parent_tool_use_id"] = parent_id + chunk.turn_id = turn_id + chunk.session_id = chunk.session_id or current_session_id + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + emitted_content = True + yield chunk except Exception as exc: if expected_trailing_error is None or str(exc) != expected_trailing_error: raise self.logger.debug(f"Ignoring Claude Code process exit after error result: {exc}") + finally: + if usage is not None: + self._record_token_usage(usage) + if result is None: + raise RuntimeError("Claude Code result message not received") + yield self._chunk( + ChunkEnum.REPLY_END, + event="end", + session_id=result["session_id"], + turn_id=turn_id, + status=result["status"], + chunk={key: value for key, value in result.items() if key != "events"}, + metadata={"error": result["error"]}, + ) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index cb29a3fb..d9b91c80 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -12,11 +12,13 @@ import os from typing import Any, TYPE_CHECKING from .base_agent_wrapper import BaseAgentWrapper +from .codex_session_history import item_chunks from ..codex_app_server import CodexAppServer from ..component_registry import R from ..job import BackgroundJob, StreamJob from ...enumeration import ChunkEnum -from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage +from ...schema import AgentReply, SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage +from ...schema.stream_chunk import stream_status if TYPE_CHECKING: from openai_codex import AsyncThread, RunInput @@ -76,12 +78,12 @@ class CodexAgentWrapper(BaseAgentWrapper): if invalid: raise TypeError(f"Configure Codex process options on codex_app_server: {', '.join(invalid)}") - def _reply_kwargs(self, kwargs: dict[str, Any], *, stream: bool = False) -> dict[str, Any]: + def _reply_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: self._reject_server_options(kwargs) invalid = sorted({"app_server", "cwd"}.intersection(kwargs)) if invalid: raise TypeError(f"Configure Codex wrapper options on the component: {', '.join(invalid)}") - return self._merged_stream_kwargs(kwargs) if stream else self._merged_kwargs(kwargs) + return self._merged_kwargs(kwargs) def _ensure_skills(self, skills: list[str] | str | None) -> None: """Expose selected project skills through Codex's repo-level directory.""" @@ -187,12 +189,6 @@ class CodexAgentWrapper(BaseAgentWrapper): except Exception as exc: self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") - @staticmethod - def _serialize(value: Any) -> Any: - from pydantic_core import to_jsonable_python - - 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() @@ -210,38 +206,55 @@ class CodexAgentWrapper(BaseAgentWrapper): 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 + def _finish_reply(self, session_id, completed, items, native_usage, usage, kwargs) -> AgentReply: + """Build the native TurnResult and portable outcome from this turn's events.""" + from openai_codex import TurnResult - 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) + by_id = {item.root.id: item for item in items} + by_id.update({item.root.id: item for item in completed.items}) + items = list(by_id.values()) + final_response = None + fallback = None + for wrapped in reversed(items): + item = wrapped.root + if item.type != "agentMessage": + continue + if item.phase is not None and item.phase.value == "final_answer": + final_response = item.text + break + if item.phase is None and fallback is None: + fallback = item.text + if final_response is None: + final_response = fallback + result = TurnResult( + id=completed.id, + status=completed.status, + error=completed.error, + started_at=completed.started_at, + completed_at=completed.completed_at, + duration_ms=completed.duration_ms, + final_response=final_response, + items=items, + usage=native_usage, + ) + status = completed.status.value + error = completed.error.message if completed.error else None + structured_output = None + if status == "completed" and kwargs.get("output_schema") is not None: try: - 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": usage.usage.model_dump() if usage.usage is not None else None, - } - if kwargs.get("output_schema") is not None: - try: - response["structured_output"] = json.loads(final_response) - except json.JSONDecodeError as exc: - raise ValueError("Codex returned invalid JSON for the requested output_schema") from exc - return response + structured_output = json.loads(final_response or "") + except json.JSONDecodeError: + status = "failed" + error = "Codex returned invalid JSON for the requested output_schema" + return self._reply_result( + session_id=session_id, + result=final_response, + structured_output=structured_output, + usage=usage, + status=status, + error=error, + raw=self._serialize(result), + ) async def compact_session(self, session_id: str) -> None: """Hold the conversation lock until native compaction has finished.""" @@ -297,13 +310,39 @@ class CodexAgentWrapper(BaseAgentWrapper): ) -> 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) + make_chunk = partial(cls._chunk, session_id=session_id, turn_id=getattr(payload, "turn_id", None)) if method == "turn/started": - return [make_chunk(ChunkEnum.REPLY_START, metadata={"turn_id": payload.turn.id})] + return [ + cls._chunk( + ChunkEnum.REPLY_START, + event="start", + session_id=session_id, + turn_id=payload.turn.id, + status="running", + metadata={"turn_id": payload.turn.id}, + ), + ] if method == "item/agentMessage/delta": - return [make_chunk(ChunkEnum.CONTENT, block_id=payload.item_id, chunk=payload.delta)] + return [ + make_chunk( + ChunkEnum.CONTENT, + block_id=payload.item_id, + message_id=payload.item_id, + chunk=payload.delta, + ), + ] if method in {"item/reasoning/summaryTextDelta", "item/reasoning/textDelta", "item/plan/delta"}: - return [make_chunk(ChunkEnum.THINK, block_id=payload.item_id, chunk=payload.delta)] + kind = {"item/reasoning/summaryTextDelta": "summary", "item/reasoning/textDelta": "content"}.get(method) + index = getattr(payload, "summary_index" if kind == "summary" else "content_index", 0) + return [ + make_chunk( + ChunkEnum.THINK, + block_id=f"{payload.item_id}:{kind}:{index}" if kind else payload.item_id, + message_id=payload.item_id, + chunk=payload.delta, + metadata={"kind": f"reasoning_{kind}" if kind else "plan", "index": index}, + ), + ] if method in {"item/commandExecution/outputDelta", "item/fileChange/outputDelta"}: return [ make_chunk( @@ -311,6 +350,7 @@ class CodexAgentWrapper(BaseAgentWrapper): block_id=payload.item_id, tool_call_id=payload.item_id, chunk=payload.delta, + status="running", ), ] if method == "item/mcpToolCall/progress": @@ -320,6 +360,7 @@ class CodexAgentWrapper(BaseAgentWrapper): block_id=payload.item_id, tool_call_id=payload.item_id, chunk=payload.message, + status="running", ), ] if method in {"item/autoApprovalReview/started", "item/autoApprovalReview/completed"}: @@ -335,6 +376,8 @@ class CodexAgentWrapper(BaseAgentWrapper): block_id=target_item_id or review_id, tool_call_id=target_item_id, chunk=action, + event="start" if status == "started" else "end", + status="running" if status == "started" else stream_status((review or {}).get("status")), metadata={ "review_id": review_id, "status": status, @@ -345,22 +388,12 @@ class CodexAgentWrapper(BaseAgentWrapper): ), ] if method in {"item/started", "item/completed"}: - item = payload.item.root - item_type, item_id = item.type, item.id - tool_types = {"commandExecution", "fileChange", "mcpToolCall", "dynamicToolCall", "collabAgentToolCall"} - if item_type not in tool_types: - return [] - name = getattr(item, "tool", None) or item_type - chunk_type = ChunkEnum.TOOL_CALL if method == "item/started" else ChunkEnum.TOOL_RESULT - return [ - make_chunk( - chunk_type, - block_id=item_id, - tool_call_id=item_id, - tool_call_name=name, - chunk=cls._serialize(item), - ), - ] + return item_chunks( + payload.item.root, + event="start" if method == "item/started" else "end", + session_id=session_id, + turn_id=payload.turn_id, + ) if method == "thread/tokenUsage/updated": if usage is None: return [] @@ -371,7 +404,11 @@ class CodexAgentWrapper(BaseAgentWrapper): chunk=usage.model_dump(), input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, - metadata={"usage": usage.model_dump()}, + metadata={ + "usage": usage.model_dump(), + "native": cls._serialize(getattr(payload, "token_usage", None)), + }, + event="end", ), ] if method == "error": @@ -379,17 +416,25 @@ class CodexAgentWrapper(BaseAgentWrapper): make_chunk( ChunkEnum.ERROR, chunk=payload.error.message, - metadata={"will_retry": payload.will_retry}, + metadata={"will_retry": payload.will_retry, "native": cls._serialize(payload.error)}, ), ] if method == "turn/completed": turn = payload.turn - chunks = [] + chunks = [ + chunk + for item in getattr(turn, "items", []) + for chunk in item_chunks(getattr(item, "root", item), session_id=session_id, turn_id=turn.id) + ] if turn.error: chunks.append(make_chunk(ChunkEnum.ERROR, chunk=turn.error.message)) chunks.append( - make_chunk( + cls._chunk( ChunkEnum.REPLY_END, + event="end", + session_id=session_id, + turn_id=turn.id, + status=stream_status(turn.status), metadata={ "turn_id": turn.id, "status": turn.status.value, @@ -398,39 +443,60 @@ class CodexAgentWrapper(BaseAgentWrapper): ), ) return chunks - if getattr(payload, "turn_id", None): - return [ - make_chunk( - ChunkEnum.DATA, - block_id=getattr(payload, "item_id", None), - chunk=cls._serialize(payload), - metadata={"codex_method": method}, - ), - ] - return [] + diagnostic = method.startswith("hook/") or not getattr(payload, "turn_id", None) + return [ + make_chunk( + ChunkEnum.DATA, + block_id=getattr(payload, "item_id", None), + chunk=cls._serialize(payload), + metadata={"codex_method": method, **({"visibility": "diagnostic"} if diagnostic else {})}, + ), + ] 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) + """Execute once, retaining every notification and a complete terminal outcome.""" + kwargs = self._reply_kwargs(kwargs) 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 = self._turn_stream(turn, usage) - completed = False + items, native_usage, completed, terminal = [], None, None, None try: async for event in stream: - if event.method == "turn/completed": - completed = True + payload = event.payload + if event.method == "item/completed" and payload.turn_id == turn.id: + items.append(payload.item) + elif event.method == "thread/tokenUsage/updated" and payload.turn_id == turn.id: + native_usage = payload.token_usage + elif event.method == "turn/completed" and payload.turn.id == turn.id: + completed = payload.turn + update = usage.usage if usage.usage != emitted_usage else None + chunks = self._event_to_chunks(event, thread.id, usage=update) 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 + if not chunks: + chunks = [self._chunk(ChunkEnum.DATA, chunk=None, metadata={"visibility": "diagnostic"})] + for chunk in chunks: + chunk.metadata.pop("native", None) + chunks[0].metadata["native"] = {"method": event.method, "payload": self._serialize(payload)} + for chunk in chunks: + chunk.session_id = thread.id + chunk.turn_id = chunk.turn_id or turn.id + if chunk.chunk_type == ChunkEnum.REPLY_END: + terminal = chunk + else: + yield chunk + if completed is None or terminal is None: + raise RuntimeError("Codex turn completed event not received") + result = self._finish_reply(thread.id, completed, items, native_usage, usage.usage, kwargs) finally: if usage.usage is not None: self._record_token_usage(usage.usage) - if not completed: + if completed is None: await self._interrupt(turn) await stream.aclose() + terminal.chunk = {key: value for key, value in result.items() if key != "events"} + terminal.status = result["status"] + terminal.metadata["error"] = result["error"] + yield terminal diff --git a/reme/components/agent_wrapper/codex_session_history.py b/reme/components/agent_wrapper/codex_session_history.py index 7ae33dfb..8ff9459b 100644 --- a/reme/components/agent_wrapper/codex_session_history.py +++ b/reme/components/agent_wrapper/codex_session_history.py @@ -9,6 +9,7 @@ from uuid import uuid4 from ...enumeration import ChunkEnum from ...schema import SessionMessage, SessionSnapshot, SessionStatus, StreamChunk, TokenUsage +from ...schema.stream_chunk import stream_status _IDENTIFIER = re.compile(r"[A-Za-z0-9_-]{1,128}\Z") @@ -125,6 +126,70 @@ def new_record(message_id: str, content: str, thread) -> dict: } +def item_chunks(item, *, event="end", session_id=None, turn_id=None) -> list[StreamChunk]: + """Project the same native item snapshots for live streams and history.""" + payload = item.model_dump(mode="json", by_alias=True) + common = {"event": event, "session_id": session_id, "turn_id": turn_id, "message_id": item.id} + if item.type == "userMessage": + return [] + if item.type in {"agentMessage", "plan", "reasoning"}: + metadata = { + "native": {key: value for key, value in payload.items() if key not in {"text", "summary", "content"}}, + } + if item.type == "reasoning": + return [ + StreamChunk( + **common, + chunk_type=ChunkEnum.THINK, + block_id=f"{item.id}:{kind}:{index}", + chunk=text, + metadata={**metadata, "kind": f"reasoning_{kind}", "index": index}, + ) + for kind in ("summary", "content") + for index, text in enumerate(getattr(item, kind, None) or []) + ] + return [ + StreamChunk( + **common, + chunk_type=ChunkEnum.CONTENT if item.type == "agentMessage" else ChunkEnum.THINK, + block_id=item.id, + chunk=item.text, + metadata={**metadata, "kind": item.type, "phase": payload.get("phase")}, + ), + ] + if item.type in { + "commandExecution", + "fileChange", + "mcpToolCall", + "dynamicToolCall", + "collabAgentToolCall", + "webSearch", + }: + native_status = getattr(item, "status", None) + status = ( + stream_status(native_status) + if native_status is not None + else "running" if event == "start" else "completed" + ) + if event == "end": + if getattr(item, "error", None) or getattr(item, "success", None) is False: + status = "failed" + if getattr(item, "exit_code", None) not in {None, 0}: + status = "failed" + return [ + StreamChunk( + **common, + chunk_type=ChunkEnum.TOOL_CALL if event == "start" else ChunkEnum.TOOL_RESULT, + block_id=item.id, + tool_call_id=item.id, + tool_call_name=getattr(item, "tool", None) or item.type, + chunk=payload, + status=status, + ), + ] + return [StreamChunk(**common, chunk_type=ChunkEnum.DATA, block_id=item.id, chunk=payload)] + + 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]] = {} @@ -140,35 +205,17 @@ def session_snapshot(thread, records: list[dict], recovered_ids: set[str]) -> Se 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")}, - ), - ) + else: + chunks.extend(item_chunks(item, session_id=thread.id, turn_id=turn.id)) 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)) + status = stream_status(getattr(turn, "status", None)) + if chunks and status != "running": + chunks.append(StreamChunk(chunk_type=ChunkEnum.REPLY_END, event="end", turn_id=turn.id, status=status)) if chunks: messages.append( SessionMessage( diff --git a/reme/components/client/http_client.py b/reme/components/client/http_client.py index b304e1bf..ccd7ee62 100644 --- a/reme/components/client/http_client.py +++ b/reme/components/client/http_client.py @@ -54,24 +54,19 @@ class HttpClient(BaseClient): continue data_str = line[len("data:") :] if data_str.strip() == "[DONE]": + yield StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True) return - try: - data = json.loads(data_str) - except json.JSONDecodeError: - continue + data = json.loads(data_str) chunk = StreamChunk(**data) - if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True: - # Recoverable notifications remain visible; terminal errors stop the request. - raise RuntimeError(str(chunk.chunk)) - if chunk.done: - return yield chunk - else: - body = await resp.aread() - yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=body.decode()) + if chunk.done or chunk.chunk_type == ChunkEnum.DONE: + return + raise RuntimeError("Stream ended without DONE") + body = await resp.aread() + yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=body.decode()) async def stream_chunks(self, action: str, **kwargs) -> AsyncGenerator[StreamChunk, None]: - """HTTP-specific richer access: yield raw StreamChunk objects (no display formatting).""" + """Yield the complete event stream, including errors and DONE, without display formatting.""" async for chunk in self._iter_stream_chunks(action, kwargs): yield chunk @@ -115,8 +110,31 @@ class HttpClient(BaseClient): # pylint: disable=invalid-overridden-method async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: """Yield text chunks for CLI display; JSON responses are pretty-formatted.""" + text_blocks: dict[str, str] = {} async for chunk in self._iter_stream_chunks(action, payload): + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True: + raise RuntimeError(str(chunk.chunk)) + if chunk.chunk_type == ChunkEnum.DONE: + continue + if chunk.chunk_type == ChunkEnum.REPLY_END and chunk.status != "completed": + outcome = chunk.chunk if isinstance(chunk.chunk, dict) else {} + raise RuntimeError( + outcome.get("error") or chunk.metadata.get("error") or f"Agent reply ended: {chunk.status}", + ) chunk_payload = chunk.chunk + if chunk.chunk_type in {ChunkEnum.REPLY_START, ChunkEnum.REPLY_END}: + continue + if chunk_payload is None or chunk.metadata.get("visibility") == "diagnostic": + continue + if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk_payload, str) and chunk.block_id: + previous = text_blocks.get(chunk.block_id, "") + text_blocks[chunk.block_id] = previous + chunk_payload if chunk.event == "delta" else chunk_payload + if chunk.event != "delta": + chunk_payload = ( + chunk_payload[len(previous) :] if chunk_payload.startswith(previous) else "\n" + chunk_payload + ) + if not chunk_payload: + continue text = chunk_payload if isinstance(chunk_payload, str) else json.dumps(chunk_payload, ensure_ascii=False) yield self._format_for_display(text) diff --git a/reme/components/job/base_job.py b/reme/components/job/base_job.py index c0149173..08165431 100644 --- a/reme/components/job/base_job.py +++ b/reme/components/job/base_job.py @@ -131,7 +131,7 @@ class BaseJob(BaseComponent): else: if execution is not None: if not context.response.success: - status, error = "failed", str(context.response.answer) + status, error = "failed", context.stream_error or 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(): diff --git a/reme/components/job/stream_job.py b/reme/components/job/stream_job.py index 5f437c7e..662575b2 100644 --- a/reme/components/job/stream_job.py +++ b/reme/components/job/stream_job.py @@ -21,5 +21,6 @@ class StreamJob(BaseJob): context.response.success = False context.response.answer = _describe_exception(e) await context.add_stream_string(str(e), ChunkEnum.ERROR) + context.stream_error = context.response.answer # 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 10ad66f1..b90abe8e 100644 --- a/reme/components/runtime_context.py +++ b/reme/components/runtime_context.py @@ -75,6 +75,10 @@ class RuntimeContext: 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) + if chunk.chunk_type == ChunkEnum.REPLY_END and chunk.status == "failed": + self.stream_error = self.stream_error or str( + chunk.metadata.get("error") or chunk.metadata.get("finished_reason") or "Agent reply failed", + ) await self.stream_queue.put(chunk) async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": diff --git a/reme/enumeration/chunk_enum.py b/reme/enumeration/chunk_enum.py index eb855371..0c3994ea 100644 --- a/reme/enumeration/chunk_enum.py +++ b/reme/enumeration/chunk_enum.py @@ -20,14 +20,13 @@ class ChunkEnum(str, Enum): ExceedMaxItersEvent -> ERROR Claude Code SDK events -> ChunkEnum mapping: - message_start -> REPLY_START + message_start/stop -> DATA (model-message boundaries) message_delta -> USAGE - message_stop -> REPLY_END content_block_start/delta/stop (text) -> CONTENT content_block_start/delta/stop (thinking) -> THINK content_block_start/delta/stop (tool_use) -> TOOL_CALL ToolResultBlock -> TOOL_RESULT - ResultMessage -> USAGE + ResultMessage -> USAGE + REPLY_END ResultMessage.is_error -> ERROR Codex app-server notifications follow the same lifecycle categories; diff --git a/reme/schema/__init__.py b/reme/schema/__init__.py index 0699aee5..de9e5717 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_reply import AgentReply, collect_reply, replay_reply from .agent_session import SessionMessage, SessionSnapshot, SessionStatus from .compaction import CallDecision, CompactionOptions, CompactionRequest, CompactionResult from .dream import ( @@ -29,6 +30,9 @@ from .token_usage import TokenUsage from .traverse_graph import TraverseGraph, TraverseGraphEdge, TraverseGraphNode __all__ = [ + "AgentReply", + "collect_reply", + "replay_reply", "CallDecision", "CompactionOptions", "CompactionRequest", diff --git a/reme/schema/agent_reply.py b/reme/schema/agent_reply.py new file mode 100644 index 00000000..a4b83728 --- /dev/null +++ b/reme/schema/agent_reply.py @@ -0,0 +1,61 @@ +"""Agent outcomes and lossless collection/replay of their event stream.""" + +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import aclosing +from copy import deepcopy +from typing import Any, Literal, TypedDict + +from .stream_chunk import StreamChunk +from ..enumeration import ChunkEnum + +AgentReplyStatus = Literal["completed", "failed", "interrupted", "waiting"] + + +class AgentReply(TypedDict): + """One invocation's outcome and complete, ordered JSON event trace. + + ``result`` is final reply text; ``usage`` covers this invocation only. + ``raw`` contains JSON-serializable native results, including provider + usage details and termination reasons omitted from the portable fields. + """ + + session_id: str + result: str + structured_output: Any + usage: dict[str, int] | None + status: AgentReplyStatus + error: str | None + raw: dict[str, Any] + events: list[dict[str, Any]] + + +async def collect_reply( + stream: AsyncGenerator[StreamChunk, None], + on_chunk: Callable[[StreamChunk], Awaitable[Any]] | None = None, +) -> AgentReply: + """Collect an invocation without deriving its outcome from display text. + + The terminal reply_end carries the same outcome fields as AgentReply. + Missing/duplicate outcomes fail explicitly; DONE is only a transport marker. + """ + events = [] + outcome = None + async with aclosing(stream): + async for chunk in stream: + events.append(chunk.model_dump(mode="json")) + if on_chunk is not None: + await on_chunk(chunk) + if chunk.chunk_type == ChunkEnum.REPLY_END: + required = AgentReply.__required_keys__ - {"events"} + if outcome is not None or not isinstance(chunk.chunk, dict) or not required <= chunk.chunk.keys(): + raise ValueError("Expected exactly one complete reply_end outcome") + outcome = chunk.chunk + if outcome is None: + raise ValueError("Agent stream ended without a reply_end outcome") + return AgentReply(**{**outcome, "events": events}) + + +async def replay_reply(reply: AgentReply) -> AsyncGenerator[StreamChunk, None]: + """Replay fresh chunks in their original order, including native metadata.""" + for event in reply["events"]: + yield StreamChunk.model_validate(deepcopy(event)) diff --git a/reme/schema/stream_chunk.py b/reme/schema/stream_chunk.py index e74dc6f4..dfa08112 100644 --- a/reme/schema/stream_chunk.py +++ b/reme/schema/stream_chunk.py @@ -1,43 +1,61 @@ """Stream chunk schema for incremental responses (e.g. LLM streaming).""" -from typing import Any +from typing import Any, Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, JsonValue from ..enumeration import ChunkEnum +StreamStatus = Literal["running", "completed", "failed", "interrupted", "waiting"] + + +def stream_status(value: Any) -> StreamStatus | None: + """Normalize provider execution states without treating unknown states as success.""" + return { + "inProgress": "running", + "running": "running", + "started": "running", + "completed": "completed", + "success": "completed", + "succeeded": "completed", + "error": "failed", + "failed": "failed", + "declined": "failed", + "denied": "failed", + "exceed_max_iters": "failed", + "interrupted": "interrupted", + "cancelled": "interrupted", + "waiting": "waiting", + }.get(getattr(value, "value", value)) + class StreamChunk(BaseModel): - """A single chunk in a unified streaming response sequence. + """Backend-independent display event, with provider extras in metadata. - Carries full information from AgentScope, Claude Code SDK, and Codex - backends. Optional fields are ``None`` by default so that simple - text-only streams (e.g. plain CONTENT deltas) stay lightweight. + ``start`` creates a block, ``delta`` appends content, and ``end`` settles it. + A non-null start/end payload is a complete snapshot replacing accumulated + content; null means no content update (an empty string is a valid snapshot). + Simple text producers use the default ``delta`` operation. - Fields: - chunk_type: Category of this chunk (see ChunkEnum). - chunk: Payload, typically a text delta, but can be - dict/list for structured data (tool-call JSON, - usage stats, etc.). - done: Terminal marker. True only for the final DONE chunk. - session_id: Backend session identifier (Codex uses its thread id). - block_id: Content-block identifier for matching start/delta/end - sequences (both backends assign block IDs). - tool_call_id: Tool-call identifier for correlating call deltas - with their start event and result. - tool_call_name: Name of the tool being invoked. - media_type: MIME type for DATA blocks (e.g. ``"image/png"``). - input_tokens: Prompt tokens consumed (populated on USAGE chunks). - output_tokens: Completion tokens generated (populated on USAGE chunks). - metadata: Backend-specific extras that don't warrant a dedicated - field. + Reply markers delimit one agent invocation, including its tool/model loop. + DONE closes the transport and does not imply success. IDs correlate a + session, invocation (turn), model message, content block, and tool call. + ``reply_end.chunk`` carries the complete outcome (excluding its event trace). + ``metadata.native`` preserves complete provider events; ``metadata.kind`` + distinguishes plans, reasoning summaries, and other display subtypes. """ + model_config = ConfigDict(extra="allow") + chunk_type: ChunkEnum = Field(default=ChunkEnum.CONTENT, description="Type of chunk content") - chunk: str | dict | list = Field(default="", description="Chunk payload") + event: Literal["start", "delta", "end"] = Field(default="delta", description="Block update operation") + chunk: JsonValue = Field(default="", description="Delta or complete start/end snapshot") done: bool = Field(default=False, description="Whether this is the final chunk") + status: StreamStatus | None = Field(default=None, description="Normalized execution state") session_id: str | None = Field(default=None, description="Session identifier") + turn_id: str | None = Field(default=None, description="Agent invocation identifier") + message_id: str | None = Field(default=None, description="Model message identifier") block_id: str | None = Field(default=None, description="Content block identifier") tool_call_id: str | None = Field(default=None, description="Tool call identifier") tool_call_name: str | None = Field(default=None, description="Tool call name") diff --git a/reme/steps/common/chat.py b/reme/steps/common/chat.py index c9a34f68..5bbf22cc 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -2,7 +2,7 @@ from ..base_step import BaseStep from ...components import R -from ...enumeration import ChunkEnum +from ...schema import collect_reply @R.register("chat_step") @@ -33,18 +33,11 @@ this chat intentionally provides read-only tools. Reply in the user's language." if session_id := self.context.get("session_id"): wrapper_kwargs["resume"] = str(session_id) - parts: list[str] = [] - async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - if chunk.chunk_type == ChunkEnum.REPLY_END: - answer = "".join(parts).strip() - if answer: - chunk.metadata["answer"] = answer - await self.context.add_stream_chunk(chunk) - if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): - parts.append(chunk.chunk) - if chunk.session_id: - self.context.response.metadata["session_id"] = chunk.session_id - - self.context.response.success = True - self.context.response.answer = "".join(parts).strip() + result = await collect_reply( + self.agent_wrapper.reply_stream(query, **wrapper_kwargs), + on_chunk=self.context.add_stream_chunk, + ) + self.context.response.success = result["status"] == "completed" and self.context.stream_error is None + self.context.response.answer = result["result"] + self.context.response.metadata.update(result) return self.context.response diff --git a/reme/steps/common/llm_demo.py b/reme/steps/common/llm_demo.py index a5f2e6bf..9da23180 100644 --- a/reme/steps/common/llm_demo.py +++ b/reme/steps/common/llm_demo.py @@ -42,18 +42,17 @@ class LLMDemoStep(BaseStep): result = await self.agent_wrapper.reply(query, **wrapper_kwargs) - structured_content = result.get("structured_output") - text = (result.get("result") or "").strip() + text = result["result"].strip() self.logger.info(f"[{self.name}] response: {text!r}") - self.context.response.success = True - self.context.response.answer = text + self.context.response.success = result["status"] == "completed" + self.context.response.answer = result["error"] or text self.context.response.metadata.update( { + **result, "query": query, "sys_prompt": sys_prompt, "response": text, - "structured_output": structured_content, }, ) return self.context.response diff --git a/reme/steps/common/stream_llm_demo.py b/reme/steps/common/stream_llm_demo.py index aa554acf..2c8b7428 100644 --- a/reme/steps/common/stream_llm_demo.py +++ b/reme/steps/common/stream_llm_demo.py @@ -2,7 +2,7 @@ from ..base_step import BaseStep from ...components import R -from ...enumeration import ChunkEnum +from ...schema import collect_reply @R.register("stream_llm_demo_step") @@ -39,15 +39,19 @@ class StreamLLMDemoStep(BaseStep): } if self.context.stream: - text = await self._stream_reply(query, **wrapper_kwargs) + result = await collect_reply( + self.agent_wrapper.reply_stream(query, **wrapper_kwargs), + on_chunk=self.context.add_stream_chunk, + ) else: result = await self.agent_wrapper.reply(query, **wrapper_kwargs) - text = (result.get("result") or "").strip() + text = result["result"].strip() + self.context.response.metadata.update(result) self.logger.debug(f"[{self.name}] response: {text!r}") - self.context.response.success = True - self.context.response.answer = text + self.context.response.success = result["status"] == "completed" and self.context.stream_error is None + self.context.response.answer = result["error"] or text self.context.response.metadata.update( { "query": query, @@ -56,19 +60,3 @@ class StreamLLMDemoStep(BaseStep): }, ) return self.context.response - - async def _stream_reply(self, query: str, **wrapper_kwargs) -> str: - """Stream unified chunks to the context stream queue.""" - assert self.context is not None - text_parts: list[str] = [] - - async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - await self.context.add_stream_chunk(chunk) - - if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): - text_parts.append(chunk.chunk) - - if chunk.session_id: - self.context.response.metadata["session_id"] = chunk.session_id - - return "".join(text_parts).strip() diff --git a/reme/steps/evolve/_evolve.py b/reme/steps/evolve/_evolve.py index 73a0148b..dc93ea04 100644 --- a/reme/steps/evolve/_evolve.py +++ b/reme/steps/evolve/_evolve.py @@ -5,7 +5,7 @@ import zoneinfo from agentscope.message import Msg -from ...schema import Response +from ...schema import AgentReply, Response def now(timezone: str | None = None) -> datetime.datetime: @@ -31,17 +31,12 @@ def format_history(messages: list[Msg], include_timestamp: bool = True) -> str: return "\n\n".join(lines) or "(empty)" -def agent_reply_result_text(reply_result: dict) -> str: - """Return the final user-visible text block from an agent reply result.""" - last_message = reply_result.get("last_message") or {} - content = last_message.get("content") if isinstance(last_message, dict) else None - if isinstance(content, list): - for block in reversed(content): - if isinstance(block, dict) and block.get("type") == "text": - text = str(block.get("text") or "").strip() - if text: - return text - return str(reply_result.get("result") or "").strip() +def agent_reply_result_text(reply_result: AgentReply) -> str: + """Require a completed invocation and return its backend-neutral final text.""" + status = reply_result["status"] + if status != "completed" or reply_result.get("error") is not None: + raise RuntimeError(reply_result.get("error") or f"Agent reply ended: {status}") + return reply_result["result"].strip() def passthrough_response(step, skip_key: str) -> Response: diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 630c626a..3f90a112 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -677,6 +677,7 @@ class AutoMemoryStep(BaseStep): job_tools=self.create_tools if created else self.update_tools, **reply_kwargs, ) + text = agent_reply_result_text(result) self.logger.info(f"[{self.name}] agent done path={note_path} has_result={bool(result.get('result'))}") if created: @@ -692,7 +693,7 @@ class AutoMemoryStep(BaseStep): return if note is None: self.context.response.success = True - self.context.response.answer = agent_reply_result_text(result) + self.context.response.answer = text self.context.response.metadata.update( {"date": day, "path": None, "created": False, "modified": False, "n_messages": len(messages)}, ) @@ -730,7 +731,7 @@ class AutoMemoryStep(BaseStep): source_conversation = self._session_link(session_id) self.context.response.success = True - self.context.response.answer = agent_reply_result_text(result) + self.context.response.answer = text self.context.response.metadata.update( { "date": day, diff --git a/reme/steps/evolve/auto_tag.py b/reme/steps/evolve/auto_tag.py index 823d59cb..028b7ee1 100644 --- a/reme/steps/evolve/auto_tag.py +++ b/reme/steps/evolve/auto_tag.py @@ -120,6 +120,7 @@ class AutoTagStep(BaseStep): job_tools=self.tools, injected_job_kwargs=tool_context, ) + text = agent_reply_result_text(result) path = self.workspace_path / target.path metadata = dict(frontmatter.loads(path.read_text(encoding="utf-8")).metadata or {}) @@ -139,7 +140,7 @@ class AutoTagStep(BaseStep): ) if not response.success: raise RuntimeError(str(response.answer)) - return agent_reply_result_text(result) + return text async def execute(self): assert self.context is not None diff --git a/reme/steps/evolve/base_auto_resource.py b/reme/steps/evolve/base_auto_resource.py index ff7a9275..8a1ab07e 100644 --- a/reme/steps/evolve/base_auto_resource.py +++ b/reme/steps/evolve/base_auto_resource.py @@ -414,6 +414,7 @@ class BaseAutoResourceStep(BaseStep): } try: result = await self.agent_wrapper.reply(inputs, **agent_kwargs) + text = agent_reply_result_text(result) except (Exception, asyncio.CancelledError): self.context.response.success = False try: @@ -430,7 +431,7 @@ class BaseAutoResourceStep(BaseStep): metadata=note_metadata, ) self.context.response.success = True - self.context.response.answer = agent_reply_result_text(result) + self.context.response.answer = text if note_path is None: self.logger.info(f"[{self.name}] done without note file_path={file_path} modified=False") return None diff --git a/reme/steps/evolve/dream/integrate.py b/reme/steps/evolve/dream/integrate.py index f3946318..14f6c13a 100644 --- a/reme/steps/evolve/dream/integrate.py +++ b/reme/steps/evolve/dream/integrate.py @@ -146,6 +146,7 @@ class DreamIntegrateStep(BaseStep): ), job_tools=list(_TOOLS), ) + raw_result = agent_reply_result_text(result) except Exception as e: # noqa: BLE001 after = _snapshot_digest(workspace, digest_dir) changed = _changed_digest_paths(before, after) @@ -177,7 +178,6 @@ class DreamIntegrateStep(BaseStep): return try: - raw_result = agent_reply_result_text(result) outcome = IntegrateOutcome.model_validate(parse_structured_reply(raw_result)) if not self._valid_target( workspace, diff --git a/reme/steps/evolve/proactive/agenda.py b/reme/steps/evolve/proactive/agenda.py index 919a4a93..0cc840f7 100644 --- a/reme/steps/evolve/proactive/agenda.py +++ b/reme/steps/evolve/proactive/agenda.py @@ -216,6 +216,7 @@ class ProactiveAgendaStep(BaseStep): wrapper.reply(user_message, system_prompt=system_prompt), timeout=self.llm_timeout_seconds, ) + raw = agent_reply_result_text(result) except asyncio.TimeoutError: self.logger.warning(f"[{self.name}] LLM reply timed out after {self.llm_timeout_seconds}s") return self._fallback_agenda(carded_ids) @@ -224,7 +225,6 @@ class ProactiveAgendaStep(BaseStep): f"[{self.name}] LLM reply failed ({type(e).__name__}: {e}); using fallback agenda", ) return self._fallback_agenda(carded_ids) - raw = agent_reply_result_text(result) agenda_raw, suppressed_raw = parse_agenda_reply(raw) allowed = set(carded_ids) diff --git a/reme/steps/evolve/proactive/plan.py b/reme/steps/evolve/proactive/plan.py index 8cbf4ad1..54eeb3a2 100644 --- a/reme/steps/evolve/proactive/plan.py +++ b/reme/steps/evolve/proactive/plan.py @@ -213,6 +213,7 @@ class ProactivePlanStep(BaseStep): wrapper.reply(user_message, system_prompt=system_prompt), timeout=self.llm_timeout_seconds, ) + raw = agent_reply_result_text(result) except asyncio.TimeoutError: self.logger.warning(f"[{self.name}] LLM reply timed out after {self.llm_timeout_seconds}s") return {} @@ -221,7 +222,6 @@ class ProactivePlanStep(BaseStep): f"[{self.name}] LLM reply failed ({type(e).__name__}: {e}); using fallback cards", ) return {} - raw = agent_reply_result_text(result) cards_by_id: dict[str, dict] = {} for card in parse_plan_reply(raw): cid = str(card.get("topic_id") or "").strip() diff --git a/reme/utils/common_utils.py b/reme/utils/common_utils.py index fd03eb65..0259c918 100644 --- a/reme/utils/common_utils.py +++ b/reme/utils/common_utils.py @@ -208,12 +208,24 @@ async def call_action( """POST to /{action}; return parsed JSON (dict) for JSON endpoints, raw text for SSE.""" from ..components.client.http_client import HttpClient - pieces: list[str] = [] + pieces: dict[str, str] = {} + result = None async with HttpClient(host=host, port=port, timeout=timeout) as client: async for chunk in client.stream_chunks(action, **kwargs): payload = chunk.chunk - pieces.append(payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)) - raw = "".join(pieces) + if chunk.chunk_type == "error" and chunk.metadata.get("will_retry") is not True: + raise RuntimeError(str(payload)) + if chunk.chunk_type == "reply_end" and isinstance(payload, dict): + if payload["status"] != "completed": + raise RuntimeError(payload["error"] or f"Agent reply ended: {payload['status']}") + result = payload["result"] + continue + if payload is None or chunk.metadata.get("visibility") == "diagnostic": + continue + key = f"{chunk.chunk_type}:{chunk.block_id}" if chunk.block_id else str(len(pieces)) + text = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False) + pieces[key] = pieces.get(key, "") + text if chunk.event == "delta" else text + raw = result if result is not None else "".join(pieces.values()) try: return json.loads(raw) except (ValueError, json.JSONDecodeError): diff --git a/reme_studio/app/chat-stream.ts b/reme_studio/app/chat-stream.ts index 6b30b5ac..a52b7e17 100644 --- a/reme_studio/app/chat-stream.ts +++ b/reme_studio/app/chat-stream.ts @@ -5,6 +5,7 @@ import type { StreamChunk, StreamPayload, ToolBlock, + StreamStatus, } from "./types"; const detailTypes = new Set(["think", "data", "approval", "usage"]); @@ -70,17 +71,42 @@ export async function readChatStream( } const payloadText = (payload: StreamPayload): string => - typeof payload === "string" ? payload : JSON.stringify(payload, null, 2); + payload === null + ? "" + : typeof payload === "string" + ? payload + : JSON.stringify(payload, null, 2); -const appendError = (blocks: ChatBlock[], text: string): ChatBlock[] => { +const appendError = ( + blocks: ChatBlock[], + text: string, + metadata?: Record, +): ChatBlock[] => { if (!text) return blocks; const last = blocks.at(-1); - if (last?.type === "error") { + if ( + last?.type === "error" && + last.metadata?.will_retry === metadata?.will_retry + ) { return [...blocks.slice(0, -1), { ...last, text: last.text + text }]; } - return [...blocks, { id: `error:${blocks.length}`, type: "error", text }]; + return [ + ...blocks, + { id: `error:${blocks.length}`, type: "error", text, metadata }, + ]; }; +const settled = (status: string): status is "done" | "error" | "interrupted" => + ["done", "error", "interrupted"].includes(status); +const blockStatus = (status?: StreamStatus | null) => + status === "failed" + ? "error" + : status === "interrupted" + ? "interrupted" + : status === "waiting" + ? "waiting" + : "done"; + export function finishChatMessage( message: ChatMessage, error?: string, @@ -88,33 +114,35 @@ export function finishChatMessage( ): ChatMessage { const failed = Boolean(error) || - (finishTools && message.blocks?.some((block) => block.type === "error")); + message.status === "failed" || + message.blocks?.some( + (block) => block.type === "error" && block.metadata?.will_retry !== true, + ); + const status = failed + ? "failed" + : message.status === "running" + ? "interrupted" + : message.status === "interrupted" || message.status === "waiting" + ? message.status + : "completed"; 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" - ) + if (!finishTools && block.type === "tool" && !settled(block.status)) return block; return { ...block, - status: - block.status === "error" || (failed && block.status !== "done") - ? "error" - : "done", + status: settled(block.status) ? block.status : blockStatus(status), expanded: false, }; }); if (error) blocks = appendError(blocks, error); - return { ...message, blocks }; + return { ...message, status, blocks }; } function applyContent(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { + if (chunk.chunk === null) return blocks; const text = payloadText(chunk.chunk); - if (!text) return blocks; + if (!text && (!chunk.event || chunk.event === "delta")) return blocks; const id = chunk.block_id || ""; const index = id ? blocks.findIndex( @@ -124,12 +152,21 @@ function applyContent(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { const current = blocks[index]; if (current?.type === "content" && (id || index === blocks.length - 1)) { const next = [...blocks]; - next[index] = { ...current, text: current.text + text }; + next[index] = { + ...current, + text: chunk.event && chunk.event !== "delta" ? text : current.text + text, + metadata: { ...current.metadata, ...chunk.metadata }, + }; return next; } return [ ...blocks, - { id: `content:${id || blocks.length}`, type: "content", text }, + { + id: `content:${id || blocks.length}`, + type: "content", + text, + metadata: chunk.metadata, + }, ]; } @@ -160,10 +197,14 @@ function detailId(blocks: ChatBlock[], chunk: StreamChunk): string { function applyDetail(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { const id = detailId(blocks, chunk); const index = blocks.findIndex((block) => block.id === id); - const payloads = payloadText(chunk.chunk) ? [chunk.chunk] : []; + const payloads = + chunk.chunk !== null && (payloadText(chunk.chunk) || chunk.event === "end") + ? [chunk.chunk] + : []; const done = - chunk.chunk_type === "usage" && - (chunk.input_tokens !== undefined || chunk.output_tokens !== undefined); + chunk.event === "end" || + (chunk.chunk_type === "usage" && + (chunk.input_tokens != null || chunk.output_tokens != null)); const type = detailTypes.has(chunk.chunk_type) ? (chunk.chunk_type as DetailBlock["type"]) : "unknown"; @@ -175,8 +216,11 @@ function applyDetail(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { current.type !== "error" ? { ...current, - payloads: [...current.payloads, ...payloads], - status: done ? "done" : "streaming", + payloads: + chunk.event && chunk.event !== "delta" && chunk.chunk !== null + ? payloads + : [...current.payloads, ...payloads], + status: done ? blockStatus(chunk.status) : "streaming", expanded: current.expanded, mediaType: chunk.media_type || current.mediaType, inputTokens: chunk.input_tokens ?? current.inputTokens, @@ -188,11 +232,11 @@ function applyDetail(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { type, sourceType: chunk.chunk_type, payloads, - status: done ? "done" : "streaming", + status: done ? blockStatus(chunk.status) : "streaming", expanded: true, - mediaType: chunk.media_type, - inputTokens: chunk.input_tokens, - outputTokens: chunk.output_tokens, + mediaType: chunk.media_type || undefined, + inputTokens: chunk.input_tokens ?? undefined, + outputTokens: chunk.output_tokens ?? undefined, metadata: chunk.metadata, }; if (index < 0) return [...blocks, nextBlock]; @@ -203,72 +247,55 @@ function applyDetail(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { function applyTool(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { const last = blocks.at(-1); - const activeToolId = - last?.type === "tool" && last.status !== "done" && last.status !== "error" - ? last.id.slice("tool:".length) - : undefined; const toolId = chunk.tool_call_id || chunk.block_id || - activeToolId || - `${chunk.chunk_type}:${blocks.length}`; + (last?.type === "tool" && !settled(last.status) + ? last.id.slice("tool:".length) + : `${blocks.length}`); const id = `tool:${toolId}`; const index = blocks.findIndex( (block) => block.type === "tool" && block.id === id, ); const current = index >= 0 ? (blocks[index] as ToolBlock) : undefined; - const hasPayload = payloadText(chunk.chunk).length > 0; const result = chunk.chunk_type === "tool_result"; - 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 previous = + (result ? current?.resultPayloads : current?.callPayloads) || []; + const payloads = + chunk.chunk === null + ? previous + : chunk.event && chunk.event !== "delta" + ? [chunk.chunk] + : [...previous, ...(payloadText(chunk.chunk) ? [chunk.chunk] : [])]; + // TOOL_CALL end finishes argument streaming; only TOOL_RESULT end finishes execution. const completed = - failed || - current?.status === "done" || - claudeResult !== undefined || - (state !== "" && state !== "inprogress"); + chunk.status === "completed" || + chunk.status === "failed" || + chunk.status === "interrupted" || + (result && + chunk.event === "end" && + chunk.status !== "running" && + chunk.status !== "waiting"); + const status = + current && (current.status === "error" || current.status === "interrupted") + ? current.status + : completed + ? blockStatus(chunk.status) + : current && settled(current.status) + ? current.status + : chunk.status === "waiting" + ? "waiting" + : result + ? "running" + : "calling"; const nextBlock: ToolBlock = { id, type: "tool", name: chunk.tool_call_name || current?.name || "ReMe tool", - callPayloads: result - ? current?.callPayloads || [] - : [ - ...(current?.callPayloads || []), - ...(hasPayload ? [chunk.chunk] : []), - ], - resultPayloads: result - ? [ - ...(current?.resultPayloads || []), - ...(hasPayload ? [chunk.chunk] : []), - ] - : current?.resultPayloads || [], - status: failed - ? "error" - : completed - ? "done" - : result - ? "running" - : "calling", - expanded: completed ? false : current?.expanded ?? true, + callPayloads: result ? current?.callPayloads || [] : payloads, + resultPayloads: result ? payloads : current?.resultPayloads || [], + status, + expanded: settled(status) ? false : current?.expanded ?? true, mediaType: chunk.media_type || current?.mediaType, metadata: { ...current?.metadata, ...chunk.metadata }, }; @@ -278,16 +305,23 @@ function applyTool(blocks: ChatBlock[], chunk: StreamChunk): ChatBlock[] { return next; } -export function applyStreamChunk( +function reduceStreamChunk( message: ChatMessage, chunk: StreamChunk, ): ChatMessage { if (chunk.chunk_type === "reply_end") { - const answer = - typeof chunk.metadata?.answer === "string" ? chunk.metadata.answer : ""; - const hasContent = (message.blocks || []).some( - (block) => block.type === "content" && block.text.length > 0, - ); + const outcome = + chunk.chunk && + typeof chunk.chunk === "object" && + !Array.isArray(chunk.chunk) + ? (chunk.chunk as ChatMessage["reply"]) + : undefined; + const answer = outcome?.result || ""; + const hasContent = (message.blocks || []) + .filter((block) => block.type === "content") + .map((block) => block.text) + .join("") + .endsWith(answer); const completed = answer && !hasContent ? { @@ -302,18 +336,53 @@ export function applyStreamChunk( ], } : message; - return finishChatMessage(completed, undefined, { finishTools: false }); + const failed = + chunk.status === "failed" && + !(completed.blocks || []).some( + (block) => + block.type === "error" && block.metadata?.will_retry !== true, + ); + let blocks = failed + ? appendError( + completed.blocks || [], + String( + outcome?.error || + chunk.metadata?.error || + chunk.metadata?.finished_reason || + "failed", + ), + ) + : completed.blocks; + if (outcome?.structured_output != null) { + blocks = [ + ...(blocks || []), + { + id: `structured:${message.id}`, + type: "data", + sourceType: "structured_output", + payloads: [outcome.structured_output as StreamPayload], + status: "done", + expanded: false, + }, + ]; + } + return finishChatMessage( + { + ...completed, + blocks, + status: outcome?.status || chunk.status || undefined, + reply: outcome, + }, + undefined, + { finishTools: chunk.event === "end" }, + ); } 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; + if (chunk.chunk_type === "reply_start") + return { ...message, status: "running" }; + if (chunk.metadata?.visibility === "diagnostic") return message; let blocks = message.blocks || []; if ( @@ -340,11 +409,21 @@ export function applyStreamChunk( if (chunk.chunk_type === "error") return { ...message, - blocks: appendError(blocks, payloadText(chunk.chunk)), + blocks: appendError(blocks, payloadText(chunk.chunk), chunk.metadata), }; return { ...message, blocks: applyDetail(blocks, chunk) }; } +export function applyStreamChunk( + message: ChatMessage, + chunk: StreamChunk, +): ChatMessage { + return { + ...reduceStreamChunk(message, chunk), + events: [...(message.events || []), chunk], + }; +} + export function toggleChatBlock( message: ChatMessage, blockId: string, diff --git a/reme_studio/app/i18n.ts b/reme_studio/app/i18n.ts index 6118a755..c8e8bde4 100644 --- a/reme_studio/app/i18n.ts +++ b/reme_studio/app/i18n.ts @@ -89,6 +89,8 @@ const messages = { streaming: "进行中", completed: "已完成", failed: "失败", + interrupted: "已中断", + waiting: "等待输入", askWorkspace: "询问你的工作区…", send: "发送", composerHint: "Enter 发送 · Shift + Enter 换行", @@ -247,6 +249,8 @@ const messages = { streaming: "Streaming", completed: "Completed", failed: "Failed", + interrupted: "Interrupted", + waiting: "Waiting for input", 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 index 2e2c9ced..46acda32 100644 --- a/reme_studio/app/session-sync.ts +++ b/reme_studio/app/session-sync.ts @@ -59,7 +59,19 @@ export function mergeSessionMessages( }, ]; for (const chunk of chunks) message = applyStreamChunk(message, chunk); + if (previous?.reply) { + message = applyStreamChunk(message, { + chunk_type: "reply_end", + event: "end", + chunk: previous.reply, + status: previous.reply.status, + done: false, + }); + } if (entry.completed) message = finishChatMessage(message); + if (previous?.events?.length) { + message.events = previous.events; + } const failed = message.blocks?.some( (block) => block.type === "error" || diff --git a/reme_studio/app/store.ts b/reme_studio/app/store.ts index e821b9e1..0c34657e 100644 --- a/reme_studio/app/store.ts +++ b/reme_studio/app/store.ts @@ -217,7 +217,8 @@ export const useWorkspaceStore = create()( applyStreamChunk(message, chunk), ); const turnId = - chunk.chunk_type === "reply_start" && chunk.metadata?.turn_id; + chunk.chunk_type === "reply_start" && + (chunk.turn_id || chunk.metadata?.turn_id); if (typeof turnId === "string") { const assistantIndex = updated.messages.findIndex( (message) => message.id === messageId, diff --git a/reme_studio/app/types.ts b/reme_studio/app/types.ts index ed17e1d5..14122f76 100644 --- a/reme_studio/app/types.ts +++ b/reme_studio/app/types.ts @@ -90,26 +90,58 @@ export type StreamChunkType = | "error" | "done"; -export type StreamPayload = string | Record | unknown[]; +export type StreamPayload = + | string + | number + | boolean + | Record + | unknown[] + | null; +export type StreamStatus = + | "running" + | "completed" + | "failed" + | "interrupted" + | "waiting"; export interface StreamChunk { chunk_type: StreamChunkType | (string & {}); + event?: "start" | "delta" | "end"; chunk: StreamPayload; done: boolean; - session_id?: string; - block_id?: string; - tool_call_id?: string; - tool_call_name?: string; - media_type?: string; - input_tokens?: number; - output_tokens?: number; + status?: StreamStatus | null; + session_id?: string | null; + turn_id?: string | null; + message_id?: string | null; + block_id?: string | null; + tool_call_id?: string | null; + tool_call_name?: string | null; + media_type?: string | null; + input_tokens?: number | null; + output_tokens?: number | null; metadata?: Record; } +export interface AgentReply { + session_id: string; + result: string; + structured_output: unknown; + usage: { + input_tokens: number; + output_tokens: number; + total_tokens: number; + } | null; + status: Exclude; + error: string | null; + raw: Record; + events: StreamChunk[]; +} + export interface ContentBlock { id: string; type: "content"; text: string; + metadata?: Record; } export interface DetailBlock { @@ -117,7 +149,7 @@ export interface DetailBlock { type: "think" | "data" | "approval" | "usage" | "unknown"; sourceType: string; payloads: StreamPayload[]; - status: "streaming" | "done" | "error"; + status: "streaming" | "done" | "error" | "interrupted" | "waiting"; expanded: boolean; mediaType?: string; inputTokens?: number; @@ -131,7 +163,7 @@ export interface ToolBlock { name: string; callPayloads: StreamPayload[]; resultPayloads: StreamPayload[]; - status: "calling" | "running" | "done" | "error"; + status: "calling" | "running" | "done" | "error" | "interrupted" | "waiting"; expanded: boolean; mediaType?: string; metadata?: Record; @@ -141,6 +173,7 @@ export interface ErrorBlock { id: string; type: "error"; text: string; + metadata?: Record; } export type ChatBlock = ContentBlock | DetailBlock | ToolBlock | ErrorBlock; @@ -150,6 +183,9 @@ export interface ChatMessage { role: "user" | "assistant"; content: string; blocks?: ChatBlock[]; + events?: StreamChunk[]; + reply?: Omit; + status?: StreamStatus; turnId?: string; source?: "chat" | "proactive"; createdAt?: string | null; diff --git a/reme_studio/app/workspace-chat.tsx b/reme_studio/app/workspace-chat.tsx index 216344a7..2721ba8f 100644 --- a/reme_studio/app/workspace-chat.tsx +++ b/reme_studio/app/workspace-chat.tsx @@ -79,7 +79,7 @@ function ChatBlockView({ block.status === "streaming" || block.status === "calling" || block.status === "running"; - const failed = block.status === "error"; + const failed = block.status === "error" || block.status === "interrupted"; const tool = block.type === "tool"; const call = tool ? formatStreamPayloads(block.callPayloads) : ""; const result = tool ? formatStreamPayloads(block.resultPayloads) : ""; @@ -97,7 +97,17 @@ function ChatBlockView({ {tool ? block.name : detailLabel(block, t)} - {t(failed ? "failed" : active ? "streaming" : "completed")} + {t( + block.status === "interrupted" + ? "interrupted" + : block.status === "waiting" + ? "waiting" + : failed + ? "failed" + : active + ? "streaming" + : "completed", + )} {tool ? ( @@ -231,6 +241,7 @@ export default function Chat({ const assistant: ChatMessage = { id: crypto.randomUUID(), role: "assistant", + status: "running", content: "", blocks: [], }; @@ -295,6 +306,17 @@ export default function Chat({ } /> ))} + {message.role === "assistant" && + (message.status === "interrupted" || + message.status === "waiting") && ( +

+ {t( + message.status === "interrupted" + ? "interrupted" + : "waiting", + )} +

+ )} {message.role === "assistant" && tab.streaming && !message.blocks?.some((block) => block.type !== "usage") && ( diff --git a/reme_studio/tests/chat-stream.test.mjs b/reme_studio/tests/chat-stream.test.mjs index 92d792c4..7ccf18bf 100644 --- a/reme_studio/tests/chat-stream.test.mjs +++ b/reme_studio/tests/chat-stream.test.mjs @@ -113,7 +113,7 @@ test("thinking finishes as each tool or answer phase starts", () => { message, chunk("tool_result", "00:11", { tool_call_id: "clock", - metadata: { state: "completed" }, + status: "completed", }), ); message = applyStreamChunk( @@ -192,7 +192,11 @@ test("reply completion restores the final answer when content deltas were missed ); message = applyStreamChunk( message, - chunk("reply_end", "", { metadata: { answer: "Final answer." } }), + chunk( + "reply_end", + { result: "Final answer." }, + { event: "end", status: "completed" }, + ), ); assert.equal(message.blocks.at(-1).type, "content"); @@ -232,7 +236,7 @@ test("Codex hook notifications do not become chat data blocks", () => { threadId: "thread-1", turnId: "turn-1", }, - { metadata: { codex_method: method } }, + { metadata: { codex_method: method, visibility: "diagnostic" } }, ), ); } @@ -281,6 +285,7 @@ test("Codex tool snapshots finish before the reply ends", () => { { tool_call_id: "tool-1", tool_call_name: "search", + status: "running", }, ), ); @@ -307,6 +312,8 @@ test("Codex tool snapshots finish before the reply ends", () => { { id: "tool-1", type: "mcpToolCall", status: "completed" }, { tool_call_id: "tool-1", + event: "end", + status: "completed", }, ), ); @@ -325,6 +332,8 @@ test("Codex failed and declined tools stay failed after reply completion", () => { tool_call_id: "tool-1", tool_call_name: "commandExecution", + event: "end", + status: "failed", }, ), ); @@ -362,7 +371,11 @@ test("Claude tools remain pending across message boundaries and report their act chunk( "tool_result", { tool_use_id: "tool-1", content: "Result", is_error }, - { tool_call_id: "tool-1" }, + { + tool_call_id: "tool-1", + event: "end", + status: is_error ? "failed" : "completed", + }, ), ); assert.equal(message.blocks[0].status, is_error ? "error" : "done"); @@ -378,7 +391,7 @@ test("Claude server tool results finish without an is_error field", () => { chunk( "tool_result", { tool_use_id: "server-1", content: { matches: 1 } }, - { tool_call_id: "server-1" }, + { tool_call_id: "server-1", event: "end", status: "completed" }, ), ); assert.equal(message.blocks[0].status, "done"); @@ -406,7 +419,7 @@ test("a transport failure marks pending blocks failed and preserves completed ph message, chunk("tool_result", "", { tool_call_id: "completed", - metadata: { state: "completed" }, + status: "completed", }), ); message = applyStreamChunk( @@ -444,7 +457,7 @@ test("late tool output does not reopen a completed or failed execution", () => { initial(), chunk("tool_result", "Result", { tool_call_id: "tool-1", - metadata: { state }, + status: state, }), ); message = applyStreamChunk( @@ -640,3 +653,157 @@ test("a failed stream returns its error message", async () => { assert.equal(error, "Connection lost"); }); + +test("authoritative text snapshots replace deltas, including an empty final text", () => { + let message = initial(); + for (const event of [ + chunk("content", "draft", { block_id: "answer" }), + chunk("content", "corrected", { + block_id: "answer", + event: "end", + metadata: { phase: "final_answer" }, + }), + chunk("content", "corrected", { block_id: "answer", event: "end" }), + chunk("content", null, { block_id: "answer", event: "end" }), + ]) + message = applyStreamChunk(message, event); + assert.equal(message.blocks.length, 1); + assert.equal(message.blocks[0].text, "corrected"); + assert.equal(message.blocks[0].metadata.phase, "final_answer"); + message = applyStreamChunk( + message, + chunk("content", "", { block_id: "answer", event: "end" }), + ); + assert.equal(message.blocks[0].text, ""); +}); + +test("tool snapshots replace streamed output and argument completion keeps the tool pending", () => { + let message = applyStreamChunk( + initial(), + chunk("tool_call", '{"q":', { tool_call_id: "tool" }), + ); + message = applyStreamChunk( + message, + chunk("tool_call", { q: "hello" }, { tool_call_id: "tool", event: "end" }), + ); + assert.equal(message.blocks[0].status, "calling"); + message = applyStreamChunk( + message, + chunk("tool_result", "partial", { + tool_call_id: "tool", + status: "running", + }), + ); + message = applyStreamChunk( + message, + chunk( + "tool_result", + { output: "complete" }, + { tool_call_id: "tool", event: "end", status: "completed" }, + ), + ); + assert.deepEqual(message.blocks[0].callPayloads, [{ q: "hello" }]); + assert.deepEqual(message.blocks[0].resultPayloads, [{ output: "complete" }]); + assert.equal(message.blocks[0].status, "done"); +}); + +test("retryable errors do not turn a recovered reply into a failure", () => { + let message = applyStreamChunk( + initial(), + chunk("error", "temporary", { metadata: { will_retry: true } }), + ); + message = applyStreamChunk(message, chunk("content", "recovered")); + message = applyStreamChunk( + message, + chunk("reply_end", null, { event: "end", status: "completed" }), + ); + message = applyStreamChunk(message, chunk("done", "", { done: true })); + assert.equal(message.status, "completed"); + assert.equal(message.blocks[0].metadata.will_retry, true); +}); + +test("reply interruption and input waits survive transport completion", () => { + for (const status of ["interrupted", "waiting"]) { + let message = applyStreamChunk( + initial(), + chunk("think", "working", { block_id: "thought" }), + ); + message = applyStreamChunk( + message, + chunk("tool_call", null, { tool_call_id: "tool" }), + ); + message = applyStreamChunk( + message, + chunk("reply_end", null, { event: "end", status }), + ); + message = applyStreamChunk(message, chunk("done", "", { done: true })); + assert.equal(message.status, status); + assert.equal(message.blocks[0].status, "done"); + assert.equal(message.blocks[1].status, status); + } +}); + +test("terminal outcomes retain raw data, scalar structured output and diagnostic events", () => { + const outcome = { + session_id: "session", + result: "final answer", + structured_output: false, + usage: null, + status: "waiting", + error: null, + raw: { cache: 7, signature: "signed", future: [false, null] }, + }; + const events = [ + chunk("reply_start", null, { event: "start", status: "running" }), + chunk("data", null, { + metadata: { visibility: "diagnostic", native: { future: [false, null] } }, + }), + chunk("content", "checking", { block_id: "commentary" }), + chunk("reply_end", outcome, { event: "end", status: "waiting" }), + chunk("done", "", { done: true }), + ]; + const message = events.reduce(applyStreamChunk, initial()); + assert.deepEqual(message.events, events); + assert.deepEqual(message.reply, outcome); + assert.equal(message.status, "waiting"); + assert.deepEqual( + message.blocks.map((block) => block.type), + ["content", "content", "data"], + ); + assert.equal(message.blocks[1].text, "final answer"); + assert.deepEqual(message.blocks[2].payloads, [false]); + assert.equal(formatStreamPayloads(message.blocks[2].payloads), "false"); + const replayed = JSON.parse(JSON.stringify(message.events)).reduce( + applyStreamChunk, + initial(), + ); + assert.deepEqual(replayed, message); +}); + +test("null usage fields do not falsely settle an in-progress detail block", () => { + const message = applyStreamChunk( + initial(), + chunk("usage", null, { input_tokens: null, output_tokens: null }), + ); + assert.equal(message.blocks[0].status, "streaming"); + assert.equal(message.blocks[0].inputTokens, undefined); +}); + +test("stopping a running invocation preserves interruption rather than implying success", () => { + let message = applyStreamChunk( + initial(), + chunk("reply_start", null, { status: "running" }), + ); + message = applyStreamChunk( + message, + chunk("tool_call", "{}", { tool_call_id: "pending" }), + ); + const stopped = finishChatMessage(message); + assert.equal(stopped.status, "interrupted"); + assert.equal(stopped.blocks[0].status, "interrupted"); + assert.deepEqual(stopped.events, message.events); + assert.equal( + stopped.blocks.some((block) => block.type === "error"), + false, + ); +}); diff --git a/reme_studio/tests/session-sync.test.mjs b/reme_studio/tests/session-sync.test.mjs index ec2bcc14..bf47e1d1 100644 --- a/reme_studio/tests/session-sync.test.mjs +++ b/reme_studio/tests/session-sync.test.mjs @@ -1,5 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; +import { applyStreamChunk } from "../app/chat-stream.ts"; import { mergeSessionMessages, mergeSessionTab } from "../app/session-sync.ts"; const entry = (id, role, content, extra = {}) => ({ @@ -164,6 +165,8 @@ test("history polling preserves running tools until a native completion snapshot done: false, tool_call_id: "tool-1", tool_call_name: "search", + event: "end", + status: status === "inProgress" ? "running" : status, }); const active = snapshot([ entry("a1", "assistant", "", { @@ -192,6 +195,8 @@ test("native history repairs interrupted tools and restores missed failures", () chunk: { id: "tool-1", type: "commandExecution", status }, done: false, tool_call_id: "tool-1", + event: "end", + status, }, ], }), @@ -209,3 +214,86 @@ test("native history repairs interrupted tools and restores missed failures", () const failed = mergeSessionMessages(recovered, data("failed")); assert.equal(failed[0].blocks[0].status, "error"); }); + +test("history refresh preserves the exact live event trace and structured outcome", () => { + const streamed = turn("temporary-assistant", "B"); + streamed.events = [ + { chunk_type: "data", chunk: { signature: "signed" }, done: false }, + ]; + streamed.reply = { + session_id: "session-1", + usage: null, + status: "completed", + error: null, + result: "B", + structured_output: false, + raw: { signature: "signed" }, + }; + const result = mergeSessionMessages( + [streamed], + snapshot([entry("a1", "assistant", "B")]), + ); + assert.deepEqual(result[0].events, streamed.events); + assert.deepEqual(result[0].reply, streamed.reply); +}); + +test("native completion cannot overwrite a collected application-level terminal outcome", () => { + for (const status of ["completed", "failed", "interrupted", "waiting"]) { + const outcome = { + session_id: "session-1", + result: "invalid JSON", + structured_output: false, + usage: null, + status, + error: status === "failed" ? "Codex returned invalid JSON" : null, + raw: { status: "completed", final_response: "invalid JSON" }, + }; + let live = { ...turn("temporary-assistant", "invalid JSON"), blocks: [] }; + const chunks = [ + { + chunk_type: "content", + event: "end", + block_id: "answer", + chunk: "invalid JSON", + done: false, + }, + { + chunk_type: "reply_end", + event: "end", + status, + chunk: outcome, + done: false, + }, + ]; + for (const chunk of chunks) live = applyStreamChunk(live, chunk); + const data = snapshot([ + entry("a1", "assistant", "invalid JSON", { + chunks: [ + chunks[0], + { + chunk_type: "reply_end", + event: "end", + status: "completed", + chunk: null, + done: false, + }, + ], + }), + ]); + const restored = mergeSessionMessages([live], data); + assert.equal(restored[0].status, status); + assert.equal(restored[0].reply.status, status); + assert.deepEqual(restored[0].events, live.events); + assert.deepEqual( + restored[0].blocks.filter((block) => block.type === "data")[0].payloads, + [false], + ); + assert.deepEqual( + restored[0].blocks + .filter((block) => block.type === "error") + .map((block) => block.text), + status === "failed" ? [outcome.error] : [], + ); + assert.equal(mergeSessionMessages(restored, data), restored); + } +}); diff --git a/tests/integration/test_agent_session.py b/tests/integration/test_agent_session.py index eb7deeb2..bdd6dff2 100644 --- a/tests/integration/test_agent_session.py +++ b/tests/integration/test_agent_session.py @@ -7,6 +7,7 @@ environment or a .env file at the repo root. Hits the real LLM API. import asyncio import sys from pathlib import Path +from uuid import uuid4 INTEGRATION_DIR = Path(__file__).resolve().parent sys.path.insert(0, str(INTEGRATION_DIR)) @@ -24,15 +25,16 @@ async def _run_session_persistence() -> None: try: wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] - sid = "test-persist-session" + sid = str(uuid4()) # First call: establish session - _, msg_1 = await wrapper.reply( + reply_1 = await wrapper.reply( "My favorite color is blue. Remember that.", session_id=sid, system_prompt="You are a helpful assistant. Keep answers short.", ) - text_1 = (msg_1.get_text_content() or "").strip() + assert reply_1["status"] == "completed", reply_1["error"] + text_1 = reply_1["result"].strip() print(f"\n[session_persist] reply 1: {text_1!r}") assert text_1, "Empty first reply" @@ -42,13 +44,14 @@ async def _run_session_persistence() -> None: assert len(files_after_1) == 1, f"Expected 1 session file, got {len(files_after_1)}" assert sid in files_after_1[0].name - # Second call: same session_id — agent should have previous context - _, msg_2 = await wrapper.reply( + # Resume the returned session to retain its context. + reply_2 = await wrapper.reply( "What is my favorite color?", - session_id=sid, + resume=sid, system_prompt="You are a helpful assistant. Keep answers short.", ) - text_2 = (msg_2.get_text_content() or "").strip() + assert reply_2["status"] == "completed", reply_2["error"] + text_2 = reply_2["result"].strip() print(f"[session_persist] reply 2: {text_2!r}") assert "blue" in text_2.lower(), f"Agent should recall 'blue' from session context, got: {text_2!r}" @@ -64,7 +67,7 @@ async def _run_fork_session() -> None: try: wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] - sid = "test-fork-origin" + sid = str(uuid4()) # Establish original session await wrapper.reply( @@ -76,13 +79,15 @@ async def _run_fork_session() -> None: assert len(files_before_fork) == 1 # Fork the session - forked_sid, msg_fork = await wrapper.reply( + forked = await wrapper.reply( "What is the secret number?", - session_id=sid, + resume=sid, fork_session=True, system_prompt="You are a helpful assistant. Keep answers short.", ) - text_fork = (msg_fork.get_text_content() or "").strip() + assert forked["status"] == "completed", forked["error"] + forked_sid = forked["session_id"] + text_fork = forked["result"].strip() print(f"\n[fork_session] forked reply: {text_fork!r}") assert "42" in text_fork, f"Forked session should recall '42', got: {text_fork!r}" @@ -104,31 +109,32 @@ async def _run_fork_session() -> None: await env.close_all() -async def _run_no_session_id() -> None: - """When session_id is empty, no session file should be created.""" +async def _run_auto_session_id() -> None: + """Omitting session_id creates and returns a persisted session.""" with workspace_env() as env: app = await env.make_app() try: wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] - _, msg = await wrapper.reply( + reply = await wrapper.reply( "Say hello.", system_prompt="You are a helpful assistant. Keep answers short.", ) - text = (msg.get_text_content() or "").strip() + assert reply["status"] == "completed", reply["error"] + text = reply["result"].strip() print(f"\n[no_session] reply: {text!r}") assert text, "Empty reply" files = env.session_state_files() - assert len(files) == 0, f"No session files should be created without session_id, found {files}" + assert len(files) == 1 and reply["session_id"] in files[0].name, files - print("✓ test_no_session_id passed") + print("✓ test_auto_session_id passed") finally: await env.close_all() async def _run_all() -> None: - await _run_no_session_id() + await _run_auto_session_id() await _run_session_persistence() await _run_fork_session() diff --git a/tests/integration/test_codex_agent_wrapper.py b/tests/integration/test_codex_agent_wrapper.py index 3cef6b4d..ef1c0946 100644 --- a/tests/integration/test_codex_agent_wrapper.py +++ b/tests/integration/test_codex_agent_wrapper.py @@ -98,7 +98,7 @@ async def test_live_shared_codex_sessions_tools_streaming_and_close(tmp_path): job_tools=["add_draft", "read_draft"], tool_context_id="one", ) - assert "FIRST_OK" in original["last_message"] + assert "FIRST_OK" in original["result"] resumed = await second.reply( "Call read_draft and include its exact result.", session_id=original["session_id"], @@ -106,20 +106,20 @@ async def test_live_shared_codex_sessions_tools_streaming_and_close(tmp_path): tool_context_id="one", ) assert resumed["session_id"] == original["session_id"] - assert "STATE_OK" in resumed["last_message"] + assert "STATE_OK" in resumed["result"] isolated = await second.reply( "Call read_draft. If it is empty, reply ISOLATED_OK.", job_tools=["read_draft"], tool_context_id="two", ) assert isolated["session_id"] != original["session_id"] - assert "ISOLATED_OK" in isolated["last_message"] - assert "STATE_OK" not in isolated["last_message"] + assert "ISOLATED_OK" in isolated["result"] + assert "STATE_OK" not in isolated["result"] forked = await first.reply("Reply FORK_OK.", resume=original["session_id"], fork_session=True) assert forked["session_id"] != original["session_id"] - assert "FORK_OK" in forked["last_message"] + assert "FORK_OK" in forked["result"] custom = await first.reply("Call only_custom and include its result.", job_tools=["only_custom"]) - assert "CUSTOM_JOB_OK" in custom["last_message"] + assert "CUSTOM_JOB_OK" in custom["result"] assert app.context.metadata["live_custom_calls"] >= 1 structured = await first.reply("Return marker STRUCTURED_OK.", output_schema=_StructuredResult) assert structured["structured_output"] == {"marker": "STRUCTURED_OK"} @@ -128,7 +128,7 @@ async def test_live_shared_codex_sessions_tools_streaming_and_close(tmp_path): "Read the reme_memory skill instructions and reply SKILL_OK. Do not execute its scripts.", skills=["reme_memory"], ) - assert "SKILL_OK" in skill["last_message"] + assert "SKILL_OK" in skill["result"] assert (workspace / ".agents/skills/reme_memory").is_symlink() concurrent = await asyncio.gather(first.reply("Reply ONE."), second.reply("Reply TWO.")) assert concurrent[0]["session_id"] != concurrent[1]["session_id"] diff --git a/tests/unit/auto_resource_test_support.py b/tests/unit/auto_resource_test_support.py index 604882f6..82ca39a1 100644 --- a/tests/unit/auto_resource_test_support.py +++ b/tests/unit/auto_resource_test_support.py @@ -32,7 +32,7 @@ class FakeAgentWrapper(BaseAgentWrapper): async def reply(self, inputs, **_kwargs) -> dict: """Record and accept one text-processor request.""" self.inputs = inputs - return {"result": "ok"} + return self._reply_result(session_id="stub", result="ok", raw={}) class FlakyAgentWrapper(BaseAgentWrapper): @@ -47,7 +47,7 @@ class FlakyAgentWrapper(BaseAgentWrapper): self.calls += 1 if self.calls == 1: raise RuntimeError("text provider unavailable") - return {"result": "recovered"} + return self._reply_result(session_id="stub", result="recovered", raw={}) class FakeImageAgentWrapper(AsAgentWrapper): @@ -76,7 +76,7 @@ class FakeImageAgentWrapper(AsAgentWrapper): if self.error is not None: raise self.error if not self.perform_write: - return {"result": str(self.content)} + return self._reply_result(session_id="stub", result=str(self.content), raw={}) assert isinstance(inputs, Msg) assert kwargs.get("output_schema") is None target = kwargs["injected_job_kwargs"]["_allowed_paths"] @@ -103,7 +103,7 @@ class FakeImageAgentWrapper(AsAgentWrapper): ) if self.after_write_error is not None: raise self.after_write_error - return {"result": "Saved image note", "tool_result": response} + return self._reply_result(session_id="stub", result="Saved image note", raw={"tool_result": response}) class FlakyImageAgentWrapper(FakeImageAgentWrapper): diff --git a/tests/unit/test_agent_reply.py b/tests/unit/test_agent_reply.py new file mode 100644 index 00000000..8a5a6124 --- /dev/null +++ b/tests/unit/test_agent_reply.py @@ -0,0 +1,437 @@ +"""Portable outcomes retain native result details across agent backends.""" + +# pylint: disable=redefined-outer-name + +import asyncio +import json +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from agentscope.message import AssistantMsg, Usage +from agentscope.event import ReplyStartEvent, ReplyEndEvent +from agentscope.model import StructuredResponse +from agentscope.model._model_usage import ChatUsage +from agentscope.types import ErrorInfo +from claude_agent_sdk import DeferredToolUse, ResultMessage +import pytest +from openai_codex.generated.v2_all import AgentMessageThreadItem, ItemCompletedNotification + +from reme.components.agent_wrapper import AsAgentWrapper, CcAgentWrapper +from reme.components.application_context import ApplicationContext +from reme.utils import global_counter_get_all + +from .test_codex_agent_wrapper import _Turn, _usage_event, system as _codex_system + +system = _codex_system + + +def _agentscope(tmp_path, monkeypatch, message, model=None): + """Use real SDK results with mocked model and persistence boundaries.""" + wrapper = AsAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path)), as_llm="") + + async def reply_stream(_inputs, **_kwargs): + yield ReplyStartEvent(name="agent", session_id="session-1", reply_id="reply-1") + if message.finished_reason is not None: + yield ReplyEndEvent(session_id="session-1", reply_id="reply-1", finished_reason=message.finished_reason) + yield message + + agent = SimpleNamespace( + state=SimpleNamespace(session_id="session-1", context=[message]), + reply_stream=reply_stream, + ) + wrapper.as_llm = SimpleNamespace(model=model) + monkeypatch.setattr(wrapper, "_build_agent", AsyncMock(return_value=(agent, "hello"))) + monkeypatch.setattr(wrapper, "_dump_state", AsyncMock()) + return wrapper + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reason", "status"), + [ + ("completed", "completed"), + ("error", "failed"), + ("exceed_max_iters", "failed"), + ("interrupted", "interrupted"), + (None, "waiting"), + ], +) +async def test_agentscope_outcome_preserves_native_message(tmp_path, monkeypatch, reason, status): + """Native termination, content and cache counters survive portable normalization.""" + message = AssistantMsg( + name="agent", + content="final answer", + finished_reason=reason, + usage=Usage(input_tokens=10, output_tokens=2, cache_input_tokens=7), + ) + if reason == "error": + message.error = ErrorInfo(type="upstream", message="provider unavailable") + model = SimpleNamespace(generate_structured_output=AsyncMock(return_value=StructuredResponse(content={}))) + wrapper = _agentscope(tmp_path, monkeypatch, message, model) + result = await wrapper.reply("hello", **({"output_schema": {}} if status != "completed" else {})) + + assert set(result) == {"session_id", "result", "structured_output", "usage", "status", "error", "raw", "events"} + assert result["status"] == status + assert result["result"] == "final answer" + assert result["raw"]["message"] == message.model_dump(mode="json") + assert result["usage"] == {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12} + assert result["raw"]["message"]["usage"]["cache_input_tokens"] == 7 + assert result["error"] == ( + "provider unavailable" + if reason == "error" + else "AgentScope reply ended: exceed_max_iters" if reason == "exceed_max_iters" else None + ) + model.generate_structured_output.assert_not_awaited() + await _assert_roundtrip(result) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("interrupted", [False, True]) +async def test_agentscope_preserves_extra_structured_response(tmp_path, monkeypatch, interrupted): + """The extra model call keeps its metadata and usage without overwriting the reply.""" + message = AssistantMsg( + name="agent", + content="answer", + finished_reason="completed", + usage=Usage(input_tokens=10, output_tokens=2), + ) + structured = StructuredResponse( + content={}, + usage=ChatUsage(input_tokens=3, output_tokens=1, time=0.0, cache_input_tokens=2), + metadata={"provider": "extra details"}, + finished_reason="interrupted" if interrupted else "completed", + ) + model = SimpleNamespace(generate_structured_output=AsyncMock(return_value=structured)) + wrapper = _agentscope(tmp_path, monkeypatch, message, model) + result = await wrapper.reply("hello", output_schema={}) + + assert result["status"] == ("interrupted" if interrupted else "completed") + assert result["structured_output"] == {} + assert result["usage"] == {"input_tokens": 13, "output_tokens": 3, "total_tokens": 16} + assert result["raw"]["message"] == message.model_dump(mode="json") + assert result["raw"]["structured_response"]["metadata"] == structured.metadata + assert result["raw"]["structured_response"]["usage"]["cache_input_tokens"] == 2 + await _assert_roundtrip(result) + + +@pytest.mark.asyncio +async def test_agentscope_propagates_cancellation_even_if_sdk_returns_a_message(tmp_path, monkeypatch): + """An SDK interruption fallback cannot turn caller cancellation into a normal return.""" + message = AssistantMsg( + name="agent", + content="interrupted", + finished_reason="interrupted", + usage=Usage(input_tokens=10, output_tokens=2), + ) + wrapper = _agentscope(tmp_path, monkeypatch, message) + agent = wrapper._build_agent.return_value[0] # pylint: disable=protected-access + entered = asyncio.Event() + + async def native_reply(_inputs, **_kwargs): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + yield ReplyEndEvent(session_id="session-1", reply_id="reply-1", finished_reason="interrupted") + yield message + + agent.reply_stream = native_reply + task = asyncio.create_task(wrapper.reply("hello")) + await asyncio.wait_for(entered.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + wrapper._dump_state.assert_awaited_once_with(agent.state) # pylint: disable=protected-access + counters = global_counter_get_all(wrapper.app_context.metadata, ["__token_counter", wrapper.name]) + assert counters["children"]["total_tokens"]["value"] == 12 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("overrides", "status"), + [ + ({}, "completed"), + ({"is_error": True, "subtype": "error_max_turns", "errors": ["turn limit"]}, "failed"), + ({"terminal_reason": "aborted_tools"}, "interrupted"), + ({"deferred_tool_use": DeferredToolUse(id="tool-1", name="external", input={})}, "waiting"), + ], +) +async def test_claude_outcome_preserves_native_result(tmp_path, monkeypatch, overrides, status): + """Fees, cache usage, permission denials and terminal reasons survive the common shape.""" + message = ResultMessage( + subtype="success", + duration_ms=20, + duration_api_ms=10, + is_error=False, + num_turns=2, + session_id="session-1", + result="answer", + structured_output=False, + total_cost_usd=0.01, + usage={"input_tokens": 10, "output_tokens": 2, "cache_read_input_tokens": 7}, + permission_denials=[{"tool_name": "Write", "tool_use_id": "tool-1"}], + ) + message = replace(message, **overrides) + + async def query(**_kwargs): + yield message + if message.is_error: + raise RuntimeError("Claude Code returned an error result: turn limit") + + monkeypatch.setattr("claude_agent_sdk.query", query) + wrapper = CcAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + result = await wrapper.reply("hello") + + assert set(result) == {"session_id", "result", "structured_output", "usage", "status", "error", "raw", "events"} + assert result["status"] == status + assert result["structured_output"] is False + assert result["error"] == ("turn limit" if status == "failed" else None) + assert result["raw"]["total_cost_usd"] == 0.01 + assert result["raw"]["usage"]["cache_read_input_tokens"] == 7 + assert result["raw"]["permission_denials"] == message.permission_denials + assert result["raw"]["terminal_reason"] == message.terminal_reason + await _assert_roundtrip(result) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_claude_does_not_hide_transport_errors_or_cancellation(tmp_path, monkeypatch, cancelled): + """An unrelated failure after a result is still an exception, including cancellation.""" + + async def query(**_kwargs): + yield ResultMessage( + subtype="error_during_execution", + duration_ms=1, + duration_api_ms=1, + is_error=True, + num_turns=1, + session_id="session-1", + errors=["native error"], + ) + if cancelled: + raise asyncio.CancelledError() + raise RuntimeError("transport disconnected") + + monkeypatch.setattr("claude_agent_sdk.query", query) + wrapper = CcAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + with pytest.raises(asyncio.CancelledError if cancelled else RuntimeError, match=None if cancelled else "transport"): + await wrapper.reply("hello") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted"]) +async def test_reply_preserves_native_outcomes_items_and_usage(system, monkeypatch, status): + """Failed and interrupted turns retain completed items and reported usage.""" + from openai_codex.generated.v2_all import CommandExecutionThreadItem, TurnError, TurnStatus + + first, _, _, _ = system + original = _Turn.stream + command = CommandExecutionThreadItem( + type="commandExecution", + id="command-1", + command="pwd", + command_actions=[], + source="agent", + cwd="/tmp", + status="completed", + aggregated_output="/tmp", + exit_code=0, + ) + + async def stream(turn): + yield _usage_event((10, 2), (10, 2), turn_id=turn.id) + yield SimpleNamespace( + method="item/completed", + payload=ItemCompletedNotification(threadId=turn.thread_id, turnId=turn.id, item=command, completedAtMs=0), + ) + async for event in original(turn): + if event.method == "turn/completed": + event.payload.turn.status = TurnStatus(status) + if status == "failed": + event.payload.turn.error = TurnError(message="native failure", additional_details="details") + yield event + + monkeypatch.setattr(_Turn, "stream", stream) + result = await first.reply("answer") + + assert set(result) == {"session_id", "result", "structured_output", "usage", "status", "error", "raw", "events"} + assert result["status"] == status + assert result["result"] == "answer" + assert result["usage"] == {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12} + assert result["error"] == ("native failure" if status == "failed" else None) + assert result["raw"]["status"] == status + assert result["raw"]["items"][0]["aggregatedOutput"] == "/tmp" + assert result["raw"]["usage"]["total"]["inputTokens"] == 10 + if status == "failed": + assert result["raw"]["error"]["additionalDetails"] == "details" + await _assert_roundtrip(result) + + +@pytest.mark.asyncio +async def test_reply_selects_final_answer_and_preserves_other_messages(system, monkeypatch): + """Codex message phases select final text without discarding other messages.""" + first, _, _, _ = system + original = _Turn.stream + + async def stream(turn): + for index, phase in enumerate((None, "final_answer", "commentary")): + yield SimpleNamespace( + method="item/completed", + payload=ItemCompletedNotification( + threadId=turn.thread_id, + turnId=turn.id, + completedAtMs=0, + item=AgentMessageThreadItem(type="agentMessage", id=f"item-{index}", text=str(phase), phase=phase), + ), + ) + async for event in original(turn): + yield event + + monkeypatch.setattr(_Turn, "stream", stream) + result = await first.reply("legacy fallback") + + assert result["result"] == "final_answer" + assert len(result["raw"]["items"]) == 4 + + +@pytest.mark.asyncio +async def test_invalid_structured_reply_retains_native_result(system): + """A JSON parsing failure preserves the provider's successful native turn.""" + first, _, _, _ = system + result = await first.reply("invalid JSON", output_schema={"type": "object"}) + + assert result["status"] == "failed" + assert "invalid JSON" in result["error"] + assert result["result"] == "invalid JSON" + assert result["raw"]["status"] == "completed" + assert result["raw"]["final_response"] == "invalid JSON" + + +async def _assert_roundtrip(reply): + """JSON storage, replay and collection preserve every event and outcome field.""" + from reme.schema import collect_reply, replay_reply + + stored = json.loads(json.dumps(reply)) + replayed = [chunk async for chunk in replay_reply(stored)] + assert [chunk.model_dump(mode="json") for chunk in replayed] == stored["events"] + + async def stream(): + for chunk in replayed: + yield chunk + + assert await collect_reply(stream()) == stored + replayed[0].metadata["mutated"] = True + assert "mutated" not in stored["events"][0]["metadata"] + + +@pytest.mark.asyncio +async def test_claude_preserves_unprojected_native_messages_and_final_outcome(tmp_path, monkeypatch): + """Unknown blocks, signatures and tool results stay available beyond the display projection.""" + from claude_agent_sdk import AssistantMessage, SystemMessage, ThinkingBlock, ToolResultBlock, UserMessage + + native = [ + SystemMessage(subtype="future_event", data={"future": [1, False, None]}), + AssistantMessage(content=[ThinkingBlock(thinking="check", signature="signature")], model="claude"), + UserMessage(content=[ToolResultBlock(tool_use_id="hidden", content=[{"type": "text", "text": "result"}])]), + ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session", + result="final", + structured_output=[], + ), + ] + calls = [] + + async def query(**kwargs): + calls.append(kwargs) + for message in native: + yield message + + monkeypatch.setattr("claude_agent_sdk.query", query) + wrapper = CcAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + reply = await wrapper.reply("hello", output_schema={}) + assert len(calls) == 1 + assert [event["metadata"]["native"] for event in reply["events"] if "native" in event["metadata"]] == [ + wrapper._serialize(message) for message in native # pylint: disable=protected-access + ] + assert reply["structured_output"] == [] + await _assert_roundtrip(reply) + + +@pytest.mark.asyncio +async def test_collector_preserves_empty_snapshots_unknown_fields_and_transport_marker(): + """The trace keeps block operations and future extension fields exactly, including DONE.""" + from reme.enumeration import ChunkEnum + from reme.schema import StreamChunk, collect_reply + + events = [ + StreamChunk(chunk="draft", block_id="text", metadata={"native": {"signature": "signed"}}), + StreamChunk(chunk="", event="end", block_id="text", future_field={"x": [False, None]}), + StreamChunk(chunk=None, event="end", block_id="text"), + StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status="waiting", + chunk={ + "session_id": "session", + "result": "", + "structured_output": False, + "usage": None, + "status": "waiting", + "error": None, + "raw": {"future": 1}, + }, + ), + StreamChunk(chunk_type=ChunkEnum.DONE, done=True), + ] + + async def stream(): + for event in events: + yield event + + reply = await collect_reply(stream()) + assert reply["events"] == [event.model_dump(mode="json") for event in events] + await _assert_roundtrip(reply) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("duplicate", [False, True]) +async def test_collector_rejects_missing_or_duplicate_outcomes(duplicate): + """Transport termination cannot manufacture a successful reply.""" + from reme.enumeration import ChunkEnum + from reme.schema import StreamChunk, collect_reply + + async def stream(): + if duplicate: + for _ in range(2): + yield StreamChunk(chunk_type=ChunkEnum.REPLY_END, chunk={"result": "answer"}) + yield StreamChunk(chunk_type=ChunkEnum.DONE, done=True) + + with pytest.raises(ValueError, match="reply_end"): + await collect_reply(stream()) + + +@pytest.mark.asyncio +async def test_codex_terminal_snapshot_supplies_items_missing_from_notifications(system, monkeypatch): + """The final native turn can complete or correct the item notifications.""" + first, _, _, _ = system + original = _Turn.stream + + async def stream(turn): + async for event in original(turn): + if event.method == "item/completed": + continue + if event.method == "turn/completed": + event.payload.turn.items[0].root.text = "authoritative final answer" + yield event + + monkeypatch.setattr(_Turn, "stream", stream) + reply = await first.reply("draft") + assert reply["result"] == "authoritative final answer" + assert reply["raw"]["items"][0]["text"] == "authoritative final answer" + await _assert_roundtrip(reply) diff --git a/tests/unit/test_agent_reply_consumers.py b/tests/unit/test_agent_reply_consumers.py new file mode 100644 index 00000000..ccb47b0f --- /dev/null +++ b/tests/unit/test_agent_reply_consumers.py @@ -0,0 +1,163 @@ +"""Plugin answer steps consume the complete agent outcome in both return modes.""" + +# pylint: disable=protected-access + +import asyncio +from importlib import import_module +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from reme.components.agent_wrapper import BaseAgentWrapper +from reme.components.application_context import ApplicationContext +from reme.enumeration import ChunkEnum +from reme.schema import StreamChunk +from reme.steps.index._dedup import _ToolContextDedupMixin + + +class _Agent(BaseAgentWrapper): + """Emit a text delta followed by its authoritative final snapshot.""" + + def __init__(self, status, **kwargs): + super().__init__(**kwargs) + self.result = self._reply_result( + session_id="session", + result="answer", + raw={"signature": "signed"}, + status=status, + error="provider failed" if status == "failed" else None, + ) + self.chunks = [ + StreamChunk(chunk="answer", block_id="text", turn_id="turn"), + StreamChunk(chunk="answer", event="end", block_id="text", turn_id="turn"), + StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status=status, + session_id="session", + turn_id="turn", + chunk={key: value for key, value in self.result.items() if key != "events"}, + ), + ] + + def _add_tool_context(self, kwargs): + if "tool_context_id" not in kwargs: + return + self.app_context.metadata.setdefault(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {})[ + kwargs["tool_context_id"] + ] = {"query": "question"} + + async def reply(self, inputs, **kwargs): + """Return the same outcome as the stream.""" + self._add_tool_context(kwargs) + return self.result + + async def reply_stream(self, inputs, **kwargs): + """Yield original chunks, including IDs and snapshot operations.""" + self._add_tool_context(kwargs) + for chunk in self.chunks: + yield chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("plugin", "class_name"), [("lme", "LmeAgenticAnswerStep"), ("beam", "BeamAgenticAnswerStep")]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_answer_plugins_preserve_outcome_and_stream_chunks( + tmp_path, + monkeypatch, + plugin, + class_name, + streaming, + status, +): + """Final text is not doubled, incomplete replies fail, and metadata survives forwarding.""" + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins" / plugin / "src")) + step_type = getattr(import_module(f"reme_{plugin}.agentic_answer"), class_name) + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = _Agent(status, app_context=context) + queue = asyncio.Queue() if streaming else None + response = await step_type(app_context=context, agent_wrapper=wrapper)(query="question", stream_queue=queue) + + assert response.answer == ("provider failed" if status == "failed" else "answer") + assert response.success is (status == "completed") + assert response.metadata["status"] == status + assert response.metadata["raw"] == wrapper.result["raw"] + assert not context.metadata[_ToolContextDedupMixin.TOOL_CONTEXTS_KEY] + if streaming: + assert [queue.get_nowait() for _ in wrapper.chunks] == wrapper.chunks + assert response.metadata["events"] == [chunk.model_dump(mode="json") for chunk in wrapper.chunks] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("plugin", ["lme", "beam"]) +async def test_stream_failure_releases_temporary_query_state(tmp_path, monkeypatch, plugin): + """A failed collector/transport leaves no query dedup state alive.""" + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins" / plugin / "src")) + step_type = getattr(import_module(f"reme_{plugin}.agentic_answer"), f"{plugin.title()}AgenticAnswerStep") + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = _Agent("completed", app_context=context) + + async def stream(_inputs, **kwargs): + wrapper._add_tool_context(kwargs) + yield wrapper.chunks[0] + raise RuntimeError("transport disconnected") + + monkeypatch.setattr(wrapper, "reply_stream", stream) + queue = asyncio.Queue() + with pytest.raises(RuntimeError, match="transport disconnected"): + await step_type(app_context=context, agent_wrapper=wrapper)(query="question", stream_queue=queue) + assert not context.metadata[_ToolContextDedupMixin.TOOL_CONTEXTS_KEY] + assert queue.get_nowait() == wrapper.chunks[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("plugin", ["lme", "beam"]) +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_judges_do_not_score_incomplete_agent_replies(tmp_path, monkeypatch, plugin, status): + """A parseable partial judgement cannot become a valid benchmark score.""" + if plugin == "beam": + pytest.importorskip("scipy") + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins" / f"{plugin}-judge" / "src")) + module = import_module(f"judge_{plugin}.llm_judge") + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = _Agent(status, app_context=context) + wrapper.reply = AsyncMock(return_value=wrapper.result) + if plugin == "lme": + step = module.LmeAnswerJudgeStep(app_context=context, agent_wrapper=wrapper) + wrapper.result["result"] = "YES" + args = {"query": "question", "agent_answer": "answer", "golden_answer": "answer"} + else: + step = module.BeamRubricJudgeStep(app_context=context, agent_wrapper=wrapper) + wrapper.result["result"] = '{"score": 1.0, "reason": "correct"}' + args = {"llm_response": "answer", "rubric": ["correct answer"]} + if status == "completed": + assert (await step(**args)).success + else: + with pytest.raises(RuntimeError, match="provider failed" if status == "failed" else status): + await step(**args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("plugin", ["auto-fin", "daily_paper"]) +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +async def test_structured_consumers_reject_incomplete_agent_outputs(tmp_path, monkeypatch, plugin, status): + """A valid-looking structured payload still requires a completed native invocation.""" + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins" / plugin / "src")) + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = _Agent(status, app_context=context) + wrapper.reply = AsyncMock(return_value=wrapper.result) + wrapper.result["structured_output"] = {"title": "report", "description": "summary", "body": "body"} + error = "provider failed" if status == "failed" else status + if plugin == "auto-fin": + module = import_module("reme_auto_fin.research") + schema = import_module("reme_auto_fin.schema").AutoFinReportOutput + step = module.AutoFinResearchStep(app_context=context, agent_wrapper=wrapper, prompt_dict={"test": "question"}) + with pytest.raises(RuntimeError, match=error): + await step._reply("test", schema) + else: + module = import_module("reme_daily_paper.base") + schema = import_module("reme_daily_paper").DailyPaperMarkdownOutput + with pytest.raises(RuntimeError, match=error): + module.structured_output(wrapper.result, schema) diff --git a/tests/unit/test_agent_stream.py b/tests/unit/test_agent_stream.py new file mode 100644 index 00000000..94f53143 --- /dev/null +++ b/tests/unit/test_agent_stream.py @@ -0,0 +1,270 @@ +"""Regression coverage for information preserved across agent stream adapters.""" + +# pylint: disable=missing-function-docstring,protected-access + +import asyncio +from types import SimpleNamespace + +import pytest +from agentscope.event import DataBlockDeltaEvent, ReplyEndEvent, TextBlockEndEvent, TextBlockStartEvent +from claude_agent_sdk import AssistantMessage, ResultMessage, StreamEvent, TextBlock +from openai_codex.generated.v2_all import ( + AgentMessageThreadItem, + ItemCompletedNotification, + PlanThreadItem, + ReasoningThreadItem, + WebSearchThreadItem, +) + +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.agent_wrapper.cc_agent_wrapper import CcAgentWrapper +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper +from reme.components.agent_wrapper.codex_session_history import item_chunks +from reme.components.application_context import ApplicationContext +from reme.components.client.http_client import HttpClient +from reme.components.runtime_context import RuntimeContext +from reme.enumeration import ChunkEnum +from reme.schema import StreamChunk +from reme.steps.common.chat import ChatStep +from reme.utils.common_utils import call_action + + +@pytest.mark.parametrize( + "item", + [ + AgentMessageThreadItem(type="agentMessage", id="answer", text="corrected", phase="final_answer"), + PlanThreadItem(type="plan", id="plan", text="updated plan"), + ReasoningThreadItem(type="reasoning", id="thought", summary=["one", "two"], content=["raw"]), + WebSearchThreadItem(type="webSearch", id="search", query="ReMe"), + ], +) +def test_codex_completed_items_match_history_and_retain_final_snapshots(item): + payload = ItemCompletedNotification(threadId="session", turnId="turn", item=item, completedAtMs=0) + chunks = CodexAgentWrapper._event_to_chunks(SimpleNamespace(method="item/completed", payload=payload), "session") + assert chunks == item_chunks(item, session_id="session", turn_id="turn") + assert chunks and all(chunk.event == "end" and chunk.turn_id == "turn" for chunk in chunks) + if item.type == "agentMessage": + assert chunks[0].chunk == "corrected" + assert chunks[0].metadata["phase"] == "final_answer" + elif item.type == "webSearch": + assert chunks[0].tool_call_id == "search" + assert chunks[0].status == "completed" + assert chunks[0].chunk["query"] == "ReMe" + elif item.type == "reasoning": + assert [chunk.block_id for chunk in chunks] == ["thought:summary:0", "thought:summary:1", "thought:content:0"] + assert [chunk.chunk for chunk in chunks] == ["one", "two", "raw"] + + +def test_codex_reasoning_deltas_correlate_with_their_final_sections(): + chunks = [] + for method, index in [("item/reasoning/summaryTextDelta", 1), ("item/reasoning/textDelta", 0)]: + payload = SimpleNamespace( + item_id="thought", + turn_id="turn", + delta="text", + summary_index=index, + content_index=index, + ) + chunks.extend(CodexAgentWrapper._event_to_chunks(SimpleNamespace(method=method, payload=payload), "session")) + assert [chunk.block_id for chunk in chunks] == ["thought:summary:1", "thought:content:0"] + assert chunks[0].metadata["kind"] != chunks[1].metadata["kind"] + + +def test_agentscope_block_boundaries_final_correction_and_url_data(): + start = AsAgentWrapper._event_to_chunk(TextBlockStartEvent(reply_id="reply", block_id="block")) + end = AsAgentWrapper._event_to_chunk(TextBlockEndEvent(reply_id="reply", block_id="block", text="corrected")) + url = AsAgentWrapper._event_to_chunk( + DataBlockDeltaEvent( + reply_id="reply", + block_id="image", + media_type="image/png", + url="https://example.com/a.png", + ), + ) + assert start.event == "start" and start.chunk is None + assert end.event == "end" and end.chunk == "corrected" + assert end.turn_id == "reply" + assert url.chunk == {"url": "https://example.com/a.png"} + assert url.media_type == "image/png" + interrupted = AsAgentWrapper._event_to_chunk( + ReplyEndEvent(session_id="session", reply_id="reply", finished_reason="interrupted"), + ) + assert interrupted.status == "interrupted" + assert interrupted.metadata["finished_reason"] == "interrupted" + + +def test_claude_indices_are_scoped_by_message_and_signatures_are_retained(): + states = {} + chunks = [ + CcAgentWrapper._raw_event_to_chunk( + {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}, + message_id="message", + block_states=states, + ) + for index in [0, 1] + ] + assert [chunk.block_id for chunk in chunks] == ["message:0", "message:1"] + signature = CcAgentWrapper._raw_event_to_chunk( + {"type": "content_block_delta", "index": 1, "delta": {"type": "signature_delta", "signature": "signed"}}, + message_id="message", + block_states=states, + ) + assert signature.chunk["delta"]["signature"] == "signed" + assert signature.chunk_type == ChunkEnum.DATA + unknown = CcAgentWrapper._raw_event_to_chunk( + {"type": "content_block_start", "index": 2, "content_block": {"type": "future_type", "value": 1}}, + message_id="message", + block_states=states, + ) + assert unknown.chunk["content_block"]["value"] == 1 + + +@pytest.mark.asyncio +async def test_claude_multiple_model_messages_have_one_invocation_end(tmp_path, monkeypatch): + async def query(**_kwargs): + for message_id, text in [("first", "checking"), ("second", "answer")]: + for event in [ + {"type": "message_start", "message": {"id": message_id}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_stop"}, + ]: + yield StreamEvent(uuid=message_id, session_id="session", event=event) + yield AssistantMessage(content=[TextBlock(text=text + "!")], model="claude", message_id=message_id) + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=2, + session_id="session", + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + wrapper = CcAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + chunks = [chunk async for chunk in wrapper.reply_stream("hello")] + assert sum(chunk.chunk_type == ChunkEnum.REPLY_START for chunk in chunks) == 1 + assert sum(chunk.chunk_type == ChunkEnum.REPLY_END for chunk in chunks) == 1 + assert chunks[-1].chunk_type == ChunkEnum.REPLY_END + assert chunks[-1].status == "completed" + assert len({chunk.turn_id for chunk in chunks}) == 1 + snapshots = [ + chunk + for chunk in chunks + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.event == "end" and chunk.chunk is not None + ] + assert [(chunk.block_id, chunk.chunk) for chunk in snapshots] == [("first:0", "checking!"), ("second:0", "answer!")] + + +@pytest.mark.asyncio +async def test_claude_interleaved_child_events_keep_their_own_block_state(tmp_path, monkeypatch): + async def query(**_kwargs): + for parent, event in [ + (None, {"type": "message_start", "message": {"id": "parent"}}), + (None, {"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}), + ("tool", {"type": "message_start", "message": {"id": "child"}}), + ("tool", {"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}), + (None, {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "parent"}}), + ("tool", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "child"}}), + ]: + yield StreamEvent(uuid="event", session_id="session", parent_tool_use_id=parent, event=event) + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session", + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + wrapper = CcAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + chunks = [chunk async for chunk in wrapper.reply_stream("hello")] + deltas = [chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT and chunk.event == "delta"] + assert [(chunk.block_id, chunk.chunk) for chunk in deltas] == [("parent:0", "parent"), ("child:0", "child")] + assert deltas[1].metadata["parent_tool_use_id"] == "tool" + + +@pytest.mark.asyncio +async def test_chat_aggregates_snapshots_without_duplicate_text_and_records_terminal_failure(tmp_path, monkeypatch): + async def reply_stream(_self, _inputs, **_kwargs): + for event in [ + StreamChunk(block_id="answer", chunk="draft"), + StreamChunk(block_id="answer", event="end", chunk="corrected"), + StreamChunk(block_id="answer", event="end", chunk="corrected"), + StreamChunk(block_id="removed", chunk="discard this"), + StreamChunk(block_id="removed", event="end", chunk=""), + StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status="failed", + metadata={"error": "upstream failed"}, + chunk={ + key: value + for key, value in _self._reply_result( + session_id="session", + result="corrected", + raw={}, + status="failed", + error="upstream failed", + ).items() + if key != "events" + }, + ), + ]: + yield event + + monkeypatch.setattr(AsAgentWrapper, "reply_stream", reply_stream) + context = ApplicationContext(workspace_dir=str(tmp_path)) + agent = AsAgentWrapper(as_llm="", app_context=context) + runtime = RuntimeContext(query="hello", stream_queue=asyncio.Queue()) + response = await ChatStep(agent_wrapper=agent, app_context=context)(runtime) + assert response.answer == "corrected" + assert not response.success + assert runtime.stream_error == "upstream failed" + + +@pytest.mark.asyncio +async def test_http_text_consumers_reconcile_snapshots_and_skip_null_or_diagnostic_events(monkeypatch): + async def stream(_self, _action, _payload): + for chunk in [ + StreamChunk(block_id="answer", chunk="draft"), + StreamChunk(block_id="answer", event="end", chunk="corrected"), + StreamChunk(block_id="answer", event="end", chunk="corrected"), + StreamChunk(block_id="answer", event="end", chunk=None), + StreamChunk(chunk_type=ChunkEnum.DATA, chunk={"internal": True}, metadata={"visibility": "diagnostic"}), + ]: + yield chunk + + monkeypatch.setattr(HttpClient, "_iter_stream_chunks", stream) + assert await call_action("chat") == "corrected" + display = "".join([text async for text in HttpClient()._execute("chat", {})]) + assert display == "draft\ncorrected" + + +@pytest.mark.asyncio +async def test_agentscope_closing_stream_records_reported_usage_and_closes_sdk(tmp_path, monkeypatch): + from agentscope.event import ModelCallEndEvent + from unittest.mock import AsyncMock + from reme.utils import global_counter_get_all + + closed = [] + + async def native_stream(_inputs, **_kwargs): + try: + yield ModelCallEndEvent(reply_id="reply", input_tokens=3, output_tokens=2) + await asyncio.Event().wait() + finally: + closed.append(True) + + agent = SimpleNamespace(state=SimpleNamespace(session_id="session"), reply_stream=native_stream) + wrapper = AsAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path)), as_llm="") + monkeypatch.setattr(wrapper, "_build_agent", AsyncMock(return_value=(agent, "hello"))) + monkeypatch.setattr(wrapper, "_dump_state", AsyncMock()) + stream = wrapper.reply_stream("hello") + await anext(stream) + await stream.aclose() + assert closed == [True] + counters = global_counter_get_all(wrapper.app_context.metadata, ["__token_counter", wrapper.name]) + assert counters["children"]["total_tokens"]["value"] == 5 diff --git a/tests/unit/test_auto_dream.py b/tests/unit/test_auto_dream.py index f4d1a0e1..e94dde52 100644 --- a/tests/unit/test_auto_dream.py +++ b/tests/unit/test_auto_dream.py @@ -118,7 +118,7 @@ class _ReplyAgent(BaseAgentWrapper): self.on_reply() if self.error: raise self.error - return self.result + return {**self._reply_result(session_id="stub", result="", raw={}), **self.result} class _SequenceAgent(BaseAgentWrapper): @@ -138,7 +138,7 @@ class _SequenceAgent(BaseAgentWrapper): outcome = outcome() if isinstance(outcome, Exception): raise outcome - return outcome + return {**self._reply_result(session_id="stub", result="", raw={}), **outcome} def test_scan_day_files_includes_only_markdown_day_files(): @@ -541,7 +541,7 @@ async def test_integrate_failure_only_blocks_affected_sources(tmp_path, monkeypa def integrate_good(): _touch(tmp_path / target, "integrated memory") - return {"result": json.dumps({"action": "CREATE", "target_path": target})} + return {"status": "completed", "result": json.dumps({"action": "CREATE", "target_path": target})} agent = _SequenceAgent(integrate_good, {"result": "{}"}, {"result": "{}"}) app_context = ApplicationContext(workspace_dir=str(tmp_path)) @@ -718,12 +718,12 @@ async def test_integrate_emits_actual_changes_once_across_units(tmp_path, monkey def create(): _touch(tmp_path / created, "new") - return {"result": json.dumps({"action": "CREATE", "target_path": created})} + return {"status": "completed", "result": json.dumps({"action": "CREATE", "target_path": created})} def update(): _touch(tmp_path / created, "new with more evidence") _touch(tmp_path / updated, "updated existing memory") - return {"result": json.dumps({"action": "REFINE", "target_path": created})} + return {"status": "completed", "result": json.dumps({"action": "REFINE", "target_path": created})} agent = _SequenceAgent( create, @@ -818,7 +818,7 @@ async def test_dream_jobs_tag_outputs_and_preserve_checkpoint_results( async def reply(_message, **kwargs): if kwargs["job_tools"] == ["read"]: - return {"result": json.dumps({"units": [unit]})} + return {"status": "completed", "result": json.dumps({"units": [unit]})} if "list_tags" in kwargs["job_tools"]: injected = kwargs["injected_job_kwargs"] path = injected["_allowed_paths"][0] @@ -831,15 +831,15 @@ async def test_dream_jobs_tag_outputs_and_preserve_checkpoint_results( ) await FrontmatterUpdateStep(file_store=store)(update_context) assert update_context.response.success - return {"result": "Tagged ReMe"} + return {"status": "completed", "result": "Tagged ReMe"} for path in targets: previous = (tmp_path / path).read_text(encoding="utf-8") if (tmp_path / path).exists() else "" _touch(tmp_path / path, previous + "ReMe evidence\n") if integration_failure == "exception": raise RuntimeError("integration unavailable") if integration_failure == "invalid_receipt": - return {"result": "{}"} - return {"result": json.dumps({"action": "CREATE", "target_path": targets[0]})} + return {"status": "completed", "result": "{}"} + return {"status": "completed", "result": json.dumps({"action": "CREATE", "target_path": targets[0]})} agent = _ReplyAgent() monkeypatch.setattr(agent, "reply", AsyncMock(side_effect=reply)) @@ -1278,7 +1278,7 @@ class _BatchAwareAgent(BaseAgentWrapper): } for index in range(count) ) - return {"result": json.dumps({"units": units})} + return {"status": "completed", "result": json.dumps({"units": units})} @pytest.mark.asyncio diff --git a/tests/unit/test_auto_memory_cc.py b/tests/unit/test_auto_memory_cc.py index 9ae05c1c..cf5745dc 100644 --- a/tests/unit/test_auto_memory_cc.py +++ b/tests/unit/test_auto_memory_cc.py @@ -67,14 +67,14 @@ async def test_cc_job_tags_created_updated_and_renamed_notes_only(tmp_path, monk **injected, ) assert result.success - return {"result": "Tagged ReMe"} + return {"status": "completed", "result": "Tagged ReMe"} memory_calls.append(mode) if mode == "create": target = tmp_path / f"daily/{day}/note.md" target.parent.mkdir(parents=True) post = frontmatter.Post("ReMe memory", name="note", session_id=session_id, source_conversation=source) elif mode == "unchanged": - return {"result": "No new facts"} + return {"status": "completed", "result": "No new facts"} else: target = tmp_path / kwargs["injected_job_kwargs"]["_allowed_paths"][0] post = frontmatter.loads(target.read_text(encoding="utf-8")) @@ -82,7 +82,7 @@ async def test_cc_job_tags_created_updated_and_renamed_notes_only(tmp_path, monk if mode == "rename": post.metadata["name"] = "renamed" target.write_text(frontmatter.dumps(post), encoding="utf-8") - return {"result": "Recorded memory"} + return {"status": "completed", "result": "Recorded memory"} agent = AsyncMock(spec=BaseAgentWrapper) agent.reply.side_effect = reply diff --git a/tests/unit/test_auto_memory_direct_images.py b/tests/unit/test_auto_memory_direct_images.py index e53b483f..175b3845 100644 --- a/tests/unit/test_auto_memory_direct_images.py +++ b/tests/unit/test_auto_memory_direct_images.py @@ -69,7 +69,7 @@ def _saved_line(message): def memory_setup(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) wrapper = AsAgentWrapper(backend="agentscope", as_llm="") - wrapper.reply = AsyncMock(return_value={"result": "ok"}) + wrapper.reply = AsyncMock(return_value={"status": "completed", "result": "ok"}) store = LocalFileStore(embedding_store="") app = SimpleNamespace( registry=R, @@ -112,7 +112,7 @@ def test_only_include_images_is_exposed_and_disabled_by_default(): async def test_text_path_keeps_main_input_kwargs_metadata_and_jsonl(setup, monkeypatch, options, images): step, _, path = setup wrapper = CcAgentWrapper(backend="claude_code") - wrapper.reply = AsyncMock(return_value={"result": "ok"}) + wrapper.reply = AsyncMock(return_value={"status": "completed", "result": "ok"}) wrapper.kwargs["context_config"] = {"max_image_num": 0} step.kwargs["agent_wrapper"] = wrapper message = _message(images=images) @@ -380,7 +380,7 @@ async def test_configured_backend_and_job_switch(tmp_path, monkeypatch, backend, wrapper = app.context.components["agent_wrapper"]["default"] assert isinstance(wrapper, AsAgentWrapper if backend == "agentscope" else CcAgentWrapper) assert wrapper.backend == backend - wrapper.reply = AsyncMock(return_value={"result": "ok"}) + wrapper.reply = AsyncMock(return_value={"status": "completed", "result": "ok"}) monkeypatch.setattr(AutoMemoryStep, "_list_session_note", AsyncMock(return_value=None)) job = app.context.jobs["auto_memory"] await job.start() @@ -439,7 +439,7 @@ async def test_default_job_still_passes_memory_changes_to_auto_tag(setup, monkey async def write_memory(*_args, **_kwargs): _write_note(target) - return {"result": "Memory written."} + return {"status": "completed", "result": "Memory written."} monkeypatch.setattr(AutoMemoryStep, "_list_session_note", find_note) wrapper.reply.side_effect = write_memory @@ -502,7 +502,7 @@ async def test_image_links_survive_native_full_rewrite_and_keep_context(setup, m # Native write replaces frontmatter as well as the body. Auto Memory # must restore previous image sources without rewriting this body. await writer(path=note_path, name="image-memory", description="Images", content=body) - return {"result": "Updated."} + return {"status": "completed", "result": "Updated."} wrapper.reply.side_effect = full_rewrite await _run(step, [first, second], include_images=True) diff --git a/tests/unit/test_auto_resource_review_regressions.py b/tests/unit/test_auto_resource_review_regressions.py index 481f4faa..6c861185 100644 --- a/tests/unit/test_auto_resource_review_regressions.py +++ b/tests/unit/test_auto_resource_review_regressions.py @@ -384,7 +384,7 @@ async def test_loose_root_text_updates_and_deletes_original_daily_card(routed, a metadata={"source_resource": "[[resource/report.txt]]"}, ) assert response.success is True - return {"result": "Updated original report."} + return {"status": "completed", "result": "Updated original report."} env.app_context.registry = R step_cls = AutoResourceStep if routed else AutoTextResourceStep diff --git a/tests/unit/test_auto_tag.py b/tests/unit/test_auto_tag.py index 69d6a6f1..b3ce0757 100644 --- a/tests/unit/test_auto_tag.py +++ b/tests/unit/test_auto_tag.py @@ -42,7 +42,7 @@ class _TaggingWrapper(BaseAgentWrapper): post = frontmatter.loads(target.read_text(encoding="utf-8")) post.metadata[self.tag_key] = self.tags target.write_text(frontmatter.dumps(post), encoding="utf-8") - return {"result": f"tagged {path}", "last_message": {}} + return {"result": f"tagged {path}", "status": "completed"} def _write_note(path: Path) -> None: @@ -63,7 +63,7 @@ async def test_auto_tag_reports_missing_agent_write_without_inventing_empty_tags _write_note(note) original = note.read_bytes() wrapper = _TaggingWrapper(tmp_path) - monkeypatch.setattr(wrapper, "reply", AsyncMock(return_value={"result": "", "last_message": {}})) + monkeypatch.setattr(wrapper, "reply", AsyncMock(return_value={"result": "", "status": "completed"})) step = AutoTagStep(agent_wrapper=wrapper, file_store=_file_store()) response = await step(RuntimeContext(changes=[{"change": "added", "path": "memory/note.md"}])) diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 1be377ba..a91cf58a 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -82,7 +82,7 @@ class _FakeAgentWrapper(BaseAgentWrapper): self.kwargs = kwargs if self.on_reply is not None: self.on_reply(inputs, kwargs) - return {"result": "ok"} + return self._reply_result(session_id="stub", result="ok", raw={}) class _StepJob: diff --git a/tests/unit/test_base_agent_wrapper.py b/tests/unit/test_base_agent_wrapper.py index 13235156..e806ff9b 100644 --- a/tests/unit/test_base_agent_wrapper.py +++ b/tests/unit/test_base_agent_wrapper.py @@ -212,22 +212,9 @@ async def test_agentscope_structured_output_uses_model_tool_choice_policy(tmp_pa async def generate_structured_output(self, **kwargs): """Return a successful structured response.""" observed.update(kwargs) - return SimpleNamespace(content={"ok": True}, usage=None) + from agentscope.model import StructuredResponse - class FakeMessage: - """Minimal AgentScope reply message.""" - - usage = None - - @staticmethod - def model_dump(): - """Return serialized message content.""" - return {"content": "answer"} - - @staticmethod - def get_text_content(): - """Return plain text content.""" - return "answer" + return StructuredResponse(content={"ok": True}) class FakeAgent: """Minimal AgentScope agent.""" @@ -235,14 +222,13 @@ async def test_agentscope_structured_output_uses_model_tool_choice_policy(tmp_pa state = SimpleNamespace(session_id="session-1", context=["context"]) @staticmethod - async def observe(_inputs): - """Accept an input message.""" - return None + async def reply_stream(_inputs, **_kwargs): + """Stream real SDK termination and message shapes.""" + from agentscope.event import ReplyEndEvent + from agentscope.message import AssistantMsg - @staticmethod - async def reply(): - """Return the fake reply message.""" - return FakeMessage() + yield ReplyEndEvent(session_id="session-1", reply_id="reply", finished_reason="completed") + yield AssistantMsg(name="agent", content="answer", finished_reason="completed") wrapper = AsAgentWrapper(app_context=ApplicationContext(workspace_dir=str(tmp_path)), as_llm="") wrapper.as_llm = SimpleNamespace(model=FakeModel()) diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index c123faa7..e65071d2 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -351,7 +351,7 @@ async def test_compact_session_uses_claude_command(tmp_path, monkeypatch): async def reply(inputs, **kwargs): calls.append((inputs, kwargs)) - return {"last_message": {"is_error": False}} + return {"status": "completed", "error": None} monkeypatch.setattr(wrapper, "reply", reply) @@ -483,7 +483,9 @@ async def test_reply_stream_emits_one_reply_end_for_normal_sdk_lifecycle(tmp_pat delta_usage = next(chunk for chunk in chunks if chunk.metadata.get("stop_reason") == "end_turn") assert delta_usage.chunk_type == ChunkEnum.USAGE assert delta_usage.output_tokens == 2 - assert not recorded_usages + assert [usage.model_dump() for usage in recorded_usages] == [ + {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3}, + ] @pytest.mark.asyncio @@ -597,17 +599,23 @@ async def test_reply_stream_reports_session_mirror_errors(tmp_path, monkeypatch) @pytest.mark.asyncio -@pytest.mark.parametrize( - "wrapper_factory", - [ - _wrapper, - lambda tmp_path: AsAgentWrapper(as_llm="", app_context=ApplicationContext(workspace_dir=str(tmp_path))), - ], -) @pytest.mark.parametrize("schema", [{}, {"type": "object"}]) -async def test_reply_stream_rejects_output_schema(tmp_path, wrapper_factory, schema): - """Streaming wrappers reject structured-output schemas consistently.""" - wrapper = wrapper_factory(tmp_path) +async def test_reply_stream_supports_output_schema(tmp_path, monkeypatch, schema): + """Structured output uses the same execution path for both return modes.""" + from claude_agent_sdk import ResultMessage - with pytest.raises(NotImplementedError, match="Structured output is not supported"): - await anext(wrapper.reply_stream("hello", output_schema=schema)) + async def query(**kwargs): + assert kwargs["options"].output_format == {"type": "json_schema", "schema": schema} + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session", + structured_output={}, + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello", output_schema=schema)] + assert chunks[-1].chunk["structured_output"] == {} diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index 3dd9ac22..7e08a259 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -1,6 +1,7 @@ """Shared Codex runtime, isolated conversations, and same-Application MCP tools.""" # pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,redefined-outer-name +# pylint: disable=too-many-lines import asyncio import json @@ -17,6 +18,8 @@ from fastmcp.client import StreamableHttpTransport from fastmcp.exceptions import ToolError from openai_codex.generated.v2_all import ( AgentMessageThreadItem, + AgentMessageDeltaNotification, + TurnStartedNotification, ItemCompletedNotification, ThreadTokenUsage, ThreadTokenUsageUpdatedNotification, @@ -70,14 +73,25 @@ class _Turn: async def stream(self): try: - yield SimpleNamespace(method="turn/started", payload=SimpleNamespace(turn=SimpleNamespace(id=self.id))) + yield SimpleNamespace( + method="turn/started", + payload=TurnStartedNotification( + threadId=self.thread_id, + turn=Turn(id=self.id, items=[], status="inProgress", error=None), + ), + ) 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), + payload=AgentMessageDeltaNotification( + threadId=self.thread_id, + turnId=self.id, + itemId="item", + delta=str(self.inputs), + ), ) item = AgentMessageThreadItem(type="agentMessage", id="item", text=str(self.inputs)) yield SimpleNamespace( @@ -808,8 +822,8 @@ async def test_structured_reply_and_native_inputs(system): inputs = [TextInput("describe"), LocalImageInput("image.png")] await first.reply(inputs) assert _Codex.instances[0].threads["thread-2"].inputs[0] is inputs - with pytest.raises(NotImplementedError): - await anext(first.reply_stream("no", output_schema=_StructuredModel)) + chunks = [chunk async for chunk in first.reply_stream('{"ok": false}', output_schema=_StructuredModel)] + assert chunks[-1].chunk["structured_output"] == {"ok": False} @pytest.mark.asyncio diff --git a/tests/unit/test_common_steps.py b/tests/unit/test_common_steps.py index 9738a36f..4f51131e 100644 --- a/tests/unit/test_common_steps.py +++ b/tests/unit/test_common_steps.py @@ -6,6 +6,7 @@ from contextlib import chdir as _temp_chdir import asyncio import tempfile import warnings +from unittest.mock import AsyncMock import pytest @@ -13,7 +14,7 @@ from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext from reme.components.file_store import LocalFileStore from reme.enumeration import ChunkEnum -from reme.schema import FileFrontMatter, FileLink, FileNode, StreamChunk, TraverseGraph +from reme.schema import FileFrontMatter, FileLink, FileNode, StreamChunk, TokenUsage, TraverseGraph from reme.steps.common.add import AddStep from reme.steps.common.health_check import _file_graph_status from reme.steps.common.llm_demo import LLMDemoStep @@ -154,7 +155,7 @@ class _FakeAgentWrapper(BaseAgentWrapper): async def reply(self, inputs, **kwargs) -> dict: self.last_kwargs = kwargs - return {"result": "ok"} + return self._reply_result(session_id="session-1", result="ok", raw={}) async def reply_stream(self, inputs, **kwargs): self.last_kwargs = kwargs @@ -174,6 +175,29 @@ async def test_llm_demo_always_registers_add_tool(): assert "job_tools" not in resp.metadata +@pytest.mark.asyncio +@pytest.mark.parametrize("step_class", [LLMDemoStep, StreamLLMDemoStep]) +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_nonstream_demo_forwards_outcome_and_native_details(step_class, status): + """HTTP/tool results retain native details and never report incomplete calls as successes.""" + wrapper = _FakeAgentWrapper() + reply = wrapper._reply_result( + session_id="session-1", + result="partial or final", + status=status, + error="provider failure" if status == "failed" else None, + structured_output=False, + usage=TokenUsage(input_tokens=10, output_tokens=2), + raw={"native": "preserved"}, + ) + wrapper.reply = AsyncMock(return_value=reply) + response = await step_class()(query="hello", agent_wrapper=wrapper) + + assert response.success is (status == "completed") + assert response.answer == ("provider failure" if status == "failed" else "partial or final") + assert all(response.metadata[key] == value for key, value in reply.items()) + + @pytest.mark.asyncio async def test_stream_llm_demo_preserves_backend_chunk_metadata(): """Forward retry flags, tool identities, usage, and completion status without rebuilding chunks.""" @@ -183,7 +207,16 @@ async def test_stream_llm_demo_preserves_backend_chunk_metadata(): StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, tool_call_id="call-1", tool_call_name="add", chunk={"a": 1}), StreamChunk(chunk_type=ChunkEnum.CONTENT, session_id="session-1", block_id="text-1", chunk="hello"), StreamChunk(chunk_type=ChunkEnum.USAGE, input_tokens=3, output_tokens=1), - StreamChunk(chunk_type=ChunkEnum.REPLY_END, metadata={"status": "completed"}), + StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status="completed", + chunk={ + key: value + for key, value in wrapper._reply_result(session_id="session-1", result="hello", raw={}).items() + if key != "events" + }, + ), ] queue = asyncio.Queue() response = await StreamLLMDemoStep()(query="hello", stream_queue=queue, agent_wrapper=wrapper) diff --git a/tests/unit/test_dream_checkpoint_versions.py b/tests/unit/test_dream_checkpoint_versions.py index ec97e112..e735a5ad 100644 --- a/tests/unit/test_dream_checkpoint_versions.py +++ b/tests/unit/test_dream_checkpoint_versions.py @@ -55,7 +55,7 @@ async def test_dream_reprocesses_concurrent_source_edits_after_restart(tmp_path, def agent_for(app, blocked=False): async def reply(message, **kwargs): if "list_tags" in kwargs["job_tools"]: - return {"result": "No tags added by the controlled provider"} + return {"status": "completed", "result": "No tags added by the controlled provider"} phase = "extract" if kwargs["job_tools"] == ["read"] else "integrate" prompt = str(message) calls.append((phase, prompt)) @@ -66,11 +66,11 @@ async def test_dream_reprocesses_concurrent_source_edits_after_restart(tmp_path, await release.wait() if phase == "extract": unit = {"name": "deployment", "bucket": "wiki", "summary": fact, "paths": [SOURCE]} - return {"result": json.dumps({"units": [unit]})} + return {"status": "completed", "result": json.dumps({"units": [unit]})} action = "CORRECT" if (tmp_path / TARGET).exists() else "CREATE" response = await app.run_job("write", path=TARGET, content=fact) assert response.success, response.answer - return {"result": json.dumps({"action": action, "target_path": TARGET})} + return {"status": "completed", "result": json.dumps({"action": action, "target_path": TARGET})} agent = AsyncMock(spec=BaseAgentWrapper) agent.reply.side_effect = reply @@ -158,16 +158,16 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older async def reply(_message, **kwargs): if "list_tags" in kwargs["job_tools"]: - return {"result": "No tags added by the controlled provider"} + return {"status": "completed", "result": "No tags added by the controlled provider"} if kwargs["job_tools"] == ["read"]: calls["extract"] += 1 unit = {"name": "deployment", "bucket": "wiki", "summary": OLD, "paths": [SOURCE, index]} - return {"result": json.dumps({"units": [unit]})} + return {"status": "completed", "result": json.dumps({"units": [unit]})} calls["integrate"] += 1 if calls["integrate"] == write_attempt: written = await app.run_job("write", path=TARGET, content=OLD) assert written.success, written.answer - return {"result": "{}"} + return {"status": "completed", "result": "{}"} agent = AsyncMock(spec=BaseAgentWrapper) agent.reply.side_effect = reply diff --git a/tests/unit/test_evolve_utils.py b/tests/unit/test_evolve_utils.py index 85d5ab69..f1d38dcf 100644 --- a/tests/unit/test_evolve_utils.py +++ b/tests/unit/test_evolve_utils.py @@ -12,11 +12,12 @@ from reme.steps.evolve._evolve import agent_reply_result_text, format_history from reme.steps.evolve.auto_memory import AutoMemoryStep, _sanitize_msg_for_save -def test_agent_reply_result_text_uses_last_text_block(): - """Agent reply text should be the last text block in the final message.""" +def test_agent_reply_result_text_uses_portable_result(): + """Callers use final text selected by the wrapper, without inspecting native data.""" reply = { - "result": "intermediate text\nfinal text", - "last_message": { + "result": "final text", + "status": "completed", + "raw": { "content": [ {"type": "thinking", "thinking": "hidden"}, {"type": "text", "text": "intermediate text"}, @@ -29,13 +30,26 @@ def test_agent_reply_result_text_uses_last_text_block(): assert agent_reply_result_text(reply) == "final text" -def test_agent_reply_result_text_falls_back_to_result(): - """Agent reply text falls back to result when no text block exists.""" - reply = {"result": " fallback text \n", "last_message": {"content": [{"type": "thinking"}]}} +def test_agent_reply_result_text_trims_result(): + """Text formatting does not depend on the presence of native content blocks.""" + reply = {"result": " fallback text \n", "status": "completed", "raw": {}} assert agent_reply_result_text(reply) == "fallback text" +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +def test_agent_reply_result_text_rejects_incomplete_replies(status): + """Partial text must not turn an incomplete invocation into a success.""" + with pytest.raises(RuntimeError, match=status): + agent_reply_result_text({"result": "partial text", "status": status, "error": None}) + + +def test_agent_reply_result_text_preserves_error_message(): + """Provider failures remain actionable to the caller.""" + with pytest.raises(RuntimeError, match="rate limit"): + agent_reply_result_text({"result": "", "status": "failed", "error": "rate limit"}) + + def test_sanitize_msg_for_save_drops_tool_results_and_base64_data(): """Saved history should keep conversation context without tool output pollution.""" msg = Msg.model_validate( diff --git a/tests/unit/test_http_client.py b/tests/unit/test_http_client.py index 7b477f46..dd6a05bb 100644 --- a/tests/unit/test_http_client.py +++ b/tests/unit/test_http_client.py @@ -44,7 +44,7 @@ def test_format_for_display_shows_metadata_when_requested(): @pytest.mark.asyncio @pytest.mark.parametrize("will_retry", [True, False, None, "false"]) -async def test_stream_retry_notifications_continue_but_terminal_errors_raise(will_retry): +async def test_raw_stream_preserves_errors_and_done_but_cli_raises_terminal_errors(will_retry): """A temporary Codex reconnect must not prevent CLI callers from receiving the answer.""" events = [ {"chunk_type": "error", "chunk": "upstream error", "metadata": {"will_retry": will_retry}}, @@ -59,11 +59,95 @@ async def test_stream_retry_notifications_continue_but_terminal_errors_raise(wil client = HttpClient() client.client = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) async with client: + chunks = [chunk async for chunk in client.stream_chunks("stream_llm_demo")] + assert [chunk.chunk_type for chunk in chunks] == [ChunkEnum.ERROR, ChunkEnum.CONTENT, ChunkEnum.DONE] + assert chunks[0].metadata == {"will_retry": will_retry} + assert chunks[1].chunk == "recovered answer" if will_retry is True: - chunks = [chunk async for chunk in client.stream_chunks("stream_llm_demo")] - assert [chunk.chunk_type for chunk in chunks] == [ChunkEnum.ERROR, ChunkEnum.CONTENT] - assert chunks[0].metadata == {"will_retry": True} - assert chunks[-1].chunk == "recovered answer" + assert "recovered answer" in "".join([text async for text in client._execute("stream_llm_demo", {})]) else: with pytest.raises(RuntimeError, match="upstream error"): - _ = [chunk async for chunk in client.stream_chunks("stream_llm_demo")] + _ = [text async for text in client._execute("stream_llm_demo", {})] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_http_collection_and_replay_preserve_complete_outcomes(status): + """Raw HTTP consumption preserves failures and native metadata all the way through DONE.""" + from reme.schema import StreamChunk, collect_reply, replay_reply + + events = [ + StreamChunk(chunk_type=ChunkEnum.REPLY_START, status="running", chunk=None), + StreamChunk(chunk="draft", block_id="text"), + StreamChunk(chunk="corrected", event="end", block_id="text", metadata={"native": {"signature": "signed"}}), + StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status=status, + chunk={ + "session_id": "session", + "result": "corrected", + "structured_output": [], + "usage": None, + "status": status, + "error": "failure" if status == "failed" else None, + "raw": {"cache": 7, "future": False}, + }, + ), + StreamChunk(chunk_type=ChunkEnum.DONE, done=True), + ] + if status == "failed": + events.insert(1, StreamChunk(chunk_type=ChunkEnum.ERROR, chunk="failure")) + body = "".join(f"data: {event.model_dump_json()}\n\n" for event in events) + + def handle(_request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + client = HttpClient() + client.client = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + async with client: + reply = await collect_reply(client.stream_chunks("chat")) + assert reply["status"] == status + assert reply["events"] == [event.model_dump(mode="json") for event in events] + assert await collect_reply(replay_reply(json.loads(json.dumps(reply)))) == reply + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", ["data: invalid\n\n", 'data: {"chunk":"partial"}\n\n']) +async def test_raw_http_stream_rejects_corruption_or_missing_done(body): + """Broken streams must not silently lose events or become successful replies.""" + + def handle(_request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + client = HttpClient() + client.client = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + async with client: + with pytest.raises((ValueError, RuntimeError)): + _ = [chunk async for chunk in client.stream_chunks("chat")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +@pytest.mark.parametrize("error", [None, "provider refused the request"]) +async def test_cli_reports_incomplete_terminal_outcomes(monkeypatch, status, error): + """A partial answer followed by DONE must not hide an incomplete invocation.""" + from reme.schema import StreamChunk + + async def stream(_action, _payload): + yield StreamChunk(chunk="partial answer") + yield StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status=status, + chunk={"status": status, "error": error}, + ) + yield StreamChunk(chunk_type=ChunkEnum.DONE, done=True) + + client = HttpClient() + monkeypatch.setattr(client, "_iter_stream_chunks", stream) + output = [] + with pytest.raises(RuntimeError, match=error or status): + async for text in client._execute("chat", {}): + output.append(text) + assert output == ["partial answer"] diff --git a/tests/unit/test_injected_job_kwargs.py b/tests/unit/test_injected_job_kwargs.py index 35cd82f4..4d41e458 100644 --- a/tests/unit/test_injected_job_kwargs.py +++ b/tests/unit/test_injected_job_kwargs.py @@ -166,7 +166,7 @@ class _RecordingWrapper(BaseAgentWrapper): async def reply(self, _inputs, **kwargs) -> dict: self.calls.append(kwargs) - return {"session_id": "s-1", "last_message": {}, "result": "ok"} + return {"session_id": "s-1", "status": "completed", "result": "ok"} async def _auto_memory_step(wrapper) -> AutoMemoryStep: diff --git a/tests/unit/test_job_status.py b/tests/unit/test_job_status.py index 4ee06ed3..a0c342c6 100644 --- a/tests/unit/test_job_status.py +++ b/tests/unit/test_job_status.py @@ -233,7 +233,22 @@ async def test_chat_error_chunks_match_execution_status_without_changing_stream( 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) + yield StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status="completed" if will_retry else "failed", + chunk={ + key: value + for key, value in self._reply_result( + session_id="stub", + result="partial or recovered", + raw={}, + status="completed" if will_retry else "failed", + error=None if will_retry else "provider failed", + ).items() + if key != "events" + }, + ) context.components["agent_wrapper"] = {"default": ErrorAgent()} job = _job(context, StreamJob, track_execution=True, steps=[{"backend": "chat_step"}]) diff --git a/tests/unit/test_lme_judge.py b/tests/unit/test_lme_judge.py index 680fb11a..9bebdad5 100644 --- a/tests/unit/test_lme_judge.py +++ b/tests/unit/test_lme_judge.py @@ -25,7 +25,7 @@ def fixture_judge(monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("raw,expected", [("YES", "yes"), (" No, incorrect.", "no")]) async def test_valid_judgement_does_not_retry(judge, raw, expected): - judge.agent_wrapper.reply.return_value = {"result": raw} + judge.agent_wrapper.reply.return_value = {"result": raw, "status": "completed", "error": None} response = await judge.execute() assert response.success assert response.answer == expected @@ -35,7 +35,10 @@ async def test_valid_judgement_does_not_retry(judge, raw, expected): @pytest.mark.asyncio async def test_invalid_output_retries_once_and_recovers(judge): - judge.agent_wrapper.reply.side_effect = [{"result": "mostly correct"}, {"result": "yes"}] + judge.agent_wrapper.reply.side_effect = [ + {"result": "mostly correct", "status": "completed", "error": None}, + {"result": "yes", "status": "completed", "error": None}, + ] response = await judge.execute() assert response.success and response.answer == "yes" assert response.metadata["raw_answer_judgement"] == "yes" @@ -46,7 +49,7 @@ async def test_invalid_output_retries_once_and_recovers(judge): @pytest.mark.asyncio @pytest.mark.parametrize("raw", ["", "maybe", "yesterday", "correct"]) async def test_two_invalid_outputs_fail_without_a_score(judge, raw): - judge.agent_wrapper.reply.return_value = {"result": raw} + judge.agent_wrapper.reply.return_value = {"result": raw, "status": "completed", "error": None} with pytest.raises(ValueError, match="invalid yes/no judgement after 2 attempts"): await judge.execute() assert judge.agent_wrapper.reply.await_count == 2 @@ -59,3 +62,13 @@ async def test_transport_errors_are_not_invalid_output_retries(judge): with pytest.raises(TimeoutError): await judge.execute() assert judge.agent_wrapper.reply.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["failed", "interrupted", "waiting"]) +async def test_incomplete_replies_are_neither_scored_nor_retried(judge, status): + judge.agent_wrapper.reply.return_value = {"result": "yes", "status": status, "error": None} + with pytest.raises(RuntimeError, match=status): + await judge.execute() + assert judge.agent_wrapper.reply.await_count == 1 + assert "answer_judgement" not in judge.context.response.metadata diff --git a/tests/unit/test_memlens_benchmark.py b/tests/unit/test_memlens_benchmark.py index 1a461ca5..d30ecff2 100644 --- a/tests/unit/test_memlens_benchmark.py +++ b/tests/unit/test_memlens_benchmark.py @@ -111,7 +111,12 @@ def fixture_candidate(tmp_path, monkeypatch): components={}, ) wrapper = AsAgentWrapper(backend="agentscope", as_llm="", app_context=app) - wrapper.reply = AsyncMock(return_value={"result": "answer", "session_id": "new", "usage": {"input_tokens": 7}}) + wrapper.reply = AsyncMock( + return_value={ + **wrapper._reply_result(session_id="new", result="answer", raw={}), + "usage": {"input_tokens": 7}, + }, + ) store = LocalFileStore(embedding_store="", app_context=app) return memory, answer, app, wrapper, store, tmp_path @@ -320,16 +325,39 @@ async def test_image_only_pipeline_preserves_sessions_and_official_query(candida @pytest.mark.asyncio -async def test_observation_preserves_native_reply_and_rejects_only_interrupted(): +async def test_observation_preserves_completed_reply_and_rejects_incomplete(): from benchmark.memlens.run import observe_usage - result = {"result": "", "last_message": {"finished_reason": "exceed_max_iters", "error": None}} + result = {"result": "answer", "status": "completed", "error": None} reply = AsyncMock(return_value=result) wrapper = SimpleNamespace(reply=reply) app = SimpleNamespace(context=SimpleNamespace(components={"agent_wrapper": {"default": wrapper}})) observe_usage(app) assert await wrapper.reply("Original input", react_config={"max_iters": 30}) is result reply.assert_awaited_once_with("Original input", react_config={"max_iters": 30}) - reply.return_value = {"last_message": {"finished_reason": "interrupted"}} - with pytest.raises(RuntimeError, match="interrupted"): - await wrapper.reply("Original input") + for status in ("interrupted", "failed", "waiting"): + reply.return_value = {"status": status, "error": None} + with pytest.raises(RuntimeError, match=status): + await wrapper.reply("Original input") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_answer_preserves_native_usage_and_rejects_incomplete_outcomes(candidate, monkeypatch, status): + """Standalone answer Jobs retain the canonical outcome and native cache counters.""" + _, answer, app, wrapper, store, _ = candidate + native = {"message": {"usage": {"input_tokens": 7, "cache_input_tokens": 5}}} + result = wrapper._reply_result( + session_id="new", + result="answer", + raw=native, + status=status, + error="provider failed" if status == "failed" else None, + ) + wrapper.reply.return_value = result + monkeypatch.setattr(wrapper, "_resolve_job_tools", lambda _names: []) + step = answer.MemLensAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + response = await step(query="question") + assert response.success is (status == "completed") + assert response.answer == ("provider failed" if status == "failed" else "answer") + assert all(response.metadata[key] == value for key, value in result.items()) diff --git a/tests/unit/test_mobilemem_benchmark.py b/tests/unit/test_mobilemem_benchmark.py index 54fc256f..e6292bfa 100644 --- a/tests/unit/test_mobilemem_benchmark.py +++ b/tests/unit/test_mobilemem_benchmark.py @@ -114,7 +114,12 @@ def fixture_candidate(tmp_path, monkeypatch): components={}, ) wrapper = AsAgentWrapper(backend="agentscope", as_llm="", app_context=app) - wrapper.reply = AsyncMock(return_value={"result": "answer", "session_id": "new", "usage": {"input_tokens": 7}}) + wrapper.reply = AsyncMock( + return_value={ + **wrapper._reply_result(session_id="new", result="answer", raw={}), + "usage": {"input_tokens": 7}, + }, + ) store = LocalFileStore(embedding_store="", app_context=app) return name, memory, answer, prefix, app, wrapper, store, tmp_path @@ -183,7 +188,7 @@ async def test_candidate_and_judge_start_with_native_settings(tmp_path, monkeypa monkeypatch.setattr(MobileMemAutoMemoryStep, "_list_session_note", AsyncMock(return_value=None)) wrappers = candidate_app.context.components["agent_wrapper"] for wrapper in wrappers.values(): - wrapper.reply = AsyncMock(return_value={"result": "answer"}) + wrapper.reply = AsyncMock(return_value=wrapper._reply_result(session_id="new", result="answer", raw={})) _, scene = mobile_scene messages = mobilemem.session_messages(scene["sessions"][0], scene["image_root"]) for images, kwargs in ((True, {}), (False, {"include_images": False})): @@ -396,12 +401,35 @@ async def test_image_only_pipeline_keeps_all_sessions_and_gold_is_judge_only(mob ] == 70 ) - raw = {"last_message": {"finished_reason": "exceed_max_iters"}, "result": "native answer"} + raw = {"status": "completed", "error": None, "result": "native answer"} native = AsyncMock(return_value=raw) wrapper = SimpleNamespace(reply=native) observe_completion(SimpleNamespace(context=SimpleNamespace(components={"agent_wrapper": {"bench": wrapper}}))) assert await wrapper.reply("unchanged input", max_iters=30) is raw native.assert_awaited_once_with("unchanged input", max_iters=30) - native.return_value = {"last_message": {"finished_reason": "interrupted"}} - with pytest.raises(RuntimeError, match="interrupted"): - await wrapper.reply("unchanged input") + for status in ("interrupted", "failed", "waiting"): + native.return_value = {"status": status, "error": None} + with pytest.raises(RuntimeError, match=status): + await wrapper.reply("unchanged input") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "failed", "interrupted", "waiting"]) +async def test_answer_preserves_native_usage_and_rejects_incomplete_outcomes(candidate, monkeypatch, status): + """Standalone answer Jobs retain the canonical outcome and native cache counters.""" + _, _, answer, _, app, wrapper, store, _ = candidate + native = {"message": {"usage": {"input_tokens": 7, "cache_input_tokens": 5}}} + result = wrapper._reply_result( + session_id="new", + result="answer", + raw=native, + status=status, + error="provider failed" if status == "failed" else None, + ) + wrapper.reply.return_value = result + monkeypatch.setattr(wrapper, "_resolve_job_tools", lambda _names: []) + step = answer.MobileMemAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + response = await step(query="question") + assert response.success is (status == "completed") + assert response.answer == ("provider failed" if status == "failed" else "answer") + assert all(response.metadata[key] == value for key, value in result.items()) diff --git a/tests/unit/test_proactive_refresh.py b/tests/unit/test_proactive_refresh.py index 72139007..733f4480 100644 --- a/tests/unit/test_proactive_refresh.py +++ b/tests/unit/test_proactive_refresh.py @@ -147,11 +147,11 @@ class _AgentWrapper(BaseAgentWrapper): return { "session_id": "stub", "result": text, - "last_message": {"content": [{"type": "text", "text": text}]}, + "status": "completed", } -class _DummyModel(ChatModelBase): +class _DummyModel(ChatModelBase): # pylint: disable=abstract-method """Placeholder model satisfying dream extract's llm_available check.""" def __init__(self): diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index cd505815..3b6d8ca9 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -123,105 +123,47 @@ def test_token_counter_is_a_per_agent_metric_tree(tmp_path): @pytest.mark.asyncio -async def test_agentscope_reply_records_final_message_usage(tmp_path, monkeypatch): - """AgentScope 2.0.4.post1 reports aggregate usage on the final message.""" +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("reported", [False, True]) +async def test_agentscope_records_usage_in_both_return_modes(tmp_path, monkeypatch, streaming, reported): + """Both APIs use the same final native message and record usage exactly once.""" + from agentscope.event import ReplyEndEvent + from agentscope.message import AssistantMsg, Usage + context = ApplicationContext(workspace_dir=str(tmp_path)) wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) - message = SimpleNamespace( - # Match the AgentScope 2.0.4.post1 accumulation test: 10/20 + 5/8. - usage=SimpleNamespace(input_tokens=15, output_tokens=28), - model_dump=lambda: {"text": "answer"}, - get_text_content=lambda: "answer", + message = AssistantMsg( + name="agent", + content="answer", + finished_reason="completed", + usage=Usage(input_tokens=15, output_tokens=28) if reported else None, ) + + async def native_stream(_inputs, **_kwargs): + yield ReplyEndEvent(session_id="session-1", reply_id="reply", finished_reason="completed") + yield message + agent = SimpleNamespace( state=SimpleNamespace(session_id="session-1", context=[message]), - observe=lambda _inputs: _async_none(), - reply=lambda: _async_value(message), - reply_stream=_unexpected_reply_stream, + reply_stream=native_stream, ) async def build_agent(inputs, **_kwargs): return agent, inputs - monkeypatch.setattr(wrapper, "_build_agent", build_agent) - monkeypatch.setattr(wrapper, "_dump_state", _async_none) - - result = await wrapper.reply("hello") - - assert result["usage"] == {"input_tokens": 15, "output_tokens": 28, "total_tokens": 43} - assert ( - global_counter_get_all(context.metadata, ["__token_counter", "research"])["children"]["total_tokens"]["value"] - == 43 - ) - - -@pytest.mark.asyncio -async def test_agentscope_reply_without_usage_leaves_token_counters_unset(tmp_path, monkeypatch): - """Replies without any usage information remain visibly unavailable.""" - context = ApplicationContext(workspace_dir=str(tmp_path)) - wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) - message = SimpleNamespace( - usage=None, - model_dump=lambda: {"text": "answer"}, - get_text_content=lambda: "answer", - ) - agent = SimpleNamespace( - state=SimpleNamespace(session_id="session-1", context=[message]), - observe=lambda _inputs: _async_none(), - reply=lambda: _async_value(message), - reply_stream=_unexpected_reply_stream, - ) - - async def build_agent(inputs, **_kwargs): - return agent, inputs - - monkeypatch.setattr(wrapper, "_build_agent", build_agent) - monkeypatch.setattr(wrapper, "_dump_state", _async_none) - - result = await wrapper.reply("hello") - - assert result["usage"] is None - assert "__token_counter" not in context.metadata - - -async def _async_none(*_args, **_kwargs): - return None - - -async def _async_value(value): - return value - - -def _unexpected_reply_stream(*_args, **_kwargs): - raise AssertionError("Non-streaming replies must use Agent.reply()") - - -@pytest.mark.asyncio -async def test_agentscope_stream_reply_does_not_record_token_usage(tmp_path, monkeypatch): - """Only non-streaming AgentScope replies contribute to token accounting.""" - context = ApplicationContext(workspace_dir=str(tmp_path)) - wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) - - class FakeAgent: - """Minimal AgentScope stream double.""" - - state = type("State", (), {"session_id": "session-1"})() - - async def reply_stream(self, inputs): - """Yield no events for the supplied input.""" - if inputs is None: - yield None - - async def build_agent(inputs, **_kwargs): - """Build the minimal stream double.""" - return FakeAgent(), inputs - async def dump_state(_state): - """Avoid durable state writes in this accounting test.""" return None monkeypatch.setattr(wrapper, "_build_agent", build_agent) monkeypatch.setattr(wrapper, "_dump_state", dump_state) - - assert [chunk async for chunk in wrapper.reply_stream("hello")] == [] - assert "__token_counter" not in context.metadata + if streaming: + chunks = [chunk async for chunk in wrapper.reply_stream("hello")] + result = chunks[-1].chunk + else: + result = await wrapper.reply("hello") + assert result["usage"] == ({"input_tokens": 15, "output_tokens": 28, "total_tokens": 43} if reported else None) + counters = global_counter_get_all(context.metadata, ["__token_counter", "research"]) + if reported: + assert counters["children"]["total_tokens"]["value"] == 43 + else: + assert counters is None diff --git a/tests/unit/test_workspace_web_steps.py b/tests/unit/test_workspace_web_steps.py index febf3768..3673b525 100644 --- a/tests/unit/test_workspace_web_steps.py +++ b/tests/unit/test_workspace_web_steps.py @@ -32,7 +32,17 @@ class _StreamingAgent(BaseAgentWrapper): self.reply_kwargs = kwargs yield StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, chunk="{}", tool_call_id="tool-1", tool_call_name="search") yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk="hello", session_id="session-new") - yield StreamChunk(chunk_type=ChunkEnum.REPLY_END, chunk="", session_id="session-new") + yield StreamChunk( + chunk_type=ChunkEnum.REPLY_END, + event="end", + status="completed", + session_id="session-new", + chunk={ + key: value + for key, value in self._reply_result(session_id="session-new", result="hello", raw={}).items() + if key != "events" + }, + ) @pytest.mark.asyncio @@ -125,7 +135,7 @@ async def test_chat_step_streams_rich_chunks_and_resumes_session(tmp_path): assert response.answer == "hello" assert chunks[0].tool_call_name == "search" assert chunks[1].session_id == "session-new" - assert chunks[2].metadata["answer"] == "hello" + assert chunks[2].chunk["result"] == "hello" assert agent.reply_kwargs["resume"] == "session-old" assert agent.reply_kwargs["job_tools"] == [ "search",