feat(plan): reuse standard tool confirmation for AskUser tool (#17864)

Co-authored-by: jacob314 <jacob314@gmail.com>

packages/core/src/tools/ask-user.test.ts

@@ -6,12 +6,9 @@
 
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { AskUserTool } from './ask-user.js';
-import {
-  MessageBusType,
-  QuestionType,
-  type Question,
-} from '../confirmation-bus/types.js';
+import { QuestionType, type Question } from '../confirmation-bus/types.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
+import { ToolConfirmationOutcome } from './tools.js';
 
 describe('AskUserTool', () => {
   let mockMessageBus: MessageBus;
@@ -213,142 +210,183 @@ describe('AskUserTool', () => {
     });
   });
 
-  it('should publish ASK_USER_REQUEST and wait for response', async () => {
-    const questions = [
-      {
-        question: 'How should we proceed with this task?',
-        header: 'Approach',
-        options: [
-          {
-            label: 'Quick fix (Recommended)',
-            description:
-              'Apply the most direct solution to resolve the immediate issue.',
-          },
-          {
-            label: 'Comprehensive refactor',
-            description:
-              'Restructure the affected code for better long-term maintainability.',
-          },
-        ],
-        multiSelect: false,
-      },
-    ];
-
-    const invocation = tool.build({ questions });
-    const executePromise = invocation.execute(new AbortController().signal);
-
-    // Verify publish called with normalized questions (type defaults to CHOICE)
-    expect(mockMessageBus.publish).toHaveBeenCalledWith(
-      expect.objectContaining({
-        type: MessageBusType.ASK_USER_REQUEST,
-        questions: questions.map((q) => ({
-          ...q,
-          type: QuestionType.CHOICE,
-        })),
-      }),
-    );
-
-    // Get the correlation ID from the published message
-    const publishCall = vi.mocked(mockMessageBus.publish).mock.calls[0][0] as {
-      correlationId: string;
-    };
-    const correlationId = publishCall.correlationId;
-    expect(correlationId).toBeDefined();
-
-    // Verify subscribe called
-    expect(mockMessageBus.subscribe).toHaveBeenCalledWith(
-      MessageBusType.ASK_USER_RESPONSE,
-      expect.any(Function),
-    );
-
-    // Simulate response
-    const subscribeCall = vi
-      .mocked(mockMessageBus.subscribe)
-      .mock.calls.find((call) => call[0] === MessageBusType.ASK_USER_RESPONSE);
-    const handler = subscribeCall![1];
-
-    const answers = { '0': 'Quick fix (Recommended)' };
-    handler({
-      type: MessageBusType.ASK_USER_RESPONSE,
-      correlationId,
-      answers,
-    });
-
-    const result = await executePromise;
-    expect(result.returnDisplay).toContain('User answered:');
-    expect(result.returnDisplay).toContain(
-      '  Approach → Quick fix (Recommended)',
-    );
-    expect(JSON.parse(result.llmContent as string)).toEqual({ answers });
-  });
-
-  it('should display message when user submits without answering', async () => {
-    const questions = [
-      {
-        question: 'Which approach?',
-        header: 'Approach',
-        options: [
-          { label: 'Option A', description: 'First option' },
-          { label: 'Option B', description: 'Second option' },
-        ],
-      },
-    ];
-
-    const invocation = tool.build({ questions });
-    const executePromise = invocation.execute(new AbortController().signal);
-
-    // Get the correlation ID from the published message
-    const publishCall = vi.mocked(mockMessageBus.publish).mock.calls[0][0] as {
-      correlationId: string;
-    };
-    const correlationId = publishCall.correlationId;
-
-    // Simulate response with empty answers
-    const subscribeCall = vi
-      .mocked(mockMessageBus.subscribe)
-      .mock.calls.find((call) => call[0] === MessageBusType.ASK_USER_RESPONSE);
-    const handler = subscribeCall![1];
-
-    handler({
-      type: MessageBusType.ASK_USER_RESPONSE,
-      correlationId,
-      answers: {},
-    });
-
-    const result = await executePromise;
-    expect(result.returnDisplay).toBe(
-      'User submitted without answering questions.',
-    );
-    expect(JSON.parse(result.llmContent as string)).toEqual({ answers: {} });
-  });
-
-  it('should handle cancellation', async () => {
-    const invocation = tool.build({
-      questions: [
+  describe('shouldConfirmExecute', () => {
+    it('should return confirmation details with normalized questions', async () => {
+      const questions = [
         {
-          question: 'Which sections of the documentation should be updated?',
-          header: 'Docs',
+          question: 'How should we proceed with this task?',
+          header: 'Approach',
           options: [
             {
-              label: 'User Guide',
-              description: 'Update the main user-facing documentation.',
+              label: 'Quick fix (Recommended)',
+              description:
+                'Apply the most direct solution to resolve the immediate issue.',
             },
             {
-              label: 'API Reference',
-              description: 'Update the detailed API documentation.',
+              label: 'Comprehensive refactor',
+              description:
+                'Restructure the affected code for better long-term maintainability.',
             },
           ],
-          multiSelect: true,
+          multiSelect: false,
         },
-      ],
+      ];
+
+      const invocation = tool.build({ questions });
+      const details = await invocation.shouldConfirmExecute(
+        new AbortController().signal,
+      );
+
+      expect(details).not.toBe(false);
+      if (details && details.type === 'ask_user') {
+        expect(details.title).toBe('Ask User');
+        expect(details.questions).toEqual(
+          questions.map((q) => ({
+            ...q,
+            type: QuestionType.CHOICE,
+          })),
+        );
+        expect(typeof details.onConfirm).toBe('function');
+      } else {
+        // Type guard for TypeScript
+        expect(details).toBeTruthy();
+      }
     });
 
-    const controller = new AbortController();
-    const executePromise = invocation.execute(controller.signal);
+    it('should normalize question type to CHOICE when omitted', async () => {
+      const questions = [
+        {
+          question: 'Which approach?',
+          header: 'Approach',
+          options: [
+            { label: 'Option A', description: 'First option' },
+            { label: 'Option B', description: 'Second option' },
+          ],
+        },
+      ];
 
-    controller.abort();
+      const invocation = tool.build({ questions });
+      const details = await invocation.shouldConfirmExecute(
+        new AbortController().signal,
+      );
 
-    const result = await executePromise;
-    expect(result.error?.message).toBe('Cancelled');
+      if (details && details.type === 'ask_user') {
+        expect(details.questions[0].type).toBe(QuestionType.CHOICE);
+      }
+    });
+  });
+
+  describe('execute', () => {
+    it('should return user answers after confirmation', async () => {
+      const questions = [
+        {
+          question: 'How should we proceed with this task?',
+          header: 'Approach',
+          options: [
+            {
+              label: 'Quick fix (Recommended)',
+              description:
+                'Apply the most direct solution to resolve the immediate issue.',
+            },
+            {
+              label: 'Comprehensive refactor',
+              description:
+                'Restructure the affected code for better long-term maintainability.',
+            },
+          ],
+          multiSelect: false,
+        },
+      ];
+
+      const invocation = tool.build({ questions });
+      const details = await invocation.shouldConfirmExecute(
+        new AbortController().signal,
+      );
+
+      // Simulate confirmation with answers
+      if (details && 'onConfirm' in details) {
+        const answers = { '0': 'Quick fix (Recommended)' };
+        await details.onConfirm(ToolConfirmationOutcome.ProceedOnce, {
+          answers,
+        });
+      }
+
+      const result = await invocation.execute(new AbortController().signal);
+      expect(result.returnDisplay).toContain('User answered:');
+      expect(result.returnDisplay).toContain(
+        '  Approach → Quick fix (Recommended)',
+      );
+      expect(JSON.parse(result.llmContent as string)).toEqual({
+        answers: { '0': 'Quick fix (Recommended)' },
+      });
+    });
+
+    it('should display message when user submits without answering', async () => {
+      const questions = [
+        {
+          question: 'Which approach?',
+          header: 'Approach',
+          options: [
+            { label: 'Option A', description: 'First option' },
+            { label: 'Option B', description: 'Second option' },
+          ],
+        },
+      ];
+
+      const invocation = tool.build({ questions });
+      const details = await invocation.shouldConfirmExecute(
+        new AbortController().signal,
+      );
+
+      // Simulate confirmation with empty answers
+      if (details && 'onConfirm' in details) {
+        await details.onConfirm(ToolConfirmationOutcome.ProceedOnce, {
+          answers: {},
+        });
+      }
+
+      const result = await invocation.execute(new AbortController().signal);
+      expect(result.returnDisplay).toBe(
+        'User submitted without answering questions.',
+      );
+      expect(JSON.parse(result.llmContent as string)).toEqual({ answers: {} });
+    });
+
+    it('should handle cancellation', async () => {
+      const invocation = tool.build({
+        questions: [
+          {
+            question: 'Which sections of the documentation should be updated?',
+            header: 'Docs',
+            options: [
+              {
+                label: 'User Guide',
+                description: 'Update the main user-facing documentation.',
+              },
+              {
+                label: 'API Reference',
+                description: 'Update the detailed API documentation.',
+              },
+            ],
+            multiSelect: true,
+          },
+        ],
+      });
+
+      const details = await invocation.shouldConfirmExecute(
+        new AbortController().signal,
+      );
+
+      // Simulate cancellation
+      if (details && 'onConfirm' in details) {
+        await details.onConfirm(ToolConfirmationOutcome.Cancel);
+      }
+
+      const result = await invocation.execute(new AbortController().signal);
+      expect(result.returnDisplay).toBe('User dismissed dialog');
+      expect(result.llmContent).toBe(
+        'User dismissed ask_user dialog without answering.',
+      );
+    });
   });
 });

packages/core/src/tools/ask-user.ts

@@ -9,17 +9,12 @@ import {
   BaseToolInvocation,
   type ToolResult,
   Kind,
-  type ToolCallConfirmationDetails,
+  type ToolAskUserConfirmationDetails,
+  type ToolConfirmationPayload,
+  ToolConfirmationOutcome,
 } from './tools.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
-import {
-  MessageBusType,
-  QuestionType,
-  type Question,
-  type AskUserRequest,
-  type AskUserResponse,
-} from '../confirmation-bus/types.js';
-import { randomUUID } from 'node:crypto';
+import { QuestionType, type Question } from '../confirmation-bus/types.js';
 import { ASK_USER_TOOL_NAME, ASK_USER_DISPLAY_NAME } from './tool-names.js';
 
 export interface AskUserParams {
@@ -165,100 +160,61 @@ export class AskUserInvocation extends BaseToolInvocation<
   AskUserParams,
   ToolResult
 > {
+  private confirmationOutcome: ToolConfirmationOutcome | null = null;
+  private userAnswers: { [questionIndex: string]: string } = {};
+
   override async shouldConfirmExecute(
     _abortSignal: AbortSignal,
-  ): Promise<ToolCallConfirmationDetails | false> {
-    return false;
+  ): Promise<ToolAskUserConfirmationDetails | false> {
+    const normalizedQuestions = this.params.questions.map((q) => ({
+      ...q,
+      type: q.type ?? QuestionType.CHOICE,
+    }));
+
+    return {
+      type: 'ask_user',
+      title: 'Ask User',
+      questions: normalizedQuestions,
+      onConfirm: async (
+        outcome: ToolConfirmationOutcome,
+        payload?: ToolConfirmationPayload,
+      ) => {
+        this.confirmationOutcome = outcome;
+        if (payload?.answers) {
+          this.userAnswers = payload.answers;
+        }
+      },
+    };
   }
 
   getDescription(): string {
     return `Asking user: ${this.params.questions.map((q) => q.question).join(', ')}`;
   }
 
-  async execute(signal: AbortSignal): Promise<ToolResult> {
-    const correlationId = randomUUID();
+  async execute(_signal: AbortSignal): Promise<ToolResult> {
+    if (this.confirmationOutcome === ToolConfirmationOutcome.Cancel) {
+      return {
+        llmContent: 'User dismissed ask_user dialog without answering.',
+        returnDisplay: 'User dismissed dialog',
+      };
+    }
 
-    const request: AskUserRequest = {
-      type: MessageBusType.ASK_USER_REQUEST,
-      questions: this.params.questions.map((q) => ({
-        ...q,
-        type: q.type ?? QuestionType.CHOICE,
-      })),
-      correlationId,
+    const answerEntries = Object.entries(this.userAnswers);
+    const hasAnswers = answerEntries.length > 0;
+
+    const returnDisplay = hasAnswers
+      ? `**User answered:**\n${answerEntries
+          .map(([index, answer]) => {
+            const question = this.params.questions[parseInt(index, 10)];
+            const category = question?.header ?? `Q${index}`;
+            return `  ${category} → ${answer}`;
+          })
+          .join('\n')}`
+      : 'User submitted without answering questions.';
+
+    return {
+      llmContent: JSON.stringify({ answers: this.userAnswers }),
+      returnDisplay,
     };
-
-    return new Promise<ToolResult>((resolve, reject) => {
-      const responseHandler = (response: AskUserResponse): void => {
-        if (response.correlationId === correlationId) {
-          cleanup();
-
-          // Handle user cancellation
-          if (response.cancelled) {
-            resolve({
-              llmContent: 'User dismissed ask user dialog without answering.',
-              returnDisplay: 'User dismissed dialog',
-            });
-            return;
-          }
-
-          // Build formatted key-value display
-          const answerEntries = Object.entries(response.answers);
-          const hasAnswers = answerEntries.length > 0;
-
-          const returnDisplay = hasAnswers
-            ? `**User answered:**\n${answerEntries
-                .map(([index, answer]) => {
-                  const question = this.params.questions[parseInt(index, 10)];
-                  const category = question?.header ?? `Q${index}`;
-                  return `  ${category} → ${answer}`;
-                })
-                .join('\n')}`
-            : 'User submitted without answering questions.';
-
-          resolve({
-            llmContent: JSON.stringify({ answers: response.answers }),
-            returnDisplay,
-          });
-        }
-      };
-
-      const cleanup = () => {
-        if (responseHandler) {
-          this.messageBus.unsubscribe(
-            MessageBusType.ASK_USER_RESPONSE,
-            responseHandler,
-          );
-        }
-        signal.removeEventListener('abort', abortHandler);
-      };
-
-      const abortHandler = () => {
-        cleanup();
-        resolve({
-          llmContent: 'Tool execution cancelled by user.',
-          returnDisplay: 'Cancelled',
-          error: {
-            message: 'Cancelled',
-          },
-        });
-      };
-
-      if (signal.aborted) {
-        abortHandler();
-        return;
-      }
-
-      signal.addEventListener('abort', abortHandler);
-      this.messageBus.subscribe(
-        MessageBusType.ASK_USER_RESPONSE,
-        responseHandler,
-      );
-
-      // Publish request
-      this.messageBus.publish(request).catch((err) => {
-        cleanup();
-        reject(err);
-      });
-    });
   }
 }

packages/core/src/tools/tools.ts

@@ -16,6 +16,7 @@ import {
   MessageBusType,
   type ToolConfirmationRequest,
   type ToolConfirmationResponse,
+  type Question,
 } from '../confirmation-bus/types.js';
 
 /**
@@ -695,7 +696,9 @@ export interface ToolEditConfirmationDetails {
 export interface ToolConfirmationPayload {
   // used to override `modifiedProposedContent` for modifiable tools in the
   // inline modify flow
-  newContent: string;
+  newContent?: string;
+  // used for askuser tool to return user's answers
+  answers?: { [questionIndex: string]: string };
 }
 
 export interface ToolExecuteConfirmationDetails {
@@ -725,11 +728,22 @@ export interface ToolInfoConfirmationDetails {
   urls?: string[];
 }
 
+export interface ToolAskUserConfirmationDetails {
+  type: 'ask_user';
+  title: string;
+  questions: Question[];
+  onConfirm: (
+    outcome: ToolConfirmationOutcome,
+    payload?: ToolConfirmationPayload,
+  ) => Promise<void>;
+}
+
 export type ToolCallConfirmationDetails =
   | ToolEditConfirmationDetails
   | ToolExecuteConfirmationDetails
   | ToolMcpConfirmationDetails
-  | ToolInfoConfirmationDetails;
+  | ToolInfoConfirmationDetails
+  | ToolAskUserConfirmationDetails;
 
 export enum ToolConfirmationOutcome {
   ProceedOnce = 'proceed_once',