From 104587bae8326e7ad1e2d739b7e14c7665e10b2b Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Thu, 26 Mar 2026 23:43:39 -0400 Subject: [PATCH] feat(core): subagent isolation and cleanup hardening (#23903) --- packages/cli/src/ui/AppContainer.tsx | 2 +- .../cli/src/ui/hooks/useSessionBrowser.ts | 4 +- packages/cli/src/utils/sessionCleanup.test.ts | 3 + packages/cli/src/utils/sessionCleanup.ts | 44 +----- packages/cli/src/utils/sessions.ts | 2 +- .../core/src/agents/local-executor.test.ts | 112 +++++++++++-- packages/core/src/agents/local-executor.ts | 14 +- .../core/src/config/agent-loop-context.ts | 3 + packages/core/src/index.ts | 1 + .../src/services/chatRecordingService.test.ts | 89 +++++++++-- .../core/src/services/chatRecordingService.ts | 109 ++++++------- .../core/src/utils/sessionOperations.test.ts | 148 ++++++++++++++++++ packages/core/src/utils/sessionOperations.ts | 122 +++++++++++++++ 13 files changed, 520 insertions(+), 133 deletions(-) create mode 100644 packages/core/src/utils/sessionOperations.test.ts create mode 100644 packages/core/src/utils/sessionOperations.ts diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index d58ed45d89..d5b34915bc 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -726,7 +726,7 @@ export const AppContainer = (props: AppContainerProps) => { // Wrap handleDeleteSession to return a Promise for UIActions interface const handleDeleteSession = useCallback( async (session: SessionInfo): Promise => { - handleDeleteSessionSync(session); + await handleDeleteSessionSync(session); }, [handleDeleteSessionSync], ); diff --git a/packages/cli/src/ui/hooks/useSessionBrowser.ts b/packages/cli/src/ui/hooks/useSessionBrowser.ts index 9a34f68e0b..4e86c2d92e 100644 --- a/packages/cli/src/ui/hooks/useSessionBrowser.ts +++ b/packages/cli/src/ui/hooks/useSessionBrowser.ts @@ -98,7 +98,7 @@ export const useSessionBrowser = ( * Deletes a session by ID using the ChatRecordingService. */ handleDeleteSession: useCallback( - (session: SessionInfo) => { + async (session: SessionInfo) => { // Note: Chat sessions are stored on disk using a filename derived from // the session, e.g. "session--.json". // The ChatRecordingService.deleteSession API expects this file basename @@ -108,7 +108,7 @@ export const useSessionBrowser = ( .getGeminiClient() ?.getChatRecordingService(); if (chatRecordingService) { - chatRecordingService.deleteSession(session.file); + await chatRecordingService.deleteSession(session.file); } } catch (error) { coreEvents.emitFeedback('error', 'Error deleting session:', error); diff --git a/packages/cli/src/utils/sessionCleanup.test.ts b/packages/cli/src/utils/sessionCleanup.test.ts index b014159e08..eddf4c3460 100644 --- a/packages/cli/src/utils/sessionCleanup.test.ts +++ b/packages/cli/src/utils/sessionCleanup.test.ts @@ -106,6 +106,8 @@ describe('Session Cleanup (Refactored)', () => { ); // Session directory await fs.mkdir(path.join(testTempDir, sessionId), { recursive: true }); + // Subagent chats directory + await fs.mkdir(path.join(chatsDir, sessionId), { recursive: true }); } async function seedSessions() { @@ -274,6 +276,7 @@ describe('Session Cleanup (Refactored)', () => { existsSync(path.join(toolOutputsDir, `session-${sessions[1].id}`)), ).toBe(false); expect(existsSync(path.join(testTempDir, sessions[1].id))).toBe(false); // Session directory should be deleted + expect(existsSync(path.join(chatsDir, sessions[1].id))).toBe(false); // Subagent chats directory should be deleted }); it('should NOT delete sessions within the cutoff date', async () => { diff --git a/packages/cli/src/utils/sessionCleanup.ts b/packages/cli/src/utils/sessionCleanup.ts index 5ed4547604..dde926674c 100644 --- a/packages/cli/src/utils/sessionCleanup.ts +++ b/packages/cli/src/utils/sessionCleanup.ts @@ -13,6 +13,8 @@ import { Storage, TOOL_OUTPUTS_DIR, type Config, + deleteSessionArtifactsAsync, + deleteSubagentSessionDirAndArtifactsAsync, } from '@google/gemini-cli-core'; import type { Settings, SessionRetentionSettings } from '../config/settings.js'; import { getAllSessionFiles, type SessionFileEntry } from './sessionUtils.js'; @@ -59,48 +61,18 @@ function deriveShortIdFromFileName(fileName: string): string | null { return null; } -/** - * Gets the log path for a session ID. - */ -function getSessionLogPath(tempDir: string, safeSessionId: string): string { - return path.join(tempDir, 'logs', `session-${safeSessionId}.jsonl`); -} - /** * Cleans up associated artifacts (logs, tool-outputs, directory) for a session. */ -async function deleteSessionArtifactsAsync( +async function cleanupSessionAndSubagentsAsync( sessionId: string, config: Config, ): Promise { const tempDir = config.storage.getProjectTempDir(); + const chatsDir = path.join(tempDir, 'chats'); - // Cleanup logs - const logsDir = path.join(tempDir, 'logs'); - const safeSessionId = sanitizeFilenamePart(sessionId); - const logPath = getSessionLogPath(tempDir, safeSessionId); - if (logPath.startsWith(logsDir)) { - await fs.unlink(logPath).catch(() => {}); - } - - // Cleanup tool outputs - const toolOutputDir = path.join( - tempDir, - TOOL_OUTPUTS_DIR, - `session-${safeSessionId}`, - ); - const toolOutputsBase = path.join(tempDir, TOOL_OUTPUTS_DIR); - if (toolOutputDir.startsWith(toolOutputsBase)) { - await fs - .rm(toolOutputDir, { recursive: true, force: true }) - .catch(() => {}); - } - - // Cleanup session directory - const sessionDir = path.join(tempDir, safeSessionId); - if (safeSessionId && sessionDir.startsWith(tempDir + path.sep)) { - await fs.rm(sessionDir, { recursive: true, force: true }).catch(() => {}); - } + await deleteSessionArtifactsAsync(sessionId, tempDir); + await deleteSubagentSessionDirAndArtifactsAsync(sessionId, chatsDir, tempDir); } /** @@ -201,7 +173,7 @@ export async function cleanupExpiredSessions( await fs.unlink(filePath); if (fullSessionId) { - await deleteSessionArtifactsAsync(fullSessionId, config); + await cleanupSessionAndSubagentsAsync(fullSessionId, config); } result.deleted++; } else { @@ -230,7 +202,7 @@ export async function cleanupExpiredSessions( const sessionId = sessionToDelete.sessionInfo?.id; if (sessionId) { - await deleteSessionArtifactsAsync(sessionId, config); + await cleanupSessionAndSubagentsAsync(sessionId, config); } if (config.getDebugMode()) { diff --git a/packages/cli/src/utils/sessions.ts b/packages/cli/src/utils/sessions.ts index 56f9f06a6a..9a4def4995 100644 --- a/packages/cli/src/utils/sessions.ts +++ b/packages/cli/src/utils/sessions.ts @@ -97,7 +97,7 @@ export async function deleteSession( try { // Use ChatRecordingService to delete the session const chatRecordingService = new ChatRecordingService(config); - chatRecordingService.deleteSession(sessionToDelete.file); + await chatRecordingService.deleteSession(sessionToDelete.file); const time = formatRelativeTime(sessionToDelete.lastUpdated); writeToStdout( diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index fb21e1093d..32499bbaf1 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -69,6 +69,10 @@ import { type FunctionDeclaration, } from '@google/genai'; import type { Config } from '../config/config.js'; +import type { AgentLoopContext } from '../config/agent-loop-context.js'; +import type { GeminiClient } from '../core/client.js'; +import type { SandboxManager } from '../services/sandboxManager.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { MockTool } from '../test-utils/mock-tool.js'; import { getDirectoryContextString } from '../utils/environmentContext.js'; import { z } from 'zod'; @@ -377,10 +381,8 @@ describe('LocalAgentExecutor', () => { describe('create (Initialization and Validation)', () => { it('should explicitly map execution context properties to prevent unintended propagation', async () => { const definition = createTestDefinition([LS_TOOL_NAME]); - const mockGeminiClient = - {} as unknown as import('../core/client.js').GeminiClient; - const mockSandboxManager = - {} as unknown as import('../services/sandboxManager.js').SandboxManager; + const mockGeminiClient = {} as unknown as GeminiClient; + const mockSandboxManager = {} as unknown as SandboxManager; const extendedContext = { config: mockConfig, promptId: mockConfig.promptId, @@ -391,7 +393,7 @@ describe('LocalAgentExecutor', () => { geminiClient: mockGeminiClient, sandboxManager: mockSandboxManager, unintendedProperty: 'should not be here', - } as unknown as import('../config/agent-loop-context.js').AgentLoopContext; + } as unknown as AgentLoopContext; const executor = await LocalAgentExecutor.create( definition, @@ -414,7 +416,7 @@ describe('LocalAgentExecutor', () => { expect(executionContext).toBeDefined(); expect(executionContext.config).toBe(extendedContext.config); - expect(executionContext.promptId).toBe(extendedContext.promptId); + expect(executionContext.promptId).toBeDefined(); expect(executionContext.geminiClient).toBe(extendedContext.geminiClient); expect(executionContext.sandboxManager).toBe( extendedContext.sandboxManager, @@ -445,7 +447,99 @@ describe('LocalAgentExecutor', () => { expect(executionContext.messageBus).not.toBe(extendedContext.messageBus); }); - it('should create successfully with allowed tools', async () => { + it('should propagate parentSessionId from context when creating executionContext', async () => { + const parentSessionId = 'top-level-session-id'; + const currentPromptId = 'subagent-a-id'; + const mockGeminiClient = {} as unknown as GeminiClient; + const mockSandboxManager = {} as unknown as SandboxManager; + const mockMessageBus = { + derive: () => ({}), + } as unknown as MessageBus; + const mockToolRegistry = { + getMessageBus: () => mockMessageBus, + getAllToolNames: () => [], + sortTools: () => {}, + } as unknown as ToolRegistry; + + const context = { + config: mockConfig, + promptId: currentPromptId, + parentSessionId, + toolRegistry: mockToolRegistry, + promptRegistry: {} as unknown as PromptRegistry, + resourceRegistry: {} as unknown as ResourceRegistry, + geminiClient: mockGeminiClient, + sandboxManager: mockSandboxManager, + messageBus: mockMessageBus, + } as unknown as AgentLoopContext; + + const definition = createTestDefinition([]); + const executor = await LocalAgentExecutor.create(definition, context); + + mockModelResponse([ + { + name: TASK_COMPLETE_TOOL_NAME, + args: { finalResult: 'done' }, + id: 'call1', + }, + ]); + + await executor.run({ goal: 'test' }, signal); + + const chatConstructorArgs = + MockedGeminiChat.mock.calls[MockedGeminiChat.mock.calls.length - 1]; + const executionContext = chatConstructorArgs[0]; + + expect(executionContext.parentSessionId).toBe(parentSessionId); + expect(executionContext.promptId).toBe(executor['agentId']); + }); + + it('should fall back to promptId if parentSessionId is missing (top-level subagent)', async () => { + const rootSessionId = 'root-session-id'; + const mockGeminiClient = {} as unknown as GeminiClient; + const mockSandboxManager = {} as unknown as SandboxManager; + const mockMessageBus = { + derive: () => ({}), + } as unknown as MessageBus; + const mockToolRegistry = { + getMessageBus: () => mockMessageBus, + getAllToolNames: () => [], + sortTools: () => {}, + } as unknown as ToolRegistry; + + const context = { + config: mockConfig, + promptId: rootSessionId, + // parentSessionId is undefined + toolRegistry: mockToolRegistry, + promptRegistry: {} as unknown as PromptRegistry, + resourceRegistry: {} as unknown as ResourceRegistry, + geminiClient: mockGeminiClient, + sandboxManager: mockSandboxManager, + messageBus: mockMessageBus, + } as unknown as AgentLoopContext; + + const definition = createTestDefinition([]); + const executor = await LocalAgentExecutor.create(definition, context); + + mockModelResponse([ + { + name: TASK_COMPLETE_TOOL_NAME, + args: { finalResult: 'done' }, + id: 'call1', + }, + ]); + + await executor.run({ goal: 'test' }, signal); + + const chatConstructorArgs = + MockedGeminiChat.mock.calls[MockedGeminiChat.mock.calls.length - 1]; + const executionContext = chatConstructorArgs[0]; + + expect(executionContext.parentSessionId).toBe(rootSessionId); + expect(executionContext.promptId).toBe(executor['agentId']); + }); + it('should successfully with allowed tools', async () => { const definition = createTestDefinition([LS_TOOL_NAME]); const executor = await LocalAgentExecutor.create( definition, @@ -500,9 +594,7 @@ describe('LocalAgentExecutor', () => { onActivity, ); - expect(executor['agentId']).toMatch( - new RegExp(`^${parentId}-${definition.name}-`), - ); + expect(executor['agentId']).toBeDefined(); }); it('should correctly apply templates to initialMessages', async () => { diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 2a47036486..c9e4341f03 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -121,7 +121,8 @@ export class LocalAgentExecutor { private get executionContext(): AgentLoopContext { return { config: this.context.config, - promptId: this.context.promptId, + promptId: this.agentId, + parentSessionId: this.context.parentSessionId || this.context.promptId, // Always preserve the main agent session ID geminiClient: this.context.geminiClient, sandboxManager: this.context.sandboxManager, toolRegistry: this.toolRegistry, @@ -255,9 +256,6 @@ export class LocalAgentExecutor { agentToolRegistry.sortTools(); - // Get the parent prompt ID from context - const parentPromptId = context.promptId; - // Get the parent tool call ID from context const toolContext = getToolCallContext(); const parentCallId = toolContext?.callId; @@ -265,7 +263,6 @@ export class LocalAgentExecutor { return new LocalAgentExecutor( definition, context, - parentPromptId, agentToolRegistry, agentPromptRegistry, agentResourceRegistry, @@ -283,7 +280,6 @@ export class LocalAgentExecutor { private constructor( definition: LocalAgentDefinition, context: AgentLoopContext, - parentPromptId: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, resourceRegistry: ResourceRegistry, @@ -299,11 +295,7 @@ export class LocalAgentExecutor { this.compressionService = new ChatCompressionService(); this.parentCallId = parentCallId; - const randomIdPart = Math.random().toString(36).slice(2, 8); - // parentPromptId will be undefined if this agent is invoked directly - // (top-level), rather than as a sub-agent. - const parentPrefix = parentPromptId ? `${parentPromptId}-` : ''; - this.agentId = `${parentPrefix}${this.definition.name}-${randomIdPart}`; + this.agentId = Math.random().toString(36).slice(2, 8); } /** diff --git a/packages/core/src/config/agent-loop-context.ts b/packages/core/src/config/agent-loop-context.ts index b16326a7ce..7325fc0b73 100644 --- a/packages/core/src/config/agent-loop-context.ts +++ b/packages/core/src/config/agent-loop-context.ts @@ -23,6 +23,9 @@ export interface AgentLoopContext { /** The unique ID for the current user turn or agent thought loop. */ readonly promptId: string; + /** The unique ID for the parent session if this is a subagent. */ + readonly parentSessionId?: string; + /** The registry of tools available to the agent in this context. */ readonly toolRegistry: ToolRegistry; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 9b98a1bbe2..0edb8b3462 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -84,6 +84,7 @@ export * from './utils/authConsent.js'; export * from './utils/googleQuotaErrors.js'; export * from './utils/googleErrors.js'; export * from './utils/fileUtils.js'; +export * from './utils/sessionOperations.js'; export * from './utils/planUtils.js'; export * from './utils/approvalModeUtils.js'; export * from './utils/fileDiffUtils.js'; diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 6b395b92e0..b84f387e1f 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -108,6 +108,30 @@ describe('ChatRecordingService', () => { expect(conversation.kind).toBe('subagent'); }); + it('should create a subdirectory for subagents if parentSessionId is present', () => { + const parentSessionId = 'test-parent-uuid'; + Object.defineProperty(mockConfig, 'parentSessionId', { + value: parentSessionId, + writable: true, + configurable: true, + }); + + chatRecordingService.initialize(undefined, 'subagent'); + chatRecordingService.recordMessage({ + type: 'user', + content: 'ping', + model: 'm', + }); + + const chatsDir = path.join(testTempDir, 'chats'); + const subagentDir = path.join(chatsDir, parentSessionId); + expect(fs.existsSync(subagentDir)).toBe(true); + + const files = fs.readdirSync(subagentDir); + expect(files.length).toBeGreaterThan(0); + expect(files[0]).toBe('test-session-id.json'); + }); + it('should resume from an existing session if provided', () => { const chatsDir = path.join(testTempDir, 'chats'); fs.mkdirSync(chatsDir, { recursive: true }); @@ -437,7 +461,7 @@ describe('ChatRecordingService', () => { }); describe('deleteSession', () => { - it('should delete the session file, tool outputs, session directory, and logs if they exist', () => { + it('should delete the session file, tool outputs, session directory, and logs if they exist', async () => { const sessionId = 'test-session-id'; const shortId = '12345678'; const chatsDir = path.join(testTempDir, 'chats'); @@ -464,7 +488,7 @@ describe('ChatRecordingService', () => { fs.mkdirSync(toolOutputDir, { recursive: true }); // Call with shortId - chatRecordingService.deleteSession(shortId); + await chatRecordingService.deleteSession(shortId); expect(fs.existsSync(sessionFile)).toBe(false); expect(fs.existsSync(logFile)).toBe(false); @@ -472,7 +496,7 @@ describe('ChatRecordingService', () => { expect(fs.existsSync(sessionDir)).toBe(false); }); - it('should delete subagent files and their logs when parent is deleted', () => { + it('should delete subagent files and their logs when parent is deleted', async () => { const parentSessionId = '12345678-session-id'; const shortId = '12345678'; const subagentSessionId = 'subagent-session-id'; @@ -494,11 +518,10 @@ describe('ChatRecordingService', () => { JSON.stringify({ sessionId: parentSessionId }), ); - // Create subagent session file - const subagentFile = path.join( - chatsDir, - `session-2023-01-01T00-01-${shortId}.json`, - ); + // Create subagent session file in subdirectory + const subagentDir = path.join(chatsDir, parentSessionId); + fs.mkdirSync(subagentDir, { recursive: true }); + const subagentFile = path.join(subagentDir, `${subagentSessionId}.json`); fs.writeFileSync( subagentFile, JSON.stringify({ sessionId: subagentSessionId, kind: 'subagent' }), @@ -526,17 +549,55 @@ describe('ChatRecordingService', () => { fs.mkdirSync(subagentToolOutputDir, { recursive: true }); // Call with parent sessionId - chatRecordingService.deleteSession(parentSessionId); + await chatRecordingService.deleteSession(parentSessionId); expect(fs.existsSync(parentFile)).toBe(false); expect(fs.existsSync(subagentFile)).toBe(false); + expect(fs.existsSync(subagentDir)).toBe(false); // Subagent directory should be deleted expect(fs.existsSync(parentLog)).toBe(false); expect(fs.existsSync(subagentLog)).toBe(false); expect(fs.existsSync(parentToolOutputDir)).toBe(false); expect(fs.existsSync(subagentToolOutputDir)).toBe(false); }); - it('should delete by basename', () => { + it('should delete subagent files and their logs when parent is deleted (legacy flat structure)', async () => { + const parentSessionId = '12345678-session-id'; + const shortId = '12345678'; + const subagentSessionId = 'subagent-session-id'; + const chatsDir = path.join(testTempDir, 'chats'); + const logsDir = path.join(testTempDir, 'logs'); + + fs.mkdirSync(chatsDir, { recursive: true }); + fs.mkdirSync(logsDir, { recursive: true }); + + // Create parent session file + const parentFile = path.join( + chatsDir, + `session-2023-01-01T00-00-${shortId}.json`, + ); + fs.writeFileSync( + parentFile, + JSON.stringify({ sessionId: parentSessionId }), + ); + + // Create legacy subagent session file (flat in chatsDir) + const subagentFile = path.join( + chatsDir, + `session-2023-01-01T00-01-${shortId}.json`, + ); + fs.writeFileSync( + subagentFile, + JSON.stringify({ sessionId: subagentSessionId, kind: 'subagent' }), + ); + + // Call with parent sessionId + await chatRecordingService.deleteSession(parentSessionId); + + expect(fs.existsSync(parentFile)).toBe(false); + expect(fs.existsSync(subagentFile)).toBe(false); + }); + + it('should delete by basename', async () => { const sessionId = 'test-session-id'; const shortId = '12345678'; const chatsDir = path.join(testTempDir, 'chats'); @@ -553,16 +614,16 @@ describe('ChatRecordingService', () => { fs.writeFileSync(logFile, '{}'); // Call with basename - chatRecordingService.deleteSession(basename); + await chatRecordingService.deleteSession(basename); expect(fs.existsSync(sessionFile)).toBe(false); expect(fs.existsSync(logFile)).toBe(false); }); - it('should not throw if session file does not exist', () => { - expect(() => + it('should not throw if session file does not exist', async () => { + await expect( chatRecordingService.deleteSession('non-existent'), - ).not.toThrow(); + ).resolves.not.toThrow(); }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index a161b7da80..f4aea75fd0 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -7,9 +7,13 @@ import { type Status } from '../scheduler/types.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 { sanitizeFilenamePart } from '../utils/fileUtils.js'; +import { + deleteSessionArtifactsAsync, + deleteSubagentSessionDirAndArtifactsAsync, +} from '../utils/sessionOperations.js'; import { randomUUID } from 'node:crypto'; import type { Content, @@ -172,20 +176,46 @@ export class ChatRecordingService { } else { // Create new session this.sessionId = this.context.promptId; - const chatsDir = path.join( + let chatsDir = path.join( this.context.config.storage.getProjectTempDir(), 'chats', ); + + // subagents are nested under the complete parent session id + if (this.kind === 'subagent' && this.context.parentSessionId) { + const safeParentId = sanitizeFilenamePart( + this.context.parentSessionId, + ); + if (!safeParentId) { + throw new Error( + `Invalid parentSessionId after sanitization: ${this.context.parentSessionId}`, + ); + } + chatsDir = path.join(chatsDir, safeParentId); + } + fs.mkdirSync(chatsDir, { recursive: true }); const timestamp = new Date() .toISOString() .slice(0, 16) .replace(/:/g, '-'); - const filename = `${SESSION_FILE_PREFIX}${timestamp}-${this.sessionId.slice( - 0, - 8, - )}.json`; + const safeSessionId = sanitizeFilenamePart(this.sessionId); + if (!safeSessionId) { + throw new Error( + `Invalid sessionId after sanitization: ${this.sessionId}`, + ); + } + + let filename: string; + if (this.kind === 'subagent') { + filename = `${safeSessionId}.json`; + } else { + filename = `${SESSION_FILE_PREFIX}${timestamp}-${safeSessionId.slice( + 0, + 8, + )}.json`; + } this.conversationFile = path.join(chatsDir, filename); this.writeConversation({ @@ -596,21 +626,22 @@ export class ChatRecordingService { * * @throws {Error} If shortId validation fails. */ - deleteSession(sessionIdOrBasename: string): void { + async deleteSession(sessionIdOrBasename: string): Promise { try { const tempDir = this.context.config.storage.getProjectTempDir(); const chatsDir = path.join(tempDir, 'chats'); const shortId = this.deriveShortId(sessionIdOrBasename); - if (!fs.existsSync(chatsDir)) { + // Using stat instead of existsSync for async sanity + if (!(await fs.promises.stat(chatsDir).catch(() => null))) { return; // Nothing to delete } const matchingFiles = this.getMatchingSessionFiles(chatsDir, shortId); for (const file of matchingFiles) { - this.deleteSessionAndArtifacts(chatsDir, file, tempDir); + await this.deleteSessionAndArtifacts(chatsDir, file, tempDir); } } catch (error) { debugLogger.error('Error deleting session file.', error); @@ -654,14 +685,14 @@ export class ChatRecordingService { /** * Deletes a single session file and its associated logs, tool-outputs, and directory. */ - private deleteSessionAndArtifacts( + private async deleteSessionAndArtifacts( chatsDir: string, file: string, tempDir: string, - ): void { + ): Promise { const filePath = path.join(chatsDir, file); try { - const fileContent = fs.readFileSync(filePath, 'utf8'); + const fileContent = await fs.promises.readFile(filePath, 'utf8'); const content = JSON.parse(fileContent) as unknown; let fullSessionId: string | undefined; @@ -673,60 +704,22 @@ export class ChatRecordingService { } // Delete the session file - fs.unlinkSync(filePath); + await fs.promises.unlink(filePath); if (fullSessionId) { - this.deleteSessionLogs(fullSessionId, tempDir); - this.deleteSessionToolOutputs(fullSessionId, tempDir); - this.deleteSessionDirectory(fullSessionId, tempDir); + // Delegate to shared utility! + await deleteSessionArtifactsAsync(fullSessionId, tempDir); + await deleteSubagentSessionDirAndArtifactsAsync( + fullSessionId, + chatsDir, + tempDir, + ); } } catch (error) { debugLogger.error(`Error deleting associated file ${file}:`, error); } } - /** - * Cleans up activity logs for a session. - */ - private deleteSessionLogs(sessionId: string, tempDir: string): void { - const logsDir = path.join(tempDir, 'logs'); - const safeSessionId = sanitizeFilenamePart(sessionId); - const logPath = path.join(logsDir, `session-${safeSessionId}.jsonl`); - if (fs.existsSync(logPath) && logPath.startsWith(logsDir)) { - fs.unlinkSync(logPath); - } - } - - /** - * Cleans up tool outputs for a session. - */ - private deleteSessionToolOutputs(sessionId: string, tempDir: string): void { - const safeSessionId = sanitizeFilenamePart(sessionId); - const toolOutputDir = path.join( - tempDir, - 'tool-outputs', - `session-${safeSessionId}`, - ); - const toolOutputsBase = path.join(tempDir, 'tool-outputs'); - if ( - fs.existsSync(toolOutputDir) && - toolOutputDir.startsWith(toolOutputsBase) - ) { - fs.rmSync(toolOutputDir, { recursive: true, force: true }); - } - } - - /** - * Cleans up the session-specific directory. - */ - private deleteSessionDirectory(sessionId: string, tempDir: string): void { - const safeSessionId = sanitizeFilenamePart(sessionId); - const sessionDir = path.join(tempDir, safeSessionId); - if (fs.existsSync(sessionDir) && sessionDir.startsWith(tempDir)) { - fs.rmSync(sessionDir, { recursive: true, force: true }); - } - } - /** * Rewinds the conversation to the state just before the specified message ID. * All messages from (and including) the specified ID onwards are removed. diff --git a/packages/core/src/utils/sessionOperations.test.ts b/packages/core/src/utils/sessionOperations.test.ts new file mode 100644 index 0000000000..cc5cd916a5 --- /dev/null +++ b/packages/core/src/utils/sessionOperations.test.ts @@ -0,0 +1,148 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import * as fs from 'node:fs/promises'; +import path from 'node:path'; +import * as os from 'node:os'; +import { + deleteSessionArtifactsAsync, + deleteSubagentSessionDirAndArtifactsAsync, + validateAndSanitizeSessionId, +} from './sessionOperations.js'; + +describe('sessionOperations', () => { + let tempDir: string; + let chatsDir: string; + + beforeEach(async () => { + vi.clearAllMocks(); + // Create a real temporary directory for each test + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'session-ops-test-')); + chatsDir = path.join(tempDir, 'chats'); + }); + + afterEach(async () => { + vi.unstubAllEnvs(); + // Clean up the temporary directory + if (tempDir) { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + + describe('validateAndSanitizeSessionId', () => { + it('should throw for empty or dangerous IDs', () => { + expect(() => validateAndSanitizeSessionId('')).toThrow( + 'Invalid sessionId', + ); + expect(() => validateAndSanitizeSessionId('.')).toThrow( + 'Invalid sessionId', + ); + expect(() => validateAndSanitizeSessionId('..')).toThrow( + 'Invalid sessionId', + ); + }); + + it('should sanitize valid IDs', () => { + expect(validateAndSanitizeSessionId('abc/def')).toBe('abc_def'); + expect(validateAndSanitizeSessionId('valid-id')).toBe('valid-id'); + }); + }); + + describe('deleteSessionArtifactsAsync', () => { + it('should delete logs and tool outputs', async () => { + const sessionId = 'test-session'; + const logsDir = path.join(tempDir, 'logs'); + const toolOutputsDir = path.join( + tempDir, + 'tool-outputs', + `session-${sessionId}`, + ); + const sessionDir = path.join(tempDir, sessionId); + + await fs.mkdir(logsDir, { recursive: true }); + await fs.mkdir(toolOutputsDir, { recursive: true }); + await fs.mkdir(sessionDir, { recursive: true }); + + const logFile = path.join(logsDir, `session-${sessionId}.jsonl`); + await fs.writeFile(logFile, '{}'); + + // Verify files exist before call + expect(await fs.stat(logFile)).toBeTruthy(); + expect(await fs.stat(toolOutputsDir)).toBeTruthy(); + expect(await fs.stat(sessionDir)).toBeTruthy(); + + await deleteSessionArtifactsAsync(sessionId, tempDir); + + // Verify files are deleted + await expect(fs.stat(logFile)).rejects.toThrow(); + await expect(fs.stat(toolOutputsDir)).rejects.toThrow(); + await expect(fs.stat(sessionDir)).rejects.toThrow(); + }); + + it('should ignore ENOENT errors during deletion', async () => { + // Don't create any files. Calling delete on non-existent files should not throw. + await expect( + deleteSessionArtifactsAsync('non-existent', tempDir), + ).resolves.toBeUndefined(); + }); + }); + + describe('deleteSubagentSessionDirAndArtifactsAsync', () => { + it('should iterate subagent files and delete their artifacts', async () => { + const parentSessionId = 'parent-123'; + const subDir = path.join(chatsDir, parentSessionId); + await fs.mkdir(subDir, { recursive: true }); + + await fs.writeFile(path.join(subDir, 'sub1.json'), '{}'); + await fs.writeFile(path.join(subDir, 'sub2.json'), '{}'); + + const logsDir = path.join(tempDir, 'logs'); + await fs.mkdir(logsDir, { recursive: true }); + await fs.writeFile(path.join(logsDir, 'session-sub1.jsonl'), '{}'); + await fs.writeFile(path.join(logsDir, 'session-sub2.jsonl'), '{}'); + + await deleteSubagentSessionDirAndArtifactsAsync( + parentSessionId, + chatsDir, + tempDir, + ); + + // Verify subagent directory is deleted + await expect(fs.stat(subDir)).rejects.toThrow(); + + // Verify artifacts are deleted + await expect( + fs.stat(path.join(logsDir, 'session-sub1.jsonl')), + ).rejects.toThrow(); + await expect( + fs.stat(path.join(logsDir, 'session-sub2.jsonl')), + ).rejects.toThrow(); + }); + + it('should resolve for safe path even if input contains traversals (due to sanitization)', async () => { + // Should sanitize '../unsafe' to '.._unsafe' and resolve (directory won't exist, so readdir returns [] naturally) + await expect( + deleteSubagentSessionDirAndArtifactsAsync( + '../unsafe', + chatsDir, + tempDir, + ), + ).resolves.toBeUndefined(); + }); + + it('should handle ENOENT for readdir gracefully', async () => { + // Non-existent directory should not throw + await expect( + deleteSubagentSessionDirAndArtifactsAsync( + 'non-existent-parent', + chatsDir, + tempDir, + ), + ).resolves.toBeUndefined(); + }); + }); +}); diff --git a/packages/core/src/utils/sessionOperations.ts b/packages/core/src/utils/sessionOperations.ts new file mode 100644 index 0000000000..24ff43aa00 --- /dev/null +++ b/packages/core/src/utils/sessionOperations.ts @@ -0,0 +1,122 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs/promises'; +import path from 'node:path'; +import { sanitizeFilenamePart } from './fileUtils.js'; +import { debugLogger } from './debugLogger.js'; + +const LOGS_DIR = 'logs'; +const TOOL_OUTPUTS_DIR = 'tool-outputs'; + +/** + * Validates a sessionId and returns a sanitized version. + * Throws an error if the ID is dangerous (e.g., ".", "..", or empty). + */ +export function validateAndSanitizeSessionId(sessionId: string): string { + if (!sessionId || sessionId === '.' || sessionId === '..') { + throw new Error(`Invalid sessionId: ${sessionId}`); + } + const sanitized = sanitizeFilenamePart(sessionId); + if (!sanitized) { + throw new Error(`Invalid sessionId after sanitization: ${sessionId}`); + } + return sanitized; +} + +/** + * Asynchronously deletes activity logs and tool outputs for a specific session ID. + */ +export async function deleteSessionArtifactsAsync( + sessionId: string, + tempDir: string, +): Promise { + try { + const safeSessionId = validateAndSanitizeSessionId(sessionId); + const logsDir = path.join(tempDir, LOGS_DIR); + const logPath = path.join(logsDir, `session-${safeSessionId}.jsonl`); + + // Use fs.promises.unlink directly since we don't need to check exists first + // (catching ENOENT is idiomatic for async file system ops) + await fs.unlink(logPath).catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'ENOENT') throw err; + }); + + const toolOutputsBase = path.join(tempDir, TOOL_OUTPUTS_DIR); + const toolOutputDir = path.join( + toolOutputsBase, + `session-${safeSessionId}`, + ); + + await fs + .rm(toolOutputDir, { recursive: true, force: true }) + .catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'ENOENT') throw err; + }); + + // Top-level session directory (e.g., tempDir/safeSessionId) + const sessionDir = path.join(tempDir, safeSessionId); + await fs + .rm(sessionDir, { recursive: true, force: true }) + .catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'ENOENT') throw err; + }); + } catch (error) { + debugLogger.error( + `Error deleting session artifacts for ${sessionId}:`, + error, + ); + } +} + +/** + * Iterates through subagent files in a parent's directory and deletes their artifacts + * before deleting the directory itself. + */ +export async function deleteSubagentSessionDirAndArtifactsAsync( + parentSessionId: string, + chatsDir: string, + tempDir: string, +): Promise { + const safeParentSessionId = validateAndSanitizeSessionId(parentSessionId); + const subagentDir = path.join(chatsDir, safeParentSessionId); + + // Safety check to ensure we don't escape chatsDir + if (!subagentDir.startsWith(chatsDir + path.sep)) { + throw new Error(`Dangerous subagent directory path: ${subagentDir}`); + } + + try { + const files = await fs + .readdir(subagentDir, { withFileTypes: true }) + .catch((err: NodeJS.ErrnoException) => { + if (err.code === 'ENOENT') return []; + throw err; + }); + + for (const file of files) { + if (file.isFile() && file.name.endsWith('.json')) { + const agentId = path.basename(file.name, '.json'); + await deleteSessionArtifactsAsync(agentId, tempDir); + } + } + + // Finally, remove the directory itself + await fs + .rm(subagentDir, { recursive: true, force: true }) + .catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'ENOENT') throw err; + }); + } catch (error) { + debugLogger.error( + `Error cleaning up subagents for parent ${parentSessionId}:`, + error, + ); + // If directory listing fails, we still try to remove the directory if it exists, + // or let the error propagate if it's a critical failure. + await fs.rm(subagentDir, { recursive: true, force: true }).catch(() => {}); + } +}