fix(ui): keep following a chat turn when its stream breaks

ui/src/api/chat/a2aGrpcChatClient.test.ts

@@ -13,7 +13,7 @@
  */
 
 import { afterEach, describe, expect, it } from "vitest";
-import { createRouterTransport } from "@connectrpc/connect";
+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,102 @@ 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")] }],
+      },
+    },
+  };
+
+  async function run(
+    broken: () => AsyncGenerator<unknown>,
+    subscribe: () => AsyncGenerator<unknown> = async function* () {},
+  ) {
+    const subscribed: unknown[] = [];
+    serve(({ service }) => {
+      service(A2AService, {
+        sendStreamingMessage: broken as never,
+        subscribeToTask: async function* (request) {
+          subscribed.push({ tenant: request.tenant, id: request.id });
+          for await (const frame of subscribe()) yield frame as never;
+        },
+      });
+    });
+    const events: ChatEvent[] = [];
+    let error: unknown;
+    try {
+      for await (const event of new A2AGrpcChatClient().send({
+        conversation: CONVERSATION,
+        text: "why is checkout crashlooping?",
+      })) {
+        events.push(event);
+      }
+    } catch (cause) {
+      error = cause;
+    }
+    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);
+  });
+});

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,24 @@ const SHARE_HEADER = "X-Share-Token";
  */
 const HISTORY_PAGE_LIMIT = 50;
 
+/** Consecutive broken streams, with no frame in between, before a turn is given up. */
+const REATTACH_LIMIT = 3;
+
+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 +547,56 @@ export class A2AGrpcChatClient implements ChatClient {
      */
     const artifacts = new Map<string, string>();
 
-    let stream: AsyncIterable<{ payload: { case?: string; value?: unknown } }>;
+    const options = this.callOptions(conversation, signal);
+    let stream: AsyncIterable<StreamFrame>;
     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.
+     */
+    let taskId = "";
+    let settled = false;
+    const follow = async function* (): AsyncGenerator<StreamFrame> {
+      let source = stream;
+      for (let failures = 0; ; ) {
+        try {
+          for await (const frame of source) {
+            failures = 0;
+            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) {
+          failures += 1;
+          const reattach = !signal?.aborted && taskId !== "" && !settled
+            && failures <= REATTACH_LIMIT && isBrokenTransport(error);
+          if (!reattach) throw error;
+          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") {