From 68b079cc5e874a9090e03c0bb94f185482a6fe64 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Fri, 17 Jul 2026 11:11:35 -0700 Subject: [PATCH] [eric] ws: slim status frames, cross-socket dedupe, REST-seeded resume cursor, closed-only history --- backend/apps/agents/agents.py | 10 +- backend/apps/agents/core/ws_manager.py | 30 +++- .../manager/session/SessionLifecycle.py | 4 + backend/tests/test_ws_efficiency.py | 155 ++++++++++++++++++ .../src/app/pages/AgentChat/AgentChat.tsx | 9 +- frontend/src/shared/state/agentsSlice.ts | 8 +- frontend/src/shared/ws/WebSocketManager.ts | 29 ++++ 7 files changed, 238 insertions(+), 7 deletions(-) create mode 100644 backend/tests/test_ws_efficiency.py diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 07c059d0..8c48f640 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -9,6 +9,7 @@ from typeguard import typechecked from backend.apps.agents.agent_manager import agent_manager from backend.apps.agents.core.models import AgentConfig, AgentSession, ApprovalResponse +from backend.apps.agents.core.seq_log import seq_log from backend.apps.agents.manager.session.history_compaction import estimate_post_compact_input from backend.config.Apps import SubApp @@ -92,7 +93,11 @@ async def get_session(session_id: str): session = await agent_manager.resume_session(session_id) except ValueError: raise HTTPException(status_code=404, detail="Session not found") - return session.model_dump(mode="json") + # Seq read before the dump (no await between = atomic): the client seeds its WS resume cursor from this, so a REST hydrate isn't followed by a full from-zero replay of everything it just received. + event_seq = seq_log.current_seq(session_id) + payload = session.model_dump(mode="json") + payload["event_seq"] = event_seq + return payload @agents.router.post("/launch") async def launch_agent(config: AgentConfig): @@ -273,10 +278,11 @@ async def delete_session(session_id: str): return {"ok": True} @agents.router.get("/history") -async def get_history(q: str = "", limit: int = 20, offset: int = 0, dashboard_id: str = ""): +async def get_history(q: str = "", limit: int = 20, offset: int = 0, dashboard_id: str = "", closed_only: int = 0): return agent_manager.get_history( q=q, limit=limit, offset=offset, dashboard_id=dashboard_id or None, + closed_only=bool(closed_only), ) @agents.router.get("/sessions/{session_id}/browser-agents") diff --git a/backend/apps/agents/core/ws_manager.py b/backend/apps/agents/core/ws_manager.py index 5b8022ff..6996e5cc 100644 --- a/backend/apps/agents/core/ws_manager.py +++ b/backend/apps/agents/core/ws_manager.py @@ -23,6 +23,29 @@ BROWSER_CMD_REBROADCAST_S = 3.0 P_WS_RECONNECT_WAIT_S = 8.0 +def slim_status_data(event: str, data: dict) -> dict: + """agent:status frames carry session METADATA, never the transcript: every message already + reaches clients as its own agent:message event (and the stream), so re-shipping full history + per status flip was pure duplication, and replayed stale copies rolled clients backwards. + Preview fields mirror p_session_list_item so collapsed-card previews keep working.""" + if event != "agent:status": + return data + sess = data.get("session") + if not isinstance(sess, dict) or not sess.get("messages"): + return data + messages = sess["messages"] + last = messages[-1].get("content", "") + first_user = next((m.get("content") for m in messages if m.get("role") == "user"), "") + slim = dict(sess) + slim["messages"] = [] + slim["last_message_preview"] = last[:120] if isinstance(last, str) else "" + slim["first_user_message"] = first_user[:200] if isinstance(first_user, str) else "" + slim["message_count"] = len(messages) + out = dict(data) + out["session"] = slim + return out + + async def await_reconnect(has_conn) -> bool: """Poll up to P_WS_RECONNECT_WAIT_S for a dashboard socket to (re)appear. `has_conn` is a 0-arg callable returning truthy when connected.""" @@ -94,6 +117,7 @@ class ConnectionManager: async def send_to_session(self, session_id: str, event: str, data: dict): """Broadcast a session event with monotonic sequencing; terminal statuses also persist to disk.""" + data = slim_status_data(event, data) async with seq_log.stamp(session_id, event, data) as (seq, payload_str): for ws in list(self.connections.get(session_id, [])): try: @@ -163,6 +187,10 @@ class ConnectionManager: "to_seq": newest, } + # Live log and the client is at (or past) the top: caught up, nothing to send. Without this, every cursor-seeded reconnect got the persisted terminal frame re-sent as "replay". + if newest is not None and last_seq >= newest and newest > 0: + return {"ok": True, "replayed": 0, "current_seq": newest} + terminal = seq_log.load_terminal(session_id) if terminal is not None: try: @@ -225,7 +253,7 @@ class ConnectionManager: async def broadcast_global(self, event: str, data: dict): """Send to all dashboard connections; bypasses seq_log (dashboard resumes via full state refetch).""" - payload = json.dumps({"event": event, "data": data}) + payload = json.dumps({"event": event, "data": slim_status_data(event, data)}) dead: list[WebSocket] = [] for ws in list(self.global_connections): try: diff --git a/backend/apps/agents/manager/session/SessionLifecycle.py b/backend/apps/agents/manager/session/SessionLifecycle.py index 19ba4654..60eec353 100644 --- a/backend/apps/agents/manager/session/SessionLifecycle.py +++ b/backend/apps/agents/manager/session/SessionLifecycle.py @@ -153,6 +153,7 @@ class SessionLifecycle(AgentManagerProtocol): limit: int = 20, offset: int = 0, dashboard_id: Optional[str] = None, + closed_only: bool = False, ) -> Dict: """Return paginated, optionally filtered summaries of closed sessions.""" all_data = load_all_session_data() @@ -161,6 +162,9 @@ class SessionLifecycle(AgentManagerProtocol): q_lower = q.strip().lower() history = [] for sid, data in all_data: + # The boot fetch wants CLOSED sessions only: open ones landing in the client's history map made its resurrection gate swallow their terminal frames. Search keeps the full pool (open sessions on other dashboards are reachable nowhere else). + if closed_only and not data.get("closed_at"): + continue if dashboard_id and data.get("dashboard_id") != dashboard_id: continue if q_lower: diff --git a/backend/tests/test_ws_efficiency.py b/backend/tests/test_ws_efficiency.py new file mode 100644 index 00000000..2e1b749f --- /dev/null +++ b/backend/tests/test_ws_efficiency.py @@ -0,0 +1,155 @@ +"""WS efficiency batch: agent:status frames are slimmed (metadata + previews, never the +transcript), GET /sessions returns the WS seq cursor for resume seeding, and /history's +closed_only filter keeps open sessions out of the client's resurrection gate. Each pins a +live-proven failure mode: full transcripts on status frames were replayed stale and rolled +clients backwards, and an open session in the history map had its terminal frame swallowed.""" + +import asyncio + +import pytest +from fastapi.testclient import TestClient + +from backend.apps.agents.core.ws_manager import ws_manager, slim_status_data +from backend.apps.agents.core.seq_log import seq_log +from backend.apps.agents.agent_manager import agent_manager +from backend.apps.agents.core.models import AgentSession, Message +from backend.main import app + + +def p_client() -> TestClient: + import backend.auth as auth_mod + if not auth_mod.TOKEN: + import secrets + auth_mod.TOKEN = secrets.token_urlsafe(32) + return TestClient(app, headers={"Authorization": f"Bearer {auth_mod.TOKEN}"}) + + +def p_status_data(n_msgs: int = 3) -> dict: + session = { + "id": "s1", + "status": "completed", + "messages": ( + [{"role": "user", "content": "first user question here"}] + + [{"role": "assistant", "content": f"reply {i} " + "x" * 200} for i in range(n_msgs - 1)] + ), + "name": "t", + } + return {"session_id": "s1", "status": "completed", "session": session} + + +def test_slim_drops_transcript_and_adds_previews(): + out = slim_status_data("agent:status", p_status_data(3)) + sess = out["session"] + assert sess["messages"] == [] + assert sess["message_count"] == 3 + assert sess["first_user_message"].startswith("first user question") + assert sess["last_message_preview"].startswith("reply 1") + assert len(sess["last_message_preview"]) <= 120 + # Original input is not mutated (callers may reuse their dicts). + assert p_status_data(3)["session"]["messages"] != [] + + +def test_slim_leaves_non_status_and_messageless_frames_alone(): + msg_data = {"session_id": "s1", "message": {"role": "user", "content": "hello"}} + assert slim_status_data("agent:message", msg_data) is msg_data + no_sess = {"session_id": "s1", "status": "running"} + assert slim_status_data("agent:status", no_sess) is no_sess + + +class p_FakeWs: + def __init__(self) -> None: + self.frames: list = [] + + async def send_text(self, s: str) -> None: + import json + self.frames.append(json.loads(s)) + + +def test_send_to_session_slims_status_for_both_socket_kinds_and_the_replay_buffer(): + sess_ws, dash_ws = p_FakeWs(), p_FakeWs() + sid = "slimtest-session" + ws_manager.connections[sid] = [sess_ws] + ws_manager.global_connections.append(dash_ws) + try: + asyncio.run(ws_manager.send_to_session(sid, "agent:status", p_status_data(4))) + asyncio.run(ws_manager.send_to_session(sid, "agent:message", { + "session_id": sid, "message": {"role": "assistant", "content": "full text stays"}, + })) + for ws in (sess_ws, dash_ws): + status = ws.frames[0] + assert status["data"]["session"]["messages"] == [] + assert status["data"]["session"]["message_count"] == 4 + msg = ws.frames[1] + assert msg["data"]["message"]["content"] == "full text stays" + # Both sockets got the SAME stamped seq (the frontend dedupes on it). + assert sess_ws.frames[0]["seq"] == dash_ws.frames[0]["seq"] + # The ring buffer stores the slim frame, so replays are slim too. + _, _, events = seq_log.replay(sid, 0) + import json + assert json.loads(events[0])["data"]["session"]["messages"] == [] + finally: + ws_manager.connections.pop(sid, None) + ws_manager.global_connections.remove(dash_ws) + seq_log.clear(sid) + + +def test_get_session_returns_event_seq_cursor(): + s = AgentSession(name="t", model="sonnet") + s.messages = [Message(role="user", content="hi")] + agent_manager.sessions[s.id] = s + try: + asyncio.run(ws_manager.send_to_session(s.id, "agent:status", {"session_id": s.id, "status": "running"})) + asyncio.run(ws_manager.send_to_session(s.id, "agent:status", {"session_id": s.id, "status": "completed"})) + res = p_client().get(f"/api/agents/sessions/{s.id}") + assert res.status_code == 200 + body = res.json() + assert body["event_seq"] == seq_log.current_seq(s.id) == 2 + assert body["messages"][0]["content"] == "hi" + finally: + agent_manager.sessions.pop(s.id, None) + seq_log.clear(s.id) + + +def test_replay_caught_up_client_gets_nothing_not_the_terminal_frame(): + sid = "caughtup-session" + ws = p_FakeWs() + try: + asyncio.run(ws_manager.send_to_session(sid, "agent:status", {"session_id": sid, "status": "running"})) + asyncio.run(ws_manager.send_to_session(sid, "agent:status", {"session_id": sid, "status": "completed"})) + top = seq_log.current_seq(sid) + ack = asyncio.run(ws_manager.replay_to(sid, ws, top)) + assert ack == {"ok": True, "replayed": 0, "current_seq": top} + assert ws.frames == [] + # A behind client still gets the real replay. + ack2 = asyncio.run(ws_manager.replay_to(sid, ws, top - 1)) + assert ack2["replayed"] == 1 + finally: + seq_log.clear(sid) + + +def test_history_closed_only_filters_open_sessions(monkeypatch): + rows = [ + ("open1", {"id": "open1", "name": "open chat", "closed_at": None, "dashboard_id": None}), + ("closed1", {"id": "closed1", "name": "closed chat", "closed_at": "2026-07-01T00:00:00", "dashboard_id": None}), + ] + import backend.apps.agents.manager.session.SessionLifecycle as lifecycle_mod + monkeypatch.setattr(lifecycle_mod, "load_all_session_data", lambda: list(rows)) + closed = agent_manager.get_history(closed_only=True) + assert [s["id"] for s in closed["sessions"]] == ["closed1"] + # Search keeps the full pool: open sessions on other dashboards are reachable nowhere else. + everything = agent_manager.get_history() + assert {s["id"] for s in everything["sessions"]} == {"open1", "closed1"} + + +def test_history_route_threads_closed_only(monkeypatch): + seen = {} + + def p_spy(**kwargs): + seen.update(kwargs) + return {"sessions": [], "total": 0, "has_more": False} + + monkeypatch.setattr(agent_manager, "get_history", p_spy) + assert p_client().get("/api/agents/history?closed_only=1").status_code == 200 + assert seen["closed_only"] is True + assert p_client().get("/api/agents/history").status_code == 200 + assert seen["closed_only"] is False diff --git a/frontend/src/app/pages/AgentChat/AgentChat.tsx b/frontend/src/app/pages/AgentChat/AgentChat.tsx index 097f6dd5..a04f548c 100644 --- a/frontend/src/app/pages/AgentChat/AgentChat.tsx +++ b/frontend/src/app/pages/AgentChat/AgentChat.tsx @@ -47,7 +47,7 @@ import { displayChatTitle, isLegacyAutoName } from '@/shared/state/sessionDispla import { Typewriter } from '@/app/components/feedback/Animated'; import { store } from '@/shared/state/store'; import { fetchModes } from '@/shared/state/modesSlice'; -import { createSessionWs, acquireSessionWs, releaseSessionWs } from '@/shared/ws/WebSocketManager'; +import { createSessionWs, acquireSessionWs, releaseSessionWs, seedSessionSeq } from '@/shared/ws/WebSocketManager'; import StreamingBubble from './bubbles/StreamingBubble'; import WelcomeQuickReplies from './WelcomeQuickReplies'; import { useWelcomeGreeting } from './useWelcomeGreeting'; @@ -374,7 +374,12 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose dispatch(fetchSession(id)); } else { try { - await dispatch(fetchSession(id)); + const action = await dispatch(fetchSession(id)); + // Seed the resume cursor from the snapshot's seq so the connect below doesn't replay the whole ring buffer we just hydrated over REST. + if (fetchSession.fulfilled.match(action)) { + const seq = (action.payload as { event_seq?: number }).event_seq; + if (typeof seq === 'number') seedSessionSeq(id, seq); + } } catch { // Even if the REST hydrate fails, still connect, the WS resume protocol can hydrate from buffered events as a fallback. } diff --git a/frontend/src/shared/state/agentsSlice.ts b/frontend/src/shared/state/agentsSlice.ts index 6d21dcc6..29f51123 100644 --- a/frontend/src/shared/state/agentsSlice.ts +++ b/frontend/src/shared/state/agentsSlice.ts @@ -89,6 +89,8 @@ export interface AgentSession { last_message_preview?: string; first_user_message?: string; message_count?: number; + /** WS seq high-water at snapshot time (GET /sessions only); seeds the resume cursor so connect skips replaying what REST just delivered. */ + event_seq?: number; pending_approvals: ApprovalRequest[]; branches: Record; active_branch_id: string; @@ -530,7 +532,8 @@ export const deleteSession = createAsyncThunk( export const fetchHistory = createAsyncThunk( 'agents/fetchHistory', async ({ dashboardId }: { dashboardId?: string } = {}) => { - const params = new URLSearchParams({ limit: '10000' }); + // closed_only: an OPEN session landing in state.history made updateSession's resurrection gate swallow its terminal frames (card stuck running, final answer invisible). Search (searchHistory) keeps the full pool. + const params = new URLSearchParams({ limit: '10000', closed_only: '1' }); if (dashboardId) params.set('dashboard_id', dashboardId); const res = await fetch(`${AGENTS_API}/history?${params}`); const data = await res.json(); @@ -717,7 +720,8 @@ const agentsSlice = createSlice({ if (state.history[action.payload.id]) { if (action.payload.status === 'running' || action.payload.mode === 'browser-agent') { delete state.history[action.payload.id]; - } else { + } else if (!state.sessions[action.payload.id]) { + // Gate only truly-closed sessions (no live card): a late frame must not resurrect them. A LIVE session that leaked into history used to have its completed frame swallowed here, leaving the card stuck running. return; } } diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts index 1211a751..320262f0 100644 --- a/frontend/src/shared/ws/WebSocketManager.ts +++ b/frontend/src/shared/ws/WebSocketManager.ts @@ -324,6 +324,20 @@ class WebSocketManager { } } + // Cross-socket dedupe: the backend fans every session frame out to BOTH the dashboard socket and the chat's own socket (same stamped seq), so an expanded chat parsed and reduced everything twice, and whichever copy landed second could be a replayed stale one. Time-windowed rather than a high-water mark so a deliberate later replay (gap recovery resets lastSeq to 0) is never starved. + if (typeof msg.seq === 'number' && session_id) { + const key = `${session_id}:${msg.seq}`; + const now = Date.now(); + const seen = _recentFrameTimes.get(key); + if (seen !== undefined && now - seen < FRAME_DEDUPE_WINDOW_MS) return; + _recentFrameTimes.set(key, now); + if (_recentFrameTimes.size > 4000) { + for (const [k, t] of _recentFrameTimes) { + if (now - t >= FRAME_DEDUPE_WINDOW_MS) _recentFrameTimes.delete(k); + } + } + } + // ----- Connection-scoped frames (no business-logic side effects) ----- if (event === 'server:pong') { @@ -351,6 +365,10 @@ class WebSocketManager { // Reset lastSeq, the REST refetch is the new authoritative baseline; subsequent server events with seq numbers will re-establish the high-water mark. Also wipe the cross-mount persistent map so a remount during this gap window doesn't resurrect the stale value. this.lastSeq = 0; _sessionLastSeq.delete(session_id); + // The recovery replay re-delivers seqs possibly seen moments ago; drop them from the dedupe window so it's never starved. + for (const k of _recentFrameTimes.keys()) { + if (k.startsWith(`${session_id}:`)) _recentFrameTimes.delete(k); + } } return; } @@ -947,6 +965,17 @@ export const dashboardWs = new WebSocketManager(`${WS_BASE}/ws/dashboard`, { ski // Per-session high-water mark for the resume protocol. Survives across AgentChat mounts/unmounts so reopening a chat doesn't re-trigger a full replay from the server's ring buffer. Why this exists: AgentChat uses `key={session.id}` on the embedded instance inside AgentCard, so every expand/collapse remounts the component, which constructs a fresh WebSocketManager. Without this persistent map, each fresh manager starts at last_seq=0 and asks the server for the entire buffered history. The server faithfully replays it, the client renders the typewriter animation again, and the user sees their completed chat "type itself out" on every reopen. Lifetime: tied to the JS module load, which means the page tab. Lost on full app reload (intentional, that should re-hydrate from REST). On backend restart the buffers are wiped anyway, so a stale lastSeq pointing past the buffer top falls into the "fresh client" path on the server (last_seq>0 but no buffer) which short-circuits to a no-op replay. Safe. const _sessionLastSeq: Map = new Map(); +// (session_id:seq) -> arrival time; entries older than the window are prunable. Bounded by event rate x window, not session count. +const FRAME_DEDUPE_WINDOW_MS = 5_000; +const _recentFrameTimes: Map = new Map(); + +/** Seed the resume cursor from a REST hydrate (GET /sessions returns event_seq), so the follow-up WS connect replays only what happened AFTER the snapshot instead of the whole ring buffer the client just received as JSON. Never lowers an existing high-water mark. */ +export function seedSessionSeq(sessionId: string, seq: number): void { + if (typeof seq !== 'number' || seq <= 0) return; + const cur = _sessionLastSeq.get(sessionId) ?? 0; + if (seq > cur) _sessionLastSeq.set(sessionId, seq); +} + export function createSessionWs(sessionId: string): WebSocketManager { return new WebSocketManager(`${WS_BASE}/ws/agents/${sessionId}`, { sessionId }); }