Skip to content

Commit 1d72a12

Browse files
kschaabgemini-code-assist[bot]
kschaab
and
gemini-code-assist[bot]
authored
fix(a2a-server): resolve tool approval race condition and improve status reporting (#26479)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 8f0edcd commit 1d72a12

4 files changed

Lines changed: 359 additions & 54 deletions

File tree

‎packages/a2a-server/src/agent/task-event-driven.test.ts‎

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ describe('Task Event-Driven Scheduler', () => {
6666
handler({
6767
type: MessageBusType.TOOL_CALLS_UPDATE,
6868
toolCalls: [toolCall],
69+
schedulerId: 'task-id',
6970
});
7071

7172
expect(mockEventBus.publish).toHaveBeenCalledWith(
@@ -106,6 +107,7 @@ describe('Task Event-Driven Scheduler', () => {
106107
handler({
107108
type: MessageBusType.TOOL_CALLS_UPDATE,
108109
toolCalls: [toolCall],
110+
schedulerId: 'task-id',
109111
});
110112

111113
// Simulate A2A client confirmation
@@ -148,7 +150,11 @@ describe('Task Event-Driven Scheduler', () => {
148150
const handler = (messageBus.subscribe as Mock).mock.calls.find(
149151
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
150152
)?.[1];
151-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
153+
handler({
154+
type: MessageBusType.TOOL_CALLS_UPDATE,
155+
toolCalls: [toolCall],
156+
schedulerId: 'task-id',
157+
});
152158

153159
// Simulate Rejection (Cancel)
154160
const handled = await (
@@ -174,7 +180,11 @@ describe('Task Event-Driven Scheduler', () => {
174180
correlationId: 'corr-2',
175181
confirmationDetails: { type: 'info', title: 'test', prompt: 'test' },
176182
};
177-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall2] });
183+
handler({
184+
type: MessageBusType.TOOL_CALLS_UPDATE,
185+
toolCalls: [toolCall2],
186+
schedulerId: 'task-id',
187+
});
178188

179189
// Simulate ModifyWithEditor
180190
const handled2 = await (
@@ -215,7 +225,11 @@ describe('Task Event-Driven Scheduler', () => {
215225
const handler = (messageBus.subscribe as Mock).mock.calls.find(
216226
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
217227
)?.[1];
218-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
228+
handler({
229+
type: MessageBusType.TOOL_CALLS_UPDATE,
230+
toolCalls: [toolCall],
231+
schedulerId: 'task-id',
232+
});
219233

220234
// Simulate ProceedOnce for MCP
221235
const handled = await (
@@ -255,7 +269,11 @@ describe('Task Event-Driven Scheduler', () => {
255269
const handler = (messageBus.subscribe as Mock).mock.calls.find(
256270
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
257271
)?.[1];
258-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
272+
handler({
273+
type: MessageBusType.TOOL_CALLS_UPDATE,
274+
toolCalls: [toolCall],
275+
schedulerId: 'task-id',
276+
});
259277

260278
const handled = await (
261279
task as unknown as {
@@ -294,7 +312,11 @@ describe('Task Event-Driven Scheduler', () => {
294312
const handler = (messageBus.subscribe as Mock).mock.calls.find(
295313
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
296314
)?.[1];
297-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
315+
handler({
316+
type: MessageBusType.TOOL_CALLS_UPDATE,
317+
toolCalls: [toolCall],
318+
schedulerId: 'task-id',
319+
});
298320

299321
const handled = await (
300322
task as unknown as {
@@ -333,7 +355,11 @@ describe('Task Event-Driven Scheduler', () => {
333355
const handler = (messageBus.subscribe as Mock).mock.calls.find(
334356
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
335357
)?.[1];
336-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
358+
handler({
359+
type: MessageBusType.TOOL_CALLS_UPDATE,
360+
toolCalls: [toolCall],
361+
schedulerId: 'task-id',
362+
});
337363

338364
const handled = await (
339365
task as unknown as {
@@ -376,7 +402,11 @@ describe('Task Event-Driven Scheduler', () => {
376402
const handler = (yoloMessageBus.subscribe as Mock).mock.calls.find(
377403
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
378404
)?.[1];
379-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
405+
handler({
406+
type: MessageBusType.TOOL_CALLS_UPDATE,
407+
toolCalls: [toolCall],
408+
schedulerId: 'task-id',
409+
});
380410

381411
// Should NOT auto-publish ProceedOnce anymore, because PolicyEngine handles it directly
382412
expect(yoloMessageBus.publish).not.toHaveBeenCalledWith(
@@ -419,6 +449,7 @@ describe('Task Event-Driven Scheduler', () => {
419449
handler({
420450
type: MessageBusType.TOOL_CALLS_UPDATE,
421451
toolCalls: [toolCall],
452+
schedulerId: 'task-id',
422453
});
423454

424455
// Should publish artifact update for output
@@ -453,7 +484,11 @@ describe('Task Event-Driven Scheduler', () => {
453484
const handler = (messageBus.subscribe as Mock).mock.calls.find(
454485
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
455486
)?.[1];
456-
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
487+
handler({
488+
type: MessageBusType.TOOL_CALLS_UPDATE,
489+
toolCalls: [toolCall],
490+
schedulerId: 'task-id',
491+
});
457492

458493
// The tool should be complete and registered appropriately, eventually
459494
// triggering the toolCompletionPromise resolution when all clear.
@@ -533,6 +568,7 @@ describe('Task Event-Driven Scheduler', () => {
533568
handler({
534569
type: MessageBusType.TOOL_CALLS_UPDATE,
535570
toolCalls: [toolCall1, toolCall2],
571+
schedulerId: 'task-id',
536572
});
537573

538574
// Confirm first tool call
@@ -600,6 +636,7 @@ describe('Task Event-Driven Scheduler', () => {
600636
handler({
601637
type: MessageBusType.TOOL_CALLS_UPDATE,
602638
toolCalls: [toolCall1, toolCall2],
639+
schedulerId: 'task-id',
603640
});
604641

605642
// Should NOT transition to input-required yet
@@ -621,6 +658,7 @@ describe('Task Event-Driven Scheduler', () => {
621658
handler({
622659
type: MessageBusType.TOOL_CALLS_UPDATE,
623660
toolCalls: [toolCall1Complete, toolCall2],
661+
schedulerId: 'task-id',
624662
});
625663

626664
// Now it should transition

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

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ import {
1212
type ToolCallRequestInfo,
1313
type GitService,
1414
type CompletedToolCall,
15+
type ToolCall,
16+
type ToolCallsUpdateMessage,
17+
MessageBusType,
1518
} from '@google/gemini-cli-core';
1619
import { createMockConfig } from '../utils/testing_utils.js';
1720
import type { ExecutionEventBus, RequestContext } from '@a2a-js/sdk/server';
@@ -460,4 +463,204 @@ describe('Task', () => {
460463
expect(task.currentPromptId).toBe(expectedPromptId2);
461464
});
462465
});
466+
467+
describe('Race Condition Fix', () => {
468+
const mockConfig = createMockConfig();
469+
const mockEventBus: ExecutionEventBus = {
470+
publish: vi.fn(),
471+
on: vi.fn(),
472+
off: vi.fn(),
473+
once: vi.fn(),
474+
removeAllListeners: vi.fn(),
475+
finished: vi.fn(),
476+
};
477+
478+
beforeEach(() => {
479+
vi.clearAllMocks();
480+
});
481+
482+
it('should NOT transition to input-required if a tool is still validating', async () => {
483+
// @ts-expect-error - Calling private constructor
484+
const task = new Task(
485+
'task-id',
486+
'context-id',
487+
mockConfig as Config,
488+
mockEventBus,
489+
);
490+
491+
// Manually register two tool calls
492+
task['_registerToolCall']('tool-1', 'awaiting_approval');
493+
task['_registerToolCall']('tool-2', 'validating');
494+
495+
// Call checkInputRequiredState (private)
496+
task['checkInputRequiredState']();
497+
498+
// Verify task state did NOT change to input-required
499+
expect(task.taskState).not.toBe('input-required');
500+
expect(mockEventBus.publish).not.toHaveBeenCalledWith(
501+
expect.objectContaining({
502+
status: expect.objectContaining({ state: 'input-required' }),
503+
}),
504+
);
505+
});
506+
507+
it('should transition to input-required if all active tools are awaiting approval', async () => {
508+
// @ts-expect-error - Calling private constructor
509+
const task = new Task(
510+
'task-id',
511+
'context-id',
512+
mockConfig as Config,
513+
mockEventBus,
514+
);
515+
516+
// Transition from submitted to working first to simulate normal flow
517+
task.taskState = 'working';
518+
519+
// Manually register tool calls
520+
task['_registerToolCall']('tool-1', 'awaiting_approval');
521+
522+
// Call checkInputRequiredState
523+
task['checkInputRequiredState']();
524+
525+
// Verify task state changed to input-required
526+
expect(task.taskState).toBe('input-required');
527+
expect(mockEventBus.publish).toHaveBeenCalledWith(
528+
expect.objectContaining({
529+
status: expect.objectContaining({ state: 'input-required' }),
530+
}),
531+
);
532+
});
533+
534+
it('handleEventDrivenToolCallsUpdate should ignore events for other schedulers', async () => {
535+
// @ts-expect-error - Calling private constructor
536+
const task = new Task(
537+
'task-id',
538+
'context-id',
539+
mockConfig as Config,
540+
mockEventBus,
541+
);
542+
543+
const handleEventDrivenToolCallSpy = vi.spyOn(
544+
task as unknown as {
545+
handleEventDrivenToolCall: Task['handleEventDrivenToolCall'];
546+
},
547+
'handleEventDrivenToolCall',
548+
);
549+
550+
const otherEvent: ToolCallsUpdateMessage = {
551+
type: MessageBusType.TOOL_CALLS_UPDATE,
552+
toolCalls: [
553+
{ request: { callId: '1' }, status: 'executing' } as ToolCall,
554+
],
555+
schedulerId: 'other-task-id',
556+
};
557+
558+
task['handleEventDrivenToolCallsUpdate'](otherEvent);
559+
560+
expect(handleEventDrivenToolCallSpy).not.toHaveBeenCalled();
561+
562+
const ownEvent: ToolCallsUpdateMessage = {
563+
type: MessageBusType.TOOL_CALLS_UPDATE,
564+
toolCalls: [
565+
{ request: { callId: '1' }, status: 'executing' } as ToolCall,
566+
],
567+
schedulerId: 'task-id',
568+
};
569+
570+
task['handleEventDrivenToolCallsUpdate'](ownEvent);
571+
572+
expect(handleEventDrivenToolCallSpy).toHaveBeenCalled();
573+
});
574+
});
575+
576+
describe('Serialization and Mapping', () => {
577+
it('should map internal "validating" status to "scheduled" for the client and include outcome', async () => {
578+
const mockConfig = createMockConfig();
579+
const mockEventBus: ExecutionEventBus = {
580+
publish: vi.fn(),
581+
on: vi.fn(),
582+
off: vi.fn(),
583+
once: vi.fn(),
584+
removeAllListeners: vi.fn(),
585+
finished: vi.fn(),
586+
};
587+
588+
// @ts-expect-error - Calling private constructor
589+
const task = new Task(
590+
'task-id',
591+
'context-id',
592+
mockConfig as Config,
593+
mockEventBus,
594+
);
595+
596+
const mockToolCall = {
597+
request: { callId: 'tool-1' },
598+
status: 'validating',
599+
outcome: 'accepted',
600+
tool: { name: 'test-tool' },
601+
};
602+
603+
const message = task['toolStatusMessage'](
604+
mockToolCall as unknown as ToolCall,
605+
'task-id',
606+
'context-id',
607+
);
608+
const serialized = (
609+
message.parts![0] as {
610+
data: { status: string; outcome: string };
611+
}
612+
).data;
613+
614+
expect(serialized.status).toBe('scheduled');
615+
expect(serialized.outcome).toBe('accepted');
616+
});
617+
618+
it('should correctly detect changes when status or outcome changes', async () => {
619+
const mockConfig = createMockConfig();
620+
const mockEventBus: ExecutionEventBus = {
621+
publish: vi.fn(),
622+
on: vi.fn(),
623+
off: vi.fn(),
624+
once: vi.fn(),
625+
removeAllListeners: vi.fn(),
626+
finished: vi.fn(),
627+
};
628+
629+
// @ts-expect-error - Calling private constructor
630+
const task = new Task(
631+
'task-id',
632+
'context-id',
633+
mockConfig as Config,
634+
mockEventBus,
635+
);
636+
637+
const toolCall1 = {
638+
request: { callId: 'tool-1' },
639+
status: 'awaiting_approval',
640+
};
641+
642+
// First update - should trigger change
643+
const changed1 = task['handleEventDrivenToolCall'](
644+
toolCall1 as unknown as ToolCall,
645+
);
646+
expect(changed1).toBe(true);
647+
648+
// Second update with same status - should NOT trigger change
649+
const changed2 = task['handleEventDrivenToolCall'](
650+
toolCall1 as unknown as ToolCall,
651+
);
652+
expect(changed2).toBe(false);
653+
654+
// Update with new outcome - SHOULD trigger change
655+
const toolCall2 = {
656+
request: { callId: 'tool-1' },
657+
status: 'awaiting_approval',
658+
outcome: 'accepted',
659+
};
660+
const changed3 = task['handleEventDrivenToolCall'](
661+
toolCall2 as unknown as ToolCall,
662+
);
663+
expect(changed3).toBe(true);
664+
});
665+
});
463666
});

0 commit comments

Comments
 (0)