feat(claude-code-hooks): improve session handling and add tests

🤖 Generated with assistance of OhMyOpenCode

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)
+    },
   }
 }

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()
+  })
+})

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<string, unknown> | 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()
+}

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()
+}

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<typeof setInterval>
+    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()
+  })
+})

src/hooks/claude-code-hooks/tool-input-cache.ts

@@ -11,12 +11,36 @@ const cache = new Map<string, CacheEntry>()
 
 const CACHE_TTL = 60000 // 1 minute
 
+let cleanupInterval: ReturnType<typeof setInterval> | 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<string, unknown>
 ): 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
 }

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
 }