[Bug] Add checkpoint/resume mechanism for long-running Agent tasks — run_stream() state is lost on process crash

Self check

• I'm on the latest version and searched existing issues (incl. closed) — no duplicate.

Environment

• Version: v2.1.9 ( latest main )
• OS: macOS 15 (also affects Linux/Windows — the issue is architecture-level, not OS-dependent)
• Python: 3.12
• Install: source
• Model & channel: any multi-turn Agent task (e.g. "read 20 files and write a summary" via terminal channel)

What happened?

What happened?

When the Agent process crashes (or is killed, or reaches max_turns) mid-task, all execution progress is lost. The user must re-send the same prompt and the Agent re-executes every tool call from scratch. For a 30-turn task that crashes at turn 25, this means ~83% wasted work.

Root cause

run_stream() (agent/protocol/agent_stream.py:673) maintains all execution state in memory:

# agent/protocol/agent_stream.py:256
self.messages = messages if messages is not None else []  # ← in-memory only

# agent/protocol/agent_stream.py:732
turn = 0  # ← local variable, lost on crash

The ReAct loop (while turn < self.max_turns, line 767) accumulates tool results and LLM reasoning into self.messages, but there is no persistence mechanism anywhere in the loop. When the process exits:

• self.messages — gone
• turn counter — gone
• All tool call results accumulated over N turns — gone

The existing cancel_event (line 251) handles graceful user-initiated cancellation, but not unexpected termination (OOM kill, power loss, deploy restart).

Proposed change (minimal)

Add an opt-in CheckpointManager that serializes {messages, turn} to a JSON file after each turn. On restart with the same session_id, the Agent resumes from the saved turn instead of starting over.

New file: agent/protocol/checkpoint.py (~60 lines)

import json
from pathlib import Path
from datetime import datetime, timezone

class CheckpointManager:
    """Atomic checkpoint persistence for Agent run_stream() execution."""

    def __init__(self, checkpoint_dir: str):
        self.dir = Path(checkpoint_dir)
        self.dir.mkdir(parents=True, exist_ok=True)

    def save(self, session_id: str, messages: list, turn: int):
        """Serialize state to JSON using atomic write (temp + rename)."""
        data = {
            "version": 1,
            "turn": turn,
            "saved_at": datetime.now(timezone.utc).isoformat(),
            "messages": messages,
        }
        path = self.dir / f"{session_id}.json"
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(data, ensure_ascii=False, default=str),
                       encoding="utf-8")
        tmp.rename(path)  # Atomic on POSIX and Windows

    def load(self, session_id: str) -> Optional[tuple]:
        """Load checkpoint. Returns (messages, turn) or None."""
        path = self.dir / f"{session_id}.json"
        if not path.exists():
            return None
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            return data["messages"], data["turn"]
        except (json.JSONDecodeError, KeyError):
            return None  # Corrupted checkpoint → safe fallback to fresh start

    def clear(self, session_id: str):
        """Remove checkpoint after successful completion."""
        (self.dir / f"{session_id}.json").unlink(missing_ok=True)

Integration in run_stream() (3 changes, ~20 lines):

# 1. Method signature: add two optional parameters
def run_stream(self, user_message: str,
               session_id: str = "",
               checkpoint_dir: Optional[str] = None) -> str:

    # 2. Initialize checkpoint manager (before the while loop)
    cp = CheckpointManager(checkpoint_dir) if checkpoint_dir else None
    if cp:
        saved = cp.load(session_id)
        if saved:
            self.messages, turn = saved  # Resume from checkpoint
            logger.info(f"Resumed from checkpoint at turn {turn}")
        else:
            self._append_user_message(user_message)
            turn = 0

    # ... (existing ReAct loop, unchanged)

    while turn < self.max_turns:
        # ... (existing logic)
        # 3a. Save checkpoint after each turn
        if cp:
            cp.save(session_id, self.messages, turn)

    # 3b. Clean up on successful completion
    if cp:
        cp.clear(session_id)
    return final_response

What stays the same

• run_stream() without checkpoint_dir → behavior identical to current code (zero overhead)
• All tool execution, LLM calling, steering, and cancel logic — unchanged
• _trim_messages(), _validate_and_fix_messages() — called before cp.save() to ensure consistent state
• _smart_compact_to_budget() — still works as the reactive overflow handler

Design decisions

Decision	Choice	Why
Atomic write	temp file + rename	Prevents partial file if process dies mid-write
Serialization	JSON (stdlib)	No new dependency; messages are already plain dicts
Trigger frequency	Every turn	Agent turns take 1-30s; 5ms disk write is negligible
Corruption handling	Return None → fresh start	Safe degradation over silent failure
What to persist	{messages, turn} only	These two fields are sufficient to resume; tool_failure_history and _retrieved_mcp_names are safe to rebuild

I'm happy to open a PR implementing the above (new checkpoint.py ~60 lines + 3 integration points in agent_stream.py ~20 lines + ~100 lines of tests). Just let me know if this direction works for you, or if you'd prefer a different persistence approach.

---

中文简述

run_stream() 的所有执行状态（消息列表 + 轮次计数器）仅存在于内存中。进程崩溃后无法恢复，长任务必须从头重跑。

建议增加可选的 Checkpoint 机制：

• 新增 CheckpointManager 类，使用 JSON 文件 + 原子写入（temp + rename）保存每 turn 的 {messages, turn}；
• run_stream() 新增 checkpoint_dir 可选参数（默认 None 时零开销，不影响现有调用方）；
• 启动时检测 checkpoint 文件，自动从断点恢复执行。

改动量：新增 checkpoint.py ~60 行 + 修改 agent_stream.py 3 处 ~20 行 + 测试 ~100 行。可直接提 PR。

Logs

N/A — This issue was identified through source code analysis of scheduler_service.py, not through runtime observation. The root cause is determinable from the code path alone (see Root cause section above).