Skip to content

Commit 28c8f4e

Browse files
TheGreatAxiosTheGreatAxios
committed
Keep live tool success when archive write fails
1 parent 0ee6a94 commit 28c8f4e

4 files changed

Lines changed: 200 additions & 18 deletions

File tree

‎src/agent/archive-tool-result.test.ts‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,108 @@ describe("withAuthorizedArchiveResult", () => {
7878
});
7979
expect(withAuthorizedArchiveResult(inner, undefined)).toBe(inner);
8080
});
81+
82+
test("returns live success when archive write throws after the tool succeeds", async () => {
83+
const archive = {
84+
recordAuthorizedPayload: async () => {
85+
throw new Error("disk full");
86+
},
87+
} as unknown as CompactionArchive;
88+
const tool = withAuthorizedArchiveResult(
89+
stringTool({
90+
definition: {
91+
name: "manage_tasks",
92+
description: "tasks",
93+
inputSchema: { type: "object", properties: {} },
94+
},
95+
handler: async () => "listed 1 task",
96+
}),
97+
() => archive,
98+
);
99+
expect(tool.kind).toBe("full");
100+
if (tool.kind !== "full") return;
101+
const result = await tool.handler(
102+
{
103+
id: "call-mt",
104+
name: "manage_tasks",
105+
arguments: { action: "create", tasks: [] },
106+
},
107+
new AbortController().signal,
108+
);
109+
expect(result.content).toBe("listed 1 task");
110+
expect(result.isError).not.toBe(true);
111+
});
112+
113+
test("abort after the string tool fulfills does not rewrite success as isError", async () => {
114+
const archive = memoryArchive("sess-abort-after");
115+
let resolveInner!: (value: string) => void;
116+
const innerPending = new Promise<string>((resolve) => {
117+
resolveInner = resolve;
118+
});
119+
const tool = withAuthorizedArchiveResult(
120+
stringTool({
121+
definition: {
122+
name: "manage_tasks",
123+
description: "tasks",
124+
inputSchema: { type: "object", properties: {} },
125+
},
126+
handler: async () => innerPending,
127+
}),
128+
() => archive,
129+
);
130+
expect(tool.kind).toBe("full");
131+
if (tool.kind !== "full") return;
132+
const ac = new AbortController();
133+
const resultPromise = tool.handler(
134+
{
135+
id: "call-abort-after",
136+
name: "manage_tasks",
137+
arguments: {},
138+
},
139+
ac.signal,
140+
);
141+
await Promise.resolve();
142+
resolveInner("listed 1 task");
143+
ac.abort();
144+
const result = await resultPromise;
145+
expect(result.content).toBe("listed 1 task");
146+
expect(result.isError).not.toBe(true);
147+
const occs = await archive.listOccurrences();
148+
const hits = occs.filter(
149+
(occ) => occ.kind === "tool_result" && occ.callId === "call-abort-after",
150+
);
151+
expect(hits).toHaveLength(1);
152+
expect(hits[0]?.provenance).toBe("agent:post-policy");
153+
});
154+
155+
test("already-aborted signal does not return live success", async () => {
156+
const archive = memoryArchive("sess-abort-entry");
157+
const tool = withAuthorizedArchiveResult(
158+
stringTool({
159+
definition: {
160+
name: "manage_tasks",
161+
description: "tasks",
162+
inputSchema: { type: "object", properties: {} },
163+
},
164+
handler: async () => "listed 1 task",
165+
}),
166+
() => archive,
167+
);
168+
expect(tool.kind).toBe("full");
169+
if (tool.kind !== "full") return;
170+
const ac = new AbortController();
171+
ac.abort();
172+
const result = await tool.handler(
173+
{
174+
id: "call-abort-entry",
175+
name: "manage_tasks",
176+
arguments: {},
177+
},
178+
ac.signal,
179+
);
180+
expect(result.content).toBe("aborted");
181+
expect(result.isError).toBe(true);
182+
});
81183
});
82184

