fix: drop replay tool calls outside allowlist

This commit is contained in:
Josh Lehman
2026-03-20 07:10:46 -07:00
parent 1da3ca52f5
commit cfbdb1ffce
2 changed files with 52 additions and 5 deletions

View File

@@ -1040,6 +1040,46 @@ describe("wrapStreamFnSanitizeMalformedToolCalls", () => {
},
]);
});
it("drops replayed tool calls that are no longer allowlisted", async () => {
const messages = [
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "write", arguments: {} }],
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "write",
content: [{ type: "text", text: "stale result" }],
isError: false,
},
{
role: "user",
content: [{ type: "text", text: "retry" }],
},
];
const baseFn = vi.fn((_model, _context) =>
createFakeStream({ events: [], resultMessage: { role: "assistant", content: [] } }),
);
const wrapped = wrapStreamFnSanitizeMalformedToolCalls(baseFn as never, new Set(["read"]));
const stream = wrapped({} as never, { messages } as never, {} as never) as
| FakeWrappedStream
| Promise<FakeWrappedStream>;
await Promise.resolve(stream);
expect(baseFn).toHaveBeenCalledTimes(1);
const seenContext = baseFn.mock.calls[0]?.[1] as {
messages: Array<{ role?: string }>;
};
expect(seenContext.messages).toEqual([
{
role: "user",
content: [{ type: "text", text: "retry" }],
},
]);
});
});
describe("wrapStreamFnRepairMalformedToolCallArguments", () => {

View File

@@ -689,7 +689,13 @@ function resolveReplayToolCallName(
if (!trimmed || trimmed.length > REPLAY_TOOL_CALL_NAME_MAX_CHARS || /\s/.test(trimmed)) {
return null;
}
return trimmed;
if (!allowedToolNames || allowedToolNames.size === 0) {
return trimmed;
}
return (
resolveExactAllowedToolName(trimmed, allowedToolNames) ??
resolveStructuredAllowedToolName(trimmed, allowedToolNames)
);
}
function sanitizeReplayToolCallInputs(
@@ -717,22 +723,23 @@ function sanitizeReplayToolCallInputs(
nextContent.push(block);
continue;
}
const replayBlock = block as ReplayToolCallBlock;
if (!replayToolCallHasInput(block) || !replayToolCallNonEmptyString(block.id)) {
if (!replayToolCallHasInput(replayBlock) || !replayToolCallNonEmptyString(replayBlock.id)) {
changed = true;
messageChanged = true;
continue;
}
const rawName = typeof block.name === "string" ? block.name : "";
const resolvedName = resolveReplayToolCallName(rawName, block.id, allowedToolNames);
const rawName = typeof replayBlock.name === "string" ? replayBlock.name : "";
const resolvedName = resolveReplayToolCallName(rawName, replayBlock.id, allowedToolNames);
if (!resolvedName) {
changed = true;
messageChanged = true;
continue;
}
if (block.name !== resolvedName) {
if (replayBlock.name !== resolvedName) {
nextContent.push({ ...(block as object), name: resolvedName } as typeof block);
changed = true;
messageChanged = true;