Przeglądaj źródła

fix: guard session replay restoration

Problem:
Stale replay loads could overwrite the selected session, while incomplete tool protocol could hide visible assistant text or render whitespace-only records.

Risk:
Low. Changes are limited to replay restoration guards and transcript repair, with focused, full, and Node verification.
zhenyu.hu 2 tygodni temu
rodzic
commit
792dcb5b29

+ 2 - 0
src/agent_lab/application/runtime.py

@@ -1325,6 +1325,8 @@ class DebugRuntime:
             reply_ids = {reply.tool_call_id for reply in replies}
             if len(replies) == len(expected_ids) and reply_ids == expected_ids:
                 repaired.extend([message, *replies])
+            elif message.content.strip():
+                repaired.append(message.model_copy(update={"tool_calls": []}))
             index = cursor
         return repaired
 

+ 12 - 1
src/agent_lab/presentation/static/app.js

@@ -37,6 +37,7 @@ let hasBackendRoundStats = false;
 let pendingEnabledTools = null;
 let availableTools = [];
 let currentSessionId = null;
+let sessionReplayLoadToken = 0;
 let sessions = [];
 let auditEntries = [];
 let selectedAuditKey = null;
@@ -205,6 +206,10 @@ async function loadSelectedSession() {
 }
 
 async function loadSessionReplay(sessionId, options = {}) {
+  const loadToken = ++sessionReplayLoadToken;
+  const isCurrentLoad = () => (
+    loadToken === sessionReplayLoadToken && sessionId === currentSessionId
+  );
   try {
     const [session, messages, audit, usage] = await Promise.all([
       fetchJson(`/api/sessions/${sessionId}`),
@@ -212,6 +217,9 @@ async function loadSessionReplay(sessionId, options = {}) {
       fetchJson(`/api/sessions/${sessionId}/audit`),
       fetchJson(`/api/sessions/${sessionId}/usage`),
     ]);
+    if (!isCurrentLoad()) {
+      return;
+    }
     restoreWorkspaceSnapshot(session.config || {});
     renderReplayMessages(messages);
     renderAuditReplay(audit);
@@ -222,6 +230,9 @@ async function loadSessionReplay(sessionId, options = {}) {
       setSessionStatus("Session loaded");
     }
   } catch (error) {
+    if (!isCurrentLoad()) {
+      return;
+    }
     if (!options.preserveStatus) {
       setSessionStatus("Failed to load session");
     }
@@ -994,7 +1005,7 @@ function renderReplayMessages(messages) {
   messages.forEach((message) => {
     if (
       message.role === "assistant" &&
-      !message.content &&
+      !message.content.trim() &&
       Array.isArray(message.tool_calls) &&
       message.tool_calls.length > 0
     ) {

+ 18 - 2
tests/test_debug_runtime.py

@@ -2197,8 +2197,22 @@ async def test_direct_restore_drops_invalid_protocol_fragments_and_keeps_visible
             ChatMessage(role="assistant", content="visible after orphan"),
             ChatMessage(
                 role="assistant",
-                content="",
-                tool_calls=[_tool_call("incomplete-call")],
+                content="visible incomplete tool preface",
+                tool_calls=[
+                    _tool_call("incomplete-visible-1"),
+                    _tool_call("incomplete-visible-2"),
+                ],
+            ),
+            ChatMessage(
+                role="tool",
+                content='{"partial":true}',
+                name="safe_tool",
+                tool_call_id="incomplete-visible-1",
+            ),
+            ChatMessage(
+                role="assistant",
+                content="  \n\t",
+                tool_calls=[_tool_call("incomplete-whitespace")],
             ),
         ],
     )
@@ -2234,9 +2248,11 @@ async def test_direct_restore_drops_invalid_protocol_fragments_and_keeps_visible
         ("tool", '{"id":2}'),
         ("assistant", "visible answer"),
         ("assistant", "visible after orphan"),
+        ("assistant", "visible incomplete tool preface"),
         ("user", "current user"),
     ]
     assert transcript[2].tool_calls == complete_calls
+    assert transcript[-2].tool_calls == []
 
 
 @pytest.mark.asyncio

+ 181 - 7
tests/test_websocket_api.py

@@ -1,6 +1,7 @@
 import json
 import asyncio
 import logging
+import subprocess
 from collections.abc import AsyncIterator
 from pathlib import Path
 
@@ -26,6 +27,86 @@ def _css_rule(css: str, selector: str) -> str:
     return css[start:end]
 
 
+def _run_node_json(source: str) -> object:
+    completed = subprocess.run(
+        ["node", "--input-type=module", "-e", source],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert completed.returncode == 0, completed.stderr
+    return json.loads(completed.stdout)
+
+
+def _run_load_session_replay_scenario(scenario: str) -> object:
+    js = Path("src/agent_lab/presentation/static/app.js").read_text()
+    load_replay = js[
+        js.index("async function loadSessionReplay(") :
+        js.index("function setSessionStatus(")
+    ]
+    harness = """
+let currentSessionId = null;
+let sessionReplayLoadToken = 0;
+const sessionTitle = { value: "" };
+const sessionList = { value: "" };
+const mutations = {
+  workspace: [],
+  transcript: [],
+  audit: [],
+  usage: [],
+  status: [],
+};
+const pending = new Map();
+
+function fetchJson(url) {
+  let resolve;
+  let reject;
+  const promise = new Promise((resolvePromise, rejectPromise) => {
+    resolve = resolvePromise;
+    reject = rejectPromise;
+  });
+  pending.set(url, { resolve, reject });
+  return promise;
+}
+
+function resolveReplay(sessionId) {
+  pending.get(`/api/sessions/${sessionId}`).resolve({
+    title: `${sessionId} title`,
+    config: { marker: sessionId },
+  });
+  pending.get(`/api/sessions/${sessionId}/messages`).resolve([
+    { content: sessionId },
+  ]);
+  pending.get(`/api/sessions/${sessionId}/audit`).resolve([
+    { event: sessionId },
+  ]);
+  pending.get(`/api/sessions/${sessionId}/usage`).resolve({
+    session: { marker: sessionId },
+  });
+}
+
+function restoreWorkspaceSnapshot(config) {
+  mutations.workspace.push(config.marker);
+}
+function renderReplayMessages(messages) {
+  mutations.transcript.push(messages[0].content);
+}
+function renderAuditReplay(audit) {
+  mutations.audit.push(audit[0].event);
+}
+function renderSessionUsage(usage) {
+  mutations.usage.push(usage.session.marker);
+}
+function setSessionStatus(status) {
+  mutations.status.push(status);
+}
+function isSocketOpen() {
+  return false;
+}
+"""
+    return _run_node_json("\n".join([harness, load_replay, scenario]))
+
+
 class FakeRuntime:
     def __init__(self) -> None:
         self.requests: list[DebugRunRequest] = []
@@ -1534,6 +1615,14 @@ def test_static_app_reuses_websocket_for_session_turns():
 
 def test_static_app_loads_session_replay_and_usage():
     js = Path("src/agent_lab/presentation/static/app.js").read_text()
+    create_session = js[
+        js.index("async function createSession()") :
+        js.index("async function loadSelectedSession()")
+    ]
+    load_session = js[
+        js.index("async function loadSelectedSession()") :
+        js.index("async function loadSessionReplay(")
+    ]
 
     assert "function loadSessions(" in js
     assert "function createSession(" in js
@@ -1552,21 +1641,106 @@ def test_static_app_loads_session_replay_and_usage():
     assert "function renderSessionUsage(" in js
     assert "sessionUsageSummary" in js
     assert "Finish current turn before switching sessions" in js
+    assert "await loadSessionReplay(session.id, { preserveStatus: true });" in create_session
+    assert "await loadSessionReplay(sessionId);" in load_session
+
+
+def test_static_session_replay_ignores_superseded_success():
+    js = Path("src/agent_lab/presentation/static/app.js").read_text()
+    result = _run_load_session_replay_scenario(
+        """
+currentSessionId = "old";
+const oldLoad = loadSessionReplay("old");
+currentSessionId = "new";
+const newLoad = loadSessionReplay("new");
+resolveReplay("new");
+await newLoad;
+resolveReplay("old");
+await oldLoad;
+console.log(JSON.stringify({
+  ...mutations,
+  title: sessionTitle.value,
+  selected: sessionList.value,
+}));
+"""
+    )
+
+    assert result == {
+        "workspace": ["new"],
+        "transcript": ["new"],
+        "audit": ["new"],
+        "usage": ["new"],
+        "status": ["Session loaded"],
+        "title": "new title",
+        "selected": "new",
+    }
+    assert "let sessionReplayLoadToken = 0;" in js
+
+
+def test_static_session_replay_ignores_superseded_failure_status():
+    result = _run_load_session_replay_scenario(
+        """
+currentSessionId = "old";
+const oldLoad = loadSessionReplay("old");
+currentSessionId = "new";
+const newLoad = loadSessionReplay("new");
+resolveReplay("new");
+await newLoad;
+pending.get("/api/sessions/old").reject(new Error("stale failure"));
+await oldLoad;
+console.log(JSON.stringify(mutations.status));
+"""
+    )
+
+    assert result == ["Session loaded"]
 
 
-def test_static_replay_hides_empty_assistant_tool_protocol_records():
+def test_static_replay_hides_whitespace_assistant_tool_protocol_records():
     js = Path("src/agent_lab/presentation/static/app.js").read_text()
     render_replay = js[
         js.index("function renderReplayMessages(") :
         js.index("function appendAuditEntry(")
     ]
+    result = _run_node_json(
+        "\n".join(
+            [
+                """
+let activeAssistant = null;
+const messagesEl = {
+  children: [],
+  scrollHeight: 0,
+  scrollTop: 0,
+  set textContent(value) {
+    if (value === "") {
+      this.children = [];
+    }
+  },
+  append(item) {
+    this.children.push(item);
+  },
+};
+const document = {
+  createElement() {
+    return { className: "", textContent: "" };
+  },
+};
+function appendLog(kind, content) {
+  messagesEl.append({ className: `message ${kind}`, textContent: content });
+}
+""",
+                render_replay,
+                """
+renderReplayMessages([
+  { role: "assistant", content: "  \\n\\t", tool_calls: [{ id: "call-1" }] },
+  { role: "assistant", content: "visible", tool_calls: [] },
+]);
+console.log(JSON.stringify(messagesEl.children.map((item) => item.textContent)));
+""",
+            ]
+        )
+    )
 
-    assert 'message.role === "assistant"' in render_replay
-    assert "!message.content" in render_replay
-    assert "Array.isArray(message.tool_calls)" in render_replay
-    assert "message.tool_calls.length > 0" in render_replay
-    assert "return;" in render_replay
-    assert '["user", "assistant", "tool"].includes(message.role)' in render_replay
+    assert result == ["visible"]
 
 
 def test_static_session_switch_guard_uses_active_turn_not_open_socket():