fix(frontend): restore user input after stream reconnect (#5428)

* fix(frontend): restore input on stream reconnect

* fix(frontend): close reconnect review feedback

* fix(frontend): address reconnect review feedback
This commit is contained in:
Xuehao Xu 2026-09-15 08:16:33 +08:00 committed by GitHub
parent b46fb476ed
commit d1f77fc1f8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 303 additions and 18 deletions

View File

@ -86,6 +86,101 @@ type StreamPart = {
data: unknown;
};
type ReconnectInputSnapshot = Record<string, unknown> & {
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<ReturnType<LangGraphClient["runs"]["get"]>>,
durableValues?: unknown,
signal?: AbortSignal,
): Promise<ReconnectInputSnapshot | undefined> {
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<boolean> {
signal?: AbortSignal,
): Promise<Awaited<ReturnType<LangGraphClient["runs"]["get"]>> | 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<StreamPart>;
resume: (runId: string, lastEventId?: string) => AsyncIterable<StreamPart>;
signal?: AbortSignal;
reconnectRun?: Awaited<ReturnType<LangGraphClient["runs"]["get"]>>;
}): AsyncGenerator<StreamPart> {
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,

View File

@ -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 },
]);