Skip to content

Commit b660dbf

Browse files
authored
Merge pull request #123 from Lellansin/dev/memory-leak-fixes-origin-main
fix(session): 修复会话清理内存泄漏并补充回归测试
2 parents 135046f + 7de1366 commit b660dbf

4 files changed

Lines changed: 305 additions & 10 deletions

File tree

src/common/state.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,17 @@ const snippetsBySession = new Map<string, Map<string, FileSnippet>>();
2929
const snippetCountersBySession = new Map<string, number>();
3030
const fileVersionsBySession = new Map<string, Map<string, number>>();
3131

32+
export function clearSessionState(sessionId: string): void {
33+
if (!sessionId) {
34+
return;
35+
}
36+
37+
fileStatesBySession.delete(sessionId);
38+
snippetsBySession.delete(sessionId);
39+
snippetCountersBySession.delete(sessionId);
40+
fileVersionsBySession.delete(sessionId);
41+
}
42+
3243
export function normalizeFilePath(filePath: string, platform: NodeJS.Platform = process.platform): string {
3344
const nativePath = normalizeNativeFilePath(filePath, platform);
3445
return platform === "win32" ? path.win32.normalize(nativePath) : path.normalize(nativePath);

src/session.ts

Lines changed: 67 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ import {
4444
type PermissionToolCall,
4545
type UserToolPermission,
4646
} from "./common/permissions";
47+
import { clearSessionWorkingDir } from "./tools/bash-handler";
48+
import { clearSessionState } from "./common/state";
4749

4850
export type { PermissionScope } from "./settings";
4951
export type {
@@ -337,6 +339,18 @@ export class SessionManager {
337339
}
338340

339341
dispose(): void {
342+
const controller = this.activePromptController;
343+
if (controller && !controller.signal.aborted) {
344+
controller.abort();
345+
}
346+
this.activePromptController = null;
347+
for (const sessionController of this.sessionControllers.values()) {
348+
if (!sessionController.signal.aborted) {
349+
sessionController.abort();
350+
}
351+
}
352+
this.sessionControllers.clear();
353+
this.processTimeoutControls.clear();
340354
this.mcpManager.disconnect();
341355
}
342356

@@ -970,7 +984,12 @@ The candidate skills are as follows:\n\n`;
970984
const droppedEntries = sortedEntries.filter((item) => !keptIds.has(item.id));
971985
index.entries = keptEntries;
972986
this.saveSessionsIndex(index);
973-
this.removeSessionMessages(droppedEntries.map((item) => item.id));
987+
for (const dropped of droppedEntries) {
988+
this.cleanupSessionResources(dropped.id, {
989+
removeMessages: true,
990+
processIds: this.getProcessIds(dropped.processes ?? null),
991+
});
992+
}
974993

975994
const promptToolOptions = this.getPromptToolOptions();
976995
const systemPrompt = getSystemPrompt(this.projectRoot, promptToolOptions);
@@ -1581,23 +1600,25 @@ ${skillMd}
15811600

15821601
/**
15831602
* Delete a session by its ID.
1584-
* Removes the session entry from the index and deletes the associated messages file.
1603+
* Removes the session entry from the index and cleans up associated resources
1604+
* such as message files, in-memory state caches, working directory state,
1605+
* session controllers, and tracked process timeout controls.
15851606
* Returns true if the session was found and deleted, false otherwise.
15861607
*/
15871608
deleteSession(sessionId: string): boolean {
15881609
const index = this.loadSessionsIndex();
1589-
const entryIndex = index.entries.findIndex((entry) => entry.id === sessionId);
1590-
if (entryIndex === -1) {
1610+
const targetEntry = index.entries.find((entry) => entry.id === sessionId) ?? null;
1611+
const nextEntries = index.entries.filter((entry) => entry.id !== sessionId);
1612+
if (nextEntries.length === index.entries.length) {
15911613
return false;
15921614
}
15931615

1594-
// Remove from index
1595-
index.entries.splice(entryIndex, 1);
1616+
index.entries = nextEntries;
15961617
this.saveSessionsIndex(index);
1597-
1598-
// Remove messages file
1599-
this.removeSessionMessages([sessionId]);
1600-
1618+
this.cleanupSessionResources(sessionId, {
1619+
removeMessages: true,
1620+
processIds: this.getProcessIds(targetEntry?.processes ?? null),
1621+
});
16011622
return true;
16021623
}
16031624

@@ -1844,6 +1865,42 @@ ${skillMd}
18441865
}
18451866
}
18461867

1868+
private cleanupSessionResources(
1869+
sessionId: string,
1870+
options: { removeMessages: boolean; processIds?: number[] }
1871+
): void {
1872+
const processIds = options.processIds ?? [];
1873+
for (const pid of processIds) {
1874+
const processControlKey = this.getProcessControlKey(sessionId, pid);
1875+
if (!this.processTimeoutControls.has(processControlKey)) {
1876+
continue;
1877+
}
1878+
1879+
const killedGroup = killProcessTree(pid, "SIGKILL");
1880+
if (killedGroup) {
1881+
this.processTimeoutControls.delete(processControlKey);
1882+
continue;
1883+
}
1884+
try {
1885+
process.kill(pid, "SIGKILL");
1886+
} catch {
1887+
// ignore process-kill failures during cleanup
1888+
}
1889+
this.processTimeoutControls.delete(processControlKey);
1890+
}
1891+
1892+
clearSessionState(sessionId);
1893+
clearSessionWorkingDir(sessionId);
1894+
const controller = this.sessionControllers.get(sessionId);
1895+
if (controller && !controller.signal.aborted) {
1896+
controller.abort();
1897+
}
1898+
this.sessionControllers.delete(sessionId);
1899+
if (options.removeMessages) {
1900+
this.removeSessionMessages([sessionId]);
1901+
}
1902+
}
1903+
18471904
private appendSessionMessage(sessionId: string, message: SessionMessage): void {
18481905
this.ensureProjectDir();
18491906
const messagePath = this.getSessionMessagesPath(sessionId);

src/tests/memory-leak.test.ts

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
import { afterEach, test } from "node:test";
2+
import assert from "node:assert/strict";
3+
import * as fs from "fs";
4+
import * as os from "os";
5+
import * as path from "path";
6+
import { SessionManager } from "../session";
7+
import { handleBashTool } from "../tools/bash-handler";
8+
import * as state from "../common/state";
9+
import type { ToolExecutionContext } from "../tools/executor";
10+
11+
const originalHome = process.env.HOME;
12+
const originalUserProfile = process.env.USERPROFILE;
13+
const tempDirs: string[] = [];
14+
15+
function createTempDir(prefix: string): string {
16+
const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix));
17+
tempDirs.push(dir);
18+
return dir;
19+
}
20+
21+
function setHomeDir(dir: string): void {
22+
process.env.HOME = dir;
23+
if (process.platform === "win32") {
24+
process.env.USERPROFILE = dir;
25+
}
26+
}
27+
28+
function createSessionManager(projectRoot: string): SessionManager {
29+
return new SessionManager({
30+
projectRoot,
31+
createOpenAIClient: () => ({
32+
client: null,
33+
model: "test",
34+
baseURL: "https://api.test.com",
35+
thinkingEnabled: false,
36+
reasoningEffort: "high",
37+
debugLogEnabled: false,
38+
env: {},
39+
}),
40+
getResolvedSettings: () => ({ model: "test" }),
41+
renderMarkdown: (text: string) => text,
42+
onAssistantMessage: () => {},
43+
});
44+
}
45+
46+
afterEach(() => {
47+
if (originalHome === undefined) {
48+
delete process.env.HOME;
49+
} else {
50+
process.env.HOME = originalHome;
51+
}
52+
if (originalUserProfile === undefined) {
53+
delete process.env.USERPROFILE;
54+
} else {
55+
process.env.USERPROFILE = originalUserProfile;
56+
}
57+
58+
while (tempDirs.length > 0) {
59+
const dir = tempDirs.pop();
60+
if (dir) {
61+
fs.rmSync(dir, { recursive: true, force: true });
62+
}
63+
}
64+
});
65+
66+
test("SessionManager.deleteSession clears state cache for that session", async () => {
67+
const home = createTempDir("deepcode-mem-home-");
68+
const projectRoot = createTempDir("deepcode-mem-workspace-");
69+
setHomeDir(home);
70+
const manager = createSessionManager(projectRoot);
71+
72+
const sessionId = await manager.createSession({ text: "seed" });
73+
const filePath = path.join(projectRoot, "a.txt");
74+
fs.writeFileSync(filePath, "hello");
75+
state.recordFileState(sessionId, { filePath, content: "hello", timestamp: Date.now() }, { incrementVersion: true });
76+
const snippet = state.createSnippet(sessionId, filePath, 1, 1, "hello");
77+
const fileVersionBeforeDelete = state.getFileVersion(sessionId, filePath);
78+
79+
assert.ok(state.wasFileRead(sessionId, filePath));
80+
assert.ok(snippet);
81+
assert.ok(state.getSnippet(sessionId, snippet!.id));
82+
assert.equal(fileVersionBeforeDelete, 1);
83+
84+
assert.equal(manager.deleteSession(sessionId), true);
85+
assert.equal(state.wasFileRead(sessionId, filePath), false);
86+
assert.equal(state.getSnippet(sessionId, snippet!.id), null);
87+
assert.equal(state.getFileVersion(sessionId, filePath), 0);
88+
});
89+
90+
test("SessionManager.createSession auto-prune clears dropped session state cache", async () => {
91+
const home = createTempDir("deepcode-mem-home-");
92+
const projectRoot = createTempDir("deepcode-mem-workspace-");
93+
setHomeDir(home);
94+
const manager = createSessionManager(projectRoot);
95+
96+
const firstSession = await manager.createSession({ text: "first" });
97+
const filePath = path.join(projectRoot, "first.txt");
98+
fs.writeFileSync(filePath, "first");
99+
state.recordFileState(firstSession, { filePath, content: "first", timestamp: Date.now() });
100+
assert.equal(state.wasFileRead(firstSession, filePath), true);
101+
102+
for (let i = 0; i < 60; i += 1) {
103+
await manager.createSession({ text: `session-${i}` });
104+
}
105+
106+
const remaining = manager.listSessions().map((entry) => entry.id);
107+
assert.equal(remaining.includes(firstSession), false);
108+
assert.equal(state.wasFileRead(firstSession, filePath), false);
109+
});
110+
111+
test("SessionManager.deleteSession clears controller map entry", async () => {
112+
const home = createTempDir("deepcode-mem-home-");
113+
const projectRoot = createTempDir("deepcode-mem-workspace-");
114+
setHomeDir(home);
115+
const manager = createSessionManager(projectRoot);
116+
117+
const sessionId = await manager.createSession({ text: "seed" });
118+
const controllers = (manager as unknown as { sessionControllers: Map<string, AbortController> }).sessionControllers;
119+
controllers.set(sessionId, new AbortController());
120+
assert.equal(controllers.has(sessionId), true);
121+
122+
assert.equal(manager.deleteSession(sessionId), true);
123+
assert.equal(controllers.has(sessionId), false);
124+
});
125+
126+
test("SessionManager.dispose aborts and clears controllers", () => {
127+
const projectRoot = createTempDir("deepcode-mem-workspace-");
128+
const manager = createSessionManager(projectRoot);
129+
const controllers = (manager as unknown as { sessionControllers: Map<string, AbortController> }).sessionControllers;
130+
131+
const controllerA = new AbortController();
132+
const controllerB = new AbortController();
133+
controllers.set("a", controllerA);
134+
controllers.set("b", controllerB);
135+
assert.equal(controllers.size, 2);
136+
137+
manager.dispose();
138+
assert.equal(controllers.size, 0);
139+
});
140+
141+
test("Deleted session id reuse should reset bash cwd to project root", async () => {
142+
const home = createTempDir("deepcode-mem-home-");
143+
const projectRoot = createTempDir("deepcode-mem-workspace-");
144+
setHomeDir(home);
145+
const manager = createSessionManager(projectRoot);
146+
147+
const sessionId = await manager.createSession({ text: "bash-session" });
148+
const sub = path.join(projectRoot, "sub");
149+
fs.mkdirSync(sub, { recursive: true });
150+
151+
const context: ToolExecutionContext = {
152+
sessionId,
153+
projectRoot,
154+
toolCall: { id: "call-1", type: "function", function: { name: "bash", arguments: "{}" } },
155+
createOpenAIClient: () => ({
156+
client: null,
157+
model: "test",
158+
baseURL: "",
159+
thinkingEnabled: false,
160+
reasoningEffort: "high",
161+
debugLogEnabled: false,
162+
env: {},
163+
}),
164+
};
165+
166+
const first = await handleBashTool({ command: `cd "${sub}" && pwd` }, context);
167+
assert.equal(first.ok, true);
168+
169+
assert.equal(manager.deleteSession(sessionId), true);
170+
171+
const second = await handleBashTool({ command: "pwd" }, context);
172+
assert.equal(second.ok, true);
173+
174+
const output = (second.output ?? "").trim();
175+
const metadataCwd =
176+
second.metadata && typeof second.metadata.cwd === "string" ? (second.metadata.cwd as string) : null;
177+
const observedCwd = (metadataCwd ?? output).replace(/\\/g, "/").replace(/\/+$/, "");
178+
const normalizedSub = fs.realpathSync(sub).replace(/\\/g, "/").replace(/\/+$/, "");
179+
assert.notEqual(
180+
observedCwd,
181+
normalizedSub,
182+
`expected cwd not to stay on deleted session subdir ${normalizedSub}, got output=${output}, metadata.cwd=${String(metadataCwd)}`
183+
);
184+
});
185+
186+
test("deleteSession should not kill untracked stale persisted pids", async () => {
187+
const home = createTempDir("deepcode-mem-home-");
188+
const projectRoot = createTempDir("deepcode-mem-workspace-");
189+
setHomeDir(home);
190+
const manager = createSessionManager(projectRoot);
191+
const sessionId = await manager.createSession({ text: "stale-pid" });
192+
193+
const privateManager = manager as unknown as {
194+
updateSessionEntry: (
195+
sessionId: string,
196+
updater: (entry: { processes: Map<string, { startTime: string; command: string }> | null }) => {
197+
processes: Map<string, { startTime: string; command: string }> | null;
198+
}
199+
) => unknown;
200+
};
201+
privateManager.updateSessionEntry(sessionId, (entry) => ({
202+
...entry,
203+
processes: new Map([["999999", { startTime: new Date().toISOString(), command: "sleep 999" }]]),
204+
}));
205+
206+
const originalKill = process.kill;
207+
let killCalls = 0;
208+
const mockedKill = ((pid: number, signal?: NodeJS.Signals | number) => {
209+
killCalls += 1;
210+
return originalKill(pid, signal);
211+
}) as typeof process.kill;
212+
process.kill = mockedKill;
213+
try {
214+
assert.equal(manager.deleteSession(sessionId), true);
215+
} finally {
216+
process.kill = originalKill;
217+
}
218+
219+
assert.equal(killCalls, 0);
220+
});

src/tools/bash-handler.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,13 @@ const MAX_OUTPUT_CHARS = 30000;
1515
const MAX_CAPTURE_CHARS = 10 * 1024 * 1024;
1616
const sessionWorkingDirs = new Map<string, string>();
1717

18+
export function clearSessionWorkingDir(sessionId: string): void {
19+
if (!sessionId) {
20+
return;
21+
}
22+
sessionWorkingDirs.delete(sessionId);
23+
}
24+
1825
type ToolCommandResult = {
1926
ok: boolean;
2027
output: string;

0 commit comments

Comments
 (0)