diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 2a0b19e627e..552191d16aa 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -11315,6 +11315,7 @@ describe('CoreToolScheduler telemetry spans', () => { shouldAvoidPermissionPrompts?: boolean; experimentalZedIntegration?: boolean; approvalMode?: ApprovalMode; + ideMode?: boolean; includeSensitiveSpanAttributes?: boolean; sensitiveSpanAttributeMaxLength?: number; onToolCallsUpdate?: ReturnType; @@ -11386,7 +11387,7 @@ describe('CoreToolScheduler telemetry spans', () => { getInputFormat: () => options.inputFormat ?? InputFormat.TEXT, getExperimentalZedIntegration: () => options.experimentalZedIntegration ?? false, - getIdeMode: () => false, + getIdeMode: () => options.ideMode ?? false, getShouldAvoidPermissionPrompts: () => options.shouldAvoidPermissionPrompts ?? false, getTelemetryIncludeSensitiveSpanAttributes: () => @@ -12975,7 +12976,7 @@ describe('CoreToolScheduler telemetry spans', () => { experimentalZedIntegration?: boolean; args?: Record; abortController?: AbortController; - tools?: MockTool[]; + tools?: AnyDeclarativeTool[]; }): Promise<{ scheduler: CoreToolScheduler; onAllToolCallsComplete: ReturnType; @@ -12989,7 +12990,7 @@ describe('CoreToolScheduler telemetry spans', () => { [ { callId: 'ask-call', - name: 'mockTool', + name: options.tools?.[0]?.name ?? 'mockTool', args: options.args ?? { input: 'x' }, isClientInitiated: false, prompt_id: 'prompt-ask', @@ -13084,6 +13085,178 @@ describe('CoreToolScheduler telemetry spans', () => { expect(blocked[0].ended).toBe(true); }); + it('shows the edit diff when a PreToolUse ask requires approval', async () => { + const messageBus = askMessageBus('review protected file'); + const { onToolCallsUpdate } = await scheduleWithAsk({ + messageBus, + tools: [new MockEditTool()], + }); + + const waiting = (await waitForStatus( + onToolCallsUpdate, + 'awaiting_approval', + )) as WaitingToolCall; + expect(waiting.confirmationDetails).toMatchObject({ + type: 'edit', + fileName: 'test.txt', + newContent: 'new content', + fileDiff: + '--- test.txt\n+++ test.txt\n@@ -1,1 +1,1 @@\n-old content\n+new content', + hideAlwaysAllow: true, + hideModify: true, + warnings: ['review protected file'], + }); + }); + + it('forwards the host denial reason when a bounced edit confirmation is cancelled', async () => { + const execute = vi.fn(); + const messageBus = askMessageBus('review protected file'); + const { onToolCallsUpdate, onAllToolCallsComplete } = await scheduleWithAsk( + { + messageBus, + execute, + tools: [new MockEditTool()], + }, + ); + + const waiting = (await waitForStatus( + onToolCallsUpdate, + 'awaiting_approval', + )) as WaitingToolCall; + expect(waiting.confirmationDetails.type).toBe('edit'); + + // stream-json hosts deny with a reason: permissionController calls + // onConfirm(Cancel, { cancelMessage }). The bounced edit wrapper must + // forward the payload like the info fallback (and the pre-PR synthetic + // prompt) instead of dropping it. + await waiting.confirmationDetails.onConfirm( + ToolConfirmationOutcome.Cancel, + { cancelMessage: 'host policy: no edits' }, + ); + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + const completed = onAllToolCallsComplete.mock.calls.at( + -1, + )?.[0] as ToolCall[]; + expect(completed[0].status).toBe('cancelled'); + expect(execute).not.toHaveBeenCalled(); + // The host's reason — not the generic 'User did not allow tool call' — + // must reach the model via the cancelled response. + expect( + JSON.stringify((completed[0] as { response?: unknown }).response), + ).toContain('host policy: no edits'); + }); + + it('refuses a stale round-1 IDE resolution for a bounced edit confirmation', async () => { + // Round-1 edit confirmation (DEFAULT mode, IDE diffing on) opens the + // IDE diff. The user approves via the scheduler path; the PreToolUse + // hook returns 'ask' and the call bounces back to awaiting_approval. + // Only THEN does the round-1 openDiff resolve — with edited panel + // content — mirroring ToolConfirmationMessage.handleConfirm, which + // fires onConfirm before awaiting resolveDiffFromCli. That stale + // resolution must not answer the bounced confirmation: without the + // bouncedAwaitingApproval guard its content would flow through + // _applyInlineModify and execute with the hook never re-consulted. + let resolveIdeDiff!: (r: { status: 'accepted'; content: string }) => void; + const ideDiffResolution = new Promise<{ + status: 'accepted'; + content: string; + }>((resolve) => { + resolveIdeDiff = resolve; + }); + vi.mocked(IdeClient.getInstance).mockResolvedValue( + mockIdeClient as unknown as IdeClient, + ); + mockIdeClient.isDiffingEnabled.mockReturnValue(true); + mockIdeClient.openDiff.mockReset(); + mockIdeClient.openDiff.mockReturnValue(ideDiffResolution); + + const execute = vi.fn().mockResolvedValue({ + llmContent: 'ok', + returnDisplay: 'ok', + }); + const messageBus = askMessageBus('review protected file'); + const awaitingSnapshots: ToolCallConfirmationDetails[] = []; + const onToolCallsUpdate = vi.fn((calls: ToolCall[]) => { + for (const call of calls) { + if ( + call.request.callId === 'stale-ide-bounce' && + call.status === 'awaiting_approval' + ) { + awaitingSnapshots.push(call.confirmationDetails); + } + } + }); + const { scheduler, onAllToolCallsComplete } = buildScheduler({ + tools: [new MockEditTool(execute)], + messageBus, + disableHooks: false, + approvalMode: ApprovalMode.DEFAULT, + ideMode: true, + onToolCallsUpdate, + }); + + await scheduler.schedule( + [ + { + callId: 'stale-ide-bounce', + name: 'mockEditTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-stale-ide-bounce', + }, + ], + new AbortController().signal, + ); + + // Round-1 confirmation opened the IDE diff. + await vi.waitFor(() => { + expect(awaitingSnapshots).toHaveLength(1); + }); + expect(mockIdeClient.openDiff).toHaveBeenCalledTimes(1); + const round1 = awaitingSnapshots[0]; + + // Approve round-1 via the scheduler path; the hook ask bounces the + // call back to awaiting_approval with its own edit confirmation. + await round1.onConfirm(ToolConfirmationOutcome.ProceedOnce); + await vi.waitFor(() => { + expect(awaitingSnapshots).toHaveLength(2); + }); + const bounced = awaitingSnapshots[1]; + expect(bounced).toMatchObject({ type: 'edit', hideModify: true }); + + // The stale round-1 IDE diff now resolves as accepted with edited + // panel content. It must be refused — the call stays parked on the + // bounced confirmation with the hook-reviewed content untouched. + resolveIdeDiff({ status: 'accepted', content: 'STALE-PANEL-CONTENT' }); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(execute).not.toHaveBeenCalled(); + expect(awaitingSnapshots).toHaveLength(2); + expect( + (awaitingSnapshots.at(-1) as { newContent?: string }).newContent, + ).toBe('new content'); + + // The bounced confirmation is the only valid resolver: approving it + // executes exactly once (and the hook does not re-fire). + await bounced.onConfirm(ToolConfirmationOutcome.ProceedOnce); + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + const completed = onAllToolCallsComplete.mock.calls.at( + -1, + )?.[0] as ToolCall[]; + expect(completed[0].status).toBe('success'); + expect(execute).toHaveBeenCalledTimes(1); + expect(preToolUseCallCount(messageBus)).toBe(1); + + // Leave the module-level IDE mocks the way this test found them. + mockIdeClient.openDiff.mockReset(); + mockIdeClient.isDiffingEnabled.mockReset(); + vi.mocked(IdeClient.getInstance).mockReset(); + }); + it('creates new confirmation details when a tool bounces after approval', async () => { const execute = vi.fn().mockResolvedValue({ llmContent: 'ok', diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 3a6a9b2b95f..ea3073ca2cd 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -4110,6 +4110,17 @@ export class CoreToolScheduler { ); if (!still) return; + // Guard: a PreToolUse-'ask' bounce re-enters awaiting_approval, so the + // guard above alone would let this stale round-1 resolution answer the + // BOUNCED confirmation. The accept path would flow resolution.content + // through _applyInlineModify (bounced edit details are type 'edit', and + // that path does not check hideModify) and execute IDE-panel content + // the hook never reviewed on the hook-skipping re-execution; the + // reject path would cancel a prompt the user never answered. Only the + // bounce's own confirmation may resolve a bounced call — the round-1 + // diff is closed by resolveDiffFromCli regardless. + if (this.bouncedAwaitingApproval.has(callId)) return; + if (resolution.status === 'accepted') { // When content is unchanged, skip the inline modify path so that // the original tool params (e.g. partial old_string for edit tool) @@ -4397,39 +4408,74 @@ export class CoreToolScheduler { /** * Bounce a tool from the EXECUTION phase back to awaiting_approval so the * user can confirm a PreToolUse 'ask' decision in the TUI. Reuses the - * standard confirmation machinery: a synthetic 'info' confirmation whose - * onConfirm routes through handleConfirmationResponse (ProceedOnce → - * re-execute, Cancel → cancelled). `hideAlwaysAllow` is set because the - * hook re-evaluates on every call, so an "always allow" rule is - * meaningless. The callId is added to `bouncedAwaitingApproval` BEFORE - * the status change so executeSingleToolCall's finally keeps the tool - * span open across the bounce and the re-execution skips the hook + - * prelude (see `_executeToolCallBody`). + * standard confirmation machinery, including the existing diff view for + * edit tools. `hideAlwaysAllow` is set because the hook re-evaluates on + * every call, so an "always allow" rule is meaningless. The callId is + * added to `bouncedAwaitingApproval` BEFORE the status change so + * executeSingleToolCall's finally keeps the tool span open across the + * bounce and the re-execution skips the hook + prelude (see + * `_executeToolCallBody`). */ - private bounceToAwaitingApprovalForAsk( + private async bounceToAwaitingApprovalForAsk( scheduledCall: ScheduledToolCall, reason: string | undefined, toolSpan: Span, signal: AbortSignal, - ): void { + ): Promise { const { callId, name: toolName } = scheduledCall.request; const canonicalName = canonicalToolName(toolName); + const hookReason = + reason || + `A PreToolUse hook requested confirmation before running ${toolName}.`; + + let confirmationDetails: ToolCallConfirmationDetails | undefined; + if (scheduledCall.tool.kind === Kind.Edit) { + try { + const editDetails = + await scheduledCall.invocation.getConfirmationDetails(signal); + if (editDetails.type === 'edit') { + confirmationDetails = { + ...editDetails, + hideAlwaysAllow: true, + hideModify: true, + warnings: [hookReason, ...(editDetails.warnings ?? [])], + onConfirm: (outcome, payload) => + this.handleConfirmationResponse( + callId, + editDetails.onConfirm, + outcome, + signal, + // Forward the host's denial reason (the stream-json + // permissionController sends { cancelMessage } on deny) but + // keep the modify channel closed: hideModify is set above, + // so a payload's newContent must not rewrite the + // hook-reviewed content on a bounce. + payload?.cancelMessage + ? { cancelMessage: payload.cancelMessage } + : undefined, + ), + }; + } + } catch (error) { + debugLogger.warn( + `Failed to prepare edit confirmation for ${toolName}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } + + if (this.cancelPreExecutionIfAborted(callId, signal, toolSpan)) return; this.bouncedAwaitingApproval.add(callId); - const confirmationDetails: ToolCallConfirmationDetails = { + confirmationDetails ??= { type: 'info', title: `Hook requested confirmation to run ${toolName}`, - prompt: - reason || - `A PreToolUse hook requested confirmation before running ${toolName}.`, + prompt: hookReason, renderPromptAsPlainText: true, hideAlwaysAllow: true, onConfirm: (outcome, payload) => this.handleConfirmationResponse( callId, - // No real tool onConfirm — for this synthetic prompt all of the - // approve/deny handling lives in handleConfirmationResponse. async () => {}, outcome, signal, @@ -4616,7 +4662,7 @@ export class CoreToolScheduler { // Preserve the tool_use_id so the post-approval re-execution // reuses it (see the toolUseId comment above). this.bouncedToolUseId.set(callId, toolUseId); - this.bounceToAwaitingApprovalForAsk( + await this.bounceToAwaitingApprovalForAsk( scheduledCall, preHookResult.blockReason, span,