Skip to content

Commit 9d5a46f

Browse files
committed
Keep archive search on primary prompts, payloads, and the search budget
Workers have no archive plugin, so advertising archive:/// sent grep onto the filesystem. Matching the formatted URI line also made pattern archive hit every occurrence. Sit archive search after shell-guard and abort the walk so the same 10s budget applies.
1 parent 673d314 commit 9d5a46f

6 files changed

Lines changed: 223 additions & 15 deletions

‎src/agent/posix-tool-plugins.test.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,34 @@ describe("buildCorePosixToolPlugins", () => {
378378
}
379379
});
380380

381+
test("evidence archive search sits after shell-guard so grep/search inherit the 10s budget", async () => {
382+
const cwd = await mkdtemp(join(tmpdir(), "ic-posix-plugins-"));
383+
try {
384+
const gate = createPermissionGate({
385+
approvals: [],
386+
interactive: false,
387+
skipPermissions: true,
388+
reactorGated: false,
389+
cwd,
390+
});
391+
const plugins = buildCorePosixToolPlugins({
392+
cwd,
393+
permissionGate: gate,
394+
getEvidenceArchive: () => undefined,
395+
});
396+
const shellGuardIndex = findMiddlewareIndex(plugins, "[command timed out after");
397+
const archiveIndex = findMiddlewareIndex(
398+
plugins,
399+
"evidence archive is not available in this session",
400+
);
401+
expect(shellGuardIndex).toBeGreaterThanOrEqual(0);
402+
expect(archiveIndex).toBeGreaterThanOrEqual(0);
403+
expect(archiveIndex).toBeGreaterThan(shellGuardIndex);
404+
} finally {
405+
await rm(cwd, { recursive: true, force: true });
406+
}
407+
});
408+
381409
test("a real line-range edit_file call verifies as success through the wired plugin chain", async () => {
382410
const cwd = await mkdtemp(join(tmpdir(), "ic-posix-plugins-"));
383411
try {

‎src/agent/posix-tool-plugins.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,6 @@ export function buildCorePosixToolPlugins(
106106
resultTruncationPlugin(truncationOptions),
107107
toolResultSecretScrubPlugin(),
108108
pathEscapePlugin(cwd, createWorktreeRootsProvider(cwd), { allowOutside }),
109-
...(getEvidenceArchive !== undefined ? [evidenceArchiveSearchPlugin(getEvidenceArchive)] : []),
110109
evidenceArchivePathGuardPlugin(),
111110
deleteFilePlugin(cwd, { allowOutside }),
112111
toolOutputUriPlugin(),
@@ -119,6 +118,9 @@ export function buildCorePosixToolPlugins(
119118
? { getBackgroundShellRegistry }
120119
: {}),
121120
}),
121+
...(getEvidenceArchive !== undefined
122+
? [evidenceArchiveSearchPlugin(getEvidenceArchive)]
123+
: []),
122124
readFileGuardPlugin(cwd, readFileGuard),
123125
ripgrepPlugin(cwd),
124126
// Verify wraps the line-range short-circuit (composeMiddleware runs plugins

‎src/agent/prompts.ts‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -226,16 +226,16 @@ export function buildPromptDisciplineBlock(
226226

227227
const TOOL_SUMMARIES: Record<string, string> = {
228228
read_file:
229-
"read a file, tool-output:///{callId} from a prior tool result, or archive:///{occurrenceId} (prefer over cat/head/tail in the shell). Only read_file a tool-output:// URI if the truncation notice named one",
229+
"read a file or tool-output:///{callId} from a prior tool result (prefer over cat/head/tail in the shell). Only read_file a tool-output:// URI if the truncation notice named one",
230230
write_file: "create or overwrite a file (never shell redirects or heredocs)",
231231
edit_file:
232232
"make a surgical edit (exact old_string match, or start_line/end_line line-range mode; never include read_file's NNNNNN\\t line prefix; substring failures include nearby file text; prefer over sed/awk in the shell)",
233233
delete_file: "delete one file with an explicit outcome (never shell rm)",
234234
run_shell:
235235
"run a shell command (builds, tests, git; pass timeout ms to bound long commands; never to read/write/delete files, search trees, or talk to the user)",
236236
search_files:
237-
"find files by name or pattern (bounded; timeout + output caps — safer than open-ended shell find); path archive:/// lists evidence-archive refs",
238-
grep: "search file contents (bounded; timeout + output caps — safer than open-ended shell grep -r/rg); path archive:/// searches this session's evidence archive",
237+
"find files by name or pattern (bounded; timeout + output caps — safer than open-ended shell find)",
238+
grep: "search file contents (bounded; timeout + output caps — safer than open-ended shell grep -r/rg)",
239239
list_dir: "list a directory's entries (bounded listing)",
240240
lsp: "resolve symbols — goToDefinition, findReferences, hover (prefer before reading huge files)",
241241
web_search: "search the web (use instead of curl or wget)",
@@ -263,11 +263,24 @@ const TOOL_SUMMARIES: Record<string, string> = {
263263
"look up skill descriptions by capability (catalog — call directly, do not tool_search for this)",
264264
};
265265

266+
const ARCHIVE_TOOL_SUMMARIES: Partial<Record<string, string>> = {
267+
read_file:
268+
"read a file, tool-output:///{callId} from a prior tool result, or archive:///{occurrenceId} (prefer over cat/head/tail in the shell). Only read_file a tool-output:// URI if the truncation notice named one",
269+
search_files:
270+
"find files by name or pattern (bounded; timeout + output caps — safer than open-ended shell find); path archive:/// lists evidence-archive refs",
271+
grep: "search file contents (bounded; timeout + output caps — safer than open-ended shell grep -r/rg); path archive:/// searches this session's evidence archive",
272+
};
273+
266274
export function buildAvailableTools(
267275
tools: readonly string[] = CORE_TOOL_NAMES,
276+
opts: { advertiseArchive?: boolean } = {},
268277
): string {
278+
const summaries =
279+
opts.advertiseArchive === true
280+
? { ...TOOL_SUMMARIES, ...ARCHIVE_TOOL_SUMMARIES }
281+
: TOOL_SUMMARIES;
269282
const lines = tools.map(
270-
(tool) => `- ${tool}: ${TOOL_SUMMARIES[tool] ?? "available"}`,
283+
(tool) => `- ${tool}: ${summaries[tool] ?? "available"}`,
271284
);
272285
return ["Tools:", ...lines].join("\n");
273286
}
@@ -365,6 +378,7 @@ export function buildChatSystemPrompt(
365378
baseSection(baseOverride, sessionMode),
366379
buildAvailableTools(
367380
coreToolNamesForSessionMode(sessionMode, toolAvailability),
381+
{ advertiseArchive: true },
368382
),
369383
];
370384
if (skills.length > 0) sections.push(buildSkillsSection(skills));

‎src/plugins/evidence-archive-search-plugin.test.ts‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,87 @@ describe("evidenceArchiveSearchPlugin", () => {
210210
);
211211
expect(result.content).toBe("passthrough:grep");
212212
});
213+
214+
test("grep pattern archive searches payloads and metadata, not the URI line", async () => {
215+
const archive = memoryArchive("sess-uri");
216+
const occ = await archive.recordAuthorizedPayload({
217+
kind: "user_message",
218+
payload: "plain-payload-without-scheme",
219+
});
220+
const plugin = evidenceArchiveSearchPlugin(() => archive);
221+
const handler = plugin.middleware ? plugin.middleware(nextHandler) : nextHandler;
222+
223+
const uriHits = await handler(
224+
makeCall("grep", { pattern: "archive", path: "archive:///" }),
225+
new AbortController().signal,
226+
);
227+
expect(String(uriHits.content)).toContain("no matches");
228+
expect(String(uriHits.content)).not.toContain(formatArchiveRef(occ.occurrenceId));
229+
230+
const payloadHits = await handler(
231+
makeCall("grep", { pattern: "plain-payload-without-scheme", path: "archive:///" }),
232+
new AbortController().signal,
233+
);
234+
expect(String(payloadHits.content)).toContain(formatArchiveRef(occ.occurrenceId));
235+
expect(String(payloadHits.content)).toContain("plain-payload-without-scheme");
236+
});
237+
238+
test("archive grep and search_files honor abort before loading remaining payloads", async () => {
239+
const archive = memoryArchive("sess-abort");
240+
const first = await archive.recordAuthorizedPayload({
241+
kind: "user_message",
242+
payload: "abort-first",
243+
});
244+
await archive.recordAuthorizedPayload({
245+
kind: "assistant_text",
246+
payload: "abort-second",
247+
});
248+
const plugin = evidenceArchiveSearchPlugin(() => archive);
249+
const handler = plugin.middleware ? plugin.middleware(nextHandler) : nextHandler;
250+
const controller = new AbortController();
251+
const reads: string[] = [];
252+
const orig = archive.readAuthorizedPayload.bind(archive);
253+
archive.readAuthorizedPayload = async (occurrenceId) => {
254+
reads.push(occurrenceId);
255+
if (reads.length === 1) controller.abort();
256+
return orig(occurrenceId);
257+
};
258+
259+
const grepResult = await handler(
260+
makeCall("grep", { pattern: "abort-second", path: "archive:///" }),
261+
controller.signal,
262+
);
263+
expect(grepResult.isError).toBe(true);
264+
expect(String(grepResult.content).toLowerCase()).toMatch(/abort/);
265+
expect(reads).toEqual([first.occurrenceId]);
266+
267+
const searchController = new AbortController();
268+
searchController.abort();
269+
const searchResult = await handler(
270+
makeCall("search_files", { pattern: "*", path: "archive:///" }),
271+
searchController.signal,
272+
);
273+
expect(searchResult.isError).toBe(true);
274+
expect(String(searchResult.content).toLowerCase()).toMatch(/abort/);
275+
});
276+
277+
test("archive grep includes posix-style context around payload matches", async () => {
278+
const archive = memoryArchive("sess-context");
279+
const occ = await archive.recordAuthorizedPayload({
280+
kind: "tool_result",
281+
payload: "alpha\nbeta-hit\ngamma",
282+
});
283+
const plugin = evidenceArchiveSearchPlugin(() => archive);
284+
const handler = plugin.middleware ? plugin.middleware(nextHandler) : nextHandler;
285+
const hits = await handler(
286+
makeCall("grep", { pattern: "beta-hit", path: "archive:///", context: 1 }),
287+
new AbortController().signal,
288+
);
289+
const ref = formatArchiveRef(occ.occurrenceId);
290+
expect(String(hits.content)).toContain(`${ref}-1-alpha`);
291+
expect(String(hits.content)).toContain(`${ref}:2:beta-hit`);
292+
expect(String(hits.content)).toContain(`${ref}-3-gamma`);
293+
});
213294
});
214295

215296
describe("createAgentToolset archive mount", () => {
@@ -233,6 +314,14 @@ describe("createAgentToolset archive mount", () => {
233314
.currentDefinitions()
234315
.find((d) => d.name === "read_file");
235316
expect(workerRead?.description ?? "").not.toContain("archive:///");
317+
const workerGrep = worker.dynamicRunner.currentDefinitions().find((d) => d.name === "grep");
318+
expect(workerGrep?.description ?? "").not.toContain("archive:///");
319+
expect(JSON.stringify(workerGrep?.inputSchema ?? {})).not.toContain("archive:///");
320+
const workerSearch = worker.dynamicRunner
321+
.currentDefinitions()
322+
.find((d) => d.name === "search_files");
323+
expect(workerSearch?.description ?? "").not.toContain("archive:///");
324+
expect(JSON.stringify(workerSearch?.inputSchema ?? {})).not.toContain("archive:///");
236325
await worker.dispose();
237326

238327
const primary = await createAgentToolset({

‎src/plugins/evidence-archive-search-plugin.ts‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,13 @@ export function evidenceArchiveSearchPlugin(
174174
const maxResults = num(call.arguments.max_results) ?? SEARCH_DEFAULT_MAX;
175175
return {
176176
callId: call.id,
177-
content: await searchArchiveFiles(archive, pattern, target.occurrenceId, maxResults),
177+
content: await searchArchiveFiles(
178+
archive,
179+
pattern,
180+
target.occurrenceId,
181+
maxResults,
182+
signal,
183+
),
178184
};
179185
}
180186
const pattern = str(call.arguments.pattern);
@@ -187,9 +193,19 @@ export function evidenceArchiveSearchPlugin(
187193
}
188194
const maxResults = num(call.arguments.max_results) ?? GREP_DEFAULT_MAX;
189195
const glob = str(call.arguments.glob);
196+
const contextArg = num(call.arguments.context);
197+
const context = contextArg !== undefined && contextArg > 0 ? Math.floor(contextArg) : 0;
190198
return {
191199
callId: call.id,
192-
content: await grepArchive(archive, pattern, target.occurrenceId, maxResults, glob),
200+
content: await grepArchive(
201+
archive,
202+
pattern,
203+
target.occurrenceId,
204+
maxResults,
205+
glob,
206+
context,
207+
signal,
208+
),
193209
};
194210
} catch (err) {
195211
return {
@@ -232,10 +248,13 @@ async function searchArchiveFiles(
232248
pattern: string,
233249
occurrenceId: string | undefined,
234250
maxResults: number,
251+
signal: AbortSignal,
235252
): Promise<string> {
253+
signal.throwIfAborted();
236254
const occurrences = await selectOccurrences(archive, occurrenceId);
237255
const hits: string[] = [];
238256
for (const occ of occurrences) {
257+
signal.throwIfAborted();
239258
if (hits.length >= maxResults) break;
240259
if (!matchesArchiveName(pattern, occ)) continue;
241260
hits.push(formatArchiveRef(occ.occurrenceId));
@@ -246,13 +265,58 @@ async function searchArchiveFiles(
246265
return hits.join("\n");
247266
}
248267

268+
function grepPayloadHits(
269+
ref: string,
270+
lines: string[],
271+
regex: RegExp,
272+
context: number,
273+
remaining: number,
274+
): string[] {
275+
if (remaining <= 0) return [];
276+
const matchLines: number[] = [];
277+
for (let i = 0; i < lines.length; i++) {
278+
if (regex.test(lines[i] ?? "")) matchLines.push(i);
279+
if (matchLines.length >= remaining) break;
280+
}
281+
if (matchLines.length === 0) return [];
282+
if (context <= 0) {
283+
return matchLines.map((i) => `${ref}:${i + 1}:${lines[i] ?? ""}`);
284+
}
285+
const matchSet = new Set(matchLines);
286+
const ranges: { start: number; end: number }[] = [];
287+
for (const i of matchLines) {
288+
const start = Math.max(0, i - context);
289+
const end = Math.min(lines.length - 1, i + context);
290+
const prev = ranges[ranges.length - 1];
291+
if (prev !== undefined && start <= prev.end + 1) {
292+
prev.end = end;
293+
} else {
294+
ranges.push({ start, end });
295+
}
296+
}
297+
const out: string[] = [];
298+
let first = true;
299+
for (const range of ranges) {
300+
if (!first) out.push("--");
301+
first = false;
302+
for (let i = range.start; i <= range.end; i++) {
303+
const sep = matchSet.has(i) ? ":" : "-";
304+
out.push(`${ref}${sep}${i + 1}${sep}${lines[i] ?? ""}`);
305+
}
306+
}
307+
return out;
308+
}
309+
249310
async function grepArchive(
250311
archive: CompactionArchive,
251312
pattern: string,
252313
occurrenceId: string | undefined,
253314
maxResults: number,
254315
glob: string | undefined,
316+
context: number,
317+
signal: AbortSignal,
255318
): Promise<string> {
319+
signal.throwIfAborted();
256320
let regex: RegExp;
257321
try {
258322
regex = new RegExp(pattern);
@@ -262,26 +326,26 @@ async function grepArchive(
262326
const occurrences = await selectOccurrences(archive, occurrenceId);
263327
const hits: string[] = [];
264328
for (const occ of occurrences) {
329+
signal.throwIfAborted();
265330
if (hits.length >= maxResults) break;
266331
if (glob !== undefined && !matchesArchiveName(glob, occ)) continue;
267332
const ref = formatArchiveRef(occ.occurrenceId);
268-
if (regex.test(metadataBlob(occ)) || regex.test(formatHit(occ))) {
333+
if (regex.test(metadataBlob(occ))) {
269334
hits.push(`${ref}:1:${formatHit(occ)}`);
270335
continue;
271336
}
272337
if (occ.gap === true) continue;
273338
let payload: string;
274339
try {
275340
payload = await archive.readAuthorizedPayload(occ.occurrenceId);
276-
} catch {
341+
} catch (err) {
342+
if (signal.aborted) throw err;
277343
continue;
278344
}
279-
const lines = payload.split("\n");
280-
for (let i = 0; i < lines.length; i++) {
281-
if (hits.length >= maxResults) break;
282-
const line = lines[i] ?? "";
283-
if (regex.test(line)) hits.push(`${ref}:${i + 1}:${line}`);
284-
}
345+
signal.throwIfAborted();
346+
hits.push(
347+
...grepPayloadHits(ref, payload.split("\n"), regex, context, maxResults - hits.length),
348+
);
285349
}
286350
if (hits.length === 0) return `no matches for /${pattern}/`;
287351
return hits.join("\n");

‎src/prompts.test.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,17 @@ test("sub-agent prompt does not advertise tool_search (it gets the full toolset)
432432
expect(prompt).toContain("your full toolset");
433433
});
434434

435+
test("worker prompt does not advertise archive:///; primary chat prompt does", () => {
436+
const worker = buildSubAgentSystemPrompt();
437+
expect(worker).not.toContain("archive:///");
438+
const primary = buildChatSystemPrompt();
439+
expect(primary).toContain("archive:///");
440+
expect(buildAvailableTools(["read_file", "grep", "search_files"])).not.toContain("archive:///");
441+
expect(
442+
buildAvailableTools(["read_file", "grep", "search_files"], { advertiseArchive: true }),
443+
).toContain("archive:///");
444+
});
445+
435446
// Pins the appendix-last invariant for JS-plugin agents: regardless of how the
436447
// systemPromptRole is sourced (data-only markdown vs. a JS plugin's
437448
// `agentPlugin.agents[i].systemPromptRole`), `buildSubAgentSystemPrompt` is

0 commit comments

Comments
 (0)