Fix session cleanup memory leaks and add regression tests

Author: Lellansin

src/common/state.ts

@@ -29,6 +29,17 @@ const snippetsBySession = new Map<string, Map<string, FileSnippet>>();
 const snippetCountersBySession = new Map<string, number>();
 const fileVersionsBySession = new Map<string, Map<string, number>>();
 
+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);

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);
@@ -1588,25 +1607,20 @@ ${skillMd}
     return index.entries.find((entry) => entry.id === sessionId) ?? null;
   }
 
-  /**
-   * Delete a session by its ID.
-   * Removes the session entry from the index and deletes the associated messages file.
-   * 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 +1867,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);

src/tests/memory-leak.test.ts

@@ -0,0 +1,213 @@
+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<string, AbortController> }).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<string, AbortController> }).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 normalizedRoot = fs.realpathSync(projectRoot);
+  assert.ok(output.startsWith(normalizedRoot), `expected cwd to reset to ${normalizedRoot}, got ${output}`);
+});
+
+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<string, { startTime: string; command: string }> | null }) => {
+        processes: Map<string, { startTime: string; command: string }> | 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);
+});

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<string, string>();
 
+export function clearSessionWorkingDir(sessionId: string): void {
+  if (!sessionId) {
+    return;
+  }
+  sessionWorkingDirs.delete(sessionId);
+}
+
 type ToolCommandResult = {
   ok: boolean;
   output: string;