From 5566292cc83f0e30c4ec890aa25cd6ae000f51f8 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Fri, 26 Dec 2025 15:51:39 -0500 Subject: [PATCH] refactor(core): extract static concerns from CoreToolScheduler (#15589) --- packages/core/src/agents/local-executor.ts | 3 +- .../core/src/code_assist/telemetry.test.ts | 2 +- .../core/src/core/coreToolScheduler.test.ts | 579 +----------------- packages/core/src/core/coreToolScheduler.ts | 402 ++---------- packages/core/src/core/turn.ts | 29 +- packages/core/src/index.ts | 1 + packages/core/src/scheduler/types.ts | 135 ++++ .../src/telemetry/loggers.test.circular.ts | 8 +- .../core/src/utils/checkpointUtils.test.ts | 2 +- packages/core/src/utils/checkpointUtils.ts | 2 +- packages/core/src/utils/fileUtils.test.ts | 210 +++++++ packages/core/src/utils/fileUtils.ts | 65 ++ .../generateContentResponseUtilities.test.ts | 312 ++++++++++ .../utils/generateContentResponseUtilities.ts | 117 ++++ packages/core/src/utils/tool-utils.test.ts | 26 +- packages/core/src/utils/tool-utils.ts | 38 ++ 16 files changed, 983 insertions(+), 948 deletions(-) create mode 100644 packages/core/src/scheduler/types.ts diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index cadf384e69..3a713c0167 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -17,7 +17,8 @@ import type { } from '@google/genai'; import { executeToolCall } from '../core/nonInteractiveToolExecutor.js'; import { ToolRegistry } from '../tools/tool-registry.js'; -import { type ToolCallRequestInfo, CompressionStatus } from '../core/turn.js'; +import { CompressionStatus } from '../core/turn.js'; +import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { ChatCompressionService } from '../services/chatCompressionService.js'; import { getDirectoryContextString } from '../utils/environmentContext.js'; import { promptIdContext } from '../utils/promptIdContext.js'; diff --git a/packages/core/src/code_assist/telemetry.test.ts b/packages/core/src/code_assist/telemetry.test.ts index d6a821a0b0..b036679d80 100644 --- a/packages/core/src/code_assist/telemetry.test.ts +++ b/packages/core/src/code_assist/telemetry.test.ts @@ -30,7 +30,7 @@ import { type AnyToolInvocation, } from '../tools/tools.js'; import type { Config } from '../config/config.js'; -import type { ToolCallResponseInfo } from '../core/turn.js'; +import type { ToolCallResponseInfo } from '../scheduler/types.js'; function createMockResponse( candidates: GenerateContentResponse['candidates'] = [], diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 6474331adb..a1b1a90b79 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -4,18 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { describe, it, expect, vi } from 'vitest'; import type { Mock } from 'vitest'; -import { - CoreToolScheduler, - convertToFunctionResponse, - truncateAndSaveToFile, -} from './coreToolScheduler.js'; +import { CoreToolScheduler } from './coreToolScheduler.js'; import type { ToolCall, WaitingToolCall, ErroredToolCall, -} from './coreToolScheduler.js'; +} from '../scheduler/types.js'; import type { ToolCallConfirmationDetails, ToolConfirmationPayload, @@ -36,20 +32,14 @@ import { HookSystem, } from '../index.js'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; -import type { Part, PartListUnion } from '@google/genai'; 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'; vi.mock('fs/promises', () => ({ writeFile: vi.fn(), @@ -655,42 +645,6 @@ describe('CoreToolScheduler', () => { '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'], - } as unknown as ToolRegistry; - 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', () => { @@ -771,310 +725,6 @@ 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: Part = { 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: PartListUnion = [{ 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: PartListUnion = [{ 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: Part = { - 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: Part = { - 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: Part = { - 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: Part = { - 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: PartListUnion = [ - { 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: PartListUnion = [ - { 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: Part = { 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: PartListUnion = []; - 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: Part = {}; // 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< Record, ToolResult @@ -2130,224 +1780,3 @@ 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: string[] = []; - 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: string[] = []; - 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 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 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: string[] = []; - 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, - ); - }); -}); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 85632265d4..cd660da40f 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -4,44 +4,36 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { - ToolCallRequestInfo, - ToolCallResponseInfo, - ToolCallConfirmationDetails, - ToolResult, - ToolResultDisplay, - EditorType, - Config, - ToolConfirmationPayload, - AnyDeclarativeTool, - AnyToolInvocation, - AnsiOutput, -} from '../index.js'; import { + type ToolResult, + type ToolResultDisplay, + type AnyDeclarativeTool, + type AnyToolInvocation, + type ToolCallConfirmationDetails, + type ToolConfirmationPayload, ToolConfirmationOutcome, - ApprovalMode, - logToolCall, - ToolErrorType, - ToolCallEvent, - logToolOutputTruncated, - ToolOutputTruncatedEvent, - runInDevTraceSpan, -} from '../index.js'; -import { READ_FILE_TOOL_NAME, SHELL_TOOL_NAME } from '../tools/tool-names.js'; -import type { Part, PartListUnion } from '@google/genai'; -import { supportsMultimodalFunctionResponse } from '../config/models.js'; +} from '../tools/tools.js'; +import type { EditorType } from '../utils/editor.js'; +import type { Config } from '../config/config.js'; +import type { AnsiOutput } from '../utils/terminalSerializer.js'; +import { ApprovalMode } from '../policy/types.js'; +import { logToolCall, logToolOutputTruncated } from '../telemetry/loggers.js'; +import { ToolErrorType } from '../tools/tool-error.js'; +import { ToolCallEvent, ToolOutputTruncatedEvent } from '../telemetry/types.js'; +import { runInDevTraceSpan } from '../telemetry/trace.js'; +import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; import type { ModifyContext } from '../tools/modifiable-tool.js'; import { isModifiableDeclarativeTool, modifyWithEditor, } from '../tools/modifiable-tool.js'; import * as Diff from 'diff'; -import * as fs from 'node:fs/promises'; -import * as path from 'node:path'; import { SHELL_TOOL_NAMES } from '../utils/shell-utils.js'; -import { doesToolInvocationMatch } from '../utils/tool-utils.js'; +import { + doesToolInvocationMatch, + getToolSuggestion, +} from '../utils/tool-utils.js'; import { isShellInvocationAllowlisted } from '../utils/shell-permissions.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'; @@ -50,221 +42,46 @@ import { fireToolNotificationHook, executeToolWithHooks, } from './coreToolHookTriggers.js'; -import { debugLogger } from '../utils/debugLogger.js'; +import { + type ToolCall, + type ValidatingToolCall, + type ScheduledToolCall, + type ErroredToolCall, + type SuccessfulToolCall, + type ExecutingToolCall, + type CancelledToolCall, + type WaitingToolCall, + type Status, + type CompletedToolCall, + type ConfirmHandler, + type OutputUpdateHandler, + type AllToolCallsCompleteHandler, + type ToolCallsUpdateHandler, + type ToolCallRequestInfo, + type ToolCallResponseInfo, +} from '../scheduler/types.js'; +import { saveTruncatedContent } from '../utils/fileUtils.js'; +import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js'; -export type ValidatingToolCall = { - status: 'validating'; - request: ToolCallRequestInfo; - tool: AnyDeclarativeTool; - invocation: AnyToolInvocation; - startTime?: number; - outcome?: ToolConfirmationOutcome; +export type { + ToolCall, + ValidatingToolCall, + ScheduledToolCall, + ErroredToolCall, + SuccessfulToolCall, + ExecutingToolCall, + CancelledToolCall, + WaitingToolCall, + Status, + CompletedToolCall, + ConfirmHandler, + OutputUpdateHandler, + AllToolCallsCompleteHandler, + ToolCallsUpdateHandler, + ToolCallRequestInfo, + ToolCallResponseInfo, }; -export type ScheduledToolCall = { - status: 'scheduled'; - request: ToolCallRequestInfo; - tool: AnyDeclarativeTool; - invocation: AnyToolInvocation; - startTime?: number; - outcome?: ToolConfirmationOutcome; -}; - -export type ErroredToolCall = { - status: 'error'; - request: ToolCallRequestInfo; - response: ToolCallResponseInfo; - tool?: AnyDeclarativeTool; - durationMs?: number; - outcome?: ToolConfirmationOutcome; -}; - -export type SuccessfulToolCall = { - status: 'success'; - request: ToolCallRequestInfo; - tool: AnyDeclarativeTool; - response: ToolCallResponseInfo; - invocation: AnyToolInvocation; - durationMs?: number; - outcome?: ToolConfirmationOutcome; -}; - -export type ExecutingToolCall = { - status: 'executing'; - request: ToolCallRequestInfo; - tool: AnyDeclarativeTool; - invocation: AnyToolInvocation; - liveOutput?: string | AnsiOutput; - startTime?: number; - outcome?: ToolConfirmationOutcome; - pid?: number; -}; - -export type CancelledToolCall = { - status: 'cancelled'; - request: ToolCallRequestInfo; - response: ToolCallResponseInfo; - tool: AnyDeclarativeTool; - invocation: AnyToolInvocation; - durationMs?: number; - outcome?: ToolConfirmationOutcome; -}; - -export type WaitingToolCall = { - status: 'awaiting_approval'; - request: ToolCallRequestInfo; - tool: AnyDeclarativeTool; - invocation: AnyToolInvocation; - confirmationDetails: ToolCallConfirmationDetails; - startTime?: number; - outcome?: ToolConfirmationOutcome; -}; - -export type Status = ToolCall['status']; - -export type ToolCall = - | ValidatingToolCall - | ScheduledToolCall - | ErroredToolCall - | SuccessfulToolCall - | ExecutingToolCall - | CancelledToolCall - | WaitingToolCall; - -export type CompletedToolCall = - | SuccessfulToolCall - | CancelledToolCall - | ErroredToolCall; - -export type ConfirmHandler = ( - toolCall: WaitingToolCall, -) => Promise; - -export type OutputUpdateHandler = ( - toolCallId: string, - outputChunk: string | AnsiOutput, -) => void; - -export type AllToolCallsCompleteHandler = ( - completedToolCalls: CompletedToolCall[], -) => Promise; - -export type ToolCallsUpdateHandler = (toolCalls: ToolCall[]) => void; - -/** - * Formats tool output for a Gemini FunctionResponse. - */ -function createFunctionResponsePart( - callId: string, - toolName: string, - output: string, -): Part { - return { - functionResponse: { - id: callId, - name: toolName, - response: { output }, - }, - }; -} - -export function convertToFunctionResponse( - toolName: string, - callId: string, - llmContent: PartListUnion, - model: string, -): Part[] { - if (typeof llmContent === 'string') { - return [createFunctionResponsePart(callId, toolName, llmContent)]; - } - - const parts = toParts(llmContent); - - // Separate text from binary types - const textParts: string[] = []; - const inlineDataParts: Part[] = []; - const fileDataParts: Part[] = []; - - for (const part of parts) { - if (part.text !== undefined) { - textParts.push(part.text); - } else if (part.inlineData) { - inlineDataParts.push(part); - } else if (part.fileData) { - fileDataParts.push(part); - } else if (part.functionResponse) { - if (parts.length > 1) { - debugLogger.warn( - 'convertToFunctionResponse received multiple parts with a functionResponse. Only the functionResponse will be used, other parts will be ignored', - ); - } - // Handle passthrough case - return [ - { - functionResponse: { - id: callId, - name: toolName, - response: part.functionResponse.response, - }, - }, - ]; - } - // Ignore other part types - } - - // Build the primary response part - const part: Part = { - functionResponse: { - id: callId, - name: toolName, - response: textParts.length > 0 ? { output: textParts.join('\n') } : {}, - }, - }; - - const isMultimodalFRSupported = supportsMultimodalFunctionResponse(model); - const siblingParts: Part[] = [...fileDataParts]; - - if (inlineDataParts.length > 0) { - if (isMultimodalFRSupported) { - // Nest inlineData if supported by the model - (part.functionResponse as unknown as { parts: Part[] }).parts = - inlineDataParts; - } else { - // Otherwise treat as siblings - siblingParts.push(...inlineDataParts); - } - } - - // Add descriptive text if the response object is empty but we have binary content - if ( - textParts.length === 0 && - (inlineDataParts.length > 0 || fileDataParts.length > 0) - ) { - const totalBinaryItems = inlineDataParts.length + fileDataParts.length; - part.functionResponse!.response = { - output: `Binary content provided (${totalBinaryItems} item(s)).`, - }; - } - - if (siblingParts.length > 0) { - return [part, ...siblingParts]; - } - - return [part]; -} - -function toParts(input: PartListUnion): Part[] { - const parts: Part[] = []; - for (const part of Array.isArray(input) ? input : [input]) { - if (typeof part === 'string') { - parts.push({ text: part }); - } else if (part) { - parts.push(part); - } - } - return parts; -} - const createErrorResponse = ( request: ToolCallRequestInfo, error: Error, @@ -286,70 +103,6 @@ const createErrorResponse = ( contentLength: error.message.length, }); -export async function truncateAndSaveToFile( - content: string, - callId: string, - projectTempDir: string, - threshold: number, - truncateLines: number, -): Promise<{ content: string; outputFile?: string }> { - if (content.length <= threshold) { - return { content }; - } - - let lines = content.split('\n'); - let fileContent = content; - - // If the content is long but has few lines, wrap it to enable line-based truncation. - if (lines.length <= truncateLines) { - const wrapWidth = 120; // A reasonable width for wrapping. - const wrappedLines: string[] = []; - for (const line of lines) { - if (line.length > wrapWidth) { - for (let i = 0; i < line.length; i += wrapWidth) { - wrappedLines.push(line.substring(i, i + wrapWidth)); - } - } else { - wrappedLines.push(line); - } - } - lines = wrappedLines; - fileContent = lines.join('\n'); - } - - const head = Math.floor(truncateLines / 5); - const beginning = lines.slice(0, head); - const end = lines.slice(-(truncateLines - head)); - const truncatedContent = - beginning.join('\n') + '\n... [CONTENT TRUNCATED] ...\n' + end.join('\n'); - - // Sanitize callId to prevent path traversal. - const safeFileName = `${path.basename(callId)}.output`; - const outputFile = path.join(projectTempDir, safeFileName); - try { - await fs.writeFile(outputFile, fileContent); - - return { - content: `Tool output was too large and has been truncated. -The full output has been saved to: ${outputFile} -To read the complete output, use the ${READ_FILE_TOOL_NAME} tool with the absolute file path above. For large files, you can use the offset and limit parameters to read specific sections: -- ${READ_FILE_TOOL_NAME} tool with offset=0, limit=100 to see the first 100 lines -- ${READ_FILE_TOOL_NAME} tool with offset=N to skip N lines from the beginning -- ${READ_FILE_TOOL_NAME} tool with limit=M to read only M lines at a time -The truncated output below shows the beginning and end of the content. The marker '... [CONTENT TRUNCATED] ...' indicates where content was removed. -This allows you to efficiently examine different parts of the output without loading the entire file. -Truncated part of the output: -${truncatedContent}`, - outputFile, - }; - } catch (_error) { - return { - content: - truncatedContent + `\n[Note: Could not save full output to file]`, - }; - } -} - interface CoreToolSchedulerOptions { config: Config; outputUpdateHandler?: OutputUpdateHandler; @@ -656,40 +409,6 @@ export class CoreToolScheduler { } } - /** - * Generates a suggestion string for a tool name that was not found in the registry. - * It finds the closest matches based on Levenshtein distance. - * @param unknownToolName The tool name that was not found. - * @param topN The number of suggestions to return. Defaults to 3. - * @returns A suggestion string like " Did you mean 'tool'?" or " Did you mean one of: 'tool1', 'tool2'?", or an empty string if no suggestions are found. - */ - private getToolSuggestion(unknownToolName: string, topN = 3): string { - const allToolNames = this.config.getToolRegistry().getAllToolNames(); - - const matches = allToolNames.map((toolName) => ({ - name: toolName, - distance: levenshtein.get(unknownToolName, toolName), - })); - - matches.sort((a, b) => a.distance - b.distance); - - const topNResults = matches.slice(0, topN); - - if (topNResults.length === 0) { - return ''; - } - - const suggestedNames = topNResults - .map((match) => `"${match.name}"`) - .join(', '); - - if (topNResults.length > 1) { - return ` Did you mean one of: ${suggestedNames}?`; - } else { - return ` Did you mean ${suggestedNames}?`; - } - } - schedule( request: ToolCallRequestInfo | ToolCallRequestInfo[], signal: AbortSignal, @@ -784,7 +503,10 @@ export class CoreToolScheduler { .getToolRegistry() .getTool(reqInfo.name); if (!toolInstance) { - const suggestion = this.getToolSuggestion(reqInfo.name); + const suggestion = getToolSuggestion( + reqInfo.name, + this.config.getToolRegistry().getAllToolNames(), + ); const errorMessage = `Tool "${reqInfo.name}" not found in registry. Tools must use the exact names that are registered.${suggestion}`; return { status: 'error', @@ -1222,7 +944,7 @@ export class CoreToolScheduler { const threshold = this.config.getTruncateToolOutputThreshold(); const lines = this.config.getTruncateToolOutputLines(); - const truncatedResult = await truncateAndSaveToFile( + const truncatedResult = await saveTruncatedContent( content, callId, this.config.storage.getProjectTempDir(), diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 86c7fdf49c..21191b34f2 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -5,7 +5,6 @@ */ import type { - Part, PartListUnion, GenerateContentResponse, FunctionCall, @@ -16,9 +15,7 @@ import type { import type { ToolCallConfirmationDetails, ToolResult, - ToolResultDisplay, } from '../tools/tools.js'; -import type { ToolErrorType } from '../tools/tool-error.js'; import { getResponseText } from '../utils/partUtils.js'; import { reportError } from '../utils/errorReporting.js'; import { @@ -33,7 +30,11 @@ import { createUserContent } from '@google/genai'; import type { ModelConfigKey } from '../services/modelConfigService.js'; import { getCitations } from '../utils/generateContentResponseUtilities.js'; -// Define a structure for tools passed to the server +import { + type ToolCallRequestInfo, + type ToolCallResponseInfo, +} from '../scheduler/types.js'; + export interface ServerTool { name: string; schema: FunctionDeclaration; @@ -102,26 +103,6 @@ export interface GeminiFinishedEventValue { usageMetadata: GenerateContentResponseUsageMetadata | undefined; } -export interface ToolCallRequestInfo { - callId: string; - name: string; - args: Record; - isClientInitiated: boolean; - prompt_id: string; - checkpoint?: string; - traceId?: string; -} - -export interface ToolCallResponseInfo { - callId: string; - responseParts: Part[]; - resultDisplay: ToolResultDisplay | undefined; - error: Error | undefined; - errorType: ToolErrorType | undefined; - outputFile?: string | undefined; - contentLength?: number; -} - export interface ServerToolCallConfirmationDetails { request: ToolCallRequestInfo; details: ToolCallConfirmationDetails; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 35804532d8..0165fffcbf 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -35,6 +35,7 @@ export * from './core/tokenLimits.js'; export * from './core/turn.js'; export * from './core/geminiRequest.js'; export * from './core/coreToolScheduler.js'; +export * from './scheduler/types.js'; export * from './core/nonInteractiveToolExecutor.js'; export * from './core/recordingContentGenerator.js'; diff --git a/packages/core/src/scheduler/types.ts b/packages/core/src/scheduler/types.ts new file mode 100644 index 0000000000..3a43a47704 --- /dev/null +++ b/packages/core/src/scheduler/types.ts @@ -0,0 +1,135 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Part } from '@google/genai'; +import type { + AnyDeclarativeTool, + AnyToolInvocation, + ToolCallConfirmationDetails, + ToolConfirmationOutcome, + ToolResultDisplay, +} from '../tools/tools.js'; +import type { AnsiOutput } from '../utils/terminalSerializer.js'; +import type { ToolErrorType } from '../tools/tool-error.js'; + +export interface ToolCallRequestInfo { + callId: string; + name: string; + args: Record; + isClientInitiated: boolean; + prompt_id: string; + checkpoint?: string; + traceId?: string; +} + +export interface ToolCallResponseInfo { + callId: string; + responseParts: Part[]; + resultDisplay: ToolResultDisplay | undefined; + error: Error | undefined; + errorType: ToolErrorType | undefined; + outputFile?: string | undefined; + contentLength?: number; +} + +export type ValidatingToolCall = { + status: 'validating'; + request: ToolCallRequestInfo; + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + startTime?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type ScheduledToolCall = { + status: 'scheduled'; + request: ToolCallRequestInfo; + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + startTime?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type ErroredToolCall = { + status: 'error'; + request: ToolCallRequestInfo; + response: ToolCallResponseInfo; + tool?: AnyDeclarativeTool; + durationMs?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type SuccessfulToolCall = { + status: 'success'; + request: ToolCallRequestInfo; + tool: AnyDeclarativeTool; + response: ToolCallResponseInfo; + invocation: AnyToolInvocation; + durationMs?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type ExecutingToolCall = { + status: 'executing'; + request: ToolCallRequestInfo; + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + liveOutput?: string | AnsiOutput; + startTime?: number; + outcome?: ToolConfirmationOutcome; + pid?: number; +}; + +export type CancelledToolCall = { + status: 'cancelled'; + request: ToolCallRequestInfo; + response: ToolCallResponseInfo; + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + durationMs?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type WaitingToolCall = { + status: 'awaiting_approval'; + request: ToolCallRequestInfo; + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + confirmationDetails: ToolCallConfirmationDetails; + startTime?: number; + outcome?: ToolConfirmationOutcome; +}; + +export type Status = ToolCall['status']; + +export type ToolCall = + | ValidatingToolCall + | ScheduledToolCall + | ErroredToolCall + | SuccessfulToolCall + | ExecutingToolCall + | CancelledToolCall + | WaitingToolCall; + +export type CompletedToolCall = + | SuccessfulToolCall + | CancelledToolCall + | ErroredToolCall; + +export type ConfirmHandler = ( + toolCall: WaitingToolCall, +) => Promise; + +export type OutputUpdateHandler = ( + toolCallId: string, + outputChunk: string | AnsiOutput, +) => void; + +export type AllToolCallsCompleteHandler = ( + completedToolCalls: CompletedToolCall[], +) => Promise; + +export type ToolCallsUpdateHandler = (toolCalls: ToolCall[]) => void; diff --git a/packages/core/src/telemetry/loggers.test.circular.ts b/packages/core/src/telemetry/loggers.test.circular.ts index 2f82743de3..060c70ffec 100644 --- a/packages/core/src/telemetry/loggers.test.circular.ts +++ b/packages/core/src/telemetry/loggers.test.circular.ts @@ -13,10 +13,10 @@ import { logToolCall } from './loggers.js'; import { ToolCallEvent } from './types.js'; import type { Config } from '../config/config.js'; import type { CompletedToolCall } from '../core/coreToolScheduler.js'; -import type { - ToolCallRequestInfo, - ToolCallResponseInfo, -} from '../core/turn.js'; +import { + type ToolCallRequestInfo, + type ToolCallResponseInfo, +} from '../scheduler/types.js'; import { MockTool } from '../test-utils/mock-tool.js'; describe('Circular Reference Handling', () => { diff --git a/packages/core/src/utils/checkpointUtils.test.ts b/packages/core/src/utils/checkpointUtils.test.ts index 50eb9325e3..e956730571 100644 --- a/packages/core/src/utils/checkpointUtils.test.ts +++ b/packages/core/src/utils/checkpointUtils.test.ts @@ -16,7 +16,7 @@ import { } from './checkpointUtils.js'; import type { GitService } from '../services/gitService.js'; import type { GeminiClient } from '../core/client.js'; -import type { ToolCallRequestInfo } from '../core/turn.js'; +import type { ToolCallRequestInfo } from '../scheduler/types.js'; describe('checkpoint utils', () => { describe('getToolCallDataSchema', () => { diff --git a/packages/core/src/utils/checkpointUtils.ts b/packages/core/src/utils/checkpointUtils.ts index a59ef73691..5bd66d7be9 100644 --- a/packages/core/src/utils/checkpointUtils.ts +++ b/packages/core/src/utils/checkpointUtils.ts @@ -10,7 +10,7 @@ import type { GeminiClient } from '../core/client.js'; import { getErrorMessage } from './errors.js'; import { z } from 'zod'; import type { Content } from '@google/genai'; -import type { ToolCallRequestInfo } from '../core/turn.js'; +import type { ToolCallRequestInfo } from '../scheduler/types.js'; export interface ToolCallData { history?: HistoryType; diff --git a/packages/core/src/utils/fileUtils.test.ts b/packages/core/src/utils/fileUtils.test.ts index 4d51206565..750151d3a4 100644 --- a/packages/core/src/utils/fileUtils.test.ts +++ b/packages/core/src/utils/fileUtils.test.ts @@ -32,6 +32,7 @@ import { readFileWithEncoding, fileExists, readWasmBinaryFromDisk, + saveTruncatedContent, } from './fileUtils.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; @@ -1022,4 +1023,213 @@ describe('fileUtils', () => { } }); }); + + describe('saveTruncatedContent', () => { + const THRESHOLD = 40_000; + const TRUNCATE_LINES = 1000; + + it('should return content unchanged if below threshold', async () => { + const content = 'Short content'; + const callId = 'test-call-id'; + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + expect(result).toEqual({ content }); + const outputFile = path.join(tempRootDir, `${callId}.output`); + expect(await fileExists(outputFile)).toBe(false); + }); + + 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)); + const content = lines.join('\n'); + const callId = 'test-call-id'; + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + const expectedOutputFile = path.join(tempRootDir, `${callId}.output`); + expect(result.outputFile).toBe(expectedOutputFile); + + const savedContent = await fsPromises.readFile( + expectedOutputFile, + 'utf-8', + ); + expect(savedContent).toBe(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 wrapWidth = 120; + + // Manually wrap the content to generate the expected file content + const wrappedLines: string[] = []; + for (let i = 0; i < content.length; i += wrapWidth) { + wrappedLines.push(content.substring(i, i + wrapWidth)); + } + const expectedFileContent = wrappedLines.join('\n'); + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + const expectedOutputFile = path.join(tempRootDir, `${callId}.output`); + expect(result.outputFile).toBe(expectedOutputFile); + + const savedContent = await fsPromises.readFile( + expectedOutputFile, + 'utf-8', + ); + expect(savedContent).toBe(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 save to correct file path with call ID', async () => { + const content = 'a'.repeat(200_000); + const callId = 'unique-call-123'; + const wrapWidth = 120; + + // Manually wrap the content to generate the expected file content + const wrappedLines: string[] = []; + for (let i = 0; i < content.length; i += wrapWidth) { + wrappedLines.push(content.substring(i, i + wrapWidth)); + } + const expectedFileContent = wrappedLines.join('\n'); + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + const expectedPath = path.join(tempRootDir, `${callId}.output`); + expect(result.outputFile).toBe(expectedPath); + + const savedContent = await fsPromises.readFile(expectedPath, 'utf-8'); + expect(savedContent).toBe(expectedFileContent); + }); + + it('should include helpful instructions in truncated message', async () => { + const content = 'a'.repeat(200_000); + const callId = 'test-call-id'; + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + 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 sanitize callId to prevent path traversal', async () => { + const content = 'a'.repeat(200_000); + const callId = '../../../../../etc/passwd'; + const wrapWidth = 120; + + // Manually wrap the content to generate the expected file content + const wrappedLines: string[] = []; + for (let i = 0; i < content.length; i += wrapWidth) { + wrappedLines.push(content.substring(i, i + wrapWidth)); + } + const expectedFileContent = wrappedLines.join('\n'); + + await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + const expectedPath = path.join(tempRootDir, 'passwd.output'); + + const savedContent = await fsPromises.readFile(expectedPath, 'utf-8'); + expect(savedContent).toBe(expectedFileContent); + }); + + it('should handle file write errors gracefully', async () => { + const content = 'a'.repeat(50_000); + const callId = 'test-call-id-fail'; + + const writeFileSpy = vi + .spyOn(fsPromises, 'writeFile') + .mockRejectedValue(new Error('File write failed')); + + const result = await saveTruncatedContent( + content, + callId, + tempRootDir, + THRESHOLD, + TRUNCATE_LINES, + ); + + expect(result.outputFile).toBeUndefined(); + expect(result.content).toContain( + '[Note: Could not save full output to file]', + ); + expect(writeFileSpy).toHaveBeenCalled(); + + writeFileSpy.mockRestore(); + }); + }); }); diff --git a/packages/core/src/utils/fileUtils.ts b/packages/core/src/utils/fileUtils.ts index 321162f192..d91a651236 100644 --- a/packages/core/src/utils/fileUtils.ts +++ b/packages/core/src/utils/fileUtils.ts @@ -15,6 +15,7 @@ import { ToolErrorType } from '../tools/tool-error.js'; import { BINARY_EXTENSIONS } from './ignorePatterns.js'; import { createRequire as createModuleRequire } from 'node:module'; import { debugLogger } from './debugLogger.js'; +import { READ_FILE_TOOL_NAME } from '../tools/tool-names.js'; const requireModule = createModuleRequire(import.meta.url); @@ -515,3 +516,67 @@ export async function fileExists(filePath: string): Promise { return false; } } + +export async function saveTruncatedContent( + content: string, + callId: string, + projectTempDir: string, + threshold: number, + truncateLines: number, +): Promise<{ content: string; outputFile?: string }> { + if (content.length <= threshold) { + return { content }; + } + + let lines = content.split('\n'); + let fileContent = content; + + // If the content is long but has few lines, wrap it to enable line-based truncation. + if (lines.length <= truncateLines) { + const wrapWidth = 120; // A reasonable width for wrapping. + const wrappedLines: string[] = []; + for (const line of lines) { + if (line.length > wrapWidth) { + for (let i = 0; i < line.length; i += wrapWidth) { + wrappedLines.push(line.substring(i, i + wrapWidth)); + } + } else { + wrappedLines.push(line); + } + } + lines = wrappedLines; + fileContent = lines.join('\n'); + } + + const head = Math.floor(truncateLines / 5); + const beginning = lines.slice(0, head); + const end = lines.slice(-(truncateLines - head)); + const truncatedContent = + beginning.join('\n') + '\n... [CONTENT TRUNCATED] ...\n' + end.join('\n'); + + // Sanitize callId to prevent path traversal. + const safeFileName = `${path.basename(callId)}.output`; + const outputFile = path.join(projectTempDir, safeFileName); + try { + await fsPromises.writeFile(outputFile, fileContent); + + return { + content: `Tool output was too large and has been truncated. +The full output has been saved to: ${outputFile} +To read the complete output, use the ${READ_FILE_TOOL_NAME} tool with the absolute file path above. For large files, you can use the offset and limit parameters to read specific sections: +- ${READ_FILE_TOOL_NAME} tool with offset=0, limit=100 to see the first 100 lines +- ${READ_FILE_TOOL_NAME} tool with offset=N to skip N lines from the beginning +- ${READ_FILE_TOOL_NAME} tool with limit=M to read only M lines at a time +The truncated output below shows the beginning and end of the content. The marker '... [CONTENT TRUNCATED] ...' indicates where content was removed. +This allows you to efficiently examine different parts of the output without loading the entire file. +Truncated part of the output: +${truncatedContent}`, + outputFile, + }; + } catch (_error) { + return { + content: + truncatedContent + `\n[Note: Could not save full output to file]`, + }; + } +} diff --git a/packages/core/src/utils/generateContentResponseUtilities.test.ts b/packages/core/src/utils/generateContentResponseUtilities.test.ts index b4c413fcbf..0562f91888 100644 --- a/packages/core/src/utils/generateContentResponseUtilities.test.ts +++ b/packages/core/src/utils/generateContentResponseUtilities.test.ts @@ -14,14 +14,20 @@ import { getStructuredResponse, getStructuredResponseFromParts, getCitations, + convertToFunctionResponse, } from './generateContentResponseUtilities.js'; import type { GenerateContentResponse, Part, SafetyRating, CitationMetadata, + PartListUnion, } from '@google/genai'; import { FinishReason } from '@google/genai'; +import { + DEFAULT_GEMINI_MODEL, + PREVIEW_GEMINI_MODEL, +} from '../config/models.js'; const mockTextPart = (text: string): Part => ({ text }); const mockFunctionCallPart = ( @@ -72,6 +78,312 @@ const minimalMockResponse = ( }); describe('generateContentResponseUtilities', () => { + 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: Part = { 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: PartListUnion = [{ 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: PartListUnion = [{ 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: Part = { + 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: Part = { + 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: Part = { + 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: Part = { + 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: PartListUnion = [ + { 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: PartListUnion = [ + { 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: Part = { 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: PartListUnion = []; + 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: Part = {}; // An empty part object + const result = convertToFunctionResponse( + toolName, + callId, + llmContent, + PREVIEW_GEMINI_MODEL, + ); + expect(result).toEqual([ + { + functionResponse: { + name: toolName, + id: callId, + response: {}, + }, + }, + ]); + }); + }); + describe('getCitations', () => { it('should return empty array for no candidates', () => { expect(getCitations(minimalMockResponse(undefined))).toEqual([]); diff --git a/packages/core/src/utils/generateContentResponseUtilities.ts b/packages/core/src/utils/generateContentResponseUtilities.ts index 2532988533..5151da9f6d 100644 --- a/packages/core/src/utils/generateContentResponseUtilities.ts +++ b/packages/core/src/utils/generateContentResponseUtilities.ts @@ -8,8 +8,125 @@ import type { GenerateContentResponse, Part, FunctionCall, + PartListUnion, } from '@google/genai'; import { getResponseText } from './partUtils.js'; +import { supportsMultimodalFunctionResponse } from '../config/models.js'; +import { debugLogger } from './debugLogger.js'; + +/** + * Formats tool output for a Gemini FunctionResponse. + */ +function createFunctionResponsePart( + callId: string, + toolName: string, + output: string, +): Part { + return { + functionResponse: { + id: callId, + name: toolName, + response: { output }, + }, + }; +} + +function toParts(input: PartListUnion): Part[] { + const parts: Part[] = []; + for (const part of Array.isArray(input) ? input : [input]) { + if (typeof part === 'string') { + parts.push({ text: part }); + } else if (part) { + parts.push(part); + } + } + return parts; +} + +export function convertToFunctionResponse( + toolName: string, + callId: string, + llmContent: PartListUnion, + model: string, +): Part[] { + if (typeof llmContent === 'string') { + return [createFunctionResponsePart(callId, toolName, llmContent)]; + } + + const parts = toParts(llmContent); + + // Separate text from binary types + const textParts: string[] = []; + const inlineDataParts: Part[] = []; + const fileDataParts: Part[] = []; + + for (const part of parts) { + if (part.text !== undefined) { + textParts.push(part.text); + } else if (part.inlineData) { + inlineDataParts.push(part); + } else if (part.fileData) { + fileDataParts.push(part); + } else if (part.functionResponse) { + if (parts.length > 1) { + debugLogger.warn( + 'convertToFunctionResponse received multiple parts with a functionResponse. Only the functionResponse will be used, other parts will be ignored', + ); + } + // Handle passthrough case + return [ + { + functionResponse: { + id: callId, + name: toolName, + response: part.functionResponse.response, + }, + }, + ]; + } + // Ignore other part types + } + + // Build the primary response part + const part: Part = { + functionResponse: { + id: callId, + name: toolName, + response: textParts.length > 0 ? { output: textParts.join('\n') } : {}, + }, + }; + + const isMultimodalFRSupported = supportsMultimodalFunctionResponse(model); + const siblingParts: Part[] = [...fileDataParts]; + + if (inlineDataParts.length > 0) { + if (isMultimodalFRSupported) { + // Nest inlineData if supported by the model + (part.functionResponse as unknown as { parts: Part[] }).parts = + inlineDataParts; + } else { + // Otherwise treat as siblings + siblingParts.push(...inlineDataParts); + } + } + + // Add descriptive text if the response object is empty but we have binary content + if ( + textParts.length === 0 && + (inlineDataParts.length > 0 || fileDataParts.length > 0) + ) { + const totalBinaryItems = inlineDataParts.length + fileDataParts.length; + part.functionResponse!.response = { + output: `Binary content provided (${totalBinaryItems} item(s)).`, + }; + } + + if (siblingParts.length > 0) { + return [part, ...siblingParts]; + } + + return [part]; +} export function getResponseTextFromParts(parts: Part[]): string | undefined { if (!parts) { diff --git a/packages/core/src/utils/tool-utils.test.ts b/packages/core/src/utils/tool-utils.test.ts index 672bc72309..861281b2da 100644 --- a/packages/core/src/utils/tool-utils.test.ts +++ b/packages/core/src/utils/tool-utils.test.ts @@ -5,10 +5,34 @@ */ import { expect, describe, it } from 'vitest'; -import { doesToolInvocationMatch } from './tool-utils.js'; +import { doesToolInvocationMatch, getToolSuggestion } from './tool-utils.js'; import type { AnyToolInvocation, Config } from '../index.js'; import { ReadFileTool } from '../tools/read-file.js'; +describe('getToolSuggestion', () => { + it('should suggest the top N closest tool names for a typo', () => { + const allToolNames = ['list_files', 'read_file', 'write_file']; + + // Test that the right tool is selected, with only 1 result, for typos + const misspelledTool = getToolSuggestion('list_fils', allToolNames, 1); + expect(misspelledTool).toBe(' Did you mean "list_files"?'); + + // Test that the right tool is selected, with only 1 result, for prefixes + const prefixedTool = getToolSuggestion( + 'github.list_files', + allToolNames, + 1, + ); + expect(prefixedTool).toBe(' Did you mean "list_files"?'); + + // Test that the right tool is first + const suggestionMultiple = getToolSuggestion('list_fils', allToolNames); + expect(suggestionMultiple).toBe( + ' Did you mean one of: "list_files", "read_file", "write_file"?', + ); + }); +}); + describe('doesToolInvocationMatch', () => { it('should not match a partial command prefix', () => { const invocation = { diff --git a/packages/core/src/utils/tool-utils.ts b/packages/core/src/utils/tool-utils.ts index 45c6cbe665..0d2dec8625 100644 --- a/packages/core/src/utils/tool-utils.ts +++ b/packages/core/src/utils/tool-utils.ts @@ -7,6 +7,44 @@ import type { AnyDeclarativeTool, AnyToolInvocation } from '../index.js'; import { isTool } from '../index.js'; import { SHELL_TOOL_NAMES } from './shell-utils.js'; +import levenshtein from 'fast-levenshtein'; + +/** + * Generates a suggestion string for a tool name that was not found in the registry. + * It finds the closest matches based on Levenshtein distance. + * @param unknownToolName The tool name that was not found. + * @param allToolNames The list of all available tool names. + * @param topN The number of suggestions to return. Defaults to 3. + * @returns A suggestion string like " Did you mean 'tool'?" or " Did you mean one of: 'tool1', 'tool2'?", or an empty string if no suggestions are found. + */ +export function getToolSuggestion( + unknownToolName: string, + allToolNames: string[], + topN = 3, +): string { + const matches = allToolNames.map((toolName) => ({ + name: toolName, + distance: levenshtein.get(unknownToolName, toolName), + })); + + matches.sort((a, b) => a.distance - b.distance); + + const topNResults = matches.slice(0, topN); + + if (topNResults.length === 0) { + return ''; + } + + const suggestedNames = topNResults + .map((match) => `"${match.name}"`) + .join(', '); + + if (topNResults.length > 1) { + return ` Did you mean one of: ${suggestedNames}?`; + } else { + return ` Did you mean ${suggestedNames}?`; + } +} /** * Checks if a tool invocation matches any of a list of patterns.