Skip to content

Commit d909c09

Browse files
code-yeongyucode-yeongyu
authored
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)
1 parent 5c73f47 commit d909c09

File tree

4 files changed

+210
-56
lines changed

4 files changed

+210
-56
lines changed

‎src/hooks/directory-agents-injector/index.ts‎

Lines changed: 65 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,15 @@ interface ToolExecuteOutput {
2020
metadata: unknown;
2121
}
2222

23+
interface ToolExecuteBeforeOutput {
24+
args: unknown;
25+
}
26+
27+
interface BatchToolCall {
28+
tool: string;
29+
parameters: Record<string, unknown>;
30+
}
31+
2332
interface EventInput {
2433
event: {
2534
type: string;
@@ -29,6 +38,7 @@ interface EventInput {
2938

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

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

40-
function resolveFilePath(title: string): string | null {
41-
if (!title) return null;
42-
if (title.startsWith("/")) return title;
43-
return resolve(ctx.directory, title);
50+
function resolveFilePath(path: string): string | null {
51+
if (!path) return null;
52+
if (path.startsWith("/")) return path;
53+
return resolve(ctx.directory, path);
4454
}
4555

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

66-
const toolExecuteAfter = async (
67-
input: ToolExecuteInput,
76+
function processFilePathForInjection(
77+
filePath: string,
78+
sessionID: string,
6879
output: ToolExecuteOutput,
69-
) => {
70-
if (input.tool.toLowerCase() !== "read") return;
71-
72-
const filePath = resolveFilePath(output.title);
73-
if (!filePath) return;
80+
): void {
81+
const resolved = resolveFilePath(filePath);
82+
if (!resolved) return;
7483

75-
const dir = dirname(filePath);
76-
const cache = getSessionCache(input.sessionID);
84+
const dir = dirname(resolved);
85+
const cache = getSessionCache(sessionID);
7786
const agentsPaths = findAgentsMdUp(dir);
7887

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

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

92-
if (toInject.length === 0) return;
99+
saveInjectedPaths(sessionID, cache);
100+
}
101+
102+
const toolExecuteBefore = async (
103+
input: ToolExecuteInput,
104+
output: ToolExecuteBeforeOutput,
105+
) => {
106+
if (input.tool.toLowerCase() !== "batch") return;
107+
108+
const args = output.args as { tool_calls?: BatchToolCall[] } | undefined;
109+
if (!args?.tool_calls) return;
110+
111+
const readFilePaths: string[] = [];
112+
for (const call of args.tool_calls) {
113+
if (call.tool.toLowerCase() === "read" && call.parameters?.filePath) {
114+
readFilePaths.push(call.parameters.filePath as string);
115+
}
116+
}
117+
118+
if (readFilePaths.length > 0) {
119+
pendingBatchReads.set(input.callID, readFilePaths);
120+
}
121+
};
122+
123+
const toolExecuteAfter = async (
124+
input: ToolExecuteInput,
125+
output: ToolExecuteOutput,
126+
) => {
127+
const toolName = input.tool.toLowerCase();
93128

94-
for (const { path, content } of toInject) {
95-
output.output += `\n\n[Directory Context: ${path}]\n${content}`;
129+
if (toolName === "read") {
130+
processFilePathForInjection(output.title, input.sessionID, output);
131+
return;
96132
}
97133

98-
saveInjectedPaths(input.sessionID, cache);
134+
if (toolName === "batch") {
135+
const filePaths = pendingBatchReads.get(input.callID);
136+
if (filePaths) {
137+
for (const filePath of filePaths) {
138+
processFilePathForInjection(filePath, input.sessionID, output);
139+
}
140+
pendingBatchReads.delete(input.callID);
141+
}
142+
}
99143
};
100144

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

122166
return {
167+
"tool.execute.before": toolExecuteBefore,
123168
"tool.execute.after": toolExecuteAfter,
124169
event: eventHandler,
125170
};

‎src/hooks/directory-readme-injector/index.ts‎

Lines changed: 65 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,15 @@ interface ToolExecuteOutput {
2020
metadata: unknown;
2121
}
2222

23+
interface ToolExecuteBeforeOutput {
24+
args: unknown;
25+
}
26+
27+
interface BatchToolCall {
28+
tool: string;
29+
parameters: Record<string, unknown>;
30+
}
31+
2332
interface EventInput {
2433
event: {
2534
type: string;
@@ -29,6 +38,7 @@ interface EventInput {
2938

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

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

40-
function resolveFilePath(title: string): string | null {
41-
if (!title) return null;
42-
if (title.startsWith("/")) return title;
43-
return resolve(ctx.directory, title);
50+
function resolveFilePath(path: string): string | null {
51+
if (!path) return null;
52+
if (path.startsWith("/")) return path;
53+
return resolve(ctx.directory, path);
4454
}
4555

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

66-
const toolExecuteAfter = async (
67-
input: ToolExecuteInput,
76+
function processFilePathForInjection(
77+
filePath: string,
78+
sessionID: string,
6879
output: ToolExecuteOutput,
69-
) => {
70-
if (input.tool.toLowerCase() !== "read") return;
71-
72-
const filePath = resolveFilePath(output.title);
73-
if (!filePath) return;
80+
): void {
81+
const resolved = resolveFilePath(filePath);
82+
if (!resolved) return;
7483

75-
const dir = dirname(filePath);
76-
const cache = getSessionCache(input.sessionID);
84+
const dir = dirname(resolved);
85+
const cache = getSessionCache(sessionID);
7786
const readmePaths = findReadmeMdUp(dir);
7887

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

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

92-
if (toInject.length === 0) return;
99+
saveInjectedPaths(sessionID, cache);
100+
}
101+
102+
const toolExecuteBefore = async (
103+
input: ToolExecuteInput,
104+
output: ToolExecuteBeforeOutput,
105+
) => {
106+
if (input.tool.toLowerCase() !== "batch") return;
107+
108+
const args = output.args as { tool_calls?: BatchToolCall[] } | undefined;
109+
if (!args?.tool_calls) return;
110+
111+
const readFilePaths: string[] = [];
112+
for (const call of args.tool_calls) {
113+
if (call.tool.toLowerCase() === "read" && call.parameters?.filePath) {
114+
readFilePaths.push(call.parameters.filePath as string);
115+
}
116+
}
117+
118+
if (readFilePaths.length > 0) {
119+
pendingBatchReads.set(input.callID, readFilePaths);
120+
}
121+
};
122+
123+
const toolExecuteAfter = async (
124+
input: ToolExecuteInput,
125+
output: ToolExecuteOutput,
126+
) => {
127+
const toolName = input.tool.toLowerCase();
93128

94-
for (const { path, content } of toInject) {
95-
output.output += `\n\n[Project README: ${path}]\n${content}`;
129+
if (toolName === "read") {
130+
processFilePathForInjection(output.title, input.sessionID, output);
131+
return;
96132
}
97133

98-
saveInjectedPaths(input.sessionID, cache);
134+
if (toolName === "batch") {
135+
const filePaths = pendingBatchReads.get(input.callID);
136+
if (filePaths) {
137+
for (const filePath of filePaths) {
138+
processFilePathForInjection(filePath, input.sessionID, output);
139+
}
140+
pendingBatchReads.delete(input.callID);
141+
}
142+
}
99143
};
100144

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

122166
return {
167+
"tool.execute.before": toolExecuteBefore,
123168
"tool.execute.after": toolExecuteAfter,
124169
event: eventHandler,
125170
};

0 commit comments

Comments
 (0)