diff --git a/backend/api/server.py b/backend/api/server.py index 3cfdd80..16e5103 100644 --- a/backend/api/server.py +++ b/backend/api/server.py @@ -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: diff --git a/frontend/app/components/assistant-ui/threadlist-sidebar.tsx b/frontend/app/components/assistant-ui/threadlist-sidebar.tsx index 3e3cf7c..3e49337 100644 --- a/frontend/app/components/assistant-ui/threadlist-sidebar.tsx +++ b/frontend/app/components/assistant-ui/threadlist-sidebar.tsx @@ -148,7 +148,7 @@ export function ThreadListSidebar({ -
+
@@ -647,7 +647,7 @@ function AccountMenu({ {compact ? (