From 91465f620b237e2a2c511cc40f3e987767c05028 Mon Sep 17 00:00:00 2001 From: Shakker Date: Sun, 12 Apr 2026 06:52:38 +0100 Subject: [PATCH] fix: reserve preserved signed replay ids by owner --- src/agents/tool-call-id.test.ts | 56 +++++++++++++++++++++++++++++++-- src/agents/tool-call-id.ts | 35 +++++++++++---------- 2 files changed, 71 insertions(+), 20 deletions(-) diff --git a/src/agents/tool-call-id.test.ts b/src/agents/tool-call-id.test.ts index cf8864ed26a..6bf7ced8760 100644 --- a/src/agents/tool-call-id.test.ts +++ b/src/agents/tool-call-id.test.ts @@ -316,9 +316,10 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => { }); expect(out).toBe(input); - expect(((out[0] as Extract).content?.[1] as { id?: string }).id).toBe( - "call_1", - ); + expect( + ((out[0] as Extract).content?.[1] as { id?: string }) + .id, + ).toBe("call_1"); expect((out[1] as Extract).toolCallId).toBe("call_1"); }); @@ -368,6 +369,55 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => { expect((out[3] as Extract).toolCallId).toBe("call1"); }); + it("rewrites later signed turns when an earlier signed turn already owns the raw id", () => { + const input = castAgentMessages([ + { + role: "assistant", + content: [ + { type: "thinking", thinking: "internal", thinkingSignature: "sig_1" }, + { type: "toolCall", id: "call1", name: "read", arguments: {} }, + ], + }, + { + role: "toolResult", + toolCallId: "call1", + toolName: "read", + content: [{ type: "text", text: "first" }], + }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "internal", thinkingSignature: "sig_2" }, + { type: "toolCall", id: "call1", name: "read", arguments: {} }, + ], + }, + { + role: "toolResult", + toolCallId: "call1", + toolName: "read", + content: [{ type: "text", text: "second" }], + }, + ]); + + const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict", { + preserveReplaySafeThinkingToolCallIds: true, + allowedToolNames: ["read"], + }); + + expect(out).not.toBe(input); + const firstAssistant = out[0] as Extract; + const secondAssistant = out[2] as Extract; + const firstToolCall = firstAssistant.content?.[1] as { id?: string }; + const secondToolCall = secondAssistant.content?.[1] as { id?: string }; + expect(firstToolCall.id).toBe("call1"); + expect(secondToolCall.id).not.toBe("call1"); + expect(secondToolCall.id).not.toBe(firstToolCall.id); + expect((out[1] as Extract).toolCallId).toBe("call1"); + expect((out[3] as Extract).toolCallId).toBe( + secondToolCall.id, + ); + }); + it("avoids collisions with alphanumeric-only suffixes", () => { const input = buildDuplicateIdCollisionInput(); diff --git a/src/agents/tool-call-id.ts b/src/agents/tool-call-id.ts index 752bb323128..0d2d5b2a9a0 100644 --- a/src/agents/tool-call-id.ts +++ b/src/agents/tool-call-id.ts @@ -144,7 +144,7 @@ function isRedactedSessionsSpawnAttachment(item: unknown): boolean { if (!(SESSIONS_SPAWN_ATTACHMENT_METADATA_KEYS as readonly string[]).includes(key)) { return false; } - if (typeof attachment[key] !== "string" || (attachment[key] as string).trim().length === 0) { + if (typeof attachment[key] !== "string" || attachment[key].trim().length === 0) { return false; } } @@ -218,10 +218,7 @@ function isReplaySafeThinkingAssistantMessage( continue; } const typedBlock = block as ReplaySafeToolCallBlock; - if ( - typeof typedBlock.type !== "string" || - !TOOL_CALL_TYPES.has(typedBlock.type) - ) { + if (typeof typedBlock.type !== "string" || !TOOL_CALL_TYPES.has(typedBlock.type)) { continue; } sawToolCall = true; @@ -243,21 +240,28 @@ function isReplaySafeThinkingAssistantMessage( function collectReplaySafeThinkingToolIds( messages: AgentMessage[], allowedToolNames: Set | null, -): Set { +): { reservedIds: Set; preservedIndexes: Set } { const reserved = new Set(); - for (const message of messages) { + const preservedIndexes = new Set(); + for (let index = 0; index < messages.length; index += 1) { + const message = messages[index]; if (!message || typeof message !== "object" || message.role !== "assistant") { continue; } - const assistant = message as Extract; + const assistant = message; if (!isReplaySafeThinkingAssistantMessage(assistant, allowedToolNames)) { continue; } - for (const toolCall of extractToolCallsFromAssistant(assistant)) { + const toolCalls = extractToolCallsFromAssistant(assistant); + if (toolCalls.some((toolCall) => reserved.has(toolCall.id))) { + continue; + } + preservedIndexes.add(index); + for (const toolCall of toolCalls) { reserved.add(toolCall.id); } } - return reserved; + return { reservedIds: reserved, preservedIndexes }; } export function isValidCloudCodeAssistToolId(id: string, mode: ToolCallIdMode = "strict"): boolean { @@ -505,27 +509,24 @@ export function sanitizeToolCallIdsForCloudCodeAssist( const allowedToolNames = normalizeAllowedToolNames(options?.allowedToolNames); const preserveReplaySafeThinkingToolCallIds = options?.preserveReplaySafeThinkingToolCallIds === true; - const reservedIds = preserveReplaySafeThinkingToolCallIds + const replaySafeThinking = preserveReplaySafeThinkingToolCallIds ? collectReplaySafeThinkingToolIds(messages, allowedToolNames) : undefined; const { resolveAssistantId, resolveToolResultId, preserveAssistantId } = createOccurrenceAwareResolver(mode, { ...options, - reservedIds, + reservedIds: replaySafeThinking?.reservedIds, }); let changed = false; - const out = messages.map((msg) => { + const out = messages.map((msg, index) => { if (!msg || typeof msg !== "object") { return msg; } const role = (msg as { role?: unknown }).role; if (role === "assistant") { const assistant = msg as Extract; - if ( - preserveReplaySafeThinkingToolCallIds && - isReplaySafeThinkingAssistantMessage(assistant, allowedToolNames) - ) { + if (replaySafeThinking?.preservedIndexes.has(index)) { for (const toolCall of extractToolCallsFromAssistant(assistant)) { preserveAssistantId(toolCall.id); }