Skip to content

Commit 6837b08

Browse files
fix(patch): cherry-pick 02995ba to release/v0.42.0-preview.1-pr-26568 to patch version v0.42.0-preview.1 and create version 0.42.0-preview.2 (#26590)
Co-authored-by: Keith Schaab <keith.schaab@gmail.com>
1 parent 92f410d commit 6837b08

3 files changed

Lines changed: 284 additions & 39 deletions

File tree

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
/**
2+
* @license
3+
* Copyright 2026 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
8+
import { Task } from './task.js';
9+
import {
10+
MessageBusType,
11+
CoreToolCallStatus,
12+
type Config,
13+
type MessageBus,
14+
} from '@google/gemini-cli-core';
15+
import { createMockConfig } from '../utils/testing_utils.js';
16+
import type { RequestContext } from '@a2a-js/sdk/server';
17+
18+
describe('Task Race Condition', () => {
19+
let mockConfig: Config;
20+
let messageBus: MessageBus;
21+
22+
beforeEach(() => {
23+
messageBus = {
24+
subscribe: vi.fn(),
25+
unsubscribe: vi.fn(),
26+
publish: vi.fn(),
27+
} as unknown as MessageBus;
28+
mockConfig = createMockConfig({
29+
messageBus,
30+
}) as Config;
31+
});
32+
33+
it('should not hang when multiple tool confirmations are processed while waiting', async () => {
34+
// @ts-expect-error - private constructor
35+
const task = new Task('task-id', 'context-id', mockConfig);
36+
37+
// 1. Register two tools as scheduled
38+
task['_registerToolCall']('tool-1', 'scheduled');
39+
task['_registerToolCall']('tool-2', 'scheduled');
40+
41+
// 2. Both transition to awaiting_approval
42+
const updateHandler = (messageBus.subscribe as Mock).mock.calls.find(
43+
(c: unknown[]) => c[0] === MessageBusType.TOOL_CALLS_UPDATE,
44+
)?.[1];
45+
46+
updateHandler({
47+
type: MessageBusType.TOOL_CALLS_UPDATE,
48+
schedulerId: 'task-id',
49+
toolCalls: [
50+
{
51+
request: { callId: 'tool-1', name: 't1' },
52+
status: CoreToolCallStatus.AwaitingApproval,
53+
correlationId: 'corr-1',
54+
confirmationDetails: { type: 'info' },
55+
},
56+
{
57+
request: { callId: 'tool-2', name: 't2' },
58+
status: CoreToolCallStatus.AwaitingApproval,
59+
correlationId: 'corr-2',
60+
confirmationDetails: { type: 'info' },
61+
},
62+
],
63+
});
64+
65+
// 3. Confirm Tool 1. This makes isAwaitingApprovalOnly() return false.
66+
for await (const _ of task.acceptUserMessage(
67+
{
68+
userMessage: {
69+
parts: [
70+
{
71+
kind: 'data',
72+
data: { callId: 'tool-1', outcome: 'proceed_once' },
73+
},
74+
],
75+
},
76+
} as unknown as RequestContext,
77+
new AbortController().signal,
78+
)) {
79+
// consume generator
80+
}
81+
82+
// 4. Start waiting. This should now block because Tool 1 is confirmed (so we are waiting for its execution).
83+
const waitPromise = task.waitForPendingTools();
84+
85+
// 5. Confirm Tool 2 while waiting.
86+
for await (const _ of task.acceptUserMessage(
87+
{
88+
userMessage: {
89+
parts: [
90+
{
91+
kind: 'data',
92+
data: { callId: 'tool-2', outcome: 'proceed_once' },
93+
},
94+
],
95+
},
96+
} as unknown as RequestContext,
97+
new AbortController().signal,
98+
)) {
99+
// consume generator
100+
}
101+
102+
// 6. Both tools complete successfully
103+
updateHandler({
104+
type: MessageBusType.TOOL_CALLS_UPDATE,
105+
schedulerId: 'task-id',
106+
toolCalls: [
107+
{
108+
request: { callId: 'tool-1', name: 't1' },
109+
status: CoreToolCallStatus.Success,
110+
response: { responseParts: [] },
111+
},
112+
{
113+
request: { callId: 'tool-2', name: 't2' },
114+
status: CoreToolCallStatus.Success,
115+
response: { responseParts: [] },
116+
},
117+
],
118+
});
119+
120+
// 7. Verify that the original waitPromise resolves.
121+
await expect(waitPromise).resolves.toBeUndefined();
122+
});
123+
124+
it('should reject waitForPendingTools when tools are cancelled', async () => {
125+
// @ts-expect-error - private constructor
126+
const task = new Task('task-id', 'context-id', mockConfig);
127+
128+
// 1. Register a tool
129+
task['_registerToolCall']('tool-1', 'scheduled');
130+
131+
// 2. Start waiting
132+
const waitPromise = task.waitForPendingTools();
133+
134+
// 3. Cancel pending tools
135+
task.cancelPendingTools('User requested cancellation');
136+
137+
// 4. Verify waitPromise rejects with the reason
138+
await expect(waitPromise).rejects.toThrow('User requested cancellation');
139+
});
140+
141+
it('should handle concurrent tool scheduling correctly', async () => {
142+
// @ts-expect-error - private constructor
143+
const task = new Task('task-id', 'context-id', mockConfig);
144+
145+
// 1. Register a tool and start waiting
146+
task['_registerToolCall']('tool-1', 'scheduled');
147+
const waitPromise = task.waitForPendingTools();
148+
149+
// 2. Schedule another tool concurrently (e.g. from a secondary user message)
150+
// This should NOT resolve the current waitPromise until both are done
151+
await task.scheduleToolCalls(
152+
[{ callId: 'tool-2', name: 't2', args: {} }],
153+
new AbortController().signal,
154+
);
155+
156+
expect(task['pendingToolCalls'].size).toBe(2);
157+
158+
// 3. Resolve tool 1
159+
task['_resolveToolCall']('tool-1');
160+
161+
// 4. Verify waitPromise is still pending
162+
let resolved = false;
163+
waitPromise.then(() => (resolved = true));
164+
await new Promise((resolve) => setTimeout(resolve, 10));
165+
expect(resolved).toBe(false);
166+
167+
// 5. Resolve tool 2
168+
task['_resolveToolCall']('tool-2');
169+
170+
// 6. Now it should resolve
171+
await expect(waitPromise).resolves.toBeUndefined();
172+
});
173+
});

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
ApprovalMode,
1313
Scheduler,
1414
type MessageBus,
15+
type ToolLiveOutput,
1516
} from '@google/gemini-cli-core';
1617
import { createMockConfig } from '../utils/testing_utils.js';
1718
import type { ExecutionEventBus } from '@a2a-js/sdk/server';
@@ -608,6 +609,74 @@ describe('Task Event-Driven Scheduler', () => {
608609
);
609610
});
610611

612+
it('should handle multi-turn tool resolution correctly', async () => {
613+
// @ts-expect-error - Calling private constructor
614+
const task = new Task('task-id', 'context-id', mockConfig);
615+
616+
task['_registerToolCall']('1', 'scheduled');
617+
task['_registerToolCall']('2', 'scheduled');
618+
619+
const handler = (messageBus.subscribe as Mock).mock.calls.find(
620+
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
621+
)?.[1];
622+
623+
// Turn 1: Resolve tool 1
624+
handler({
625+
type: MessageBusType.TOOL_CALLS_UPDATE,
626+
toolCalls: [
627+
{
628+
request: { callId: '1', name: 't1' },
629+
status: 'success',
630+
response: { responseParts: [] },
631+
},
632+
],
633+
schedulerId: 'task-id',
634+
});
635+
636+
expect(task['pendingToolCalls'].size).toBe(1);
637+
expect(task['pendingToolCalls'].has('2')).toBe(true);
638+
639+
// Turn 2: Resolve tool 2
640+
handler({
641+
type: MessageBusType.TOOL_CALLS_UPDATE,
642+
toolCalls: [
643+
{
644+
request: { callId: '2', name: 't2' },
645+
status: 'success',
646+
response: { responseParts: [] },
647+
},
648+
],
649+
schedulerId: 'task-id',
650+
});
651+
652+
expect(task['pendingToolCalls'].size).toBe(0);
653+
});
654+
655+
it('should handle subagent progress events from the scheduler', async () => {
656+
// @ts-expect-error - Calling private constructor
657+
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
658+
659+
// Trigger _schedulerOutputUpdate with subagent progress
660+
task['_schedulerOutputUpdate']('tool-1', {
661+
isSubagentProgress: true,
662+
agentName: 'researcher',
663+
recentActivity: [],
664+
} as ToolLiveOutput);
665+
666+
expect(mockEventBus.publish).toHaveBeenCalledWith(
667+
expect.objectContaining({
668+
kind: 'artifact-update',
669+
artifact: expect.objectContaining({
670+
parts: [
671+
expect.objectContaining({
672+
text: expect.stringContaining('researcher'),
673+
}),
674+
],
675+
}),
676+
}),
677+
);
678+
});
679+
611680
it('should wait for executing tools before transitioning to input-required state', async () => {
612681
// @ts-expect-error - Calling private constructor
613682
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);

0 commit comments

Comments
 (0)