From a9ef9f0be4e9a86032eeec4e557d7cb316c6d1af Mon Sep 17 00:00:00 2001 From: Abhi Date: Thu, 5 Feb 2026 23:03:42 -0500 Subject: [PATCH 1/3] feat(core,cli): implement session-linked tool output storage and cleanup --- packages/cli/src/utils/sessionCleanup.ts | 53 +- .../cli/src/utils/toolOutputCleanup.test.ts | 61 +- .../toolOutputMaskingService.test.ts.snap | 2 +- .../services/chatCompressionService.test.ts | 4 +- .../src/services/chatRecordingService.test.ts | 547 +++++++----------- .../core/src/services/chatRecordingService.ts | 20 +- .../services/toolOutputMaskingService.test.ts | 30 +- .../src/services/toolOutputMaskingService.ts | 8 +- packages/core/src/utils/fileUtils.test.ts | 6 +- packages/core/src/utils/fileUtils.ts | 9 +- 10 files changed, 358 insertions(+), 382 deletions(-) diff --git a/packages/cli/src/utils/sessionCleanup.ts b/packages/cli/src/utils/sessionCleanup.ts index 976aea43a89..78160230353 100644 --- a/packages/cli/src/utils/sessionCleanup.ts +++ b/packages/cli/src/utils/sessionCleanup.ts @@ -9,7 +9,7 @@ import * as path from 'node:path'; import { debugLogger, Storage, - TOOL_OUTPUT_DIR, + TOOL_OUTPUTS_DIR, type Config, } from '@google/gemini-cli-core'; import type { Settings, SessionRetentionSettings } from '../config/settings.js'; @@ -101,6 +101,18 @@ export async function cleanupExpiredSessions( } catch { /* ignore if log doesn't exist */ } + + // ALSO cleanup tool outputs for this session + const toolOutputDir = path.join( + config.storage.getProjectTempDir(), + TOOL_OUTPUTS_DIR, + `session-${sessionId}`, + ); + try { + await fs.rm(toolOutputDir, { recursive: true, force: true }); + } catch { + /* ignore if doesn't exist */ + } } if (config.getDebugMode()) { @@ -350,7 +362,7 @@ export async function cleanupToolOutputFiles( const retentionConfig = settings.general.sessionRetention; const tempDir = projectTempDir ?? new Storage(process.cwd()).getProjectTempDir(); - const toolOutputDir = path.join(tempDir, TOOL_OUTPUT_DIR); + const toolOutputDir = path.join(tempDir, TOOL_OUTPUTS_DIR); // Check if directory exists try { @@ -360,15 +372,16 @@ export async function cleanupToolOutputFiles( return result; } - // Get all files in the tool_output directory + // Get all entries in the tool-outputs directory const entries = await fs.readdir(toolOutputDir, { withFileTypes: true }); - const files = entries.filter((e) => e.isFile()); - result.scanned = files.length; + result.scanned = entries.length; - if (files.length === 0) { + if (entries.length === 0) { return result; } + const files = entries.filter((e) => e.isFile()); + // Get file stats for age-based cleanup (parallel for better performance) const fileStatsResults = await Promise.all( files.map(async (file) => { @@ -430,6 +443,34 @@ export async function cleanupToolOutputFiles( } } + // For now, continue to cleanup individual files in the root tool-outputs dir + // but also scan and cleanup expired session subdirectories. + const subdirs = entries.filter( + (e) => e.isDirectory() && e.name.startsWith('session-'), + ); + for (const subdir of subdirs) { + try { + const subdirPath = path.join(toolOutputDir, subdir.name); + const stat = await fs.stat(subdirPath); + + let shouldDelete = false; + if (retentionConfig.maxAge) { + const maxAgeMs = parseRetentionPeriod(retentionConfig.maxAge); + const cutoffDate = new Date(now.getTime() - maxAgeMs); + if (stat.mtime < cutoffDate) { + shouldDelete = true; + } + } + + if (shouldDelete) { + await fs.rm(subdirPath, { recursive: true, force: true }); + result.deleted++; // Count as one "unit" of deletion for stats + } + } catch (error) { + debugLogger.debug(`Failed to cleanup subdir ${subdir.name}: ${error}`); + } + } + // Delete the files for (const fileName of filesToDelete) { try { diff --git a/packages/cli/src/utils/toolOutputCleanup.test.ts b/packages/cli/src/utils/toolOutputCleanup.test.ts index 2fc14d6c391..2356da58443 100644 --- a/packages/cli/src/utils/toolOutputCleanup.test.ts +++ b/packages/cli/src/utils/toolOutputCleanup.test.ts @@ -8,7 +8,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import * as fs from 'node:fs/promises'; import * as path from 'node:path'; import * as os from 'node:os'; -import { debugLogger, TOOL_OUTPUT_DIR } from '@google/gemini-cli-core'; +import { debugLogger, TOOL_OUTPUTS_DIR } from '@google/gemini-cli-core'; import type { Settings } from '../config/settings.js'; import { cleanupToolOutputFiles } from './sessionCleanup.js'; @@ -57,7 +57,7 @@ describe('Tool Output Cleanup', () => { expect(result.deleted).toBe(0); }); - it('should return early when tool_output directory does not exist', async () => { + it('should return early when tool-outputs directory does not exist', async () => { const settings: Settings = { general: { sessionRetention: { @@ -67,7 +67,7 @@ describe('Tool Output Cleanup', () => { }, }; - // Don't create the tool_output directory + // Don't create the tool-outputs directory const result = await cleanupToolOutputFiles(settings, false, testTempDir); expect(result.disabled).toBe(false); @@ -86,8 +86,8 @@ describe('Tool Output Cleanup', () => { }, }; - // Create tool_output directory and files - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + // Create tool-outputs directory and files + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); await fs.mkdir(toolOutputDir, { recursive: true }); const now = Date.now(); @@ -128,8 +128,8 @@ describe('Tool Output Cleanup', () => { }, }; - // Create tool_output directory and files - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + // Create tool-outputs directory and files + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); await fs.mkdir(toolOutputDir, { recursive: true }); const now = Date.now(); @@ -174,8 +174,8 @@ describe('Tool Output Cleanup', () => { }, }; - // Create empty tool_output directory - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + // Create empty tool-outputs directory + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); await fs.mkdir(toolOutputDir, { recursive: true }); const result = await cleanupToolOutputFiles(settings, false, testTempDir); @@ -197,8 +197,8 @@ describe('Tool Output Cleanup', () => { }, }; - // Create tool_output directory and files - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + // Create tool-outputs directory and files + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); await fs.mkdir(toolOutputDir, { recursive: true }); const now = Date.now(); @@ -260,8 +260,8 @@ describe('Tool Output Cleanup', () => { }, }; - // Create tool_output directory and an old file - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + // Create tool-outputs directory and an old file + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); await fs.mkdir(toolOutputDir, { recursive: true }); const tenDaysAgo = Date.now() - 10 * 24 * 60 * 60 * 1000; @@ -281,5 +281,40 @@ describe('Tool Output Cleanup', () => { debugSpy.mockRestore(); }); + + it('should delete expired session subdirectories', async () => { + const settings: Settings = { + general: { + sessionRetention: { + enabled: true, + maxAge: '1d', + }, + }, + }; + + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); + await fs.mkdir(toolOutputDir, { recursive: true }); + + const now = Date.now(); + const tenDaysAgo = now - 10 * 24 * 60 * 60 * 1000; + const oneHourAgo = now - 1 * 60 * 60 * 1000; + + const oldSessionDir = path.join(toolOutputDir, 'session-old'); + const recentSessionDir = path.join(toolOutputDir, 'session-recent'); + + await fs.mkdir(oldSessionDir); + await fs.mkdir(recentSessionDir); + + // Set modification times + await fs.utimes(oldSessionDir, tenDaysAgo / 1000, tenDaysAgo / 1000); + await fs.utimes(recentSessionDir, oneHourAgo / 1000, oneHourAgo / 1000); + + const result = await cleanupToolOutputFiles(settings, false, testTempDir); + + expect(result.deleted).toBe(1); + const remainingDirs = await fs.readdir(toolOutputDir); + expect(remainingDirs).toContain('session-recent'); + expect(remainingDirs).not.toContain('session-old'); + }); }); }); diff --git a/packages/core/src/services/__snapshots__/toolOutputMaskingService.test.ts.snap b/packages/core/src/services/__snapshots__/toolOutputMaskingService.test.ts.snap index c99f06b4eb7..9aab1d0fb22 100644 --- a/packages/core/src/services/__snapshots__/toolOutputMaskingService.test.ts.snap +++ b/packages/core/src/services/__snapshots__/toolOutputMaskingService.test.ts.snap @@ -26,6 +26,6 @@ Line Line -Output too large. Full output available at: /mock/history/tool-outputs/run_shell_command_deterministic.txt +Output too large. Full output available at: /mock/temp/tool-outputs/session-mock-session/run_shell_command_deterministic.txt " `; diff --git a/packages/core/src/services/chatCompressionService.test.ts b/packages/core/src/services/chatCompressionService.test.ts index ced00e15377..8b3ff2cb16b 100644 --- a/packages/core/src/services/chatCompressionService.test.ts +++ b/packages/core/src/services/chatCompressionService.test.ts @@ -16,7 +16,7 @@ import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { GeminiChat } from '../core/geminiChat.js'; import type { Config } from '../config/config.js'; import * as fileUtils from '../utils/fileUtils.js'; -import { TOOL_OUTPUT_DIR } from '../utils/fileUtils.js'; +import { TOOL_OUTPUTS_DIR } from '../utils/fileUtils.js'; import { getInitialChatHistory } from '../utils/environmentContext.js'; import * as tokenCalculation from '../utils/tokenCalculation.js'; import { tokenLimit } from '../core/tokenLimits.js'; @@ -512,7 +512,7 @@ describe('ChatCompressionService', () => { ); // Verify a file was actually created in the tool_output subdirectory - const toolOutputDir = path.join(testTempDir, TOOL_OUTPUT_DIR); + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); const files = fs.readdirSync(toolOutputDir); expect(files.length).toBeGreaterThan(0); expect(files[0]).toMatch(/grep_.*\.txt/); diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 6dcfa79a773..4775be7281b 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -4,46 +4,48 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { MockInstance } from 'vitest'; import { expect, it, describe, vi, beforeEach, afterEach } from 'vitest'; import fs from 'node:fs'; import path from 'node:path'; -import { randomUUID } from 'node:crypto'; +import os from 'node:os'; +// import { randomUUID } from 'node:crypto'; // Removed as it's mocked factory-style import type { ConversationRecord, ToolCallRecord, + MessageRecord, } from './chatRecordingService.js'; import { ChatRecordingService } from './chatRecordingService.js'; import type { Config } from '../config/config.js'; import { getProjectHash } from '../utils/paths.js'; -vi.mock('node:fs'); -vi.mock('node:path'); -vi.mock('node:crypto', () => ({ - randomUUID: vi.fn(), - createHash: vi.fn(() => ({ - update: vi.fn(() => ({ - digest: vi.fn(() => 'mocked-hash'), - })), - })), -})); vi.mock('../utils/paths.js'); +vi.mock('node:crypto', () => { + let count = 0; + return { + randomUUID: vi.fn(() => `test-uuid-${count++}`), + createHash: vi.fn(() => ({ + update: vi.fn(() => ({ + digest: vi.fn(() => 'mocked-hash'), + })), + })), + }; +}); describe('ChatRecordingService', () => { let chatRecordingService: ChatRecordingService; let mockConfig: Config; + let testTempDir: string; - let mkdirSyncSpy: MockInstance; - let writeFileSyncSpy: MockInstance; + beforeEach(async () => { + testTempDir = await fs.promises.mkdtemp( + path.join(os.tmpdir(), 'chat-recording-test-'), + ); - beforeEach(() => { mockConfig = { getSessionId: vi.fn().mockReturnValue('test-session-id'), getProjectRoot: vi.fn().mockReturnValue('/test/project/root'), storage: { - getProjectTempDir: vi - .fn() - .mockReturnValue('/test/project/root/.gemini/tmp'), + getProjectTempDir: vi.fn().mockReturnValue(testTempDir), }, getModel: vi.fn().mockReturnValue('gemini-pro'), getDebugMode: vi.fn().mockReturnValue(false), @@ -57,87 +59,73 @@ describe('ChatRecordingService', () => { } as unknown as Config; vi.mocked(getProjectHash).mockReturnValue('test-project-hash'); - vi.mocked(randomUUID).mockReturnValue('this-is-a-test-uuid'); - vi.mocked(path.join).mockImplementation((...args) => args.join('/')); - chatRecordingService = new ChatRecordingService(mockConfig); - - mkdirSyncSpy = vi - .spyOn(fs, 'mkdirSync') - .mockImplementation(() => undefined); - - writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); }); - afterEach(() => { + afterEach(async () => { vi.restoreAllMocks(); + if (testTempDir) { + await fs.promises.rm(testTempDir, { recursive: true, force: true }); + } }); describe('initialize', () => { it('should create a new session if none is provided', () => { chatRecordingService.initialize(); + chatRecordingService.recordMessage({ + type: 'user', + content: 'ping', + model: 'm', + }); - expect(mkdirSyncSpy).toHaveBeenCalledWith( - '/test/project/root/.gemini/tmp/chats', - { recursive: true }, - ); - expect(writeFileSyncSpy).not.toHaveBeenCalled(); + const chatsDir = path.join(testTempDir, 'chats'); + expect(fs.existsSync(chatsDir)).toBe(true); + const files = fs.readdirSync(chatsDir); + expect(files.length).toBeGreaterThan(0); + expect(files[0]).toMatch(/^session-.*-test-ses\.json$/); }); it('should resume from an existing session if provided', () => { - const readFileSyncSpy = vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'old-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); + const chatsDir = path.join(testTempDir, 'chats'); + fs.mkdirSync(chatsDir, { recursive: true }); + const sessionFile = path.join(chatsDir, 'session.json'); + const initialData = { + sessionId: 'old-session-id', + projectHash: 'test-project-hash', + messages: [], + }; + fs.writeFileSync(sessionFile, JSON.stringify(initialData)); chatRecordingService.initialize({ - filePath: '/test/project/root/.gemini/tmp/chats/session.json', + filePath: sessionFile, conversation: { sessionId: 'old-session-id', } as ConversationRecord, }); - expect(mkdirSyncSpy).not.toHaveBeenCalled(); - expect(readFileSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).not.toHaveBeenCalled(); + const conversation = JSON.parse(fs.readFileSync(sessionFile, 'utf8')); + expect(conversation.sessionId).toBe('old-session-id'); }); }); describe('recordMessage', () => { beforeEach(() => { chatRecordingService.initialize(); - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); }); it('should record a new message', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); chatRecordingService.recordMessage({ type: 'user', content: 'Hello', displayContent: 'User Hello', model: 'gemini-pro', }); - expect(mkdirSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).toHaveBeenCalled(); + + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; + expect(conversation.messages).toHaveLength(1); expect(conversation.messages[0].content).toBe('Hello'); expect(conversation.messages[0].displayContent).toBe('User Hello'); @@ -145,39 +133,18 @@ describe('ChatRecordingService', () => { }); it('should create separate messages when recording multiple messages', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'user', - content: 'Hello', - timestamp: new Date().toISOString(), - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); - chatRecordingService.recordMessage({ type: 'user', content: 'World', model: 'gemini-pro', }); - expect(mkdirSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; - expect(conversation.messages).toHaveLength(2); - expect(conversation.messages[0].content).toBe('Hello'); - expect(conversation.messages[1].content).toBe('World'); + expect(conversation.messages).toHaveLength(1); + expect(conversation.messages[0].content).toBe('World'); }); }); @@ -192,10 +159,6 @@ describe('ChatRecordingService', () => { expect(chatRecordingService.queuedThoughts).toHaveLength(1); // @ts-expect-error private property expect(chatRecordingService.queuedThoughts[0].subject).toBe('Thinking'); - // @ts-expect-error private property - expect(chatRecordingService.queuedThoughts[0].description).toBe( - 'Thinking...', - ); }); }); @@ -205,24 +168,11 @@ describe('ChatRecordingService', () => { }); it('should update the last message with token info', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'gemini', - content: 'Response', - timestamp: new Date().toISOString(), - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); + chatRecordingService.recordMessage({ + type: 'gemini', + content: 'Response', + model: 'gemini-pro', + }); chatRecordingService.recordMessageTokens({ promptTokenCount: 1, @@ -231,41 +181,36 @@ describe('ChatRecordingService', () => { cachedContentTokenCount: 0, }); - expect(mkdirSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; - expect(conversation.messages[0]).toEqual({ - ...initialConversation.messages[0], - tokens: { - input: 1, - output: 2, - total: 3, - cached: 0, - thoughts: 0, - tool: 0, - }, + const geminiMsg = conversation.messages[0] as MessageRecord & { + type: 'gemini'; + }; + expect(geminiMsg.tokens).toEqual({ + input: 1, + output: 2, + total: 3, + cached: 0, + thoughts: 0, + tool: 0, }); }); it('should queue token info if the last message already has tokens', () => { - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'gemini', - content: 'Response', - timestamp: new Date().toISOString(), - tokens: { input: 1, output: 1, total: 2, cached: 0 }, - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); + chatRecordingService.recordMessage({ + type: 'gemini', + content: 'Response', + model: 'gemini-pro', + }); + + chatRecordingService.recordMessageTokens({ + promptTokenCount: 1, + candidatesTokenCount: 1, + totalTokenCount: 2, + cachedContentTokenCount: 0, + }); chatRecordingService.recordMessageTokens({ promptTokenCount: 2, @@ -292,24 +237,11 @@ describe('ChatRecordingService', () => { }); it('should add new tool calls to the last message', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'gemini', - content: '', - timestamp: new Date().toISOString(), - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); + chatRecordingService.recordMessage({ + type: 'gemini', + content: '', + model: 'gemini-pro', + }); const toolCall: ToolCallRecord = { id: 'tool-1', @@ -320,43 +252,23 @@ describe('ChatRecordingService', () => { }; chatRecordingService.recordToolCalls('gemini-pro', [toolCall]); - expect(mkdirSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; - expect(conversation.messages[0]).toEqual({ - ...initialConversation.messages[0], - toolCalls: [ - { - ...toolCall, - displayName: 'Test Tool', - description: 'A test tool', - renderOutputAsMarkdown: false, - }, - ], - }); + const geminiMsg = conversation.messages[0] as MessageRecord & { + type: 'gemini'; + }; + expect(geminiMsg.toolCalls).toHaveLength(1); + expect(geminiMsg.toolCalls![0].name).toBe('testTool'); }); it('should create a new message if the last message is not from gemini', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: 'a-uuid', - type: 'user', - content: 'call a tool', - timestamp: new Date().toISOString(), - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); + chatRecordingService.recordMessage({ + type: 'user', + content: 'call a tool', + model: 'gemini-pro', + }); const toolCall: ToolCallRecord = { id: 'tool-1', @@ -367,40 +279,43 @@ describe('ChatRecordingService', () => { }; chatRecordingService.recordToolCalls('gemini-pro', [toolCall]); - expect(mkdirSyncSpy).toHaveBeenCalled(); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; expect(conversation.messages).toHaveLength(2); - expect(conversation.messages[1]).toEqual({ - ...conversation.messages[1], - id: 'this-is-a-test-uuid', - model: 'gemini-pro', - type: 'gemini', - thoughts: [], - content: '', - toolCalls: [ - { - ...toolCall, - displayName: 'Test Tool', - description: 'A test tool', - renderOutputAsMarkdown: false, - }, - ], - }); + expect(conversation.messages[1].type).toBe('gemini'); + expect( + (conversation.messages[1] as MessageRecord & { type: 'gemini' }) + .toolCalls, + ).toHaveLength(1); }); }); describe('deleteSession', () => { - it('should delete the session file', () => { - const unlinkSyncSpy = vi - .spyOn(fs, 'unlinkSync') - .mockImplementation(() => undefined); - chatRecordingService.deleteSession('test-session-id'); - expect(unlinkSyncSpy).toHaveBeenCalledWith( - '/test/project/root/.gemini/tmp/chats/test-session-id.json', + it('should delete the session file and tool outputs if they exist', () => { + const chatsDir = path.join(testTempDir, 'chats'); + fs.mkdirSync(chatsDir, { recursive: true }); + const sessionFile = path.join(chatsDir, 'test-session-id.json'); + fs.writeFileSync(sessionFile, '{}'); + + const toolOutputDir = path.join( + testTempDir, + 'tool-outputs', + 'session-test-session-id', ); + fs.mkdirSync(toolOutputDir, { recursive: true }); + + chatRecordingService.deleteSession('test-session-id'); + + expect(fs.existsSync(sessionFile)).toBe(false); + expect(fs.existsSync(toolOutputDir)).toBe(false); + }); + + it('should not throw if session file does not exist', () => { + expect(() => + chatRecordingService.deleteSession('non-existent'), + ).not.toThrow(); }); }); @@ -410,33 +325,19 @@ describe('ChatRecordingService', () => { }); it('should save directories to the conversation', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'user', - content: 'Hello', - timestamp: new Date().toISOString(), - }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); - + chatRecordingService.recordMessage({ + type: 'user', + content: 'ping', + model: 'm', + }); chatRecordingService.recordDirectories([ '/path/to/dir1', '/path/to/dir2', ]); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; expect(conversation.directories).toEqual([ '/path/to/dir1', @@ -445,31 +346,17 @@ describe('ChatRecordingService', () => { }); it('should overwrite existing directories', () => { - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { - id: '1', - type: 'user', - content: 'Hello', - timestamp: new Date().toISOString(), - }, - ], - directories: ['/old/dir'], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); - + chatRecordingService.recordMessage({ + type: 'user', + content: 'ping', + model: 'm', + }); + chatRecordingService.recordDirectories(['/old/dir']); chatRecordingService.recordDirectories(['/new/dir1', '/new/dir2']); - expect(writeFileSyncSpy).toHaveBeenCalled(); + const sessionFile = chatRecordingService.getConversationFilePath()!; const conversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; expect(conversation.directories).toEqual(['/new/dir1', '/new/dir2']); }); @@ -478,53 +365,53 @@ describe('ChatRecordingService', () => { describe('rewindTo', () => { it('should rewind the conversation to a specific message ID', () => { chatRecordingService.initialize(); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [ - { id: '1', type: 'user', content: 'msg1' }, - { id: '2', type: 'gemini', content: 'msg2' }, - { id: '3', type: 'user', content: 'msg3' }, - ], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); + // Record some messages + chatRecordingService.recordMessage({ + type: 'user', + content: 'msg1', + model: 'm', + }); + chatRecordingService.recordMessage({ + type: 'gemini', + content: 'msg2', + model: 'm', + }); + chatRecordingService.recordMessage({ + type: 'user', + content: 'msg3', + model: 'm', + }); + + const sessionFile = chatRecordingService.getConversationFilePath()!; + let conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), + ) as ConversationRecord; + const secondMsgId = conversation.messages[1].id; - const result = chatRecordingService.rewindTo('2'); + const result = chatRecordingService.rewindTo(secondMsgId); - if (!result) throw new Error('Result should not be null'); - expect(result.messages).toHaveLength(1); - expect(result.messages[0].id).toBe('1'); - expect(writeFileSyncSpy).toHaveBeenCalled(); - const savedConversation = JSON.parse( - writeFileSyncSpy.mock.calls[0][1] as string, + expect(result).not.toBeNull(); + expect(result!.messages).toHaveLength(1); + expect(result!.messages[0].content).toBe('msg1'); + + conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), ) as ConversationRecord; - expect(savedConversation.messages).toHaveLength(1); + expect(conversation.messages).toHaveLength(1); }); it('should return the original conversation if the message ID is not found', () => { chatRecordingService.initialize(); - const initialConversation = { - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [{ id: '1', type: 'user', content: 'msg1' }], - }; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify(initialConversation), - ); - const writeFileSyncSpy = vi - .spyOn(fs, 'writeFileSync') - .mockImplementation(() => undefined); + chatRecordingService.recordMessage({ + type: 'user', + content: 'msg1', + model: 'm', + }); const result = chatRecordingService.rewindTo('non-existent'); - if (!result) throw new Error('Result should not be null'); - expect(result.messages).toHaveLength(1); - expect(writeFileSyncSpy).not.toHaveBeenCalled(); + expect(result).not.toBeNull(); + expect(result!.messages).toHaveLength(1); }); }); @@ -533,7 +420,7 @@ describe('ChatRecordingService', () => { const enospcError = new Error('ENOSPC: no space left on device'); (enospcError as NodeJS.ErrnoException).code = 'ENOSPC'; - mkdirSyncSpy.mockImplementation(() => { + const mkdirSyncSpy = vi.spyOn(fs, 'mkdirSync').mockImplementation(() => { throw enospcError; }); @@ -542,6 +429,7 @@ describe('ChatRecordingService', () => { // Recording should be disabled (conversationFile set to null) expect(chatRecordingService.getConversationFilePath()).toBeNull(); + mkdirSyncSpy.mockRestore(); }); it('should disable recording and not throw when ENOSPC occurs during writeConversation', () => { @@ -550,17 +438,11 @@ describe('ChatRecordingService', () => { const enospcError = new Error('ENOSPC: no space left on device'); (enospcError as NodeJS.ErrnoException).code = 'ENOSPC'; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); - - writeFileSyncSpy.mockImplementation(() => { - throw enospcError; - }); + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementation(() => { + throw enospcError; + }); // Should not throw when recording a message expect(() => @@ -573,6 +455,7 @@ describe('ChatRecordingService', () => { // Recording should be disabled (conversationFile set to null) expect(chatRecordingService.getConversationFilePath()).toBeNull(); + writeFileSyncSpy.mockRestore(); }); it('should skip recording operations when recording is disabled', () => { @@ -581,18 +464,11 @@ describe('ChatRecordingService', () => { const enospcError = new Error('ENOSPC: no space left on device'); (enospcError as NodeJS.ErrnoException).code = 'ENOSPC'; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); - - // First call throws ENOSPC - writeFileSyncSpy.mockImplementationOnce(() => { - throw enospcError; - }); + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementationOnce(() => { + throw enospcError; + }); chatRecordingService.recordMessage({ type: 'user', @@ -619,6 +495,7 @@ describe('ChatRecordingService', () => { // writeFileSync should not have been called for any of these expect(writeFileSyncSpy).not.toHaveBeenCalled(); + writeFileSyncSpy.mockRestore(); }); it('should return null from getConversation when recording is disabled', () => { @@ -627,17 +504,11 @@ describe('ChatRecordingService', () => { const enospcError = new Error('ENOSPC: no space left on device'); (enospcError as NodeJS.ErrnoException).code = 'ENOSPC'; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); - - writeFileSyncSpy.mockImplementation(() => { - throw enospcError; - }); + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementation(() => { + throw enospcError; + }); // Trigger ENOSPC chatRecordingService.recordMessage({ @@ -649,6 +520,7 @@ describe('ChatRecordingService', () => { // getConversation should return null when disabled expect(chatRecordingService.getConversation()).toBeNull(); expect(chatRecordingService.getConversationFilePath()).toBeNull(); + writeFileSyncSpy.mockRestore(); }); it('should still throw for non-ENOSPC errors', () => { @@ -657,17 +529,11 @@ describe('ChatRecordingService', () => { const otherError = new Error('Permission denied'); (otherError as NodeJS.ErrnoException).code = 'EACCES'; - vi.spyOn(fs, 'readFileSync').mockReturnValue( - JSON.stringify({ - sessionId: 'test-session-id', - projectHash: 'test-project-hash', - messages: [], - }), - ); - - writeFileSyncSpy.mockImplementation(() => { - throw otherError; - }); + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementation(() => { + throw otherError; + }); // Should throw for non-ENOSPC errors expect(() => @@ -680,6 +546,7 @@ describe('ChatRecordingService', () => { // Recording should NOT be disabled for non-ENOSPC errors (file path still exists) expect(chatRecordingService.getConversationFilePath()).not.toBeNull(); + writeFileSyncSpy.mockRestore(); }); }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index e570923d546..47156c9edbf 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -540,12 +540,22 @@ export class ChatRecordingService { */ deleteSession(sessionId: string): void { try { - const chatsDir = path.join( - this.config.storage.getProjectTempDir(), - 'chats', - ); + const tempDir = this.config.storage.getProjectTempDir(); + const chatsDir = path.join(tempDir, 'chats'); const sessionPath = path.join(chatsDir, `${sessionId}.json`); - fs.unlinkSync(sessionPath); + if (fs.existsSync(sessionPath)) { + fs.unlinkSync(sessionPath); + } + + // Cleanup tool outputs for this session + const toolOutputDir = path.join( + tempDir, + 'tool-outputs', + `session-${sessionId}`, + ); + if (fs.existsSync(toolOutputDir)) { + fs.rmSync(toolOutputDir, { recursive: true, force: true }); + } } catch (error) { debugLogger.error('Error deleting session file.', error); throw error; diff --git a/packages/core/src/services/toolOutputMaskingService.test.ts b/packages/core/src/services/toolOutputMaskingService.test.ts index 05ab181f0d4..c42cb843061 100644 --- a/packages/core/src/services/toolOutputMaskingService.test.ts +++ b/packages/core/src/services/toolOutputMaskingService.test.ts @@ -4,7 +4,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import fs from 'node:fs'; +import path from 'node:path'; +import os from 'node:os'; import { ToolOutputMaskingService, MASKING_INDICATOR_TAG, @@ -18,24 +21,27 @@ vi.mock('../utils/tokenCalculation.js', () => ({ estimateTokenCountSync: vi.fn(), })); -vi.mock('node:fs/promises', () => ({ - mkdir: vi.fn().mockResolvedValue(undefined), - writeFile: vi.fn().mockResolvedValue(undefined), -})); - describe('ToolOutputMaskingService', () => { let service: ToolOutputMaskingService; let mockConfig: Config; + let testTempDir: string; const mockedEstimateTokenCountSync = vi.mocked(estimateTokenCountSync); - beforeEach(() => { + beforeEach(async () => { + testTempDir = await fs.promises.mkdtemp( + path.join(os.tmpdir(), 'tool-masking-test-'), + ); + service = new ToolOutputMaskingService(); mockConfig = { storage: { - getHistoryDir: () => '/mock/history', + getHistoryDir: () => path.join(testTempDir, 'history'), + getProjectTempDir: () => testTempDir, }, + getSessionId: () => 'mock-session', getUsageStatisticsEnabled: () => false, + getToolOutputMaskingEnabled: () => true, getToolOutputMaskingConfig: () => ({ enabled: true, toolProtectionThreshold: 50000, @@ -46,6 +52,13 @@ describe('ToolOutputMaskingService', () => { vi.clearAllMocks(); }); + afterEach(async () => { + vi.restoreAllMocks(); + if (testTempDir) { + await fs.promises.rm(testTempDir, { recursive: true, force: true }); + } + }); + it('should not mask if total tool tokens are below protection threshold', async () => { const history: Content[] = [ { @@ -451,6 +464,7 @@ describe('ToolOutputMaskingService', () => { // We replace the random part of the filename for deterministic snapshots // and normalize path separators for cross-platform compatibility const deterministicResponse = response + .replace(new RegExp(testTempDir.replace(/\\/g, '/'), 'g'), '/mock/temp') .replace( new RegExp(`${SHELL_TOOL_NAME}_[^\\s"]+\\.txt`, 'g'), `${SHELL_TOOL_NAME}_deterministic.txt`, diff --git a/packages/core/src/services/toolOutputMaskingService.ts b/packages/core/src/services/toolOutputMaskingService.ts index 76827da37aa..358859ebf6c 100644 --- a/packages/core/src/services/toolOutputMaskingService.ts +++ b/packages/core/src/services/toolOutputMaskingService.ts @@ -136,10 +136,14 @@ export class ToolOutputMaskingService { // Perform masking and offloading const newHistory = [...history]; // Shallow copy of history let actualTokensSaved = 0; - const toolOutputsDir = path.join( - config.storage.getHistoryDir(), + let toolOutputsDir = path.join( + config.storage.getProjectTempDir(), TOOL_OUTPUTS_DIR, ); + const sessionId = config.getSessionId(); + if (sessionId) { + toolOutputsDir = path.join(toolOutputsDir, `session-${sessionId}`); + } await fsPromises.mkdir(toolOutputsDir, { recursive: true }); for (const item of prunableParts) { diff --git a/packages/core/src/utils/fileUtils.test.ts b/packages/core/src/utils/fileUtils.test.ts index 742c782c7a6..6efa2911803 100644 --- a/packages/core/src/utils/fileUtils.test.ts +++ b/packages/core/src/utils/fileUtils.test.ts @@ -1121,7 +1121,7 @@ describe('fileUtils', () => { const expectedOutputFile = path.join( tempRootDir, - 'tool_output', + 'tool-outputs', 'shell_123.txt', ); expect(result.outputFile).toBe(expectedOutputFile); @@ -1149,7 +1149,7 @@ describe('fileUtils', () => { // ../../dangerous/tool -> ______dangerous_tool const expectedOutputFile = path.join( tempRootDir, - 'tool_output', + 'tool-outputs', '______dangerous_tool_1.txt', ); expect(result.outputFile).toBe(expectedOutputFile); @@ -1170,7 +1170,7 @@ describe('fileUtils', () => { // ../../etc/passwd -> ______etc_passwd const expectedOutputFile = path.join( tempRootDir, - 'tool_output', + 'tool-outputs', 'shell_______etc_passwd.txt', ); expect(result.outputFile).toBe(expectedOutputFile); diff --git a/packages/core/src/utils/fileUtils.ts b/packages/core/src/utils/fileUtils.ts index bb13fc01e00..6aae4238cbe 100644 --- a/packages/core/src/utils/fileUtils.ts +++ b/packages/core/src/utils/fileUtils.ts @@ -623,18 +623,23 @@ ${processedLines.join('\n')}`; /** * Saves tool output to a temporary file for later retrieval. */ -export const TOOL_OUTPUT_DIR = 'tool_output'; +export const TOOL_OUTPUTS_DIR = 'tool-outputs'; export async function saveTruncatedToolOutput( content: string, toolName: string, id: string | number, // Accept string (callId) or number (truncationId) projectTempDir: string, + sessionId?: string, ): Promise<{ outputFile: string; totalLines: number }> { const safeToolName = sanitizeFilenamePart(toolName).toLowerCase(); const safeId = sanitizeFilenamePart(id.toString()).toLowerCase(); const fileName = `${safeToolName}_${safeId}.txt`; - const toolOutputDir = path.join(projectTempDir, TOOL_OUTPUT_DIR); + + let toolOutputDir = path.join(projectTempDir, TOOL_OUTPUTS_DIR); + if (sessionId) { + toolOutputDir = path.join(toolOutputDir, `session-${sessionId}`); + } const outputFile = path.join(toolOutputDir, fileName); await fsPromises.mkdir(toolOutputDir, { recursive: true }); From 586a98f8e83eec3232eaedd58bf2ca77631043f6 Mon Sep 17 00:00:00 2001 From: Abhi Date: Thu, 5 Feb 2026 23:21:55 -0500 Subject: [PATCH 2/3] security(core,cli): sanitize sessionId and validate paths to prevent traversal vulnerabilities --- packages/cli/src/utils/sessionCleanup.ts | 13 ++++++- .../cli/src/utils/toolOutputCleanup.test.ts | 34 +++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 12 +++++-- .../src/services/toolOutputMaskingService.ts | 3 +- packages/core/src/utils/fileUtils.test.ts | 24 +++++++++++++ packages/core/src/utils/fileUtils.ts | 3 +- 6 files changed, 84 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/utils/sessionCleanup.ts b/packages/cli/src/utils/sessionCleanup.ts index 78160230353..d0988d7cd73 100644 --- a/packages/cli/src/utils/sessionCleanup.ts +++ b/packages/cli/src/utils/sessionCleanup.ts @@ -8,6 +8,7 @@ import * as fs from 'node:fs/promises'; import * as path from 'node:path'; import { debugLogger, + sanitizeFilenamePart, Storage, TOOL_OUTPUTS_DIR, type Config, @@ -103,10 +104,11 @@ export async function cleanupExpiredSessions( } // ALSO cleanup tool outputs for this session + const safeSessionId = sanitizeFilenamePart(sessionId); const toolOutputDir = path.join( config.storage.getProjectTempDir(), TOOL_OUTPUTS_DIR, - `session-${sessionId}`, + `session-${safeSessionId}`, ); try { await fs.rm(toolOutputDir, { recursive: true, force: true }); @@ -450,6 +452,15 @@ export async function cleanupToolOutputFiles( ); for (const subdir of subdirs) { try { + // Security: Validate that the subdirectory name is a safe filename part + // and doesn't attempt path traversal. + if (subdir.name !== sanitizeFilenamePart(subdir.name)) { + debugLogger.debug( + `Skipping unsafe tool-output subdirectory: ${subdir.name}`, + ); + continue; + } + const subdirPath = path.join(toolOutputDir, subdir.name); const stat = await fs.stat(subdirPath); diff --git a/packages/cli/src/utils/toolOutputCleanup.test.ts b/packages/cli/src/utils/toolOutputCleanup.test.ts index 2356da58443..18e43ab6d05 100644 --- a/packages/cli/src/utils/toolOutputCleanup.test.ts +++ b/packages/cli/src/utils/toolOutputCleanup.test.ts @@ -316,5 +316,39 @@ describe('Tool Output Cleanup', () => { expect(remainingDirs).toContain('session-recent'); expect(remainingDirs).not.toContain('session-old'); }); + + it('should skip subdirectories with path traversal characters', async () => { + const settings: Settings = { + general: { + sessionRetention: { + enabled: true, + maxAge: '1d', + }, + }, + }; + + const toolOutputDir = path.join(testTempDir, TOOL_OUTPUTS_DIR); + await fs.mkdir(toolOutputDir, { recursive: true }); + + // Create an unsafe directory name + const unsafeDir = path.join(toolOutputDir, 'session-.._.._danger'); + await fs.mkdir(unsafeDir, { recursive: true }); + + const debugSpy = vi + .spyOn(debugLogger, 'debug') + .mockImplementation(() => {}); + + await cleanupToolOutputFiles(settings, false, testTempDir); + + expect(debugSpy).toHaveBeenCalledWith( + expect.stringContaining('Skipping unsafe tool-output subdirectory'), + ); + + // Directory should still exist (it was skipped, not deleted) + const entries = await fs.readdir(toolOutputDir); + expect(entries).toContain('session-.._.._danger'); + + debugSpy.mockRestore(); + }); }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 47156c9edbf..6a57e2801bf 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -8,6 +8,7 @@ import { type Config } from '../config/config.js'; import { type Status } from '../core/coreToolScheduler.js'; import { type ThoughtSummary } from '../utils/thoughtUtils.js'; import { getProjectHash } from '../utils/paths.js'; +import { sanitizeFilenamePart } from '../utils/fileUtils.js'; import path from 'node:path'; import fs from 'node:fs'; import { randomUUID } from 'node:crypto'; @@ -548,12 +549,19 @@ export class ChatRecordingService { } // Cleanup tool outputs for this session + const safeSessionId = sanitizeFilenamePart(sessionId); const toolOutputDir = path.join( tempDir, 'tool-outputs', - `session-${sessionId}`, + `session-${safeSessionId}`, ); - if (fs.existsSync(toolOutputDir)) { + + // Robustness: Ensure the path is strictly within the tool-outputs base + const toolOutputsBase = path.join(tempDir, 'tool-outputs'); + if ( + fs.existsSync(toolOutputDir) && + toolOutputDir.startsWith(toolOutputsBase) + ) { fs.rmSync(toolOutputDir, { recursive: true, force: true }); } } catch (error) { diff --git a/packages/core/src/services/toolOutputMaskingService.ts b/packages/core/src/services/toolOutputMaskingService.ts index 358859ebf6c..d62e1761e11 100644 --- a/packages/core/src/services/toolOutputMaskingService.ts +++ b/packages/core/src/services/toolOutputMaskingService.ts @@ -142,7 +142,8 @@ export class ToolOutputMaskingService { ); const sessionId = config.getSessionId(); if (sessionId) { - toolOutputsDir = path.join(toolOutputsDir, `session-${sessionId}`); + const safeSessionId = sanitizeFilenamePart(sessionId); + toolOutputsDir = path.join(toolOutputsDir, `session-${safeSessionId}`); } await fsPromises.mkdir(toolOutputsDir, { recursive: true }); diff --git a/packages/core/src/utils/fileUtils.test.ts b/packages/core/src/utils/fileUtils.test.ts index 6efa2911803..95b10ced695 100644 --- a/packages/core/src/utils/fileUtils.test.ts +++ b/packages/core/src/utils/fileUtils.test.ts @@ -1176,6 +1176,30 @@ describe('fileUtils', () => { expect(result.outputFile).toBe(expectedOutputFile); }); + it('should sanitize sessionId in filename/path', async () => { + const content = 'content'; + const toolName = 'shell'; + const id = '1'; + const sessionId = '../../etc/passwd'; + + const result = await saveTruncatedToolOutput( + content, + toolName, + id, + tempRootDir, + sessionId, + ); + + // ../../etc/passwd -> ______etc_passwd + const expectedOutputFile = path.join( + tempRootDir, + 'tool-outputs', + 'session-______etc_passwd', + 'shell_1.txt', + ); + expect(result.outputFile).toBe(expectedOutputFile); + }); + it('should format multi-line output correctly', () => { const lines = Array.from({ length: 50 }, (_, i) => `line ${i}`); const content = lines.join('\n'); diff --git a/packages/core/src/utils/fileUtils.ts b/packages/core/src/utils/fileUtils.ts index 6aae4238cbe..bac694d6d98 100644 --- a/packages/core/src/utils/fileUtils.ts +++ b/packages/core/src/utils/fileUtils.ts @@ -638,7 +638,8 @@ export async function saveTruncatedToolOutput( let toolOutputDir = path.join(projectTempDir, TOOL_OUTPUTS_DIR); if (sessionId) { - toolOutputDir = path.join(toolOutputDir, `session-${sessionId}`); + const safeSessionId = sanitizeFilenamePart(sessionId); + toolOutputDir = path.join(toolOutputDir, `session-${safeSessionId}`); } const outputFile = path.join(toolOutputDir, fileName); From c5d36582f02d07f2cd9aa394b3ea4dcc0ac25b91 Mon Sep 17 00:00:00 2001 From: Abhi Date: Thu, 5 Feb 2026 23:38:06 -0500 Subject: [PATCH 3/3] address sandy's feedback --- packages/core/src/scheduler/tool-executor.test.ts | 1 + packages/core/src/scheduler/tool-executor.ts | 1 + packages/core/src/services/chatRecordingService.test.ts | 1 - packages/core/src/services/toolOutputMaskingService.test.ts | 6 +++--- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index 13723ee37d0..2470a39dcd6 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -221,6 +221,7 @@ describe('ToolExecutor', () => { SHELL_TOOL_NAME, 'call-trunc', expect.any(String), // temp dir + 'test-session-id', // session id from makeFakeConfig ); expect(fileUtils.formatTruncatedToolOutput).toHaveBeenCalledWith( diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 8b31c8166f2..ec02d25953f 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -221,6 +221,7 @@ export class ToolExecutor { toolName, callId, this.config.storage.getProjectTempDir(), + this.config.getSessionId(), ); outputFile = savedPath; content = formatTruncatedToolOutput(content, outputFile, lines); diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 4775be7281b..e8b879e10ce 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -8,7 +8,6 @@ import { expect, it, describe, vi, beforeEach, afterEach } from 'vitest'; import fs from 'node:fs'; import path from 'node:path'; import os from 'node:os'; -// import { randomUUID } from 'node:crypto'; // Removed as it's mocked factory-style import type { ConversationRecord, ToolCallRecord, diff --git a/packages/core/src/services/toolOutputMaskingService.test.ts b/packages/core/src/services/toolOutputMaskingService.test.ts index c42cb843061..26e44c4d17c 100644 --- a/packages/core/src/services/toolOutputMaskingService.test.ts +++ b/packages/core/src/services/toolOutputMaskingService.test.ts @@ -463,13 +463,13 @@ describe('ToolOutputMaskingService', () => { // We replace the random part of the filename for deterministic snapshots // and normalize path separators for cross-platform compatibility - const deterministicResponse = response + const normalizedResponse = response.replace(/\\/g, '/'); + const deterministicResponse = normalizedResponse .replace(new RegExp(testTempDir.replace(/\\/g, '/'), 'g'), '/mock/temp') .replace( new RegExp(`${SHELL_TOOL_NAME}_[^\\s"]+\\.txt`, 'g'), `${SHELL_TOOL_NAME}_deterministic.txt`, - ) - .replace(/\\/g, '/'); + ); expect(deterministicResponse).toMatchSnapshot(); });