diff --git a/ui/src/api/chat/a2aGrpcChatClient.test.ts b/ui/src/api/chat/a2aGrpcChatClient.test.ts index ee163f406c..0ffedd3768 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.test.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.test.ts @@ -12,8 +12,8 @@ * against a live controller, and none of them is implied by the protocol. */ -import { afterEach, describe, expect, it } from "vitest"; -import { createRouterTransport } from "@connectrpc/connect"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { Code, ConnectError, createRouterTransport } from "@connectrpc/connect"; import type { ConnectRouter } from "@connectrpc/connect"; import { fromJson } from "@bufbuild/protobuf"; import { ValueSchema } from "@bufbuild/protobuf/wkt"; @@ -1561,3 +1561,231 @@ describe("A2AGrpcChatClient files", () => { ]); }); }); + +/** + * A stream cut by something between the browser and the controller. + * + * Measured behind a proxy that closes streams after 60 s without data: the model + * read a long prompt silently, the browser saw `network error`, and the turn was + * shown as failed while the task finished on the server. + */ +describe("A2AGrpcChatClient.send, when the stream breaks", () => { + const completed = { + payload: { + case: "task" as const, + value: { + id: "task-1", + contextId: CONVERSATION.contextId, + status: { state: TaskState.COMPLETED }, + artifacts: [{ artifactId: "answer", parts: [text("checkout is out of memory")] }], + }, + }, + }; + + /** What `SubscribeToTask` opens with while the task is still running. */ + const snapshot = { + payload: { + case: "task" as const, + value: { + id: "task-1", + contextId: CONVERSATION.contextId, + status: { state: TaskState.WORKING }, + }, + }, + }; + + const broken = () => new ConnectError("network error", Code.Unavailable); + + // Re-attaching waits between attempts and judges a stream by how long it lasted, + // so the clock is driven by the test rather than waited out. + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + async function run( + sent: () => AsyncGenerator, + subscribe: (attempt: number) => AsyncGenerator = async function* () {}, + ) { + const subscribed: unknown[] = []; + serve(({ service }) => { + service(A2AService, { + sendStreamingMessage: sent as never, + subscribeToTask: async function* (request) { + subscribed.push({ tenant: request.tenant, id: request.id }); + for await (const frame of subscribe(subscribed.length)) yield frame as never; + }, + }); + }); + const events: ChatEvent[] = []; + let error: unknown; + let done = false; + const turn = (async () => { + try { + for await (const event of new A2AGrpcChatClient().send({ + conversation: CONVERSATION, + text: "why is checkout crashlooping?", + })) { + events.push(event); + } + } catch (cause) { + error = cause; + } finally { + done = true; + } + })(); + // Ten simulated minutes at most, so a loop that never ends fails the test + // instead of hanging it. + for (let second = 0; !done && second < 600; second += 1) { + await vi.advanceTimersByTimeAsync(1_000); + } + expect(done).toBe(true); + await turn; + return { events, error, subscribed }; + } + + it("follows the task it already named instead of failing the turn", async () => { + const { events, error, subscribed } = await run( + async function* () { + yield statusFrame({}); + throw new ConnectError("network error", Code.Unknown); + }, + async function* () { + yield completed; + }, + ); + expect(error).toBeUndefined(); + expect(subscribed).toEqual([{ tenant: CONVERSATION.agent, id: "task-1" }]); + expect(events.at(-1)).toMatchObject({ type: "status", state: "completed", taskId: "task-1" }); + expect(transcript(events).map(textOf)).toContain("checkout is out of memory"); + }); + + it("does not re-attach before the stream has named a task", async () => { + const { error, subscribed } = await run(async function* () { + yield* []; + throw new ConnectError("network error", Code.Unknown); + }); + expect(error).toBeDefined(); + expect(subscribed).toEqual([]); + }); + + it("surfaces an error the controller decided", async () => { + const { error, subscribed } = await run(async function* () { + yield statusFrame({}); + throw new ConnectError("dispatch busy", Code.FailedPrecondition); + }); + expect(error).toMatchObject({ message: "dispatch busy" }); + expect(subscribed).toEqual([]); + }); + + it("gives up after repeated breaks with nothing in between", async () => { + const { error, subscribed } = await run( + async function* () { + yield statusFrame({}); + throw new ConnectError("network error", Code.Unknown); + }, + async function* () { + yield* []; + throw new ConnectError("network error", Code.Unavailable); + }, + ); + expect(error).toBeDefined(); + expect(subscribed).toHaveLength(3); + }); + + // The server opens every re-attached stream with the task as it stands, so a + // path that breaks right after that frame must still run out of attempts. The + // completed task after twenty attempts only stops a loop that never gives up. + it("gives up when each re-attach breaks right after the task's snapshot", async () => { + const { error, subscribed } = await run( + async function* () { + yield statusFrame({}); + throw broken(); + }, + async function* (attempt) { + if (attempt > 20) { + yield completed; + return; + } + yield snapshot; + throw broken(); + }, + ); + expect(error).toBeDefined(); + expect(subscribed).toHaveLength(3); + }); + + it("keeps following a task whose re-attached streams are cut for being quiet", async () => { + const { events, error, subscribed } = await run( + async function* () { + yield statusFrame({}); + throw broken(); + }, + async function* (attempt) { + if (attempt > 4) { + yield completed; + return; + } + yield snapshot; + await new Promise((resolve) => setTimeout(resolve, 60_000)); + throw broken(); + }, + ); + expect(error).toBeUndefined(); + expect(subscribed).toHaveLength(5); + expect(events.at(-1)).toMatchObject({ type: "status", state: "completed", taskId: "task-1" }); + }); + + it("starts counting again after a re-attach that brought news", async () => { + const { events, error, subscribed } = await run( + async function* () { + yield statusFrame({}); + throw broken(); + }, + async function* (attempt) { + if (attempt > 4) { + yield completed; + return; + } + yield snapshot; + yield statusFrame({}); + throw broken(); + }, + ); + expect(error).toBeUndefined(); + expect(subscribed).toHaveLength(5); + expect(events.at(-1)).toMatchObject({ type: "status", state: "completed", taskId: "task-1" }); + }); + + it("stops waiting to re-attach when the reader presses stop", async () => { + const subscribed: unknown[] = []; + serve(({ service }) => { + service(A2AService, { + sendStreamingMessage: async function* () { + yield statusFrame({}) as never; + throw broken(); + }, + subscribeToTask: async function* () { + subscribed.push(true); + yield completed as never; + }, + }); + }); + const stop = new AbortController(); + const events: ChatEvent[] = []; + const turn = (async () => { + for await (const event of new A2AGrpcChatClient().send({ + conversation: CONVERSATION, + text: "why is checkout crashlooping?", + signal: stop.signal, + })) { + events.push(event); + } + })(); + const outcome = expect(turn).rejects.toMatchObject({ name: "AbortError" }); + // The stream has broken and the first re-attach is waiting its 500 ms. + await vi.advanceTimersByTimeAsync(100); + stop.abort(); + await vi.advanceTimersByTimeAsync(1_000); + await outcome; + expect(subscribed).toEqual([]); + }); +}); diff --git a/ui/src/api/chat/a2aGrpcChatClient.ts b/ui/src/api/chat/a2aGrpcChatClient.ts index d4f239af32..0b3642597b 100644 --- a/ui/src/api/chat/a2aGrpcChatClient.ts +++ b/ui/src/api/chat/a2aGrpcChatClient.ts @@ -47,6 +47,7 @@ import { create, toJson, type JsonObject } from "@bufbuild/protobuf"; import { ValueSchema } from "@bufbuild/protobuf/wkt"; +import { Code, ConnectError } from "@connectrpc/connect"; import { A2AService, MessageSchema, @@ -98,6 +99,56 @@ const SHARE_HEADER = "X-Share-Token"; */ const HISTORY_PAGE_LIMIT = 50; +/** Consecutive broken streams that came to nothing before a turn is given up. */ +const REATTACH_LIMIT = 3; + +/** + * How long a stream has to stay open for its break to count as a quiet stream + * being cut rather than a failure. + * + * A proxy closes an idle stream after its timeout, commonly 30 s or more, so a + * stream cut that late was working and the turn keeps following its task. One + * that breaks sooner with nothing new in it counts towards `REATTACH_LIMIT`. + */ +const REATTACH_HEALTHY_MS = 10_000; + +/** The wait before the first re-attach after a failure, doubled for each one after. */ +const REATTACH_BACKOFF_MS = 500; + +/** Resolves after `ms`, or rejects with the abort as soon as `signal` fires. */ +function pause(ms: number, signal: AbortSignal | undefined): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason); + return; + } + const onAbort = () => { + clearTimeout(timer); + reject(signal?.reason); + }; + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + signal?.addEventListener("abort", onAbort, { once: true }); + }); +} + +type StreamFrame = { payload: { case?: string; value?: unknown } }; + +/** + * Whether an error is the connection failing rather than the server answering. + * + * A proxy closing a quiet stream reaches the browser as a bare `network error` or + * a body that ends without its trailer — both `Unknown` — or as an HTTP 502/503/504, + * which is `Unavailable`. Anything the controller decided (a refusal, a missing + * permission) arrives with its own code and is not retried. + */ +function isBrokenTransport(error: unknown): boolean { + const { code } = ConnectError.from(error); + return code === Code.Unknown || code === Code.Unavailable; +} + /** Ids for the messages the wire did not name. */ let counter = 0; function nextId(prefix: string): string { @@ -528,19 +579,65 @@ export class A2AGrpcChatClient implements ChatClient { */ const artifacts = new Map(); - let stream: AsyncIterable<{ payload: { case?: string; value?: unknown } }>; + const options = this.callOptions(conversation, signal); + let stream: AsyncIterable; try { - stream = client.sendStreamingMessage( - request, - this.callOptions(conversation, signal), - ); + stream = client.sendStreamingMessage(request, options); } catch (error) { rethrowIfAborted(error, signal); throw fromConnectError(error, "A2AService/SendStreamingMessage"); } + /* + * The turn's frames, followed past a broken stream. + * + * A proxy between the browser and the controller may close a stream that has + * been quiet for a while — a model reading a long prompt sends nothing for + * minutes. The task keeps running on the server, so once the stream has named + * its task, losing the connection is not a failed turn: the task is re-attached + * with `SubscribeToTask`, which is read-only and returns the stored task once it + * has finished. The message itself is never sent twice. + * + * A re-attached stream opens with the task as it stands, so that first frame is + * not taken as a sign the path works: a path that breaks right after it would + * otherwise never run out of attempts. What resets the count is a frame after + * it, or a stream that stayed open long enough to have been cut for being quiet. + */ + let taskId = ""; + let settled = false; + const follow = async function* (): AsyncGenerator { + let source = stream; + for (let failures = 0; ; ) { + const opened = Date.now(); + let progressed = false; + try { + for await (const frame of source) { + if (frame.payload.case !== "task") progressed = true; + const event = frame.payload.value as + | { id?: string; taskId?: string; status?: TaskStatus } + | undefined; + const id = frame.payload.case === "task" ? event?.id : event?.taskId; + if (id) taskId = id; + if (event?.status) { + settled = !["working", "submitted"].includes(turnState(event.status.state)); + } + yield frame; + } + return; + } catch (error) { + if (progressed || Date.now() - opened >= REATTACH_HEALTHY_MS) failures = 0; + failures += 1; + const reattach = !signal?.aborted && taskId !== "" && !settled + && failures <= REATTACH_LIMIT && isBrokenTransport(error); + if (!reattach) throw error; + await pause(REATTACH_BACKOFF_MS * 2 ** (failures - 1), signal); + source = client.subscribeToTask({ tenant: conversation.agent, id: taskId }, options); + } + } + }; + try { - for await (const frame of stream) { + for await (const frame of follow()) { const payload = frame.payload; if (payload.case === "statusUpdate") {