|
9 | 9 | type Config, |
10 | 10 | MessageBusType, |
11 | 11 | ToolConfirmationOutcome, |
| 12 | + ApprovalMode, |
12 | 13 | Scheduler, |
13 | 14 | type MessageBus, |
14 | 15 | } from '@google/gemini-cli-core'; |
@@ -133,6 +134,103 @@ describe('Task Event-Driven Scheduler', () => { |
133 | 134 | ); |
134 | 135 | }); |
135 | 136 |
|
| 137 | + it('should handle Rejection (Cancel) and Modification (ModifyWithEditor)', async () => { |
| 138 | + // @ts-expect-error - Calling private constructor |
| 139 | + const task = new Task('task-id', 'context-id', mockConfig, mockEventBus); |
| 140 | + |
| 141 | + const toolCall = { |
| 142 | + request: { callId: '1', name: 'ls', args: {} }, |
| 143 | + status: 'awaiting_approval', |
| 144 | + correlationId: 'corr-1', |
| 145 | + confirmationDetails: { type: 'info', title: 'test', prompt: 'test' }, |
| 146 | + }; |
| 147 | + |
| 148 | + const handler = (messageBus.subscribe as Mock).mock.calls.find( |
| 149 | + (call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE, |
| 150 | + )?.[1]; |
| 151 | + handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] }); |
| 152 | + |
| 153 | + // Simulate Rejection (Cancel) |
| 154 | + let handled = await ( |
| 155 | + task as unknown as { |
| 156 | + _handleToolConfirmationPart: (part: unknown) => Promise<boolean>; |
| 157 | + } |
| 158 | + )._handleToolConfirmationPart({ |
| 159 | + kind: 'data', |
| 160 | + data: { callId: '1', outcome: 'cancel' }, |
| 161 | + }); |
| 162 | + expect(handled).toBe(true); |
| 163 | + expect(messageBus.publish).toHaveBeenCalledWith( |
| 164 | + expect.objectContaining({ |
| 165 | + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, |
| 166 | + correlationId: 'corr-1', |
| 167 | + confirmed: false, |
| 168 | + }), |
| 169 | + ); |
| 170 | + |
| 171 | + const toolCall2 = { |
| 172 | + request: { callId: '2', name: 'ls', args: {} }, |
| 173 | + status: 'awaiting_approval', |
| 174 | + correlationId: 'corr-2', |
| 175 | + confirmationDetails: { type: 'info', title: 'test', prompt: 'test' }, |
| 176 | + }; |
| 177 | + handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall2] }); |
| 178 | + |
| 179 | + // Simulate ModifyWithEditor |
| 180 | + handled = await ( |
| 181 | + task as unknown as { |
| 182 | + _handleToolConfirmationPart: (part: unknown) => Promise<boolean>; |
| 183 | + } |
| 184 | + )._handleToolConfirmationPart({ |
| 185 | + kind: 'data', |
| 186 | + data: { callId: '2', outcome: 'modify_with_editor' }, |
| 187 | + }); |
| 188 | + expect(handled).toBe(true); |
| 189 | + expect(messageBus.publish).toHaveBeenCalledWith( |
| 190 | + expect.objectContaining({ |
| 191 | + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, |
| 192 | + correlationId: 'corr-2', |
| 193 | + confirmed: false, |
| 194 | + outcome: ToolConfirmationOutcome.ModifyWithEditor, |
| 195 | + payload: undefined, |
| 196 | + }), |
| 197 | + ); |
| 198 | + }); |
| 199 | + |
| 200 | + it('Execution without Confirmation (ProceedAlways / YOLO mode)', async () => { |
| 201 | + // Enable YOLO mode |
| 202 | + const yoloConfig = createMockConfig({ |
| 203 | + isEventDrivenSchedulerEnabled: () => true, |
| 204 | + getApprovalMode: () => ApprovalMode.YOLO, |
| 205 | + }) as Config; |
| 206 | + const yoloMessageBus = yoloConfig.getMessageBus(); |
| 207 | + |
| 208 | + // @ts-expect-error - Calling private constructor |
| 209 | + const _task = new Task('task-id', 'context-id', yoloConfig, mockEventBus); |
| 210 | + |
| 211 | + const toolCall = { |
| 212 | + request: { callId: '1', name: 'ls', args: {} }, |
| 213 | + status: 'awaiting_approval', |
| 214 | + correlationId: 'corr-1', |
| 215 | + confirmationDetails: { type: 'info', title: 'test', prompt: 'test' }, |
| 216 | + }; |
| 217 | + |
| 218 | + const handler = (yoloMessageBus.subscribe as Mock).mock.calls.find( |
| 219 | + (call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE, |
| 220 | + )?.[1]; |
| 221 | + handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] }); |
| 222 | + |
| 223 | + // Should immediately auto-publish ProceedOnce without user intervention |
| 224 | + expect(yoloMessageBus.publish).toHaveBeenCalledWith( |
| 225 | + expect.objectContaining({ |
| 226 | + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, |
| 227 | + correlationId: 'corr-1', |
| 228 | + confirmed: true, |
| 229 | + outcome: ToolConfirmationOutcome.ProceedOnce, |
| 230 | + }), |
| 231 | + ); |
| 232 | + }); |
| 233 | + |
136 | 234 | it('should handle output updates via the message bus', async () => { |
137 | 235 | // @ts-expect-error - Calling private constructor |
138 | 236 | // eslint-disable-next-line @typescript-eslint/no-unused-vars |
@@ -170,4 +268,36 @@ describe('Task Event-Driven Scheduler', () => { |
170 | 268 | }), |
171 | 269 | ); |
172 | 270 | }); |
| 271 | + |
| 272 | + it('Artifact Creation should complete without hanging', async () => { |
| 273 | + // @ts-expect-error - Calling private constructor |
| 274 | + const task = new Task('task-id', 'context-id', mockConfig, mockEventBus); |
| 275 | + |
| 276 | + const toolCallId = 'create-file-123'; |
| 277 | + task['_registerToolCall'](toolCallId, 'executing'); |
| 278 | + |
| 279 | + const toolCall = { |
| 280 | + request: { |
| 281 | + callId: toolCallId, |
| 282 | + name: 'writeFile', |
| 283 | + args: { path: 'test.sh' }, |
| 284 | + }, |
| 285 | + status: 'success', |
| 286 | + result: { ok: true }, |
| 287 | + }; |
| 288 | + |
| 289 | + const handler = (messageBus.subscribe as Mock).mock.calls.find( |
| 290 | + (call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE, |
| 291 | + )?.[1]; |
| 292 | + handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] }); |
| 293 | + |
| 294 | + // The tool should be complete and registered appropriately, eventually |
| 295 | + // triggering the toolCompletionPromise resolution when all clear. |
| 296 | + const internalTask = task as unknown as { |
| 297 | + completedToolCalls: unknown[]; |
| 298 | + pendingToolCalls: Map<string, string>; |
| 299 | + }; |
| 300 | + expect(internalTask.completedToolCalls.length).toBe(1); |
| 301 | + expect(internalTask.pendingToolCalls.size).toBe(0); |
| 302 | + }); |
173 | 303 | }); |
0 commit comments