diff --git a/src/hooks/claude-code-hooks/claude-code-hooks-hook.ts b/src/hooks/claude-code-hooks/claude-code-hooks-hook.ts index b4c2a3124..bd711df12 100644 --- a/src/hooks/claude-code-hooks/claude-code-hooks-hook.ts +++ b/src/hooks/claude-code-hooks/claude-code-hooks-hook.ts @@ -3,7 +3,10 @@ import type { PluginConfig } from "./types" import type { ContextCollector } from "../../features/context-injector" import { createChatMessageHandler } from "./handlers/chat-message-handler" import { createPreCompactHandler } from "./handlers/pre-compact-handler" -import { createSessionEventHandler } from "./handlers/session-event-handler" +import { + createSessionEventHandler, + disposeSessionEventHandler, +} from "./handlers/session-event-handler" import { createToolExecuteAfterHandler } from "./handlers/tool-execute-after-handler" import { createToolExecuteBeforeHandler } from "./handlers/tool-execute-before-handler" @@ -17,6 +20,9 @@ export function createClaudeCodeHooksHook( "chat.message": createChatMessageHandler(ctx, config, contextCollector), "tool.execute.before": createToolExecuteBeforeHandler(ctx, config), "tool.execute.after": createToolExecuteAfterHandler(ctx, config), - event: createSessionEventHandler(ctx, config), + event: createSessionEventHandler(ctx, config, contextCollector), + dispose: (): void => { + disposeSessionEventHandler(contextCollector) + }, } } diff --git a/src/hooks/claude-code-hooks/handlers/session-event-handler.test.ts b/src/hooks/claude-code-hooks/handlers/session-event-handler.test.ts new file mode 100644 index 000000000..dc67b5223 --- /dev/null +++ b/src/hooks/claude-code-hooks/handlers/session-event-handler.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, test } from "bun:test" + +import { ContextCollector } from "../../../features/context-injector" +import { cacheToolInput, getToolInput, stopToolInputCacheCleanup } from "../tool-input-cache" +import { buildTranscriptFromSession, hasTranscriptCacheEntry } from "../transcript" +import { createSessionEventHandler, disposeSessionEventHandler } from "./session-event-handler" + +function createMockClient() { + return { + session: { + get: async () => ({ data: {} }), + prompt: async () => undefined, + messages: async () => ({ data: [] }), + }, + } +} + +describe("createSessionEventHandler", () => { + test("#given deleted session with retained caches #when session deleted arrives #then per-session resources are cleared", async () => { + //#given + const collector = new ContextCollector() + collector.register("ses_cleanup", { + id: "hook-context", + source: "custom", + content: "pending hook context", + }) + cacheToolInput("ses_cleanup", "Read", "call-1", { path: "/tmp/a" }) + await buildTranscriptFromSession(createMockClient(), "ses_cleanup", "/tmp", "Read", { path: "/tmp/a" }) + const handler = createSessionEventHandler(createMockClient() as never, {}, collector) + + //#when + await handler({ + event: { type: "session.deleted", properties: { info: { id: "ses_cleanup" } } }, + }) + + //#then + expect(collector.hasPending("ses_cleanup")).toBe(false) + expect(getToolInput("ses_cleanup", "Read", "call-1")).toBeNull() + expect(hasTranscriptCacheEntry("ses_cleanup")).toBe(false) + }) + + test("#given active singleton state #when dispose runs #then all shared caches are cleared", async () => { + //#given + const collector = new ContextCollector() + collector.register("ses_one", { + id: "ctx-1", + source: "custom", + content: "one", + }) + collector.register("ses_two", { + id: "ctx-2", + source: "custom", + content: "two", + }) + cacheToolInput("ses_one", "Read", "call-1", { path: "/tmp/one" }) + cacheToolInput("ses_two", "Read", "call-2", { path: "/tmp/two" }) + await buildTranscriptFromSession(createMockClient(), "ses_one", "/tmp", "Read", { path: "/tmp/one" }) + await buildTranscriptFromSession(createMockClient(), "ses_two", "/tmp", "Read", { path: "/tmp/two" }) + + //#when + disposeSessionEventHandler(collector) + + //#then + expect(collector.hasPending("ses_one")).toBe(false) + expect(collector.hasPending("ses_two")).toBe(false) + expect(getToolInput("ses_one", "Read", "call-1")).toBeNull() + expect(getToolInput("ses_two", "Read", "call-2")).toBeNull() + expect(hasTranscriptCacheEntry("ses_one")).toBe(false) + expect(hasTranscriptCacheEntry("ses_two")).toBe(false) + + stopToolInputCacheCleanup() + }) +}) diff --git a/src/hooks/claude-code-hooks/handlers/session-event-handler.ts b/src/hooks/claude-code-hooks/handlers/session-event-handler.ts index 4c845004c..71d0374b8 100644 --- a/src/hooks/claude-code-hooks/handlers/session-event-handler.ts +++ b/src/hooks/claude-code-hooks/handlers/session-event-handler.ts @@ -1,16 +1,24 @@ import type { PluginInput } from "@opencode-ai/plugin" +import type { ContextCollector } from "../../../features/context-injector" import { loadClaudeHooksConfig } from "../config" import { loadPluginExtendedConfig } from "../config-loader" import { executeStopHooks, type StopContext } from "../stop" +import { clearTranscriptCache } from "../transcript" +import { clearToolInputCache, stopToolInputCacheCleanup } from "../tool-input-cache" import type { PluginConfig } from "../types" import { createInternalAgentTextPart, isHookDisabled, log } from "../../../shared" import { + clearAllSessionHookState, clearSessionHookState, sessionErrorState, sessionInterruptState, } from "../session-hook-state" -export function createSessionEventHandler(ctx: PluginInput, config: PluginConfig) { +export function createSessionEventHandler( + ctx: PluginInput, + config: PluginConfig, + contextCollector?: ContextCollector, +) { return async (input: { event: { type: string; properties?: unknown } }) => { const { event } = input @@ -30,6 +38,9 @@ export function createSessionEventHandler(ctx: PluginInput, config: PluginConfig const props = event.properties as Record | undefined const sessionInfo = props?.info as { id?: string } | undefined if (sessionInfo?.id) { + clearTranscriptCache(sessionInfo.id) + clearToolInputCache(sessionInfo.id) + contextCollector?.clear(sessionInfo.id) clearSessionHookState(sessionInfo.id) } return @@ -109,3 +120,10 @@ export function createSessionEventHandler(ctx: PluginInput, config: PluginConfig clearSessionHookState(sessionID) } } + +export function disposeSessionEventHandler(contextCollector?: ContextCollector): void { + clearTranscriptCache() + stopToolInputCacheCleanup() + contextCollector?.clearAll() + clearAllSessionHookState() +} diff --git a/src/hooks/claude-code-hooks/session-hook-state.ts b/src/hooks/claude-code-hooks/session-hook-state.ts index 50a2887cb..a6b4024bd 100644 --- a/src/hooks/claude-code-hooks/session-hook-state.ts +++ b/src/hooks/claude-code-hooks/session-hook-state.ts @@ -9,3 +9,9 @@ export function clearSessionHookState(sessionID: string): void { sessionInterruptState.delete(sessionID) sessionFirstMessageProcessed.delete(sessionID) } + +export function clearAllSessionHookState(): void { + sessionErrorState.clear() + sessionInterruptState.clear() + sessionFirstMessageProcessed.clear() +} diff --git a/src/hooks/claude-code-hooks/tool-input-cache.test.ts b/src/hooks/claude-code-hooks/tool-input-cache.test.ts new file mode 100644 index 000000000..409c56897 --- /dev/null +++ b/src/hooks/claude-code-hooks/tool-input-cache.test.ts @@ -0,0 +1,54 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test" + +describe("tool-input-cache", () => { + const originalSetInterval = globalThis.setInterval + const originalClearInterval = globalThis.clearInterval + + beforeEach(() => { + globalThis.setInterval = originalSetInterval + globalThis.clearInterval = originalClearInterval + }) + + afterEach(async () => { + const modulePath = new URL("./tool-input-cache.ts", import.meta.url).pathname + const cacheModule = await import(`${modulePath}?cleanup=${Date.now()}`) + cacheModule.stopToolInputCacheCleanup() + }) + + test("#given cached entries from multiple sessions #when clearing one session #then only matching entries are removed", async () => { + //#given + const modulePath = new URL("./tool-input-cache.ts", import.meta.url).pathname + const cacheModule = await import(`${modulePath}?session-clear`) + + cacheModule.cacheToolInput("ses_a", "Read", "call-1", { path: "a" }) + cacheModule.cacheToolInput("ses_b", "Read", "call-2", { path: "b" }) + + //#when + cacheModule.clearToolInputCache("ses_a") + + //#then + expect(cacheModule.getToolInput("ses_a", "Read", "call-1")).toBeNull() + expect(cacheModule.getToolInput("ses_b", "Read", "call-2")).toEqual({ path: "b" }) + }) + + test("#given cleanup timer started #when stop cleanup runs #then interval is cleared and cache is emptied", async () => { + //#given + const intervalHandle = { unref: mock(() => {}) } as unknown as ReturnType + const setIntervalMock = mock(() => intervalHandle) + const clearIntervalMock = mock(() => {}) + globalThis.setInterval = setIntervalMock as unknown as typeof setInterval + globalThis.clearInterval = clearIntervalMock as unknown as typeof clearInterval + + const modulePath = new URL("./tool-input-cache.ts", import.meta.url).pathname + const cacheModule = await import(`${modulePath}?stop-clear`) + cacheModule.cacheToolInput("ses_stop", "Read", "call-stop", { path: "stop" }) + + //#when + cacheModule.stopToolInputCacheCleanup() + + //#then + expect(setIntervalMock).toHaveBeenCalledTimes(1) + expect(clearIntervalMock).toHaveBeenCalledWith(intervalHandle) + expect(cacheModule.getToolInput("ses_stop", "Read", "call-stop")).toBeNull() + }) +}) diff --git a/src/hooks/claude-code-hooks/tool-input-cache.ts b/src/hooks/claude-code-hooks/tool-input-cache.ts index 3b47317c6..0f7087707 100644 --- a/src/hooks/claude-code-hooks/tool-input-cache.ts +++ b/src/hooks/claude-code-hooks/tool-input-cache.ts @@ -11,12 +11,36 @@ const cache = new Map() const CACHE_TTL = 60000 // 1 minute +let cleanupInterval: ReturnType | null = null + +function pruneExpiredToolInputs(): void { + const now = Date.now() + for (const [key, entry] of cache.entries()) { + if (now - entry.timestamp > CACHE_TTL) { + cache.delete(key) + } + } +} + +function ensureCleanupInterval(): void { + if (cleanupInterval) return + + cleanupInterval = setInterval(() => { + pruneExpiredToolInputs() + }, CACHE_TTL) + + if (typeof cleanupInterval === "object" && "unref" in cleanupInterval) { + cleanupInterval.unref() + } +} + export function cacheToolInput( sessionId: string, toolName: string, invocationId: string, toolInput: Record ): void { + ensureCleanupInterval() const key = `${sessionId}:${toolName}:${invocationId}` cache.set(key, { toolInput, timestamp: Date.now() }) } @@ -30,22 +54,29 @@ export function getToolInput( const entry = cache.get(key) if (!entry) return null - cache.delete(key) + cache.delete(key) if (Date.now() - entry.timestamp > CACHE_TTL) return null return entry.toolInput } -// Periodic cleanup (every minute) -const cleanupInterval = setInterval(() => { - const now = Date.now() - for (const [key, entry] of cache.entries()) { - if (now - entry.timestamp > CACHE_TTL) { +export function clearToolInputCache(sessionId?: string): void { + if (!sessionId) { + cache.clear() + return + } + + const sessionPrefix = `${sessionId}:` + for (const key of cache.keys()) { + if (key.startsWith(sessionPrefix)) { cache.delete(key) } } -}, CACHE_TTL) -// Allow process to exit naturally even if interval is running -if (typeof cleanupInterval === "object" && "unref" in cleanupInterval) { - cleanupInterval.unref() +} + +export function stopToolInputCacheCleanup(): void { + clearToolInputCache() + if (!cleanupInterval) return + clearInterval(cleanupInterval) + cleanupInterval = null } diff --git a/src/hooks/claude-code-hooks/transcript.ts b/src/hooks/claude-code-hooks/transcript.ts index 3c1693db9..2c1c56723 100644 --- a/src/hooks/claude-code-hooks/transcript.ts +++ b/src/hooks/claude-code-hooks/transcript.ts @@ -97,6 +97,10 @@ export function clearTranscriptCache(sessionId?: string): void { } } +export function hasTranscriptCacheEntry(sessionId: string): boolean { + return transcriptCache.has(sessionId) +} + function isCacheValid(entry: TranscriptCacheEntry): boolean { return Date.now() - entry.createdAt < TRANSCRIPT_CACHE_TTL_MS }