From 2fac930e4e46eafd0b730251d8393c71da49e96d Mon Sep 17 00:00:00 2001 From: Hugo-DDT <141296842+Hugo-DDT@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:46:39 +0800 Subject: [PATCH] fix(a2a-server): preserve canceled streaming task status --- .changeset/a2a-stream-cancellation.md | 5 + packages/a2a-server/src/server.spec.ts | 273 +++++++++++++++++-------- packages/a2a-server/src/server.ts | 5 + 3 files changed, 195 insertions(+), 88 deletions(-) create mode 100644 .changeset/a2a-stream-cancellation.md diff --git a/.changeset/a2a-stream-cancellation.md b/.changeset/a2a-stream-cancellation.md new file mode 100644 index 000000000..5af231c5b --- /dev/null +++ b/.changeset/a2a-stream-cancellation.md @@ -0,0 +1,5 @@ +--- +"@voltagent/a2a-server": patch +--- + +Preserve canceled A2A task status when agent streams end normally or deliver a late chunk after cancellation. Avoid starting canceled streams and recheck cancellation after awaiting completion metadata. diff --git a/packages/a2a-server/src/server.spec.ts b/packages/a2a-server/src/server.spec.ts index e981f9099..2e0d866ca 100644 --- a/packages/a2a-server/src/server.spec.ts +++ b/packages/a2a-server/src/server.spec.ts @@ -263,115 +263,212 @@ describe("A2AServer", () => { }); }); - it("propagates cancellation to streaming agents and returns a canceled task", async () => { - const cancellationObserver = vi.fn(); + it.each(["error", "end", "late-chunk"])( + "preserves cancellation when the agent stream exits with %s", + async (exitMode) => { + const cancellationObserver = vi.fn(); - const streamText = vi - .fn() - .mockImplementation(async (_input: string, options?: { abortSignal?: AbortSignal }) => { - const abortSignal = options?.abortSignal; - - const textPromise = new Promise((resolve) => { - if (!abortSignal) { - resolve("final"); - return; - } - abortSignal.addEventListener( - "abort", - () => { - cancellationObserver(); - resolve(""); - }, - { once: true }, - ); - }); + const streamText = vi + .fn() + .mockImplementation(async (_input: string, options?: { abortSignal?: AbortSignal }) => { + const abortSignal = options?.abortSignal; - async function* textStream() { - yield "Chunk"; - await new Promise((_, reject) => { + const textPromise = new Promise((resolve) => { if (!abortSignal) { - reject(new Error("unexpected")); - return; - } - if (abortSignal.aborted) { - cancellationObserver(); - reject(new Error("aborted")); + resolve("final"); return; } - const handler = () => { - abortSignal.removeEventListener("abort", handler); - cancellationObserver(); - reject(new Error("aborted")); - }; - abortSignal.addEventListener("abort", handler, { once: true }); + abortSignal.addEventListener( + "abort", + () => { + cancellationObserver(); + resolve(""); + }, + { once: true }, + ); }); - yield "Should not reach"; + + async function* textStream() { + yield "Chunk"; + await new Promise((resolve, reject) => { + const onAbort = () => { + cancellationObserver(); + if (exitMode === "error") reject(new Error("aborted")); + else resolve(); + }; + if (!abortSignal) { + reject(new Error("unexpected")); + return; + } + if (abortSignal.aborted) { + onAbort(); + return; + } + const handler = () => { + abortSignal.removeEventListener("abort", handler); + onAbort(); + }; + abortSignal.addEventListener("abort", handler, { once: true }); + }); + if (exitMode === "end") return; + yield "Should not reach"; + } + + return { + text: textPromise, + finishReason: "stop", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + textStream: textStream(), + }; + }); + + const agent: StubAgent = { + id: "support-agent", + purpose: "Handle streamed cancellations", + generateText: vi.fn(), + streamText, + }; + + const taskStore = new InMemoryTaskStore(); + const server = createServer(agent, taskStore); + + const streamRequest = { + jsonrpc: "2.0" as const, + id: "stream-cancel", + method: "message/stream", + params: { + message: createMessage({ messageId: "msg-cancel-1" }), + }, + } satisfies JsonRpcRequest; + + const streamResponse = await server.handleRequest("support-agent", streamRequest); + expect("kind" in streamResponse && streamResponse.kind === "stream").toBe(true); + + const collected: JsonRpcResponse[] = []; + let cancelled = false; + + for await (const event of streamResponse.stream) { + const snapshot = clone(event); + collected.push(snapshot); + + if (!cancelled && snapshot.result?.status.message?.parts[0]?.text === "Chunk") { + const taskId = snapshot.result?.id; + expect(taskId).toBeDefined(); + expect(snapshot.result?.status.state).toBe("working"); + + const cancelResponse = (await server.handleRequest("support-agent", { + jsonrpc: "2.0", + id: "cancel-1", + method: "tasks/cancel", + params: { id: taskId }, + })) as JsonRpcResponse; + + expect(cancelResponse.error).toBeUndefined(); + expect(cancelResponse.result?.status.state).toBe("canceled"); + cancelled = true; + } + if (cancelled) { + expect(snapshot.result?.status.state).not.toBe("completed"); + expect(snapshot.result?.status.message?.parts[0]?.text).not.toBe("ChunkShould not reach"); } + } - return { - text: textPromise, - finishReason: "stop", - usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, - textStream: textStream(), - }; - }); + expect(cancellationObserver).toHaveBeenCalled(); + const finalSnapshot = collected.at(-1); + expect(finalSnapshot?.result?.status.state).toBe("canceled"); + const lastTaskId = finalSnapshot?.result?.id; + const stored = lastTaskId + ? await taskStore.load({ agentId: "support-agent", taskId: lastTaskId }) + : null; + expect(stored?.status.state).toBe("canceled"); + }, + ); + + it("does not start the agent when the task is canceled before stream consumption", async () => { const agent: StubAgent = { id: "support-agent", purpose: "Handle streamed cancellations", generateText: vi.fn(), - streamText, + streamText: vi.fn(), }; - const taskStore = new InMemoryTaskStore(); const server = createServer(agent, taskStore); - - const streamRequest = { - jsonrpc: "2.0" as const, - id: "stream-cancel", + const response = await server.handleRequest("support-agent", { + ...createSendRequest(), method: "message/stream", - params: { - message: createMessage({ messageId: "msg-cancel-1" }), - }, - } satisfies JsonRpcRequest; - - const streamResponse = await server.handleRequest("support-agent", streamRequest); - expect("kind" in streamResponse && streamResponse.kind === "stream").toBe(true); - - const collected: JsonRpcResponse[] = []; - let cancelled = false; - - for await (const event of streamResponse.stream) { - const snapshot = clone(event); - collected.push(snapshot); - - if (!cancelled) { - const taskId = snapshot.result?.id; - expect(taskId).toBeDefined(); - expect(snapshot.result?.status.state).toBe("working"); + }); + if (!("stream" in response)) throw new Error("Expected a stream"); + const first = await response.stream.next(); + const taskId = first.value.result.id; + await server.handleRequest("support-agent", { + jsonrpc: "2.0", + id: "cancel-before-start", + method: "tasks/cancel", + params: { id: taskId }, + }); - const cancelResponse = (await server.handleRequest("support-agent", { - jsonrpc: "2.0", - id: "cancel-1", - method: "tasks/cancel", - params: { id: taskId }, - })) as JsonRpcResponse; + const events = []; + for await (const event of response.stream) events.push(event); - expect(cancelResponse.error).toBeUndefined(); - expect(cancelResponse.result?.status.state).toBe("canceled"); - cancelled = true; - } - } + expect(agent.streamText).not.toHaveBeenCalled(); + expect(events.at(-1)?.result?.status.state).toBe("canceled"); + expect((await taskStore.load({ agentId: "support-agent", taskId }))?.status.state).toBe( + "canceled", + ); + }); - expect(cancellationObserver).toHaveBeenCalled(); - const finalSnapshot = collected.at(-1); - expect(finalSnapshot?.result?.status.state).toBe("canceled"); + it("preserves cancellation while awaiting stream completion metadata", async () => { + let signalMetadataWait!: () => void; + const metadataWait = new Promise((resolve) => { + signalMetadataWait = resolve; + }); + let finish!: (value: string) => void; + const finishReason = new Promise((resolve) => { + finish = resolve; + }); + const agent: StubAgent = { + id: "support-agent", + purpose: "Handle streamed cancellations", + generateText: vi.fn(), + streamText: vi.fn().mockResolvedValue({ + text: Promise.resolve("Chunk"), + get finishReason() { + signalMetadataWait(); + return finishReason; + }, + textStream: (async function* () { + yield "Chunk"; + })(), + }), + }; + const taskStore = new InMemoryTaskStore(); + const server = createServer(agent, taskStore); + const response = await server.handleRequest("support-agent", { + ...createSendRequest(), + method: "message/stream", + }); + if (!("stream" in response)) throw new Error("Expected a stream"); + const first = await response.stream.next(); + const taskId = first.value.result.id; + const events: JsonRpcResponse[] = []; + const consume = (async () => { + for await (const event of response.stream) events.push(event); + })(); + await metadataWait; + await server.handleRequest("support-agent", { + jsonrpc: "2.0", + id: "cancel-during-metadata", + method: "tasks/cancel", + params: { id: taskId }, + }); + finish("stop"); + await consume; - const lastTaskId = finalSnapshot?.result?.id; - const stored = lastTaskId - ? await taskStore.load({ agentId: "support-agent", taskId: lastTaskId }) - : null; - expect(stored?.status.state).toBe("canceled"); + expect(events.at(-1)?.result?.status.state).toBe("canceled"); + expect((await taskStore.load({ agentId: "support-agent", taskId }))?.status.state).toBe( + "canceled", + ); }); it("returns method not found for unknown requests", async () => { diff --git a/packages/a2a-server/src/server.ts b/packages/a2a-server/src/server.ts index eb801f063..6534b471c 100644 --- a/packages/a2a-server/src/server.ts +++ b/packages/a2a-server/src/server.ts @@ -402,6 +402,7 @@ export class A2AServer { let aggregatedText = ""; try { + abortController.signal.throwIfAborted(); const voltMessage = toVoltAgentMessage(normalizedMessage); const agentOptions = server.buildAgentCallOptions({ record, @@ -412,6 +413,7 @@ export class A2AServer { }); const streamResult = await agent.streamText(voltMessage.content, agentOptions); + abortController.signal.throwIfAborted(); responseMessage = fromVoltAgentMessage({ role: "assistant", @@ -426,6 +428,7 @@ export class A2AServer { yield createSuccessResponse(requestId, record); for await (const chunk of streamResult.textStream) { + abortController.signal.throwIfAborted(); if (!chunk) { continue; } @@ -441,6 +444,7 @@ export class A2AServer { yield createSuccessResponse(requestId, record); } + abortController.signal.throwIfAborted(); let finalText = aggregatedText; if (!finalText) { try { @@ -485,6 +489,7 @@ export class A2AServer { }; } + abortController.signal.throwIfAborted(); await taskStore.save({ agentId, data: record }); yield createSuccessResponse(requestId, record); } catch (error) {