feat(core): wire up UI for ASK_USER policy decisions in message bus (#10630)

allenhutchison · web-flow · commit b8df8b2ab8b6 · 2025-10-15T16:39:41.000Z
diff --git a/packages/a2a-server/src/utils/testing_utils.ts b/packages/a2a-server/src/utils/testing_utils.ts
@@ -50,6 +50,9 @@ export function createMockConfig(
     getEmbeddingModel: vi.fn().mockReturnValue('text-embedding-004'),
     getSessionId: vi.fn().mockReturnValue('test-session-id'),
     getUserTier: vi.fn(),
+    getEnableMessageBusIntegration: vi.fn().mockReturnValue(false),
+    getMessageBus: vi.fn(),
+    getPolicyEngine: vi.fn(),
     ...overrides,
   } as unknown as Config;
 
diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts
@@ -70,6 +70,9 @@ const mockConfig = {
   getUseModelRouter: () => false,
   getGeminiClient: () => null, // No client needed for these tests
   getShellExecutionConfig: () => ({ terminalWidth: 80, terminalHeight: 24 }),
+  getEnableMessageBusIntegration: () => false,
+  getMessageBus: () => null,
+  getPolicyEngine: () => null,
 } as unknown as Config;
 
 const mockTool = new MockTool({
diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts
@@ -24,6 +24,11 @@ export interface ToolConfirmationResponse {
   type: MessageBusType.TOOL_CONFIRMATION_RESPONSE;
   correlationId: string;
   confirmed: boolean;
+  /**
+   * When true, indicates that policy decision was ASK_USER and the tool should
+   * show its legacy confirmation UI instead of auto-proceeding.
+   */
+  requiresUserConfirmation?: boolean;
 }
 
 export interface ToolPolicyRejection {
diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts
@@ -255,6 +255,9 @@ describe('CoreToolScheduler', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -332,6 +335,9 @@ describe('CoreToolScheduler', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null,
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -365,15 +371,18 @@ describe('CoreToolScheduler', () => {
   describe('getToolSuggestion', () => {
     it('should suggest the top N closest tool names for a typo', () => {
       // Create mocked tool registry
+      const mockToolRegistry = {
+        getAllToolNames: () => ['list_files', 'read_file', 'write_file'],
+      } as unknown as ToolRegistry;
       const mockConfig = {
         getToolRegistry: () => mockToolRegistry,
         getUseSmartEdit: () => false,
         getUseModelRouter: () => false,
         getGeminiClient: () => null, // No client needed for these tests
+        getEnableMessageBusIntegration: () => false,
+        getMessageBus: () => null,
+        getPolicyEngine: () => null,
       } as unknown as Config;
-      const mockToolRegistry = {
-        getAllToolNames: () => ['list_files', 'read_file', 'write_file'],
-      } as unknown as ToolRegistry;
 
       // Create scheduler
       const scheduler = new CoreToolScheduler({
@@ -448,6 +457,9 @@ describe('CoreToolScheduler with payload', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -768,6 +780,9 @@ describe('CoreToolScheduler edit cancellation', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -874,6 +889,9 @@ describe('CoreToolScheduler YOLO mode', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -981,6 +999,9 @@ describe('CoreToolScheduler request queueing', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -1113,6 +1134,9 @@ describe('CoreToolScheduler request queueing', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -1215,6 +1239,9 @@ describe('CoreToolScheduler request queueing', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -1287,6 +1314,9 @@ describe('CoreToolScheduler request queueing', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     const testTool = new TestApprovalTool(mockConfig);
@@ -1475,6 +1505,8 @@ describe('CoreToolScheduler Sequential Execution', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null,
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
@@ -1595,6 +1627,8 @@ describe('CoreToolScheduler Sequential Execution', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null,
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
     } as unknown as Config;
 
     const scheduler = new CoreToolScheduler({
diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts
@@ -42,6 +42,8 @@ import * as path from 'node:path';
 import { doesToolInvocationMatch } from '../utils/tool-utils.js';
 import levenshtein from 'fast-levenshtein';
 import { ShellToolInvocation } from '../tools/shell.js';
+import type { ToolConfirmationRequest } from '../confirmation-bus/types.js';
+import { MessageBusType } from '../confirmation-bus/types.js';
 
 export type ValidatingToolCall = {
   status: 'validating';
@@ -352,6 +354,15 @@ export class CoreToolScheduler {
     this.onToolCallsUpdate = options.onToolCallsUpdate;
     this.getPreferredEditor = options.getPreferredEditor;
     this.onEditorClose = options.onEditorClose;
+
+    // Subscribe to message bus for ASK_USER policy decisions
+    if (this.config.getEnableMessageBusIntegration()) {
+      const messageBus = this.config.getMessageBus();
+      messageBus.subscribe(
+        MessageBusType.TOOL_CONFIRMATION_REQUEST,
+        this.handleToolConfirmationRequest.bind(this),
+      );
+    }
   }
 
   private setStatusInternal(
@@ -1160,6 +1171,26 @@ export class CoreToolScheduler {
     });
   }
 
+  /**
+   * Handle tool confirmation requests from the message bus when policy decision is ASK_USER.
+   * This publishes a response with requiresUserConfirmation=true to signal the tool
+   * that it should fall back to its legacy confirmation UI.
+   */
+  private handleToolConfirmationRequest(
+    request: ToolConfirmationRequest,
+  ): void {
+    // When ASK_USER policy decision is made, the message bus emits the request here.
+    // We respond with requiresUserConfirmation=true to tell the tool to use its
+    // legacy confirmation flow (which will show diffs, URLs, etc in the UI).
+    const messageBus = this.config.getMessageBus();
+    messageBus.publish({
+      type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
+      correlationId: request.correlationId,
+      confirmed: false, // Not auto-approved
+      requiresUserConfirmation: true, // Use legacy UI confirmation
+    });
+  }
+
   private async autoApproveCompatiblePendingTools(
     signal: AbortSignal,
     triggeringCallId: string,
diff --git a/packages/core/src/core/nonInteractiveToolExecutor.test.ts b/packages/core/src/core/nonInteractiveToolExecutor.test.ts
@@ -62,6 +62,9 @@ describe('executeToolCall', () => {
       getUseSmartEdit: () => false,
       getUseModelRouter: () => false,
       getGeminiClient: () => null, // No client needed for these tests
+      getEnableMessageBusIntegration: () => false,
+      getMessageBus: () => null,
+      getPolicyEngine: () => null,
     } as unknown as Config;
 
     abortController = new AbortController();
diff --git a/packages/core/src/tools/message-bus-integration.test.ts b/packages/core/src/tools/message-bus-integration.test.ts
@@ -52,6 +52,22 @@ class TestToolInvocation extends BaseToolInvocation<TestParams, TestResult> {
       testValue: this.params.testParam,
     };
   }
+
+  override async shouldConfirmExecute(
+    abortSignal: AbortSignal,
+  ): Promise<false> {
+    // This conditional is here to allow testing of the case where there is no message bus.
+    if (this.messageBus) {
+      const decision = await this.getMessageBusDecision(abortSignal);
+      if (decision === 'ALLOW') {
+        return false;
+      }
+      if (decision === 'DENY') {
+        throw new Error('Tool execution denied by policy');
+      }
+    }
+    return false;
+  }
 }
 
 class TestTool extends BaseDeclarativeTool<TestParams, TestResult> {
@@ -200,7 +216,7 @@ describe('Message Bus Integration', () => {
       abortController.abort();
 
       await expect(confirmationPromise).rejects.toThrow(
-        'Tool confirmation aborted',
+        'Tool execution denied by policy',
       );
     });
 
diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts
diff --git a/packages/core/src/tools/web-fetch.test.ts b/packages/core/src/tools/web-fetch.test.ts
diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts
