Fix all injection hooks not working with batch tool (#159)

* Fix AGENTS.md injection not working with batch tool (#141) 🤖 Generated with assistance of OhMyOpenCode (https://github.com/code-yeongyu/oh-my-opencode) * Extend batch tool support to rules-injector The rules-injector hook now captures file paths from batch tool calls, enabling it to inject rules into files read via the batch tool. This ensures all injection hooks work correctly for all file access patterns. 🤖 Generated with assistance of OhMyOpenCode (https://github.com/code-yeongyu/oh-my-opencode)
2025-12-22 01:53:15 +09:00
parent 5c73f47281
commit d909c09f84
4 changed files with 210 additions and 56 deletions
--- a/src/hooks/directory-agents-injector/index.ts
+++ b/src/hooks/directory-agents-injector/index.ts
@@ -20,6 +20,15 @@ interface ToolExecuteOutput {
  metadata: unknown;
 }

+interface ToolExecuteBeforeOutput {
+  args: unknown;
+}
+
+interface BatchToolCall {
+  tool: string;
+  parameters: Record<string, unknown>;
+}
+
 interface EventInput {
  event: {
    type: string;
@@ -29,6 +38,7 @@ interface EventInput {

 export function createDirectoryAgentsInjectorHook(ctx: PluginInput) {
  const sessionCaches = new Map<string, Set<string>>();
+  const pendingBatchReads = new Map<string, string[]>();

  function getSessionCache(sessionID: string): Set<string> {
    if (!sessionCaches.has(sessionID)) {
@@ -37,10 +47,10 @@ export function createDirectoryAgentsInjectorHook(ctx: PluginInput) {
    return sessionCaches.get(sessionID)!;
  }

-  function resolveFilePath(title: string): string | null {
-    if (!title) return null;
-    if (title.startsWith("/")) return title;
-    return resolve(ctx.directory, title);
+  function resolveFilePath(path: string): string | null {
+    if (!path) return null;
+    if (path.startsWith("/")) return path;
+    return resolve(ctx.directory, path);
  }

  function findAgentsMdUp(startDir: string): string[] {
@@ -63,39 +73,73 @@ export function createDirectoryAgentsInjectorHook(ctx: PluginInput) {
    return found.reverse();
  }

-  const toolExecuteAfter = async (
-    input: ToolExecuteInput,
+  function processFilePathForInjection(
+    filePath: string,
+    sessionID: string,
    output: ToolExecuteOutput,
-  ) => {
-    if (input.tool.toLowerCase() !== "read") return;
+  ): void {
+    const resolved = resolveFilePath(filePath);
+    if (!resolved) return;

-    const filePath = resolveFilePath(output.title);
-    if (!filePath) return;
-
-    const dir = dirname(filePath);
-    const cache = getSessionCache(input.sessionID);
+    const dir = dirname(resolved);
+    const cache = getSessionCache(sessionID);
    const agentsPaths = findAgentsMdUp(dir);

-    const toInject: { path: string; content: string }[] = [];
-
    for (const agentsPath of agentsPaths) {
      const agentsDir = dirname(agentsPath);
      if (cache.has(agentsDir)) continue;

      try {
        const content = readFileSync(agentsPath, "utf-8");
-        toInject.push({ path: agentsPath, content });
+        output.output += `\n\n[Directory Context: ${agentsPath}]\n${content}`;
        cache.add(agentsDir);
      } catch {}
    }

-    if (toInject.length === 0) return;
-
-    for (const { path, content } of toInject) {
-      output.output += `\n\n[Directory Context: ${path}]\n${content}`;
+    saveInjectedPaths(sessionID, cache);
  }

-    saveInjectedPaths(input.sessionID, cache);
+  const toolExecuteBefore = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteBeforeOutput,
+  ) => {
+    if (input.tool.toLowerCase() !== "batch") return;
+
+    const args = output.args as { tool_calls?: BatchToolCall[] } | undefined;
+    if (!args?.tool_calls) return;
+
+    const readFilePaths: string[] = [];
+    for (const call of args.tool_calls) {
+      if (call.tool.toLowerCase() === "read" && call.parameters?.filePath) {
+        readFilePaths.push(call.parameters.filePath as string);
+      }
+    }
+
+    if (readFilePaths.length > 0) {
+      pendingBatchReads.set(input.callID, readFilePaths);
+    }
+  };
+
+  const toolExecuteAfter = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteOutput,
+  ) => {
+    const toolName = input.tool.toLowerCase();
+
+    if (toolName === "read") {
+      processFilePathForInjection(output.title, input.sessionID, output);
+      return;
+    }
+
+    if (toolName === "batch") {
+      const filePaths = pendingBatchReads.get(input.callID);
+      if (filePaths) {
+        for (const filePath of filePaths) {
+          processFilePathForInjection(filePath, input.sessionID, output);
+        }
+        pendingBatchReads.delete(input.callID);
+      }
+    }
  };

  const eventHandler = async ({ event }: EventInput) => {
@@ -120,6 +164,7 @@ export function createDirectoryAgentsInjectorHook(ctx: PluginInput) {
  };

  return {
+    "tool.execute.before": toolExecuteBefore,
    "tool.execute.after": toolExecuteAfter,
    event: eventHandler,
  };
--- a/src/hooks/directory-readme-injector/index.ts
+++ b/src/hooks/directory-readme-injector/index.ts
@@ -20,6 +20,15 @@ interface ToolExecuteOutput {
  metadata: unknown;
 }

+interface ToolExecuteBeforeOutput {
+  args: unknown;
+}
+
+interface BatchToolCall {
+  tool: string;
+  parameters: Record<string, unknown>;
+}
+
 interface EventInput {
  event: {
    type: string;
@@ -29,6 +38,7 @@ interface EventInput {

 export function createDirectoryReadmeInjectorHook(ctx: PluginInput) {
  const sessionCaches = new Map<string, Set<string>>();
+  const pendingBatchReads = new Map<string, string[]>();

  function getSessionCache(sessionID: string): Set<string> {
    if (!sessionCaches.has(sessionID)) {
@@ -37,10 +47,10 @@ export function createDirectoryReadmeInjectorHook(ctx: PluginInput) {
    return sessionCaches.get(sessionID)!;
  }

-  function resolveFilePath(title: string): string | null {
-    if (!title) return null;
-    if (title.startsWith("/")) return title;
-    return resolve(ctx.directory, title);
+  function resolveFilePath(path: string): string | null {
+    if (!path) return null;
+    if (path.startsWith("/")) return path;
+    return resolve(ctx.directory, path);
  }

  function findReadmeMdUp(startDir: string): string[] {
@@ -63,39 +73,73 @@ export function createDirectoryReadmeInjectorHook(ctx: PluginInput) {
    return found.reverse();
  }

-  const toolExecuteAfter = async (
-    input: ToolExecuteInput,
+  function processFilePathForInjection(
+    filePath: string,
+    sessionID: string,
    output: ToolExecuteOutput,
-  ) => {
-    if (input.tool.toLowerCase() !== "read") return;
+  ): void {
+    const resolved = resolveFilePath(filePath);
+    if (!resolved) return;

-    const filePath = resolveFilePath(output.title);
-    if (!filePath) return;
-
-    const dir = dirname(filePath);
-    const cache = getSessionCache(input.sessionID);
+    const dir = dirname(resolved);
+    const cache = getSessionCache(sessionID);
    const readmePaths = findReadmeMdUp(dir);

-    const toInject: { path: string; content: string }[] = [];
-
    for (const readmePath of readmePaths) {
      const readmeDir = dirname(readmePath);
      if (cache.has(readmeDir)) continue;

      try {
        const content = readFileSync(readmePath, "utf-8");
-        toInject.push({ path: readmePath, content });
+        output.output += `\n\n[Project README: ${readmePath}]\n${content}`;
        cache.add(readmeDir);
      } catch {}
    }

-    if (toInject.length === 0) return;
-
-    for (const { path, content } of toInject) {
-      output.output += `\n\n[Project README: ${path}]\n${content}`;
+    saveInjectedPaths(sessionID, cache);
  }

-    saveInjectedPaths(input.sessionID, cache);
+  const toolExecuteBefore = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteBeforeOutput,
+  ) => {
+    if (input.tool.toLowerCase() !== "batch") return;
+
+    const args = output.args as { tool_calls?: BatchToolCall[] } | undefined;
+    if (!args?.tool_calls) return;
+
+    const readFilePaths: string[] = [];
+    for (const call of args.tool_calls) {
+      if (call.tool.toLowerCase() === "read" && call.parameters?.filePath) {
+        readFilePaths.push(call.parameters.filePath as string);
+      }
+    }
+
+    if (readFilePaths.length > 0) {
+      pendingBatchReads.set(input.callID, readFilePaths);
+    }
+  };
+
+  const toolExecuteAfter = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteOutput,
+  ) => {
+    const toolName = input.tool.toLowerCase();
+
+    if (toolName === "read") {
+      processFilePathForInjection(output.title, input.sessionID, output);
+      return;
+    }
+
+    if (toolName === "batch") {
+      const filePaths = pendingBatchReads.get(input.callID);
+      if (filePaths) {
+        for (const filePath of filePaths) {
+          processFilePathForInjection(filePath, input.sessionID, output);
+        }
+        pendingBatchReads.delete(input.callID);
+      }
+    }
  };

  const eventHandler = async ({ event }: EventInput) => {
@@ -120,6 +164,7 @@ export function createDirectoryReadmeInjectorHook(ctx: PluginInput) {
  };

  return {
+    "tool.execute.before": toolExecuteBefore,
    "tool.execute.after": toolExecuteAfter,
    event: eventHandler,
  };
--- a/src/hooks/rules-injector/index.ts
+++ b/src/hooks/rules-injector/index.ts
@@ -28,6 +28,15 @@ interface ToolExecuteOutput {
  metadata: unknown;
 }

+interface ToolExecuteBeforeOutput {
+  args: unknown;
+}
+
+interface BatchToolCall {
+  tool: string;
+  parameters: Record<string, unknown>;
+}
+
 interface EventInput {
  event: {
    type: string;
@@ -49,6 +58,7 @@ export function createRulesInjectorHook(ctx: PluginInput) {
    string,
    { contentHashes: Set<string>; realPaths: Set<string> }
  >();
+  const pendingBatchFiles = new Map<string, string[]>();

  function getSessionCache(sessionID: string): {
    contentHashes: Set<string>;
@@ -60,26 +70,25 @@ export function createRulesInjectorHook(ctx: PluginInput) {
    return sessionCaches.get(sessionID)!;
  }

-  function resolveFilePath(title: string): string | null {
-    if (!title) return null;
-    if (title.startsWith("/")) return title;
-    return resolve(ctx.directory, title);
+  function resolveFilePath(path: string): string | null {
+    if (!path) return null;
+    if (path.startsWith("/")) return path;
+    return resolve(ctx.directory, path);
  }

-  const toolExecuteAfter = async (
-    input: ToolExecuteInput,
+  function processFilePathForInjection(
+    filePath: string,
+    sessionID: string,
    output: ToolExecuteOutput
-  ) => {
-    if (!TRACKED_TOOLS.includes(input.tool.toLowerCase())) return;
+  ): void {
+    const resolved = resolveFilePath(filePath);
+    if (!resolved) return;

-    const filePath = resolveFilePath(output.title);
-    if (!filePath) return;
-
-    const projectRoot = findProjectRoot(filePath);
-    const cache = getSessionCache(input.sessionID);
+    const projectRoot = findProjectRoot(resolved);
+    const cache = getSessionCache(sessionID);
    const home = homedir();

-    const ruleFileCandidates = findRuleFiles(projectRoot, home, filePath);
+    const ruleFileCandidates = findRuleFiles(projectRoot, home, resolved);
    const toInject: RuleToInject[] = [];

    for (const candidate of ruleFileCandidates) {
@@ -89,7 +98,7 @@ export function createRulesInjectorHook(ctx: PluginInput) {
        const rawContent = readFileSync(candidate.path, "utf-8");
        const { metadata, body } = parseRuleFrontmatter(rawContent);

-        const matchResult = shouldApplyRule(metadata, filePath, projectRoot);
+        const matchResult = shouldApplyRule(metadata, resolved, projectRoot);
        if (!matchResult.applies) continue;

        const contentHash = createContentHash(body);
@@ -119,7 +128,58 @@ export function createRulesInjectorHook(ctx: PluginInput) {
      output.output += `\n\n[Rule: ${rule.relativePath}]\n[Match: ${rule.matchReason}]\n${rule.content}`;
    }

-    saveInjectedRules(input.sessionID, cache);
+    saveInjectedRules(sessionID, cache);
+  }
+
+  function extractFilePathFromToolCall(call: BatchToolCall): string | null {
+    const params = call.parameters;
+    return (params?.filePath ?? params?.file_path ?? params?.path) as string | null;
+  }
+
+  const toolExecuteBefore = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteBeforeOutput
+  ) => {
+    if (input.tool.toLowerCase() !== "batch") return;
+
+    const args = output.args as { tool_calls?: BatchToolCall[] } | undefined;
+    if (!args?.tool_calls) return;
+
+    const filePaths: string[] = [];
+    for (const call of args.tool_calls) {
+      if (TRACKED_TOOLS.includes(call.tool.toLowerCase())) {
+        const filePath = extractFilePathFromToolCall(call);
+        if (filePath) {
+          filePaths.push(filePath);
+        }
+      }
+    }
+
+    if (filePaths.length > 0) {
+      pendingBatchFiles.set(input.callID, filePaths);
+    }
+  };
+
+  const toolExecuteAfter = async (
+    input: ToolExecuteInput,
+    output: ToolExecuteOutput
+  ) => {
+    const toolName = input.tool.toLowerCase();
+
+    if (TRACKED_TOOLS.includes(toolName)) {
+      processFilePathForInjection(output.title, input.sessionID, output);
+      return;
+    }
+
+    if (toolName === "batch") {
+      const filePaths = pendingBatchFiles.get(input.callID);
+      if (filePaths) {
+        for (const filePath of filePaths) {
+          processFilePathForInjection(filePath, input.sessionID, output);
+        }
+        pendingBatchFiles.delete(input.callID);
+      }
+    }
  };

  const eventHandler = async ({ event }: EventInput) => {
@@ -144,6 +204,7 @@ export function createRulesInjectorHook(ctx: PluginInput) {
  };

  return {
+    "tool.execute.before": toolExecuteBefore,
    "tool.execute.after": toolExecuteAfter,
    event: eventHandler,
  };
--- a/src/index.ts
+++ b/src/index.ts
@@ -543,6 +543,9 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => {
      await claudeCodeHooks["tool.execute.before"](input, output);
      await nonInteractiveEnv?.["tool.execute.before"](input, output);
      await commentChecker?.["tool.execute.before"](input, output);
+      await directoryAgentsInjector?.["tool.execute.before"]?.(input, output);
+      await directoryReadmeInjector?.["tool.execute.before"]?.(input, output);
+      await rulesInjector?.["tool.execute.before"]?.(input, output);

      if (input.tool === "task") {
        const args = output.args as Record<string, unknown>;