diff --git a/studio/backend/routes/research_runs.py b/studio/backend/routes/research_runs.py index 62ee3cce4c..1880936ff6 100644 --- a/studio/backend/routes/research_runs.py +++ b/studio/backend/routes/research_runs.py @@ -24,6 +24,7 @@ from storage.studio_db import get_chat_message, get_chat_thread, upsert_chat_mes router = APIRouter() _SENSITIVE_KEY = re.compile(r"^(?:api.?key|secret|token|authorization|password)$", re.IGNORECASE) _MAX_PLAN_STEPS = 30 +_DELTA_ONLY_EVENTS = {"reasoning.updated", "report.updated"} class CreateResearchRun(BaseModel): @@ -389,7 +390,8 @@ async def research_events( cursor = int(event["seq"]) event_data = dict(event["data"]) event_data["createdAt"] = event["createdAt"] - event_data["run"] = snapshot + if event["type"] not in _DELTA_ONLY_EVENTS: + event_data["run"] = snapshot data = json.dumps(event_data, separators = (",", ":"), ensure_ascii = False) yield f"id: {cursor}\nevent: {event['type']}\ndata: {data}\n\n" if snapshot["status"] in db.TERMINAL_STATUSES and cursor >= int( diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 530cf556f4..accb29a297 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -1498,6 +1498,11 @@ def test_terminal_sse_event_contains_report_and_complete_snapshot(research_home) research_db.upsert_source( "run-1", 0, "https://example.com/final", "Final source", "Final evidence" ) + research_db.append_event( + "run-1", + "report.updated", + {"delta": "Draft chunk", "offset": 0, "length": 11}, + ) report = "# Durable report\n\nFinal markdown." assert ( research_db.finish("run-1", "worker-1", "completed", event_payload = {"report": report}) @@ -1525,6 +1530,11 @@ def test_terminal_sse_event_contains_report_and_complete_snapshot(research_home) return "".join(chunks) stream = asyncio.run(consume()) + delta = next(block for block in stream.split("\n\n") if "event: report.updated" in block) + delta_line = next(line for line in delta.splitlines() if line.startswith("data: ")) + delta_payload = json.loads(delta_line[6:]) + assert delta_payload["delta"] == "Draft chunk" + assert "run" not in delta_payload terminal = next(block for block in stream.split("\n\n") if "event: run.completed" in block) data_line = next(line for line in terminal.splitlines() if line.startswith("data: ")) payload = json.loads(data_line[6:]) diff --git a/studio/frontend/src/features/chat/api/research-api.ts b/studio/frontend/src/features/chat/api/research-api.ts index 278f414733..c56377b23b 100644 --- a/studio/frontend/src/features/chat/api/research-api.ts +++ b/studio/frontend/src/features/chat/api/research-api.ts @@ -8,6 +8,11 @@ import type { ResearchRun, } from "../types/research"; +type StreamResearchEvent = Omit & { + data: Omit; + run?: ResearchRun; +}; + type JsonObject = Record; const TERMINAL_RESEARCH_STATUSES = new Set([ "completed", @@ -139,7 +144,7 @@ export async function* streamResearchEvents( id: string, after: number, signal?: AbortSignal, -): AsyncGenerator { +): AsyncGenerator { const response = await authFetch( `/api/chat/research-runs/${id}/events?after=${Math.max(0, after)}`, { headers: { accept: "text/event-stream" }, signal }, @@ -176,18 +181,16 @@ export async function* streamResearchEvents( if (data.length > 0) { const parsed = camelize(JSON.parse(data.join("\n"))) as JsonObject; const candidate = parsed.run as ResearchRun | undefined; - if (candidate?.id && candidate.status) { - yield { - id: eventId, - event: event as ResearchEvent["event"], - createdAt: - typeof parsed.createdAt === "number" - ? parsed.createdAt - : candidate.updatedAt, - data: parsed as unknown as ResearchEvent["data"], - run: candidate, - }; - } + yield { + id: eventId, + event: event as ResearchEvent["event"], + createdAt: + typeof parsed.createdAt === "number" + ? parsed.createdAt + : (candidate?.updatedAt ?? Date.now()), + data: parsed as unknown as StreamResearchEvent["data"], + ...(candidate?.id && candidate.status ? { run: candidate } : {}), + }; } boundary = buffer.indexOf("\n\n"); } @@ -272,20 +275,31 @@ export async function* followResearchRun( ) { return; } + let currentRun: ResearchRun = run; let cursor = replayFrom ?? run.lastEventSeq; while (!signal?.aborted) { try { for await (const event of streamResearchEvents(id, cursor, signal)) { cursor = Math.max(cursor, event.id); - run = event.run; + const eventRun: ResearchRun = event.run ?? { + ...currentRun, + lastEventSeq: Math.max(currentRun.lastEventSeq, event.id), + updatedAt: Math.max(currentRun.updatedAt, event.createdAt), + }; + const hydratedEvent: ResearchEvent = { + ...event, + data: { ...event.data, run: eventRun }, + run: eventRun, + }; + currentRun = eventRun; failures = 0; - yield { run, event, source: "event" }; + yield { run: currentRun, event: hydratedEvent, source: "event" }; if ( - (event.event === "run.completed" || - event.event === "run.failed" || - event.event === "run.cancelled") && - TERMINAL_RESEARCH_STATUSES.has(event.run.status) && - (event.data.attempt ?? 0) === (event.run.retryCount ?? 0) + (hydratedEvent.event === "run.completed" || + hydratedEvent.event === "run.failed" || + hydratedEvent.event === "run.cancelled") && + TERMINAL_RESEARCH_STATUSES.has(eventRun.status) && + (hydratedEvent.data.attempt ?? 0) === (eventRun.retryCount ?? 0) ) { return; } @@ -306,21 +320,21 @@ export async function* followResearchRun( try { const fresh = await getResearchRun(id, signal); const changed = - fresh.lastEventSeq !== run.lastEventSeq || - fresh.updatedAt !== run.updatedAt || - fresh.status !== run.status || - fresh.report !== run.report; + fresh.lastEventSeq !== currentRun.lastEventSeq || + fresh.updatedAt !== currentRun.updatedAt || + fresh.status !== currentRun.status || + fresh.report !== currentRun.report; const needsCatchup = cursor < fresh.lastEventSeq; - run = fresh; + currentRun = fresh; if (replayFrom === undefined) { cursor = Math.max(cursor, fresh.lastEventSeq); } if (changed || needsCatchup) { - yield { run, source: "snapshot" }; + yield { run: currentRun, source: "snapshot" }; } if ( - TERMINAL_RESEARCH_STATUSES.has(run.status) && - cursor >= run.lastEventSeq + TERMINAL_RESEARCH_STATUSES.has(currentRun.status) && + cursor >= currentRun.lastEventSeq ) { return; } diff --git a/studio/frontend/src/features/chat/stores/research-run-store.ts b/studio/frontend/src/features/chat/stores/research-run-store.ts index bf239803ce..021f21736d 100644 --- a/studio/frontend/src/features/chat/stores/research-run-store.ts +++ b/studio/frontend/src/features/chat/stores/research-run-store.ts @@ -706,6 +706,9 @@ export function ingestResearchUpdate( } const pending = pendingStreamEvents.get(run.id); + if (pending && event.id <= pending.event.id) { + return; + } if (pending && canCoalesceStreamEvent(pending.event, event)) { const reasoningDelta = event.event === "reasoning.updated" diff --git a/tests/studio/test_deep_research_frontend_contract.py b/tests/studio/test_deep_research_frontend_contract.py index e6ba96ff80..767f3589a2 100644 --- a/tests/studio/test_deep_research_frontend_contract.py +++ b/tests/studio/test_deep_research_frontend_contract.py @@ -11,6 +11,7 @@ def source(path: str) -> str: def test_research_api_is_isolated_and_cursor_based() -> None: api = source("features/chat/api/research-api.ts") + store = source("features/chat/stores/research-run-store.ts") assert 'authFetch("/api/chat/research-runs"' in api assert "authFetch(`/api/chat/research-runs/active?${query}`)" in api assert "const { runs, hasRun }" in api @@ -22,11 +23,12 @@ def test_research_api_is_isolated_and_cursor_based() -> None: assert "Math.min(8_000, 500 * 2 ** (failures - 1))" in api assert "for await (const event of streamResearchEvents" in api assert 'source: "event"' in api - assert "fresh.report !== run.report" in api + assert "fresh.report !== currentRun.report" in api assert "await waitForReconnect(" in api assert "while (!(run || signal?.aborted))" in api assert "isPermanentResearchError(error)" in api assert 'yield { run, source: "snapshot" }' in api + assert "event.id <= pending.event.id" in store for action in ("cancel", "retry"): assert f'mutate(id, "{action}")' in api assert 'mutate(id, "approve", { planRevision, planHash })' in api