Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/a2a-stream-cancellation.md
Original file line number Diff line number Diff line change
@@ -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.
273 changes: 185 additions & 88 deletions packages/a2a-server/src/server.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string>((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<string>((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<void>((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<MessageSendParams>;

const streamResponse = await server.handleRequest("support-agent", streamRequest);
expect("kind" in streamResponse && streamResponse.kind === "stream").toBe(true);

const collected: JsonRpcResponse<TaskRecord>[] = [];
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<TaskRecord>;

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<MessageSendParams>;

const streamResponse = await server.handleRequest("support-agent", streamRequest);
expect("kind" in streamResponse && streamResponse.kind === "stream").toBe(true);

const collected: JsonRpcResponse<TaskRecord>[] = [];
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<TaskRecord>;
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<void>((resolve) => {
signalMetadataWait = resolve;
});
let finish!: (value: string) => void;
const finishReason = new Promise<string>((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<TaskRecord>[] = [];
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 () => {
Expand Down
5 changes: 5 additions & 0 deletions packages/a2a-server/src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,7 @@ export class A2AServer {
let aggregatedText = "";

try {
abortController.signal.throwIfAborted();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: A custom asynchronous TaskStore.save can pause after this check; cancellation during the save still reaches the unconditional completed yield. Recheck after each stream save so the abort path emits the canceled record instead of a stale event.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At packages/a2a-server/src/server.ts, line 405:

<comment>A custom asynchronous `TaskStore.save` can pause after this check; cancellation during the save still reaches the unconditional completed `yield`. Recheck after each stream save so the abort path emits the canceled record instead of a stale event.</comment>

<file context>
@@ -402,6 +402,7 @@ export class A2AServer {
       let aggregatedText = "";
 
       try {
+        abortController.signal.throwIfAborted();
         const voltMessage = toVoltAgentMessage(normalizedMessage);
         const agentOptions = server.buildAgentCallOptions({
</file context>

const voltMessage = toVoltAgentMessage(normalizedMessage);
const agentOptions = server.buildAgentCallOptions({
record,
Expand All @@ -412,6 +413,7 @@ export class A2AServer {
});

const streamResult = await agent.streamText(voltMessage.content, agentOptions);
abortController.signal.throwIfAborted();

responseMessage = fromVoltAgentMessage({
role: "assistant",
Expand All @@ -426,6 +428,7 @@ export class A2AServer {
yield createSuccessResponse(requestId, record);

for await (const chunk of streamResult.textStream) {
abortController.signal.throwIfAborted();
if (!chunk) {
continue;
}
Expand All @@ -441,6 +444,7 @@ export class A2AServer {
yield createSuccessResponse(requestId, record);
}

abortController.signal.throwIfAborted();
let finalText = aggregatedText;
if (!finalText) {
try {
Expand Down Expand Up @@ -485,6 +489,7 @@ export class A2AServer {
};
}

abortController.signal.throwIfAborted();
await taskStore.save({ agentId, data: record });
yield createSuccessResponse(requestId, record);
} catch (error) {
Expand Down