Skip to content

Commit c6fe974

Browse files
fix(a2a-server): ensure task cancellation aborts execution loop
1 parent 15a9429 commit c6fe974

2 files changed

Lines changed: 101 additions & 1 deletion

File tree

‎packages/a2a-server/src/agent/executor.test.ts‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,4 +300,65 @@ describe('CoderAgentExecutor', () => {
300300
true,
301301
);
302302
});
303+
304+
it('cancelTask should abort the active execution loop', async () => {
305+
const abortSpy = vi.spyOn(AbortController.prototype, 'abort');
306+
const taskId = 'test-task-to-cancel';
307+
const contextId = 'test-context';
308+
309+
const mockSocket = new EventEmitter();
310+
(requestStorage.getStore as Mock).mockReturnValue({
311+
req: { socket: mockSocket },
312+
});
313+
314+
const requestContext = {
315+
userMessage: {
316+
messageId: 'msg-1',
317+
taskId,
318+
contextId,
319+
parts: [{ kind: 'text', text: 'a long running prompt' }],
320+
metadata: {
321+
coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
322+
},
323+
},
324+
} as unknown as RequestContext;
325+
326+
// Don't await this, let it run in the background.
327+
// The mocked acceptUserMessage will hang until aborted.
328+
const primaryPromise = executor.execute(requestContext, mockEventBus);
329+
330+
// Give it a moment to start and register the controller.
331+
await new Promise((resolve) => setTimeout(resolve, 50));
332+
333+
const wrapper = executor.getTask(taskId);
334+
expect(wrapper).toBeDefined();
335+
const setTaskStateSpy = vi
336+
.spyOn(wrapper!.task, 'setTaskStateAndPublishUpdate')
337+
.mockImplementation((newState) => {
338+
// Make the mock realistic: actually update the state when called.
339+
wrapper!.task.taskState = newState;
340+
});
341+
342+
// Now, cancel the task.
343+
await executor.cancelTask(taskId, mockEventBus);
344+
345+
// Verify that the abort method on the controller was called and state was updated.
346+
expect(abortSpy).toHaveBeenCalledOnce();
347+
expect(setTaskStateSpy).toHaveBeenCalledWith(
348+
'canceled',
349+
expect.any(Object),
350+
'Task canceled by user request.',
351+
undefined,
352+
true,
353+
);
354+
355+
// Clean up the test by allowing the promise to resolve.
356+
// The abort call should have unblocked the acceptUserMessage generator.
357+
await primaryPromise;
358+
359+
// Verify task is evicted from cache
360+
expect(executor.getTask(taskId)).toBeUndefined();
361+
362+
abortSpy.mockRestore();
363+
});
303364
});

‎packages/a2a-server/src/agent/executor.ts‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ export class CoderAgentExecutor implements AgentExecutor {
9090
private tasks: Map<string, TaskWrapper> = new Map();
9191
// Track tasks with an active execution loop.
9292
private executingTasks = new Set<string>();
93+
private activeAbortControllers = new Map<string, AbortController>();
9394

9495
constructor(private taskStore?: TaskStore) {}
9596

@@ -187,6 +188,35 @@ export class CoderAgentExecutor implements AgentExecutor {
187188
logger.info(
188189
`[CoderAgentExecutor] Received cancel request for task ${taskId}`,
189190
);
191+
192+
// If there is an active execution loop, just abort it.
193+
// The `finally` block of the `execute` method will handle cleanup.
194+
const abortController = this.activeAbortControllers.get(taskId);
195+
if (abortController) {
196+
logger.info(
197+
`[CoderAgentExecutor] Aborting active execution loop for task ${taskId}.`,
198+
);
199+
const wrapper = this.tasks.get(taskId);
200+
if (wrapper) {
201+
const { task } = wrapper;
202+
task.cancelPendingTools('Task canceled by user request.');
203+
task.setTaskStateAndPublishUpdate(
204+
'canceled',
205+
{ kind: CoderAgentEvent.StateChangeEvent },
206+
'Task canceled by user request.',
207+
undefined,
208+
true,
209+
);
210+
}
211+
abortController.abort();
212+
return;
213+
}
214+
215+
// If there is no active execution loop, the task is idle.
216+
// We can clean it up directly.
217+
logger.info(
218+
`[CoderAgentExecutor] No active execution for task ${taskId}. Cleaning up directly.`,
219+
);
190220
const wrapper = this.tasks.get(taskId);
191221

192222
if (!wrapper) {
@@ -244,7 +274,7 @@ export class CoderAgentExecutor implements AgentExecutor {
244274

245275
try {
246276
logger.info(
247-
`[CoderAgentExecutor] Initiating cancellation for task ${taskId}.`,
277+
`[CoderAgentExecutor] Initiating cancellation for idle task ${taskId}.`,
248278
);
249279
task.cancelPendingTools('Task canceled by user request.');
250280

@@ -332,6 +362,14 @@ export class CoderAgentExecutor implements AgentExecutor {
332362
const abortController = new AbortController();
333363
const abortSignal = abortController.signal;
334364

365+
// Register the abort controller only for the primary execution loop.
366+
if (
367+
!this.executingTasks.has(taskId) &&
368+
!this.activeAbortControllers.has(taskId)
369+
) {
370+
this.activeAbortControllers.set(taskId, abortController);
371+
}
372+
335373
if (store) {
336374
// Grab the raw socket from the request object
337375
const socket = store.req.socket;
@@ -644,6 +682,7 @@ export class CoderAgentExecutor implements AgentExecutor {
644682
}
645683
} finally {
646684
if (isPrimaryExecution) {
685+
this.activeAbortControllers.delete(taskId); // Clean up the reference
647686
this.executingTasks.delete(taskId);
648687
logger.info(
649688
`[CoderAgentExecutor] Saving final state for task ${taskId}.`,

0 commit comments

Comments
 (0)