83185
describe("createAgentToolset authorized result capture", () => {

‎src/agent/archive-tool-result.ts‎

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,47 @@ async function recordAuthorizedToolResult(
1212
): Promise<void> {
1313
const archive = getArchive();
1414
if (archive === undefined) return;
15-
await archive.recordAuthorizedPayload({
16-
kind: "tool_result",
17-
payload: applyRecordingPolicyToValue(content),
18-
callId,
19-
provenance: isError ? "agent:error" : "agent:post-policy",
15+
try {
16+
await archive.recordAuthorizedPayload({
17+
kind: "tool_result",
18+
payload: applyRecordingPolicyToValue(content),
19+
callId,
20+
provenance: isError ? "agent:error" : "agent:post-policy",
21+
});
22+
} catch {
23+
// Archive write failures fail compact later; the live tool result stands.
24+
}
25+
}
26+
27+
type AbortOrPending<T> =
28+
| { status: "aborted" }
29+
| { status: "fulfilled"; value: T }
30+
| { status: "rejected"; error: unknown };
31+
32+
async function waitUntilAbortOrPending<T>(
33+
signal: AbortSignal,
34+
pending: Promise<T>,
35+
): Promise<AbortOrPending<T>> {
36+
const settled: Promise<AbortOrPending<T>> = pending.then(
37+
(value) => ({ status: "fulfilled", value }),
38+
(error) => ({ status: "rejected", error }),
39+
);
40+
if (signal.aborted) {
41+
void settled;
42+
return { status: "aborted" };
43+
}
44+
const aborted = new Promise<AbortOrPending<T>>((resolve) => {
45+
signal.addEventListener("abort", () => resolve({ status: "aborted" }), {
46+
once: true,
47+
});
2048
});
49+
const winner = await Promise.race([settled, aborted]);
50+
if (winner.status !== "aborted") return winner;
51+
// Abort won the race; if pending already fulfilled, keep that success.
52+
return Promise.race([
53+
settled,
54+
Promise.resolve({ status: "aborted" as const }),
55+
]);
2156
}
2257

2358
/** Record authorized results for non-posix AgentTools. Posix tools already record via resultTruncationPlugin. */
@@ -32,15 +67,34 @@ export function withAuthorizedArchiveResult(
3267
kind: "full",
3368
definition: tool.definition,
3469
handler: async (call, signal) => {
35-
try {
36-
const content = await inner(call.arguments, signal);
37-
await recordAuthorizedToolResult(getArchive, call.id, content, false);
38-
return { callId: call.id, content };
39-
} catch (err) {
40-
const content = err instanceof Error ? err.message : String(err);
70+
const outcome = await waitUntilAbortOrPending(
71+
signal,
72+
inner(call.arguments, signal),
73+
);
74+
if (outcome.status === "aborted") {
75+
await recordAuthorizedToolResult(
76+
getArchive,
77+
call.id,
78+
"aborted",
79+
true,
80+
);
81+
return { callId: call.id, content: "aborted", isError: true };
82+
}
83+
if (outcome.status === "rejected") {
84+
const content =
85+
outcome.error instanceof Error
86+
? outcome.error.message
87+
: String(outcome.error);
4188
await recordAuthorizedToolResult(getArchive, call.id, content, true);
4289
return { callId: call.id, content, isError: true };
4390
}
91+
await recordAuthorizedToolResult(
92+
getArchive,
93+
call.id,
94+
outcome.value,
95+
false,
96+
);
97+
return { callId: call.id, content: outcome.value };
4498
},
4599
};
46100
}

‎src/plugins/result-truncation-plugin.test.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,28 @@ describe("resultTruncationPlugin", () => {
464464
},
465465
]);
466466
});
467+
468+
test("returns the posix result when archive write throws after next()", async () => {
469+
const archive = {
470+
recordAuthorizedPayload: async () => {
471+
throw new Error("disk full");
472+
},
473+
} as unknown as CompactionArchive;
474+
const plugin = resultTruncationPlugin({
475+
getEvidenceArchive: () => archive,
476+
});
477+
if (plugin.middleware === undefined) throw new Error("expected middleware");
478+
const middleware = plugin.middleware(async (call) => ({
479+
callId: call.id,
480+
content: "hello from posix",
481+
}));
482+
const result = await middleware(
483+
{ id: "call-ok", name: "list_dir", arguments: { path: "." } },
484+
new AbortController().signal,
485+
);
486+
expect(result.content).toBe("hello from posix");
487+
expect(result.isError).not.toBe(true);
488+
});
467489
});
468490

469491
describe("scrub-before-spill", () => {

‎src/plugins/result-truncation-plugin.ts‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -237,13 +237,17 @@ async function archiveAuthorizedResult(
237237
isError: boolean | undefined,
238238
): Promise<void> {
239239
if (archive === undefined) return;
240-
await archive.recordAuthorizedPayload({
241-
kind: "tool_result",
242-
payload: content,
243-
callId,
244-
provenance:
245-
isError === true ? "posix:error" : "posix:post-policy-pre-truncation",
246-
});
240+
try {
241+
await archive.recordAuthorizedPayload({
242+
kind: "tool_result",
243+
payload: content,
244+
callId,
245+
provenance:
246+
isError === true ? "posix:error" : "posix:post-policy-pre-truncation",
247+
});
248+
} catch {
249+
// Archive write failures fail compact later; the live posix result stands.
250+
}
247251
}
248252

249253
export function resultTruncationPlugin(

0 commit comments

Comments
 (0)