Studio: slim research stream deltas
This commit is contained in:
alkinun
parent
179a16a4d9
commit
73d6e64453
5 changed files with 61 additions and 30 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:])
|
||||
|
|
|
|||
|
|
@ -8,6 +8,11 @@ import type {
|
|||
ResearchRun,
|
||||
} from "../types/research";
|
||||
|
||||
type StreamResearchEvent = Omit<ResearchEvent, "data" | "run"> & {
|
||||
data: Omit<ResearchEvent["data"], "run">;
|
||||
run?: ResearchRun;
|
||||
};
|
||||
|
||||
type JsonObject = Record<string, unknown>;
|
||||
const TERMINAL_RESEARCH_STATUSES = new Set([
|
||||
"completed",
|
||||
|
|
@ -139,7 +144,7 @@ export async function* streamResearchEvents(
|
|||
id: string,
|
||||
after: number,
|
||||
signal?: AbortSignal,
|
||||
): AsyncGenerator<ResearchEvent> {
|
||||
): AsyncGenerator<StreamResearchEvent> {
|
||||
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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue