@machina.ai/cell-cli-core 1.22.5-rc1 → 1.25.0-rc1

package/dist/src/core/coreToolScheduler.test.js

@@ -3,38 +3,36 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { CoreToolScheduler, convertToFunctionResponse, truncateAndSaveToFile, } from './coreToolScheduler.js';
-import { DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES, DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, BaseDeclarativeTool, BaseToolInvocation, ToolConfirmationOutcome, Kind, ApprovalMode, HookSystem, } from '../index.js';
+import { describe, it, expect, vi } from 'vitest';
+import { CoreToolScheduler } from './coreToolScheduler.js';
+import { DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES, DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, BaseDeclarativeTool, BaseToolInvocation, ToolConfirmationOutcome, Kind, ApprovalMode, HookSystem, PolicyDecision, } from '../index.js';
 import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
 import { MockModifiableTool, MockTool, MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, } from '../test-utils/mock-tool.js';
 import * as modifiableToolModule from '../tools/modifiable-tool.js';
-import * as fs from 'node:fs/promises';
-import * as path from 'node:path';
-import { isShellInvocationAllowlisted } from '../utils/shell-permissions.js';
-import { DEFAULT_GEMINI_MODEL, PREVIEW_GEMINI_MODEL, } from '../config/models.js';
+import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
+import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
 vi.mock('fs/promises', () => ({
     writeFile: vi.fn(),
 }));
 class TestApprovalTool extends BaseDeclarativeTool {
     config;
     static Name = 'testApprovalTool';
-    constructor(config) {
+    constructor(config, messageBus) {
         super(TestApprovalTool.Name, 'TestApprovalTool', 'A tool for testing approval logic', Kind.Edit, {
             properties: { id: { type: 'string' } },
             required: ['id'],
             type: 'object',
-        });
+        }, messageBus);
         this.config = config;
     }
-    createInvocation(params) {
-        return new TestApprovalInvocation(this.config, params);
+    createInvocation(params, messageBus, _toolName, _toolDisplayName) {
+        return new TestApprovalInvocation(this.config, params, messageBus);
     }
 }
 class TestApprovalInvocation extends BaseToolInvocation {
     config;
-    constructor(config, params) {
-        super(params);
+    constructor(config, params, messageBus) {
+        super(params, messageBus);
         this.config = config;
     }
     getDescription() {
@@ -70,8 +68,8 @@ class TestApprovalInvocation extends BaseToolInvocation {
 class AbortDuringConfirmationInvocation extends BaseToolInvocation {
     abortController;
     abortError;
-    constructor(abortController, abortError, params) {
-        super(params);
+    constructor(abortController, abortError, params, messageBus) {
+        super(params, messageBus);
         this.abortController = abortController;
         this.abortError = abortError;
     }
@@ -89,16 +87,16 @@ class AbortDuringConfirmationInvocation extends BaseToolInvocation {
 class AbortDuringConfirmationTool extends BaseDeclarativeTool {
     abortController;
     abortError;
-    constructor(abortController, abortError) {
+    constructor(abortController, abortError, messageBus) {
         super('abortDuringConfirmationTool', 'Abort During Confirmation Tool', 'A tool that aborts while confirming execution.', Kind.Other, {
             type: 'object',
             properties: {},
-        });
+        }, messageBus);
         this.abortController = abortController;
         this.abortError = abortError;
     }
-    createInvocation(params) {
-        return new AbortDuringConfirmationInvocation(this.abortController, this.abortError, params);
+    createInvocation(params, messageBus, _toolName, _toolDisplayName) {
+        return new AbortDuringConfirmationInvocation(this.abortController, this.abortError, params, messageBus);
     }
 }
 async function waitForStatus(onToolCallsUpdate, status, timeout = 5000) {
@@ -155,6 +153,11 @@ function createMockConfig(overrides = {}) {
         getShellExecutionConfig: () => ({
             terminalWidth: 90,
             terminalHeight: 30,
+            sanitizationConfig: {
+                enableEnvironmentVariableRedaction: true,
+                allowedEnvironmentVariables: [],
+                blockedEnvironmentVariables: [],
+            },
         }),
         storage: {
             getProjectTempDir: () => '/tmp',
@@ -163,15 +166,32 @@ function createMockConfig(overrides = {}) {
         getTruncateToolOutputLines: () => DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES,
         getToolRegistry: () => defaultToolRegistry,
         getActiveModel: () => DEFAULT_GEMINI_MODEL,
-        getUseSmartEdit: () => false,
         getGeminiClient: () => null,
-        getEnableMessageBusIntegration: () => false,
-        getMessageBus: () => null,
+        getMessageBus: () => createMockMessageBus(),
         getEnableHooks: () => false,
-        getPolicyEngine: () => null,
         getExperiments: () => { },
     };
-    return { ...baseConfig, ...overrides };
+    const finalConfig = { ...baseConfig, ...overrides };
+    // Patch the policy engine to use the final config if not overridden
+    if (!overrides.getPolicyEngine) {
+        finalConfig.getPolicyEngine = () => ({
+            check: async (toolCall, _serverName) => {
+                // Mock simple policy logic for tests
+                const mode = finalConfig.getApprovalMode();
+                if (mode === ApprovalMode.YOLO) {
+                    return { decision: PolicyDecision.ALLOW };
+                }
+                const allowed = finalConfig.getAllowedTools();
+                if (allowed &&
+                    (allowed.includes(toolCall.name) ||
+                        allowed.some((p) => toolCall.name.startsWith(p)))) {
+                    return { decision: PolicyDecision.ALLOW };
+                }
+                return { decision: PolicyDecision.ASK_USER };
+            },
+        });
+    }
+    return finalConfig;
 }
 describe('CoreToolScheduler', () => {
     it('should cancel a tool call if the signal is aborted before confirmation', async () => {
@@ -401,7 +421,7 @@ describe('CoreToolScheduler', () => {
     it('should mark tool call as cancelled when abort happens during confirmation error', async () => {
         const abortController = new AbortController();
         const abortError = new Error('Abort requested during confirmation');
-        const declarativeTool = new AbortDuringConfirmationTool(abortController, abortError);
+        const declarativeTool = new AbortDuringConfirmationTool(abortController, abortError, createMockMessageBus());
         const mockToolRegistry = {
             getTool: () => declarativeTool,
             getFunctionDeclarations: () => [],
@@ -419,7 +439,7 @@ describe('CoreToolScheduler', () => {
         const onToolCallsUpdate = vi.fn();
         const mockConfig = createMockConfig({
             getToolRegistry: () => mockToolRegistry,
-            isInteractive: () => false,
+            isInteractive: () => true,
         });
         const scheduler = new CoreToolScheduler({
             config: mockConfig,
@@ -493,35 +513,6 @@ describe('CoreToolScheduler', () => {
         const errorMessage = errorParts[0].functionResponse.response.error;
         expect(errorMessage).toContain('Tool execution for "mockTool" requires user confirmation, which is not supported in non-interactive mode.');
     });
-    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'],
-            };
-            const mockConfig = createMockConfig({
-                getToolRegistry: () => mockToolRegistry,
-                isInteractive: () => false,
-            });
-            // Create scheduler
-            const scheduler = new CoreToolScheduler({
-                config: mockConfig,
-                getPreferredEditor: () => 'vscode',
-            });
-            // Test that the right tool is selected, with only 1 result, for typos
-            // @ts-expect-error accessing private method
-            const misspelledTool = scheduler.getToolSuggestion('list_fils', 1);
-            expect(misspelledTool).toBe(' Did you mean "list_files"?');
-            // Test that the right tool is selected, with only 1 result, for prefixes
-            // @ts-expect-error accessing private method
-            const prefixedTool = scheduler.getToolSuggestion('github.list_files', 1);
-            expect(prefixedTool).toBe(' Did you mean "list_files"?');
-            // Test that the right tool is first
-            // @ts-expect-error accessing private method
-            const suggestionMultiple = scheduler.getToolSuggestion('list_fils');
-            expect(suggestionMultiple).toBe(' Did you mean one of: "list_files", "read_file", "write_file"?');
-        });
-    });
 });
 describe('CoreToolScheduler with payload', () => {
     it('should update args and diff and execute tool when payload is provided', async () => {
@@ -573,6 +564,10 @@ describe('CoreToolScheduler with payload', () => {
             const payload = { newContent: 'final version' };
             await confirmationDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce, payload);
         }
+        // After internal update, the tool should be awaiting approval again with the NEW content.
+        const updatedAwaitingCall = (await waitForStatus(onToolCallsUpdate, 'awaiting_approval'));
+        // Now confirm for real to execute.
+        await updatedAwaitingCall.confirmationDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
         // Wait for the tool execution to complete
         await vi.waitFor(() => {
             expect(onAllToolCallsComplete).toHaveBeenCalled();
@@ -585,226 +580,9 @@ describe('CoreToolScheduler with payload', () => {
         });
     });
 });
-describe('convertToFunctionResponse', () => {
-    const toolName = 'testTool';
-    const callId = 'call1';
-    it('should handle simple string llmContent', () => {
-        const llmContent = 'Simple text output';
-        const result = convertToFunctionResponse(toolName, callId, llmContent, DEFAULT_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Simple text output' },
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as a single Part with text', () => {
-        const llmContent = { text: 'Text from Part object' };
-        const result = convertToFunctionResponse(toolName, callId, llmContent, DEFAULT_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Text from Part object' },
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as a PartListUnion array with a single text Part', () => {
-        const llmContent = [{ text: 'Text from array' }];
-        const result = convertToFunctionResponse(toolName, callId, llmContent, DEFAULT_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Text from array' },
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as a PartListUnion array with multiple Parts', () => {
-        const llmContent = [{ text: 'part1' }, { text: 'part2' }];
-        const result = convertToFunctionResponse(toolName, callId, llmContent, DEFAULT_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'part1\npart2' },
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent with fileData for Gemini 3 model (should be siblings)', () => {
-        const llmContent = {
-            fileData: { mimeType: 'application/pdf', fileUri: 'gs://...' },
-        };
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Binary content provided (1 item(s)).' },
-                },
-            },
-            llmContent,
-        ]);
-    });
-    it('should handle llmContent with inlineData for Gemini 3 model (should be nested)', () => {
-        const llmContent = {
-            inlineData: { mimeType: 'image/png', data: 'base64...' },
-        };
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Binary content provided (1 item(s)).' },
-                    parts: [llmContent],
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent with fileData for non-Gemini 3 models', () => {
-        const llmContent = {
-            fileData: { mimeType: 'application/pdf', fileUri: 'gs://...' },
-        };
-        const result = convertToFunctionResponse(toolName, callId, llmContent, DEFAULT_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Binary content provided (1 item(s)).' },
-                },
-            },
-            llmContent,
-        ]);
-    });
-    it('should preserve existing functionResponse metadata', () => {
-        const innerId = 'inner-call-id';
-        const innerName = 'inner-tool-name';
-        const responseMetadata = {
-            flags: ['flag1'],
-            isError: false,
-            customData: { key: 'value' },
-        };
-        const input = {
-            functionResponse: {
-                id: innerId,
-                name: innerName,
-                response: responseMetadata,
-            },
-        };
-        const result = convertToFunctionResponse(toolName, callId, input, DEFAULT_GEMINI_MODEL);
-        expect(result).toHaveLength(1);
-        expect(result[0].functionResponse).toEqual({
-            id: callId,
-            name: toolName,
-            response: responseMetadata,
-        });
-    });
-    it('should handle llmContent as an array of multiple Parts (text and inlineData)', () => {
-        const llmContent = [
-            { text: 'Some textual description' },
-            { inlineData: { mimeType: 'image/jpeg', data: 'base64data...' } },
-            { text: 'Another text part' },
-        ];
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: {
-                        output: 'Some textual description\nAnother text part',
-                    },
-                    parts: [
-                        { inlineData: { mimeType: 'image/jpeg', data: 'base64data...' } },
-                    ],
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as an array with a single inlineData Part', () => {
-        const llmContent = [
-            { inlineData: { mimeType: 'image/gif', data: 'gifdata...' } },
-        ];
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: 'Binary content provided (1 item(s)).' },
-                    parts: llmContent,
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as a generic Part (not text, inlineData, or fileData)', () => {
-        const llmContent = { functionCall: { name: 'test', args: {} } };
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: {},
-                },
-            },
-        ]);
-    });
-    it('should handle empty string llmContent', () => {
-        const llmContent = '';
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: { output: '' },
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as an empty array', () => {
-        const llmContent = [];
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: {},
-                },
-            },
-        ]);
-    });
-    it('should handle llmContent as a Part with undefined inlineData/fileData/text', () => {
-        const llmContent = {}; // An empty part object
-        const result = convertToFunctionResponse(toolName, callId, llmContent, PREVIEW_GEMINI_MODEL);
-        expect(result).toEqual([
-            {
-                functionResponse: {
-                    name: toolName,
-                    id: callId,
-                    response: {},
-                },
-            },
-        ]);
-    });
-});
 class MockEditToolInvocation extends BaseToolInvocation {
-    constructor(params) {
-        super(params);
+    constructor(params, messageBus) {
+        super(params, messageBus);
     }
     getDescription() {
         return 'A mock edit tool invocation';
@@ -829,16 +607,16 @@ class MockEditToolInvocation extends BaseToolInvocation {
     }
 }
 class MockEditTool extends BaseDeclarativeTool {
-    constructor() {
-        super('mockEditTool', 'mockEditTool', 'A mock edit tool', Kind.Edit, {});
+    constructor(messageBus) {
+        super('mockEditTool', 'mockEditTool', 'A mock edit tool', Kind.Edit, {}, messageBus);
     }
-    createInvocation(params) {
-        return new MockEditToolInvocation(params);
+    createInvocation(params, messageBus, _toolName, _toolDisplayName) {
+        return new MockEditToolInvocation(params, messageBus);
     }
 }
 describe('CoreToolScheduler edit cancellation', () => {
     it('should preserve diff when an edit is cancelled', async () => {
-        const mockEditTool = new MockEditTool();
+        const mockEditTool = new MockEditTool(createMockMessageBus());
         const mockToolRegistry = {
             getTool: () => mockEditTool,
             getFunctionDeclarations: () => [],
@@ -1109,6 +887,11 @@ describe('CoreToolScheduler request queueing', () => {
             getShellExecutionConfig: () => ({
                 terminalWidth: 80,
                 terminalHeight: 24,
+                sanitizationConfig: {
+                    enableEnvironmentVariableRedaction: true,
+                    allowedEnvironmentVariables: [],
+                    blockedEnvironmentVariables: [],
+                },
             }),
             isInteractive: () => false,
         });
@@ -1164,9 +947,6 @@ describe('CoreToolScheduler request queueing', () => {
         }
     });
     it('should require approval for a chained shell command even when prefix is allowlisted', async () => {
-        expect(isShellInvocationAllowlisted({
-            params: { command: 'git status && rm -rf /tmp/should-not-run' },
-        }, ['run_shell_command(git)'])).toBe(false);
         const executeFn = vi.fn().mockResolvedValue({
             llmContent: 'Shell command executed',
             returnDisplay: 'Shell command executed',
@@ -1202,9 +982,17 @@ describe('CoreToolScheduler request queueing', () => {
             getShellExecutionConfig: () => ({
                 terminalWidth: 80,
                 terminalHeight: 24,
+                sanitizationConfig: {
+                    enableEnvironmentVariableRedaction: true,
+                    allowedEnvironmentVariables: [],
+                    blockedEnvironmentVariables: [],
+                },
             }),
             getToolRegistry: () => toolRegistry,
             getHookSystem: () => undefined,
+            getPolicyEngine: () => ({
+                check: async () => ({ decision: PolicyDecision.ASK_USER }),
+            }),
         });
         const scheduler = new CoreToolScheduler({
             config: mockConfig,
@@ -1307,7 +1095,7 @@ describe('CoreToolScheduler request queueing', () => {
         mockConfig.getHookSystem = vi
             .fn()
             .mockReturnValue(new HookSystem(mockConfig));
-        const testTool = new TestApprovalTool(mockConfig);
+        const testTool = new TestApprovalTool(mockConfig, mockMessageBus);
         const toolRegistry = {
             getTool: () => testTool,
             getFunctionDeclarations: () => [],
@@ -1662,119 +1450,189 @@ describe('CoreToolScheduler Sequential Execution', () => {
         });
         modifyWithEditorSpy.mockRestore();
     });
-});
-describe('truncateAndSaveToFile', () => {
-    const mockWriteFile = vi.mocked(fs.writeFile);
-    const THRESHOLD = 40_000;
-    const TRUNCATE_LINES = 1000;
-    beforeEach(() => {
-        vi.clearAllMocks();
-    });
-    it('should return content unchanged if below threshold', async () => {
-        const content = 'Short content';
-        const callId = 'test-call-id';
-        const projectTempDir = '/tmp';
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        expect(result).toEqual({ content });
-        expect(mockWriteFile).not.toHaveBeenCalled();
-    });
-    it('should truncate content by lines when content has many lines', async () => {
-        // Create content that exceeds 100,000 character threshold with many lines
-        const lines = Array(2000).fill('x'.repeat(100)); // 100 chars per line * 2000 lines = 200,000 chars
-        const content = lines.join('\n');
-        const callId = 'test-call-id';
-        const projectTempDir = '/tmp';
-        mockWriteFile.mockResolvedValue(undefined);
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        expect(result.outputFile).toBe(path.join(projectTempDir, `${callId}.output`));
-        expect(mockWriteFile).toHaveBeenCalledWith(path.join(projectTempDir, `${callId}.output`), content);
-        // Should contain the first and last lines with 1/5 head and 4/5 tail
-        const head = Math.floor(TRUNCATE_LINES / 5);
-        const beginning = lines.slice(0, head);
-        const end = lines.slice(-(TRUNCATE_LINES - head));
-        const expectedTruncated = beginning.join('\n') + '\n... [CONTENT TRUNCATED] ...\n' + end.join('\n');
-        expect(result.content).toContain('Tool output was too large and has been truncated');
-        expect(result.content).toContain('Truncated part of the output:');
-        expect(result.content).toContain(expectedTruncated);
-    });
-    it('should wrap and truncate content when content has few but long lines', async () => {
-        const content = 'a'.repeat(200_000); // A single very long line
-        const callId = 'test-call-id';
-        const projectTempDir = '/tmp';
-        const wrapWidth = 120;
-        mockWriteFile.mockResolvedValue(undefined);
-        // Manually wrap the content to generate the expected file content
-        const wrappedLines = [];
-        for (let i = 0; i < content.length; i += wrapWidth) {
-            wrappedLines.push(content.substring(i, i + wrapWidth));
-        }
-        const expectedFileContent = wrappedLines.join('\n');
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        expect(result.outputFile).toBe(path.join(projectTempDir, `${callId}.output`));
-        // Check that the file was written with the wrapped content
-        expect(mockWriteFile).toHaveBeenCalledWith(path.join(projectTempDir, `${callId}.output`), expectedFileContent);
-        // Should contain the first and last lines with 1/5 head and 4/5 tail of the wrapped content
-        const head = Math.floor(TRUNCATE_LINES / 5);
-        const beginning = wrappedLines.slice(0, head);
-        const end = wrappedLines.slice(-(TRUNCATE_LINES - head));
-        const expectedTruncated = beginning.join('\n') + '\n... [CONTENT TRUNCATED] ...\n' + end.join('\n');
-        expect(result.content).toContain('Tool output was too large and has been truncated');
-        expect(result.content).toContain('Truncated part of the output:');
-        expect(result.content).toContain(expectedTruncated);
-    });
-    it('should handle file write errors gracefully', async () => {
-        const content = 'a'.repeat(2_000_000);
-        const callId = 'test-call-id';
-        const projectTempDir = '/tmp';
-        mockWriteFile.mockRejectedValue(new Error('File write failed'));
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        expect(result.outputFile).toBeUndefined();
-        expect(result.content).toContain('[Note: Could not save full output to file]');
-        expect(mockWriteFile).toHaveBeenCalled();
-    });
-    it('should save to correct file path with call ID', async () => {
-        const content = 'a'.repeat(200_000);
-        const callId = 'unique-call-123';
-        const projectTempDir = '/custom/temp/dir';
-        const wrapWidth = 120;
-        mockWriteFile.mockResolvedValue(undefined);
-        // Manually wrap the content to generate the expected file content
-        const wrappedLines = [];
-        for (let i = 0; i < content.length; i += wrapWidth) {
-            wrappedLines.push(content.substring(i, i + wrapWidth));
-        }
-        const expectedFileContent = wrappedLines.join('\n');
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        const expectedPath = path.join(projectTempDir, `${callId}.output`);
-        expect(result.outputFile).toBe(expectedPath);
-        expect(mockWriteFile).toHaveBeenCalledWith(expectedPath, expectedFileContent);
+    it('should pass serverName to policy engine for DiscoveredMCPTool', async () => {
+        const mockMcpTool = {
+            tool: async () => ({ functionDeclarations: [] }),
+            callTool: async () => [],
+        };
+        const serverName = 'test-server';
+        const toolName = 'test-tool';
+        const mcpTool = new DiscoveredMCPTool(mockMcpTool, serverName, toolName, 'description', { type: 'object', properties: {} }, createMockMessageBus());
+        const mockToolRegistry = {
+            getTool: () => mcpTool,
+            getFunctionDeclarations: () => [],
+            tools: new Map(),
+            discovery: {},
+            registerTool: () => { },
+            getToolByName: () => mcpTool,
+            getToolByDisplayName: () => mcpTool,
+            getTools: () => [],
+            discoverTools: async () => { },
+            getAllTools: () => [],
+            getToolsByServer: () => [],
+        };
+        const mockPolicyEngineCheck = vi.fn().mockResolvedValue({
+            decision: PolicyDecision.ALLOW,
+        });
+        const mockConfig = createMockConfig({
+            getToolRegistry: () => mockToolRegistry,
+            getPolicyEngine: () => ({
+                check: mockPolicyEngineCheck,
+            }),
+            isInteractive: () => false,
+        });
+        const scheduler = new CoreToolScheduler({
+            config: mockConfig,
+            getPreferredEditor: () => 'vscode',
+        });
+        const abortController = new AbortController();
+        const request = {
+            callId: '1',
+            name: toolName,
+            args: {},
+            isClientInitiated: false,
+            prompt_id: 'prompt-id-1',
+        };
+        await scheduler.schedule(request, abortController.signal);
+        expect(mockPolicyEngineCheck).toHaveBeenCalledWith(expect.objectContaining({ name: toolName }), serverName);
     });
-    it('should include helpful instructions in truncated message', async () => {
-        const content = 'a'.repeat(2_000_000);
-        const callId = 'test-call-id';
-        const projectTempDir = '/tmp';
-        mockWriteFile.mockResolvedValue(undefined);
-        const result = await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        expect(result.content).toContain('read_file tool with the absolute file path above');
-        expect(result.content).toContain('read_file tool with offset=0, limit=100');
-        expect(result.content).toContain('read_file tool with offset=N to skip N lines');
-        expect(result.content).toContain('read_file tool with limit=M to read only M lines');
+    it('should not double-report completed tools when concurrent completions occur', async () => {
+        // Arrange
+        const executeFn = vi.fn().mockResolvedValue({ llmContent: 'success' });
+        const mockTool = new MockTool({ name: 'mockTool', execute: executeFn });
+        const declarativeTool = mockTool;
+        const mockToolRegistry = {
+            getTool: () => declarativeTool,
+            getToolByName: () => declarativeTool,
+            getFunctionDeclarations: () => [],
+            tools: new Map(),
+            discovery: {},
+            registerTool: () => { },
+            getToolByDisplayName: () => declarativeTool,
+            getTools: () => [],
+            discoverTools: async () => { },
+            getAllTools: () => [],
+            getToolsByServer: () => [],
+        };
+        let completionCallCount = 0;
+        const onAllToolCallsComplete = vi.fn().mockImplementation(async () => {
+            completionCallCount++;
+            // Simulate slow reporting (e.g. Gemini API call)
+            await new Promise((resolve) => setTimeout(resolve, 50));
+        });
+        const mockConfig = createMockConfig({
+            getToolRegistry: () => mockToolRegistry,
+            getApprovalMode: () => ApprovalMode.YOLO,
+            isInteractive: () => false,
+        });
+        const mockMessageBus = createMockMessageBus();
+        mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus);
+        mockConfig.getEnableHooks = vi.fn().mockReturnValue(false);
+        mockConfig.getHookSystem = vi
+            .fn()
+            .mockReturnValue(new HookSystem(mockConfig));
+        const scheduler = new CoreToolScheduler({
+            config: mockConfig,
+            onAllToolCallsComplete,
+            getPreferredEditor: () => 'vscode',
+        });
+        const abortController = new AbortController();
+        const request = {
+            callId: '1',
+            name: 'mockTool',
+            args: {},
+            isClientInitiated: false,
+            prompt_id: 'prompt-1',
+        };
+        // Act
+        // 1. Start execution
+        const schedulePromise = scheduler.schedule([request], abortController.signal);
+        // 2. Wait just enough for it to finish and enter checkAndNotifyCompletion
+        // (awaiting our slow mock)
+        await vi.waitFor(() => {
+            expect(completionCallCount).toBe(1);
+        });
+        // 3. Trigger a concurrent completion event (e.g. via cancelAll)
+        scheduler.cancelAll(abortController.signal);
+        await schedulePromise;
+        // Assert
+        // Even though cancelAll was called while the first completion was in progress,
+        // it should not have triggered a SECOND completion call because the first one
+        // was still 'finalizing' and will drain any new tools.
+        expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1);
     });
-    it('should sanitize callId to prevent path traversal', async () => {
-        const content = 'a'.repeat(200_000);
-        const callId = '../../../../../etc/passwd';
-        const projectTempDir = '/tmp/safe_dir';
-        const wrapWidth = 120;
-        mockWriteFile.mockResolvedValue(undefined);
-        // Manually wrap the content to generate the expected file content
-        const wrappedLines = [];
-        for (let i = 0; i < content.length; i += wrapWidth) {
-            wrappedLines.push(content.substring(i, i + wrapWidth));
-        }
-        const expectedFileContent = wrappedLines.join('\n');
-        await truncateAndSaveToFile(content, callId, projectTempDir, THRESHOLD, TRUNCATE_LINES);
-        const expectedPath = path.join(projectTempDir, 'passwd.output');
-        expect(mockWriteFile).toHaveBeenCalledWith(expectedPath, expectedFileContent);
+    it('should complete reporting all tools even mid-callback during abort', async () => {
+        // Arrange
+        const onAllToolCallsComplete = vi.fn().mockImplementation(async () => {
+            // Simulate slow reporting
+            await new Promise((resolve) => setTimeout(resolve, 50));
+        });
+        const mockTool = new MockTool({ name: 'mockTool' });
+        const mockToolRegistry = {
+            getTool: () => mockTool,
+            getToolByName: () => mockTool,
+            getFunctionDeclarations: () => [],
+            tools: new Map(),
+            discovery: {},
+            registerTool: () => { },
+            getToolByDisplayName: () => mockTool,
+            getTools: () => [],
+            discoverTools: async () => { },
+            getAllTools: () => [],
+            getToolsByServer: () => [],
+        };
+        const mockConfig = createMockConfig({
+            getToolRegistry: () => mockToolRegistry,
+            getApprovalMode: () => ApprovalMode.YOLO,
+            isInteractive: () => false,
+        });
+        const scheduler = new CoreToolScheduler({
+            config: mockConfig,
+            onAllToolCallsComplete,
+            getPreferredEditor: () => 'vscode',
+        });
+        const abortController = new AbortController();
+        const signal = abortController.signal;
+        // Act
+        // 1. Start execution of two tools
+        const schedulePromise = scheduler.schedule([
+            {
+                callId: '1',
+                name: 'mockTool',
+                args: {},
+                isClientInitiated: false,
+                prompt_id: 'prompt-1',
+            },
+            {
+                callId: '2',
+                name: 'mockTool',
+                args: {},
+                isClientInitiated: false,
+                prompt_id: 'prompt-1',
+            },
+        ], signal);
+        // 2. Wait for reporting to start
+        await vi.waitFor(() => {
+            expect(onAllToolCallsComplete).toHaveBeenCalled();
+        });
+        // 3. Abort the signal while reporting is in progress
+        abortController.abort();
+        await schedulePromise;
+        // Assert
+        // Verify that onAllToolCallsComplete was called and processed the tools,
+        // and that the scheduler didn't just drop them because of the abort.
+        expect(onAllToolCallsComplete).toHaveBeenCalled();
+        const reportedTools = onAllToolCallsComplete.mock.calls.flatMap((call) => 
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        call[0].map((t) => t.request.callId));
+        // Both tools should have been reported exactly once with success status
+        expect(reportedTools).toContain('1');
+        expect(reportedTools).toContain('2');
+        const allStatuses = onAllToolCallsComplete.mock.calls.flatMap((call) => 
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        call[0].map((t) => t.status));
+        expect(allStatuses).toEqual(['success', 'success']);
+        expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1);
     });
 });
 //# sourceMappingURL=coreToolScheduler.test.js.map