diff --git a/frontend/src/core/api/api-client.ts b/frontend/src/core/api/api-client.ts index 7408c976c..bda91372c 100644 --- a/frontend/src/core/api/api-client.ts +++ b/frontend/src/core/api/api-client.ts @@ -86,6 +86,101 @@ type StreamPart = { data: unknown; }; +type ReconnectInputSnapshot = Record & { + messages: unknown[]; +}; + +function streamOptionSignal(options: unknown): AbortSignal | undefined { + if (typeof AbortSignal === "undefined") { + return undefined; + } + if (options instanceof AbortSignal) { + return options; + } + if (typeof options !== "object" || options === null) { + return undefined; + } + const signal = Reflect.get(options, "signal"); + return signal instanceof AbortSignal ? signal : undefined; +} + +/** + * Recover the submitted input before replaying an active run. The incremental + * chat stream intentionally omits `values`, so a page reload can otherwise + * receive the run's AI/tool chunks before its human message has reached the + * durable history feed. `runs.get` retains the original graph input in + * `kwargs.input`; merge it into the latest durable values for one synthetic + * snapshot. Any read failure is deliberately ignored so reconnect semantics + * remain unchanged for deployments without run metadata. + */ +async function loadReconnectInputSnapshot( + client: LangGraphClient, + threadId: string, + runId: string, + run?: Awaited>, + durableValues?: unknown, + signal?: AbortSignal, +): Promise { + try { + const resolvedRun = + run ?? (await client.runs.get(threadId, runId, { signal })); + const runKwargs = Reflect.get(resolvedRun, "kwargs"); + const input = + typeof runKwargs === "object" && runKwargs !== null + ? Reflect.get(runKwargs, "input") + : undefined; + const inputMessages = + typeof input === "object" && input !== null + ? Reflect.get(input, "messages") + : undefined; + if (!Array.isArray(inputMessages) || inputMessages.length === 0) { + return undefined; + } + + const resolvedDurableValues = + durableValues ?? + (await client.threads.getState(threadId, undefined, { signal })).values; + const normalizedDurableValues = + typeof resolvedDurableValues === "object" && + resolvedDurableValues !== null + ? resolvedDurableValues + : {}; + const durableMessages = Array.isArray( + Reflect.get(normalizedDurableValues, "messages"), + ) + ? (Reflect.get(normalizedDurableValues, "messages") as unknown[]) + : []; + const seenIds = new Set( + durableMessages.flatMap((message) => { + const id = + typeof message === "object" && message !== null + ? Reflect.get(message, "id") + : undefined; + return typeof id === "string" && id.length > 0 ? [id] : []; + }), + ); + const messages = [ + ...durableMessages, + ...inputMessages.filter((message) => { + const id = + typeof message === "object" && message !== null + ? Reflect.get(message, "id") + : undefined; + if (typeof id !== "string" || id.length === 0) return true; + if (seenIds.has(id)) return false; + seenIds.add(id); + return true; + }), + ]; + return { ...normalizedDurableValues, messages } as ReconnectInputSnapshot; + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + throw error; + } + return undefined; + } +} + export class StreamReplayGapError extends Error { constructor( readonly gap: StreamReplayGapData, @@ -184,26 +279,25 @@ export function isRunNotCancellableError(error: unknown): boolean { } /** - * Preflight a reconnect: if the run already reached a terminal state, there is - * nothing to rejoin. Returns ``true`` when the caller should skip the - * underlying ``joinStream`` so the SDK's ``onSuccess`` path runs and - * ``isLoading`` flips back to false — instead of blocking forever on a drained - * stream bridge. + * Preflight a reconnect and return the run record when it can be read. A + * missing record or failed request returns ``undefined`` so a legitimately + * active reconnect falls back to the original join and the terminal-state + * check remains owned by the caller. * * Any error (404 for an evicted record, network blip, auth hiccup, …) falls * back to the original join so a legitimately active reconnect is never * silently suppressed. */ -async function shouldSkipReconnect( +async function getReconnectRun( client: LangGraphClient, threadId: string, runId: string, -): Promise { + signal?: AbortSignal, +): Promise> | undefined> { try { - const run = await client.runs.get(threadId, runId); - return TERMINAL_RUN_STATUSES.has(run.status); + return await client.runs.get(threadId, runId, { signal }); } catch { - return false; + return undefined; } } @@ -243,12 +337,16 @@ async function* recoverStreamReplayGaps({ expectedRunId, initialStream, resume, + signal, + reconnectRun, }: { client: LangGraphClient; threadId: string | null | undefined; expectedRunId: () => string | undefined; initialStream: AsyncIterable; resume: (runId: string, lastEventId?: string) => AsyncIterable; + signal?: AbortSignal; + reconnectRun?: Awaited>; }): AsyncGenerator { let stream = initialStream; let recoveryAttempts = 0; @@ -289,12 +387,31 @@ async function* recoverStreamReplayGaps({ }; const durableState = await client.threads - .getState(threadId) + .getState(threadId, undefined, { signal }) .catch((error: unknown) => { + if (error instanceof Error && error.name === "AbortError") { + throw error; + } throw new StreamReplayGapError(gap, recoveryAttempts, error); }); if (durableState.values != null) { - yield { event: "values", data: durableState.values }; + // A gap can arrive after the initial hydration frame but before the + // input reaches the checkpoint. Rebuild the snapshot from run metadata + // so this recovery path cannot overwrite the rescued human message. + const recoveredSnapshot = reconnectRun + ? await loadReconnectInputSnapshot( + client, + threadId, + runId, + reconnectRun, + durableState.values, + signal, + ) + : undefined; + yield { + event: "values", + data: recoveredSnapshot ?? durableState.values, + }; } rememberReconnectRun(threadId, runId); @@ -356,6 +473,8 @@ function createCompatibleClient(isMock?: boolean): LangGraphClient { threadId, expectedRunId: () => runId, initialStream, + signal: streamOptionSignal(sanitizedPayload), + reconnectRun: undefined, resume: (resolvedRunId, lastEventId) => { // Keep the recovery run id available to the shared inactive-stream // handler even if the SDK omitted its onRunCreated callback. @@ -398,10 +517,31 @@ function createCompatibleClient(isMock?: boolean): LangGraphClient { // reload after the backend's stream bridge is reaped blocks forever on a // drained condition variable, pinning ``isLoading`` true so the first // post-reload message is routed to ``stop`` instead of ``submit``. - if (threadId && (await shouldSkipReconnect(client, threadId, runId))) { + const reconnectSignal = streamOptionSignal(options); + const reconnectRun = threadId + ? await getReconnectRun(client, threadId, runId, reconnectSignal) + : undefined; + if (reconnectRun && TERMINAL_RUN_STATUSES.has(reconnectRun.status)) { clearReconnectRun(threadId, runId); return; } + if (threadId && reconnectRun) { + const reconnectSnapshot = await loadReconnectInputSnapshot( + client, + threadId, + runId, + reconnectRun, + undefined, + reconnectSignal, + ); + if (reconnectSnapshot) { + // This is an internal hydration frame. The requested network stream + // remains incremental; the SDK receives the current input before any + // replayed messages-tuple AI/tool chunks and deduplicates it against + // later history or stream copies by message id. + yield { event: "values", data: reconnectSnapshot }; + } + } const sanitizedOptions = forceChatRunStreamOptions(options); yield* handleInactiveRunStream({ threadId, @@ -411,6 +551,8 @@ function createCompatibleClient(isMock?: boolean): LangGraphClient { threadId, expectedRunId: () => runId, initialStream: originalJoinStream(threadId, runId, sanitizedOptions), + signal: reconnectSignal, + reconnectRun, resume: (resolvedRunId, lastEventId) => originalJoinStream(threadId, resolvedRunId, { ...sanitizedOptions, diff --git a/frontend/tests/unit/core/api/api-client.test.ts b/frontend/tests/unit/core/api/api-client.test.ts index 9bb4a0771..f4aea4283 100644 --- a/frontend/tests/unit/core/api/api-client.test.ts +++ b/frontend/tests/unit/core/api/api-client.test.ts @@ -242,6 +242,115 @@ test("short-circuits reconnect to a terminal run", async () => { expect(sessionStorage.removeItem).toHaveBeenCalledWith("lg:stream:thread-1"); }); +test("hydrates the active run input before replaying an incremental stream", async () => { + const sessionStorage = makeSessionStorage(); + const fetchFn = rs.fn(async (url: string | URL) => { + const path = new URL(url.toString()).pathname; + if (path.endsWith("/runs/run-input")) { + return new Response( + JSON.stringify({ + status: "running", + kwargs: { + input: { + messages: [ + { id: "human-2", type: "human", content: "Second question" }, + ], + }, + }, + }), + { status: 200 }, + ); + } + if (path.endsWith("/threads/thread-input/state")) { + return new Response( + JSON.stringify({ + values: { + messages: [ + { id: "human-1", type: "human", content: "First question" }, + { id: "human-2", type: "human", content: "Second question" }, + ], + }, + }), + { status: 200 }, + ); + } + if (path.endsWith("/runs/run-input/stream")) { + return makeSSEResponse("event: end\ndata: null\n\n"); + } + return new Response(JSON.stringify({ detail: "unexpected request" }), { + status: 500, + }); + }); + rs.stubGlobal("window", { + location: { origin: "http://localhost:2026" }, + sessionStorage, + }); + rs.stubGlobal("fetch", fetchFn); + + const entries: Array<{ event: string; data: unknown }> = []; + for await (const entry of getAPIClient(true).runs.joinStream( + "thread-input", + "run-input", + )) { + entries.push(entry); + } + + expect(entries[0]).toMatchObject({ + event: "values", + data: { + messages: [ + { id: "human-1", content: "First question" }, + { id: "human-2", content: "Second question" }, + ], + }, + }); + expect( + (entries[0]?.data as { messages: Array<{ id: string }> }).messages, + ).toHaveLength(2); +}); + +test("continues reconnect when durable state hydration fails", async () => { + const fetchFn = rs.fn(async (url: string | URL) => { + const path = new URL(url.toString()).pathname; + if (path.endsWith("/runs/run-no-state")) { + return new Response( + JSON.stringify({ + status: "running", + kwargs: { + input: { + messages: [{ id: "human-2", type: "human", content: "Second" }], + }, + }, + }), + { status: 200 }, + ); + } + if (path.endsWith("/threads/thread-no-state/state")) { + return new Response(JSON.stringify({ detail: "state unavailable" }), { + status: 404, + }); + } + if (path.endsWith("/runs/run-no-state/stream")) { + return makeSSEResponse("event: end\ndata: null\n\n"); + } + return new Response(JSON.stringify({ detail: "unexpected request" }), { + status: 500, + }); + }); + rs.stubGlobal("fetch", fetchFn); + + const entries: Array<{ event: string; data: unknown }> = []; + for await (const entry of getAPIClient(true).runs.joinStream( + "thread-no-state", + "run-no-state", + )) { + entries.push(entry); + } + + expect(entries).toEqual([{ event: "end", data: null }]); + expect(fetchFn).toHaveBeenCalledTimes(3); +}); + test("falls back to join when preflight cannot resolve the run", async () => { const sessionStorage = makeSessionStorage(); sessionStorage.setItem("lg:stream:thread-1", "run-1"); @@ -437,9 +546,25 @@ test("recovers a join stream gap from durable state and resumes after the retain const fetchFn = rs.fn(async (url: string | URL, init?: RequestInit) => { const path = url.toString(); if (path.endsWith("/runs/run-1")) { - return new Response(JSON.stringify({ status: "running" }), { - status: 200, - }); + return new Response( + JSON.stringify({ + status: "running", + kwargs: { + input: { + messages: [ + { + id: "human-2", + type: "human", + content: "Second question", + }, + ], + }, + }, + }), + { + status: 200, + }, + ); } if (path.includes("/runs/run-1/stream")) { recoveryRequests.push(init ?? {}); @@ -451,7 +576,11 @@ test("recovers a join stream gap from durable state and resumes after the retain if (path.includes("/threads/thread-1/state")) { return new Response( JSON.stringify({ - values: { messages: [{ type: "ai", content: "durable" }] }, + values: { + messages: [ + { id: "human-1", type: "human", content: "First question" }, + ], + }, next: [], tasks: [], metadata: {}, @@ -482,13 +611,27 @@ test("recovers a join stream gap from durable state and resumes after the retain } expect(received).toEqual([ + { + event: "values", + data: { + messages: [ + { id: "human-1", type: "human", content: "First question" }, + { id: "human-2", type: "human", content: "Second question" }, + ], + }, + }, { event: "custom", data: { type: "stream_replay_gap", ...gap }, }, { event: "values", - data: { messages: [{ type: "ai", content: "durable" }] }, + data: { + messages: [ + { id: "human-1", type: "human", content: "First question" }, + { id: "human-2", type: "human", content: "Second question" }, + ], + }, }, { event: "end", data: null }, ]);