Persist completed run elapsed time

2026-06-12 19:48:45 +08:00
parent aac6383fe6
commit 4a268902bf
2 changed files with 61 additions and 27 deletions
@@ -3799,34 +3799,68 @@ def _annotate_last_assistant_elapsed(
    session_id: str,
    elapsed_ms: int,
 ) -> None:
-    path = _session_json_path(directory, session_id)
    try:
-        data = json.loads(path.read_text(encoding='utf-8'))
-    except (OSError, json.JSONDecodeError):
+        stored = load_agent_session(session_id, directory=directory)
+    except (
+        FileNotFoundError,
+        OSError,
+        json.JSONDecodeError,
+        KeyError,
+        TypeError,
+        ValueError,
+    ):
        return
-    messages = data.get('messages')
-    display_messages = data.get('display_messages')
-    if not isinstance(display_messages, list):
-        display_messages = messages
-    if not isinstance(messages, list) and not isinstance(display_messages, list):
+
+    messages, messages_changed = _annotate_last_assistant_in_messages(
+        stored.messages,
+        elapsed_ms,
+    )
+    display_messages: tuple[dict[str, Any], ...] = ()
+    display_changed = False
+    if stored.display_messages:
+        display_messages, display_changed = _annotate_last_assistant_in_messages(
+            stored.display_messages,
+            elapsed_ms,
+        )
+
+    if not messages_changed and not display_changed:
        return
-    for message_list in (messages, display_messages):
-        if not isinstance(message_list, list):
+
+    save_agent_session(
+        replace(
+            stored,
+            messages=messages if messages_changed else stored.messages,
+            display_messages=(
+                display_messages
+                if display_changed
+                else stored.display_messages
+            ),
+        ),
+        directory=directory,
+    )
+
+
+def _annotate_last_assistant_in_messages(
+    messages: tuple[dict[str, Any], ...] | tuple[Any, ...],
+    elapsed_ms: int,
+) -> tuple[tuple[dict[str, Any], ...], bool]:
+    updated = [dict(message) for message in messages if isinstance(message, dict)]
+    for index in range(len(updated) - 1, -1, -1):
+        message = updated[index]
+        if message.get('role') != 'assistant':
            continue
-        for message in reversed(message_list):
-            if not isinstance(message, dict) or message.get('role') != 'assistant':
-                continue
-            metadata = message.get('metadata')
-            if not isinstance(metadata, dict):
-                metadata = {}
-                message['metadata'] = metadata
-            metadata['elapsed_ms'] = elapsed_ms
-            break
-    try:
-        path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
-    except OSError:
-        return
-    return
+        metadata = message.get('metadata')
+        if not isinstance(metadata, dict):
+            metadata = {}
+        else:
+            metadata = dict(metadata)
+        if metadata.get('elapsed_ms') == elapsed_ms:
+            return tuple(updated), False
+        metadata['elapsed_ms'] = elapsed_ms
+        message['metadata'] = metadata
+        updated[index] = message
+        return tuple(updated), True
+    return tuple(updated), False


 def _read_session_title(directory: Path, session_id: str | None) -> str | None: