diff --git a/src/common/state.ts b/src/common/state.ts index add27f35..c3e946cc 100644 --- a/src/common/state.ts +++ b/src/common/state.ts @@ -29,6 +29,17 @@ const snippetsBySession = new Map>(); const snippetCountersBySession = new Map(); const fileVersionsBySession = new Map>(); +export function clearSessionState(sessionId: string): void { + if (!sessionId) { + return; + } + + fileStatesBySession.delete(sessionId); + snippetsBySession.delete(sessionId); + snippetCountersBySession.delete(sessionId); + fileVersionsBySession.delete(sessionId); +} + export function normalizeFilePath(filePath: string, platform: NodeJS.Platform = process.platform): string { const nativePath = normalizeNativeFilePath(filePath, platform); return platform === "win32" ? path.win32.normalize(nativePath) : path.normalize(nativePath); diff --git a/src/session.ts b/src/session.ts index a9fc39e8..b11679a3 100644 --- a/src/session.ts +++ b/src/session.ts @@ -44,6 +44,8 @@ import { type PermissionToolCall, type UserToolPermission, } from "./common/permissions"; +import { clearSessionWorkingDir } from "./tools/bash-handler"; +import { clearSessionState } from "./common/state"; export type { PermissionScope } from "./settings"; export type { @@ -346,6 +348,18 @@ export class SessionManager { } dispose(): void { + const controller = this.activePromptController; + if (controller && !controller.signal.aborted) { + controller.abort(); + } + this.activePromptController = null; + for (const sessionController of this.sessionControllers.values()) { + if (!sessionController.signal.aborted) { + sessionController.abort(); + } + } + this.sessionControllers.clear(); + this.processTimeoutControls.clear(); this.mcpManager.disconnect(); } @@ -979,7 +993,12 @@ The candidate skills are as follows:\n\n`; const droppedEntries = sortedEntries.filter((item) => !keptIds.has(item.id)); index.entries = keptEntries; this.saveSessionsIndex(index); - this.removeSessionMessages(droppedEntries.map((item) => item.id)); + for (const dropped of droppedEntries) { + this.cleanupSessionResources(dropped.id, { + removeMessages: true, + processIds: this.getProcessIds(dropped.processes ?? null), + }); + } const promptToolOptions = this.getPromptToolOptions(); const systemPrompt = getSystemPrompt(this.projectRoot, promptToolOptions); @@ -1590,23 +1609,25 @@ ${skillMd} /** * Delete a session by its ID. - * Removes the session entry from the index and deletes the associated messages file. + * Removes the session entry from the index and cleans up associated resources + * such as message files, in-memory state caches, working directory state, + * session controllers, and tracked process timeout controls. * Returns true if the session was found and deleted, false otherwise. */ deleteSession(sessionId: string): boolean { const index = this.loadSessionsIndex(); - const entryIndex = index.entries.findIndex((entry) => entry.id === sessionId); - if (entryIndex === -1) { + const targetEntry = index.entries.find((entry) => entry.id === sessionId) ?? null; + const nextEntries = index.entries.filter((entry) => entry.id !== sessionId); + if (nextEntries.length === index.entries.length) { return false; } - // Remove from index - index.entries.splice(entryIndex, 1); + index.entries = nextEntries; this.saveSessionsIndex(index); - - // Remove messages file - this.removeSessionMessages([sessionId]); - + this.cleanupSessionResources(sessionId, { + removeMessages: true, + processIds: this.getProcessIds(targetEntry?.processes ?? null), + }); return true; } @@ -1853,6 +1874,42 @@ ${skillMd} } } + private cleanupSessionResources( + sessionId: string, + options: { removeMessages: boolean; processIds?: number[] } + ): void { + const processIds = options.processIds ?? []; + for (const pid of processIds) { + const processControlKey = this.getProcessControlKey(sessionId, pid); + if (!this.processTimeoutControls.has(processControlKey)) { + continue; + } + + const killedGroup = killProcessTree(pid, "SIGKILL"); + if (killedGroup) { + this.processTimeoutControls.delete(processControlKey); + continue; + } + try { + process.kill(pid, "SIGKILL"); + } catch { + // ignore process-kill failures during cleanup + } + this.processTimeoutControls.delete(processControlKey); + } + + clearSessionState(sessionId); + clearSessionWorkingDir(sessionId); + const controller = this.sessionControllers.get(sessionId); + if (controller && !controller.signal.aborted) { + controller.abort(); + } + this.sessionControllers.delete(sessionId); + if (options.removeMessages) { + this.removeSessionMessages([sessionId]); + } + } + private appendSessionMessage(sessionId: string, message: SessionMessage): void { this.ensureProjectDir(); const messagePath = this.getSessionMessagesPath(sessionId); diff --git a/src/tests/memory-leak.test.ts b/src/tests/memory-leak.test.ts new file mode 100644 index 00000000..218bcee5 --- /dev/null +++ b/src/tests/memory-leak.test.ts @@ -0,0 +1,220 @@ +import { afterEach, test } from "node:test"; +import assert from "node:assert/strict"; +import * as fs from "fs"; +import * as os from "os"; +import * as path from "path"; +import { SessionManager } from "../session"; +import { handleBashTool } from "../tools/bash-handler"; +import * as state from "../common/state"; +import type { ToolExecutionContext } from "../tools/executor"; + +const originalHome = process.env.HOME; +const originalUserProfile = process.env.USERPROFILE; +const tempDirs: string[] = []; + +function createTempDir(prefix: string): string { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix)); + tempDirs.push(dir); + return dir; +} + +function setHomeDir(dir: string): void { + process.env.HOME = dir; + if (process.platform === "win32") { + process.env.USERPROFILE = dir; + } +} + +function createSessionManager(projectRoot: string): SessionManager { + return new SessionManager({ + projectRoot, + createOpenAIClient: () => ({ + client: null, + model: "test", + baseURL: "https://api.test.com", + thinkingEnabled: false, + reasoningEffort: "high", + debugLogEnabled: false, + env: {}, + }), + getResolvedSettings: () => ({ model: "test" }), + renderMarkdown: (text: string) => text, + onAssistantMessage: () => {}, + }); +} + +afterEach(() => { + if (originalHome === undefined) { + delete process.env.HOME; + } else { + process.env.HOME = originalHome; + } + if (originalUserProfile === undefined) { + delete process.env.USERPROFILE; + } else { + process.env.USERPROFILE = originalUserProfile; + } + + while (tempDirs.length > 0) { + const dir = tempDirs.pop(); + if (dir) { + fs.rmSync(dir, { recursive: true, force: true }); + } + } +}); + +test("SessionManager.deleteSession clears state cache for that session", async () => { + const home = createTempDir("deepcode-mem-home-"); + const projectRoot = createTempDir("deepcode-mem-workspace-"); + setHomeDir(home); + const manager = createSessionManager(projectRoot); + + const sessionId = await manager.createSession({ text: "seed" }); + const filePath = path.join(projectRoot, "a.txt"); + fs.writeFileSync(filePath, "hello"); + state.recordFileState(sessionId, { filePath, content: "hello", timestamp: Date.now() }, { incrementVersion: true }); + const snippet = state.createSnippet(sessionId, filePath, 1, 1, "hello"); + const fileVersionBeforeDelete = state.getFileVersion(sessionId, filePath); + + assert.ok(state.wasFileRead(sessionId, filePath)); + assert.ok(snippet); + assert.ok(state.getSnippet(sessionId, snippet!.id)); + assert.equal(fileVersionBeforeDelete, 1); + + assert.equal(manager.deleteSession(sessionId), true); + assert.equal(state.wasFileRead(sessionId, filePath), false); + assert.equal(state.getSnippet(sessionId, snippet!.id), null); + assert.equal(state.getFileVersion(sessionId, filePath), 0); +}); + +test("SessionManager.createSession auto-prune clears dropped session state cache", async () => { + const home = createTempDir("deepcode-mem-home-"); + const projectRoot = createTempDir("deepcode-mem-workspace-"); + setHomeDir(home); + const manager = createSessionManager(projectRoot); + + const firstSession = await manager.createSession({ text: "first" }); + const filePath = path.join(projectRoot, "first.txt"); + fs.writeFileSync(filePath, "first"); + state.recordFileState(firstSession, { filePath, content: "first", timestamp: Date.now() }); + assert.equal(state.wasFileRead(firstSession, filePath), true); + + for (let i = 0; i < 60; i += 1) { + await manager.createSession({ text: `session-${i}` }); + } + + const remaining = manager.listSessions().map((entry) => entry.id); + assert.equal(remaining.includes(firstSession), false); + assert.equal(state.wasFileRead(firstSession, filePath), false); +}); + +test("SessionManager.deleteSession clears controller map entry", async () => { + const home = createTempDir("deepcode-mem-home-"); + const projectRoot = createTempDir("deepcode-mem-workspace-"); + setHomeDir(home); + const manager = createSessionManager(projectRoot); + + const sessionId = await manager.createSession({ text: "seed" }); + const controllers = (manager as unknown as { sessionControllers: Map }).sessionControllers; + controllers.set(sessionId, new AbortController()); + assert.equal(controllers.has(sessionId), true); + + assert.equal(manager.deleteSession(sessionId), true); + assert.equal(controllers.has(sessionId), false); +}); + +test("SessionManager.dispose aborts and clears controllers", () => { + const projectRoot = createTempDir("deepcode-mem-workspace-"); + const manager = createSessionManager(projectRoot); + const controllers = (manager as unknown as { sessionControllers: Map }).sessionControllers; + + const controllerA = new AbortController(); + const controllerB = new AbortController(); + controllers.set("a", controllerA); + controllers.set("b", controllerB); + assert.equal(controllers.size, 2); + + manager.dispose(); + assert.equal(controllers.size, 0); +}); + +test("Deleted session id reuse should reset bash cwd to project root", async () => { + const home = createTempDir("deepcode-mem-home-"); + const projectRoot = createTempDir("deepcode-mem-workspace-"); + setHomeDir(home); + const manager = createSessionManager(projectRoot); + + const sessionId = await manager.createSession({ text: "bash-session" }); + const sub = path.join(projectRoot, "sub"); + fs.mkdirSync(sub, { recursive: true }); + + const context: ToolExecutionContext = { + sessionId, + projectRoot, + toolCall: { id: "call-1", type: "function", function: { name: "bash", arguments: "{}" } }, + createOpenAIClient: () => ({ + client: null, + model: "test", + baseURL: "", + thinkingEnabled: false, + reasoningEffort: "high", + debugLogEnabled: false, + env: {}, + }), + }; + + const first = await handleBashTool({ command: `cd "${sub}" && pwd` }, context); + assert.equal(first.ok, true); + + assert.equal(manager.deleteSession(sessionId), true); + + const second = await handleBashTool({ command: "pwd" }, context); + assert.equal(second.ok, true); + + const output = (second.output ?? "").trim(); + const metadataCwd = + second.metadata && typeof second.metadata.cwd === "string" ? (second.metadata.cwd as string) : null; + const observedCwd = (metadataCwd ?? output).replace(/\\/g, "/").replace(/\/+$/, ""); + const normalizedSub = fs.realpathSync(sub).replace(/\\/g, "/").replace(/\/+$/, ""); + assert.notEqual( + observedCwd, + normalizedSub, + `expected cwd not to stay on deleted session subdir ${normalizedSub}, got output=${output}, metadata.cwd=${String(metadataCwd)}` + ); +}); + +test("deleteSession should not kill untracked stale persisted pids", async () => { + const home = createTempDir("deepcode-mem-home-"); + const projectRoot = createTempDir("deepcode-mem-workspace-"); + setHomeDir(home); + const manager = createSessionManager(projectRoot); + const sessionId = await manager.createSession({ text: "stale-pid" }); + + const privateManager = manager as unknown as { + updateSessionEntry: ( + sessionId: string, + updater: (entry: { processes: Map | null }) => { + processes: Map | null; + } + ) => unknown; + }; + privateManager.updateSessionEntry(sessionId, (entry) => ({ + ...entry, + processes: new Map([["999999", { startTime: new Date().toISOString(), command: "sleep 999" }]]), + })); + + const originalKill = process.kill; + let killCalls = 0; + const mockedKill = ((pid: number, signal?: NodeJS.Signals | number) => { + killCalls += 1; + return originalKill(pid, signal); + }) as typeof process.kill; + process.kill = mockedKill; + try { + assert.equal(manager.deleteSession(sessionId), true); + } finally { + process.kill = originalKill; + } + + assert.equal(killCalls, 0); +}); diff --git a/src/tools/bash-handler.ts b/src/tools/bash-handler.ts index 42722710..7d9a3736 100644 --- a/src/tools/bash-handler.ts +++ b/src/tools/bash-handler.ts @@ -15,6 +15,13 @@ const MAX_OUTPUT_CHARS = 30000; const MAX_CAPTURE_CHARS = 10 * 1024 * 1024; const sessionWorkingDirs = new Map(); +export function clearSessionWorkingDir(sessionId: string): void { + if (!sessionId) { + return; + } + sessionWorkingDirs.delete(sessionId); +} + type ToolCommandResult = { ok: boolean; output: string;