From 3c200db6216fc6a67bbfdcc0d3fd4693c25f4d23 Mon Sep 17 00:00:00 2001
From: wuyang6 <wuyang6@xiaomi.com>
Date: Thu, 14 May 2026 14:48:21 +0800
Subject: [PATCH] Persist resume prompts before execution

---
 backend/api/server.py | 94 ++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 83 insertions(+), 11 deletions(-)

diff --git a/backend/api/server.py b/backend/api/server.py
index d64a2b2..b125373 100644
--- a/backend/api/server.py
+++ b/backend/api/server.py
@@ -1630,6 +1630,29 @@ def create_app(state: AgentState) -> FastAPI:
         account_key = state._account_key(account_id)
         snapshot = state.run_manager.snapshot_latest(account_key, safe_id)
         stored_snapshot = state.run_state_store.snapshot_latest(account_key, safe_id)
+        if (
+            snapshot is not None
+            and stored_snapshot is not None
+            and stored_snapshot.get('status') in ACTIVE_RUN_STATUSES
+            and snapshot.get('run_id') == stored_snapshot.get('run_id')
+            and snapshot.get('status') not in ACTIVE_RUN_STATUSES
+        ):
+            state.run_state_store.finish(
+                str(stored_snapshot.get('run_id') or ''),
+                status=str(snapshot.get('status') or 'interrupted'),
+                elapsed_ms=(
+                    int(snapshot['elapsed_ms'])
+                    if isinstance(snapshot.get('elapsed_ms'), (int, float))
+                    else None
+                ),
+                stage=str(snapshot.get('current_stage') or ''),
+                error=(
+                    str(snapshot.get('error'))
+                    if snapshot.get('error') is not None
+                    else None
+                ),
+            )
+            stored_snapshot = state.run_state_store.snapshot_latest(account_key, safe_id)
         if snapshot is not None:
             return _merge_run_snapshot(snapshot, stored_snapshot) or snapshot
         if stored_snapshot is not None:
@@ -1735,14 +1758,12 @@ def create_app(state: AgentState) -> FastAPI:
             session_directory,
             requested_session_id,
         )
-        fallback_title = None
-        if request.resume_session_id is None:
-            fallback_title = _save_in_progress_session(
-                directory=session_directory,
-                agent=agent,
-                session_id=requested_session_id,
-                prompt=prompt,
-            )
+        fallback_title = _save_in_progress_session(
+            directory=session_directory,
+            agent=agent,
+            session_id=requested_session_id,
+            prompt=prompt,
+        )
         if run_lock.locked():
             queued_event = {
                 'type': 'run_queued',
@@ -2137,6 +2158,9 @@ def _sanitize_stored_session_for_resume(
     stored: StoredAgentSession,
 ) -> StoredAgentSession:
     trimmed, changed = _trim_incomplete_message_tail(stored.messages)
+    if trimmed and _is_pending_user_message(trimmed[-1]):
+        trimmed = trimmed[:-1]
+        changed = True
     messages = _without_run_status_messages(trimmed)
     if len(messages) != len(trimmed):
         changed = True
@@ -2186,6 +2210,16 @@ def _is_incomplete_session_message(message: dict[str, Any]) -> bool:
     return False
 
 
+def _is_pending_user_message(message: dict[str, Any]) -> bool:
+    metadata = message.get('metadata')
+    return (
+        message.get('role') == 'user'
+        and isinstance(metadata, dict)
+        and metadata.get('placeholder') is True
+        and metadata.get('status') == 'running'
+    )
+
+
 def _assistant_has_trailing_tool_call(message: dict[str, Any]) -> bool:
     if message.get('role') != 'assistant':
         return False
@@ -2359,11 +2393,49 @@ def _save_in_progress_session(
     safe_id = _safe_session_id(session_id)
     if safe_id is None or not prompt:
         return None
-    if _session_json_path(directory, safe_id).exists():
+
+    initial_title = _derive_initial_session_title(prompt)
+    pending_metadata = {'status': 'running', 'placeholder': True}
+    session_path = _session_json_path(directory, safe_id)
+    if session_path.exists():
+        # 续聊也要先落盘用户消息。否则前端流断开或刷新时，用户刚发的内容只
+        # 存在于浏览器内存里，会出现“消息丢了但后端状态还在”的错觉。
+        try:
+            stored = load_agent_session(safe_id, directory=directory)
+        except (FileNotFoundError, OSError, json.JSONDecodeError):
+            return None
+        messages = list(stored.messages)
+        if messages and _is_pending_user_message(dict(messages[-1])):
+            messages.pop()
+        messages.append(
+            {
+                'role': 'user',
+                'content': prompt,
+                'state': 'final',
+                'metadata': pending_metadata,
+                'message_id': f'user_pending_{int(time.time() * 1000)}',
+            }
+        )
+        budget_state = (
+            dict(stored.budget_state)
+            if isinstance(stored.budget_state, dict)
+            else {}
+        )
+        budget_state['status'] = 'running'
+        try:
+            save_agent_session(
+                replace(
+                    stored,
+                    messages=tuple(messages),
+                    budget_state=budget_state,
+                ),
+                directory=directory,
+            )
+        except OSError:
+            return None
         return None
 
     # 新会话的第一轮执行可能很久。先写一个运行中占位，避免刷新页面后找不到会话。
-    initial_title = _derive_initial_session_title(prompt)
     scratchpad_directory = (
         agent.runtime_config.scratchpad_root / safe_id / 'scratchpad'
     ).resolve()
@@ -2372,7 +2444,7 @@ def _save_in_progress_session(
         session = agent.build_session(None, scratchpad_directory=scratchpad_directory)
         session.append_user(
             prompt,
-            metadata={'status': 'running', 'placeholder': True},
+            metadata=pending_metadata,
             message_id='user_pending_0',
         )
         stored = StoredAgentSession(