feat(plan): handle inconsistency in schedulers (#17813)

packages/core/src/scheduler/policy.test.ts

@@ -4,8 +4,20 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, type Mocked } from 'vitest';
-import { checkPolicy, updatePolicy } from './policy.js';
+import {
+  describe,
+  it,
+  expect,
+  vi,
+  type Mocked,
+  beforeEach,
+  afterEach,
+} from 'vitest';
+import {
+  checkPolicy,
+  updatePolicy,
+  PLAN_MODE_DENIAL_MESSAGE,
+} from './policy.js';
 import type { Config } from '../config/config.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
 import { MessageBusType } from '../confirmation-bus/types.js';
@@ -15,10 +27,20 @@ import {
   type AnyDeclarativeTool,
   type ToolMcpConfirmationDetails,
   type ToolExecuteConfirmationDetails,
+  type AnyToolInvocation,
 } from '../tools/tools.js';
-import type { ValidatingToolCall } from './types.js';
+import type {
+  ValidatingToolCall,
+  ToolCallRequestInfo,
+  CompletedToolCall,
+} from './types.js';
 import type { PolicyEngine } from '../policy/policy-engine.js';
 import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
+import { CoreToolScheduler } from '../core/coreToolScheduler.js';
+import { Scheduler } from './scheduler.js';
+import { ROOT_SCHEDULER_ID } from './types.js';
+import { ToolErrorType } from '../tools/tool-error.js';
+import type { ToolRegistry } from '../tools/tool-registry.js';
 
 describe('policy.ts', () => {
   describe('checkPolicy', () => {
@@ -420,3 +442,113 @@ describe('policy.ts', () => {
     });
   });
 });
+
+describe('Plan Mode Denial Consistency', () => {
+  let mockConfig: Mocked<Config>;
+  let mockMessageBus: Mocked<MessageBus>;
+  let mockPolicyEngine: Mocked<PolicyEngine>;
+  let mockToolRegistry: Mocked<ToolRegistry>;
+  let mockTool: AnyDeclarativeTool;
+  let mockInvocation: AnyToolInvocation;
+
+  const req: ToolCallRequestInfo = {
+    callId: 'call-1',
+    name: 'test-tool',
+    args: { foo: 'bar' },
+    isClientInitiated: false,
+    prompt_id: 'prompt-1',
+    schedulerId: ROOT_SCHEDULER_ID,
+  };
+
+  beforeEach(() => {
+    mockTool = {
+      name: 'test-tool',
+      build: vi.fn(),
+    } as unknown as AnyDeclarativeTool;
+
+    mockInvocation = {
+      shouldConfirmExecute: vi.fn(),
+    } as unknown as AnyToolInvocation;
+    vi.mocked(mockTool.build).mockReturnValue(mockInvocation);
+
+    mockPolicyEngine = {
+      check: vi.fn().mockResolvedValue({ decision: PolicyDecision.DENY }), // Default to DENY for this test
+    } as unknown as Mocked<PolicyEngine>;
+
+    mockToolRegistry = {
+      getTool: vi.fn().mockReturnValue(mockTool),
+      getAllToolNames: vi.fn().mockReturnValue(['test-tool']),
+    } as unknown as Mocked<ToolRegistry>;
+
+    mockMessageBus = {
+      publish: vi.fn(),
+      subscribe: vi.fn(),
+    } as unknown as Mocked<MessageBus>;
+
+    mockConfig = {
+      getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
+      getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
+      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
+      isInteractive: vi.fn().mockReturnValue(true),
+      getEnableHooks: vi.fn().mockReturnValue(false),
+      getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.PLAN), // Key: Plan Mode
+      setApprovalMode: vi.fn(),
+      getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
+    } as unknown as Mocked<Config>;
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe.each([
+    { enableEventDrivenScheduler: false, name: 'Legacy CoreToolScheduler' },
+    { enableEventDrivenScheduler: true, name: 'Event-Driven Scheduler' },
+  ])('$name', ({ enableEventDrivenScheduler }) => {
+    it('should return the correct Plan Mode denial message when policy denies execution', async () => {
+      let resultMessage: string | undefined;
+      let resultErrorType: ToolErrorType | undefined;
+
+      const signal = new AbortController().signal;
+
+      if (enableEventDrivenScheduler) {
+        const scheduler = new Scheduler({
+          config: mockConfig,
+          messageBus: mockMessageBus,
+          getPreferredEditor: () => undefined,
+          schedulerId: ROOT_SCHEDULER_ID,
+        });
+
+        const results = await scheduler.schedule(req, signal);
+        const result = results[0];
+
+        expect(result.status).toBe('error');
+        if (result.status === 'error') {
+          resultMessage = result.response.error?.message;
+          resultErrorType = result.response.errorType;
+        }
+      } else {
+        let capturedCalls: CompletedToolCall[] = [];
+        const scheduler = new CoreToolScheduler({
+          config: mockConfig,
+          getPreferredEditor: () => undefined,
+          onAllToolCallsComplete: async (calls) => {
+            capturedCalls = calls;
+          },
+        });
+
+        await scheduler.schedule(req, signal);
+
+        expect(capturedCalls.length).toBeGreaterThan(0);
+        const call = capturedCalls[0];
+        if (call.status === 'error') {
+          resultMessage = call.response.error?.message;
+          resultErrorType = call.response.errorType;
+        }
+      }
+
+      expect(resultMessage).toBe(PLAN_MODE_DENIAL_MESSAGE);
+      expect(resultErrorType).toBe(ToolErrorType.STOP_EXECUTION);
+    });
+  });
+});