diff --git a/src/app/api/answer/stream/route.ts b/src/app/api/answer/stream/route.ts index 843ee7153..eb280883d 100644 --- a/src/app/api/answer/stream/route.ts +++ b/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) { + 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); diff --git a/tests/private-access-routes.test.ts b/tests/private-access-routes.test.ts index 183930122..78afdd138 100644 --- a/tests/private-access-routes.test.ts +++ b/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((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.",