MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-15 00:21:09 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(core,cli): implement session-linked tool output storage and cleanup (#18416)

Browse Source
This commit is contained in:
Abhi
2026-02-06 01:36:42 -05:00
committed by GitHub
parent 8ec176e005
commit 30354892b3
12 changed files with 442 additions and 386 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
packages/cli/src/utils/sessionCleanup.ts
View File

@@ -8,8 +8,9 @@ import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import {
debugLogger,
sanitizeFilenamePart,
Storage,
TOOL_OUTPUT_DIR,
TOOL_OUTPUTS_DIR,
type Config,
} from '@google/gemini-cli-core';
import type { Settings, SessionRetentionSettings } from '../config/settings.js';
@@ -101,6 +102,19 @@ export async function cleanupExpiredSessions(
} catch {
/* ignore if log doesn't exist */
}
// ALSO cleanup tool outputs for this session
const safeSessionId = sanitizeFilenamePart(sessionId);
const toolOutputDir = path.join(
config.storage.getProjectTempDir(),
TOOL_OUTPUTS_DIR,
`session-${safeSessionId}`,
);
try {
await fs.rm(toolOutputDir, { recursive: true, force: true });
} catch {
/* ignore if doesn't exist */
}
}
if (config.getDebugMode()) {
@@ -350,7 +364,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 +374,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 +445,43 @@ 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 {
// 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);
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 {

95
packages/cli/src/utils/toolOutputCleanup.test.ts
View File

@@ -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,74 @@ 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');
});
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();
});
});
});

1
packages/core/src/scheduler/tool-executor.test.ts
View File

@@ -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(

1
packages/core/src/scheduler/tool-executor.ts
View File

@@ -221,6 +221,7 @@ export class ToolExecutor {
toolName,
callId,
this.config.storage.getProjectTempDir(),
this.config.getSessionId(),
);
outputFile = savedPath;
content = formatTruncatedToolOutput(content, outputFile, lines);

2
packages/core/src/services/__snapshots__/toolOutputMaskingService.test.ts.snap
View File

@@ -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
</tool_output_masked>"
`;

4
packages/core/src/services/chatCompressionService.test.ts
View File

@@ -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/);

548
packages/core/src/services/chatRecordingService.test.ts
View File

@@ -4,46 +4,47 @@
* 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 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<typeof fs.mkdirSync>;
let writeFileSyncSpy: MockInstance<typeof fs.writeFileSync>;
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 +58,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 +132,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 +158,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 +167,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 +180,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 +236,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 +251,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 +278,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 +324,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 +345,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 +364,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 result = chatRecordingService.rewindTo('2');
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,
const sessionFile = chatRecordingService.getConversationFilePath()!;
let conversation = JSON.parse(
fs.readFileSync(sessionFile, 'utf8'),
) as ConversationRecord;
expect(savedConversation.messages).toHaveLength(1);
const secondMsgId = conversation.messages[1].id;
const result = chatRecordingService.rewindTo(secondMsgId);
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(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 +419,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 +428,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 +437,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 +454,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 +463,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 +494,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 +503,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 +519,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 +528,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 +545,7 @@ describe('ChatRecordingService', () => {
// Recording should NOT be disabled for non-ENOSPC errors (file path still exists)
expect(chatRecordingService.getConversationFilePath()).not.toBeNull();
writeFileSyncSpy.mockRestore();
});
});
});

28
packages/core/src/services/chatRecordingService.ts
View File

@@ -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';
@@ -540,12 +541,29 @@ 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 safeSessionId = sanitizeFilenamePart(sessionId);
const toolOutputDir = path.join(
tempDir,
'tool-outputs',
`session-${safeSessionId}`,
);
// 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) {
debugLogger.error('Error deleting session file.', error);
throw error;

36
packages/core/src/services/toolOutputMaskingService.test.ts
View File

@@ -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[] = [
{
@@ -450,12 +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();
});

9
packages/core/src/services/toolOutputMaskingService.ts
View File

@@ -136,10 +136,15 @@ 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) {
const safeSessionId = sanitizeFilenamePart(sessionId);
toolOutputsDir = path.join(toolOutputsDir, `session-${safeSessionId}`);
}
await fsPromises.mkdir(toolOutputsDir, { recursive: true });
for (const item of prunableParts) {

30
packages/core/src/utils/fileUtils.test.ts
View File

@@ -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,12 +1170,36 @@ describe('fileUtils', () => {
// ../../etc/passwd -> ______etc_passwd
const expectedOutputFile = path.join(
tempRootDir,
'tool_output',
'tool-outputs',
'shell_______etc_passwd.txt',
);
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');

10
packages/core/src/utils/fileUtils.ts
View File

@@ -623,18 +623,24 @@ ${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) {
const safeSessionId = sanitizeFilenamePart(sessionId);
toolOutputDir = path.join(toolOutputDir, `session-${safeSessionId}`);
}
const outputFile = path.join(toolOutputDir, fileName);
await fsPromises.mkdir(toolOutputDir, { recursive: true });