fix: abort response-cancelled answer streams

src/app/api/answer/stream/route.ts

@@ -54,6 +54,25 @@ function documentSummaryRateLimitStream(rateLimit: ApiRateLimitResult) {
   return rateLimitJsonResponse("Too many document summary requests. Retry shortly.", rateLimit);
 }
 
+function mergeAbortSignals(signals: Array<AbortSignal | undefined>) {
+  const controller = new AbortController();
+  for (const signal of signals) {
+    if (!signal) continue;
+    if (signal.aborted) {
+      controller.abort(signal.reason);
+      break;
+    }
+    signal.addEventListener(
+      "abort",
+      () => {
+        if (!controller.signal.aborted) controller.abort(signal.reason);
+      },
+      { once: true },
+    );
+  }
+  return controller.signal;
+}
+
 function streamErrorPayload(error: unknown) {
   if (error instanceof PublicApiError) {
     return {
@@ -128,7 +147,12 @@ function buildDemoStreamAnswer(body: AnswerRequestBody, fallbackReason?: string)
   );
 }
 
-function streamAnswer(body: AnswerRequestBody, accessScope: RetrievalAccessScope, signal?: AbortSignal) {
+function streamAnswer(
+  body: AnswerRequestBody,
+  accessScope: RetrievalAccessScope,
+  signal?: AbortSignal,
+  streamAbortController?: AbortController,
+) {
   const ownerId = accessScope.ownerId;
   const encoder = new TextEncoder();
   const interactionId = randomUUID();
@@ -254,6 +278,9 @@ function streamAnswer(body: AnswerRequestBody, accessScope: RetrievalAccessScope
           }
         }
       },
+      cancel() {
+        streamAbortController?.abort();
+      },
     }),
     {
       headers: {
@@ -268,7 +295,9 @@ function streamAnswer(body: AnswerRequestBody, accessScope: RetrievalAccessScope
 export async function POST(request: Request) {
   try {
     const body = await parseJsonBody(request, answerRequestSchema, "Invalid answer request.");
-    if (isDemoMode()) return streamAnswer(body, resolveRetrievalAccessScope(), request.signal);
+    const streamAbortController = new AbortController();
+    const streamSignal = mergeAbortSignals([request.signal, streamAbortController.signal]);
+    if (isDemoMode()) return streamAnswer(body, resolveRetrievalAccessScope(), streamSignal, streamAbortController);
 
     const supabase = createAdminClient();
     const access = await publicAccessContext(request, supabase);
@@ -293,7 +322,7 @@ export async function POST(request: Request) {
       if (rateLimit.limited) return rateLimitStream(rateLimit);
     }
 
-    return streamAnswer(body, resolveRetrievalAccessScope(access.ownerId), request.signal);
+    return streamAnswer(body, resolveRetrievalAccessScope(access.ownerId), streamSignal, streamAbortController);
   } catch (error) {
     if (error instanceof AuthenticationError) {
       return unauthorizedResponse(error);

tests/private-access-routes.test.ts

@@ -4459,6 +4459,55 @@ describe("private document API access", () => {
     expect(answerQuestionWithScope).not.toHaveBeenCalled();
   });
 
+  it("aborts streamed answer generation when the response body is cancelled", async () => {
+    const answerQuestionWithScope = vi.fn(
+      async ({ signal, onProgress }: { signal?: AbortSignal; onProgress?: (event: unknown) => void }) => {
+        onProgress?.({
+          stage: "ranking",
+          resultCount: 1,
+          selectedContextCount: 0,
+          australianSourceCount: 0,
+          waSourceCount: 0,
+        });
+        await new Promise<void>((resolve) => {
+          signal?.addEventListener(
+            "abort",
+            () => {
+              resolve();
+            },
+            { once: true },
+          );
+        });
+        return {
+          answer: "Owned evidence.",
+          grounded: true,
+          confidence: "medium",
+          citations: [],
+          sources: [],
+        };
+      },
+    );
+    const client = createSupabaseMock();
+    mockRuntime(client, { answerQuestionWithScope });
+    const { POST } = await import("../src/app/api/answer/stream/route");
+
+    const response = await POST(
+      authenticatedRequest("/api/answer/stream", {
+        method: "POST",
+        body: JSON.stringify({ query: "monitoring", documentId: otherDocumentId }),
+      }),
+    );
+    const reader = response.body?.getReader();
+    await reader?.read();
+    await reader?.cancel();
+
+    expect(response.status).toBe(200);
+    expect(answerQuestionWithScope).toHaveBeenCalledTimes(1);
+    const signal = answerQuestionWithScope.mock.calls.at(0)?.[0]?.signal;
+    expect(signal).toBeInstanceOf(AbortSignal);
+    expect(signal?.aborted).toBe(true);
+  });
+
   it("uses an anonymous in-memory limiter for managed local no-auth streaming answers", async () => {
     const answerQuestionWithScope = vi.fn(async () => ({
       answer: "Owned evidence.",