Merge pull request #2485 from code-yeongyu/fix/issue-2316-tool-after-error-boundary

fix: add error boundary around extract/discard hooks in tool-execute-after

src/plugin/tool-execute-after.ts

@@ -1,10 +1,19 @@
 import { consumeToolMetadata } from "../features/tool-metadata-store"
 import type { CreatedHooks } from "../create-hooks"
+import { log } from "../shared"
 import type { PluginContext } from "./types"
 import { readState, writeState } from "../hooks/ralph-loop/storage"
 
 const VERIFICATION_ATTEMPT_PATTERN = /<ulw_verification_attempt_id>(.*?)<\/ulw_verification_attempt_id>/i
 
+function getPluginDirectory(ctx: PluginContext): string | null {
+  if (typeof ctx === "object" && ctx !== null && "directory" in ctx && typeof ctx.directory === "string") {
+    return ctx.directory
+  }
+
+  return null
+}
+
 export function createToolExecuteAfterHandler(args: {
   ctx: PluginContext
   hooks: CreatedHooks
@@ -33,47 +42,77 @@ export function createToolExecuteAfterHandler(args: {
     }
 
     if (input.tool === "task") {
+      const directory = getPluginDirectory(ctx)
       const sessionId = typeof output.metadata?.sessionId === "string" ? output.metadata.sessionId : undefined
       const agent = typeof output.metadata?.agent === "string" ? output.metadata.agent : undefined
       const prompt = typeof output.metadata?.prompt === "string" ? output.metadata.prompt : undefined
       const verificationAttemptId = prompt?.match(VERIFICATION_ATTEMPT_PATTERN)?.[1]?.trim()
-      const loopState = readState(ctx.directory)
+      const loopState = directory ? readState(directory) : null
 
       if (
         agent === "oracle"
         && sessionId
         && verificationAttemptId
+        && directory
         && loopState?.active === true
         && loopState.ultrawork === true
         && loopState.verification_pending === true
         && loopState.session_id === input.sessionID
         && loopState.verification_attempt_id === verificationAttemptId
       ) {
-        writeState(ctx.directory, {
+        writeState(directory, {
           ...loopState,
           verification_session_id: sessionId,
         })
       }
     }
 
-    await hooks.toolOutputTruncator?.["tool.execute.after"]?.(input, output)
+    const runToolExecuteAfterHooks = async (): Promise<void> => {
-    await hooks.claudeCodeHooks?.["tool.execute.after"]?.(input, output)
+      await hooks.toolOutputTruncator?.["tool.execute.after"]?.(input, output)
-    await hooks.preemptiveCompaction?.["tool.execute.after"]?.(input, output)
+      await hooks.claudeCodeHooks?.["tool.execute.after"]?.(input, output)
-    await hooks.contextWindowMonitor?.["tool.execute.after"]?.(input, output)
+      await hooks.preemptiveCompaction?.["tool.execute.after"]?.(input, output)
-    await hooks.commentChecker?.["tool.execute.after"]?.(input, output)
+      await hooks.contextWindowMonitor?.["tool.execute.after"]?.(input, output)
-    await hooks.directoryAgentsInjector?.["tool.execute.after"]?.(input, output)
+      await hooks.commentChecker?.["tool.execute.after"]?.(input, output)
-    await hooks.directoryReadmeInjector?.["tool.execute.after"]?.(input, output)
+      await hooks.directoryAgentsInjector?.["tool.execute.after"]?.(input, output)
-    await hooks.rulesInjector?.["tool.execute.after"]?.(input, output)
+      await hooks.directoryReadmeInjector?.["tool.execute.after"]?.(input, output)
-    await hooks.emptyTaskResponseDetector?.["tool.execute.after"]?.(input, output)
+      await hooks.rulesInjector?.["tool.execute.after"]?.(input, output)
-    await hooks.agentUsageReminder?.["tool.execute.after"]?.(input, output)
+      await hooks.emptyTaskResponseDetector?.["tool.execute.after"]?.(input, output)
-    await hooks.categorySkillReminder?.["tool.execute.after"]?.(input, output)
+      await hooks.agentUsageReminder?.["tool.execute.after"]?.(input, output)
-    await hooks.interactiveBashSession?.["tool.execute.after"]?.(input, output)
+      await hooks.categorySkillReminder?.["tool.execute.after"]?.(input, output)
-    await hooks.editErrorRecovery?.["tool.execute.after"]?.(input, output)
+      await hooks.interactiveBashSession?.["tool.execute.after"]?.(input, output)
-    await hooks.delegateTaskRetry?.["tool.execute.after"]?.(input, output)
+      await hooks.editErrorRecovery?.["tool.execute.after"]?.(input, output)
-    await hooks.atlasHook?.["tool.execute.after"]?.(input, output)
+      await hooks.delegateTaskRetry?.["tool.execute.after"]?.(input, output)
-    await hooks.taskResumeInfo?.["tool.execute.after"]?.(input, output)
+      await hooks.atlasHook?.["tool.execute.after"]?.(input, output)
-    await hooks.readImageResizer?.["tool.execute.after"]?.(input, output)
+      await hooks.taskResumeInfo?.["tool.execute.after"]?.(input, output)
-    await hooks.hashlineReadEnhancer?.["tool.execute.after"]?.(input, output)
+      await hooks.readImageResizer?.["tool.execute.after"]?.(input, output)
-    await hooks.jsonErrorRecovery?.["tool.execute.after"]?.(input, output)
+      await hooks.hashlineReadEnhancer?.["tool.execute.after"]?.(input, output)
+      await hooks.jsonErrorRecovery?.["tool.execute.after"]?.(input, output)
+    }
+
+    if (input.tool === "extract" || input.tool === "discard") {
+      const originalOutput = {
+        title: output.title,
+        output: output.output,
+        metadata: { ...output.metadata },
+      }
+
+      try {
+        await runToolExecuteAfterHooks()
+      } catch (error) {
+        output.title = originalOutput.title
+        output.output = originalOutput.output
+        output.metadata = originalOutput.metadata
+        log("[tool-execute-after] Failed to process extract/discard hooks", {
+          tool: input.tool,
+          sessionID: input.sessionID,
+          callID: input.callID,
+          error,
+        })
+      }
+
+      return
+    }
+
+    await runToolExecuteAfterHooks()
   }
 }