Quelle  tool-call-id.test.ts

import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
import {
  isValidCloudCodeAssistToolId,
  sanitizeToolCallId,
  sanitizeToolCallIdsForCloudCodeAssist,
} from "./tool-call-id.js";

const buildDuplicateIdCollisionInput = () =>
  castAgentMessages([
    {
      role: "assistant",
      content: [
        { type: "toolCall", id: "call_a|b", name: "read", arguments: {} },
        { type: "toolCall", id: "call_a:b", name: "read", arguments: {} },
      ],
    },
    {
      role: "toolResult",
      toolCallId: "call_a|b",
      toolName: "read",
      content: [{ type: "text", text: "one" }],
    },
    {
      role: "toolResult",
      toolCallId: "call_a:b",
      toolName: "read",
      content: [{ type: "text", text: "two" }],
    },
  ]);

const readToolCall = (id: string) => ({
  type: "toolCall" as const,
  id,
  name: "read",
  arguments: {},
});

const buildToolResult = (params: {
  toolCallId: string;
  text: string;
  toolName?: string;
  toolUseId?: string;
}) => ({
  role: "toolResult" as const,
  toolCallId: params.toolCallId,
  ...(params.toolUseId ? { toolUseId: params.toolUseId } : {}),
  toolName: params.toolName ?? "read",
  content: [{ type: "text" as const, text: params.text }],
});

const signedReadAssistant = (signature: string, id: string) => ({
  role: "assistant" as const,
  content: [
    { type: "thinking" as const, thinking: "internal", thinkingSignature: signature },
    readToolCall(id),
  ],
});

const buildRepeatedEditIdInput = (params: { includeToolUseId?: boolean } = {}) =>
  castAgentMessages([
    {
      role: "assistant",
      content: [
        { type: "toolCall", id: "edit:22", name: "edit", arguments: {} },
        { type: "toolCall", id: "edit:22", name: "edit", arguments: {} },
      ],
    },
    buildToolResult({
      toolCallId: "edit:22",
      toolName: "edit",
      ...(params.includeToolUseId ? { toolUseId: "edit:22" } : {}),
      text: "one",
    }),
    buildToolResult({
      toolCallId: "edit:22",
      toolName: "edit",
      ...(params.includeToolUseId ? { toolUseId: "edit:22" } : {}),
      text: "two",
    }),
  ]);

const buildRepeatedRawIdInput = () => buildRepeatedEditIdInput();

const buildRepeatedSharedToolResultIdInput = () =>
  buildRepeatedEditIdInput({ includeToolUseId: true });

function expectCollisionIdsRemainDistinct(
  out: AgentMessage[],
  mode: "strict" | "strict9",
): { aId: string; bId: string } {
  const assistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
  const a = assistant.content?.[0] as { id?: string };
  const b = assistant.content?.[1] as { id?: string };
  expect(typeof a.id).toBe("string");
  expect(typeof b.id).toBe("string");
  expect(a.id).not.toBe(b.id);
  expect(isValidCloudCodeAssistToolId(a.id as string, mode)).toBe(true);
  expect(isValidCloudCodeAssistToolId(b.id as string, mode)).toBe(true);

  const r1 = out[1] as Extract<AgentMessage, { role: "toolResult" }>;
  const r2 = out[2] as Extract<AgentMessage, { role: "toolResult" }>;
  expect(r1.toolCallId).toBe(a.id);
  expect(r2.toolCallId).toBe(b.id);
  return { aId: a.id as string, bId: b.id as string };
}

function expectSingleToolCallRewrite(
  out: AgentMessage[],
  expectedId: string,
  mode: "strict" | "strict9",
): void {
  const assistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
  const toolCall = assistant.content?.[0] as { id?: string };
  expect(toolCall.id).toBe(expectedId);
  expect(isValidCloudCodeAssistToolId(toolCall.id as string, mode)).toBe(true);

  const result = out[1] as Extract<AgentMessage, { role: "toolResult" }>;
  expect(result.toolCallId).toBe(toolCall.id);
}

function expectToolUseIdsFollowDistinctToolCallIds(
  out: AgentMessage[],
  mode: "strict" | "strict9",
): { aId: string; bId: string } {
  const ids = expectCollisionIdsRemainDistinct(out, mode);
  const r1 = out[1] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string };
  const r2 = out[2] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string };
  expect(r1.toolUseId).toBe(ids.aId);
  expect(r2.toolUseId).toBe(ids.bId);
  return ids;
}

function expectStrict9IdLengths(ids: { aId: string; bId: string }) {
  expect(ids.aId.length).toBe(9);
  expect(ids.bId.length).toBe(9);
}

function expectDistinctStrict9Ids(out: AgentMessage[], input: AgentMessage[]) {
  expect(out).not.toBe(input);
  const ids = expectCollisionIdsRemainDistinct(out, "strict9");
  expectStrict9IdLengths(ids);
}

function expectReplaySafeSignedTurnOwnership(params: {
  input: AgentMessage[];
  preservedTurn: "first" | "second";
  firstToolCallIndex: number;
}) {
  const out = sanitizeToolCallIdsForCloudCodeAssist(params.input, "strict", {
    preserveReplaySafeThinkingToolCallIds: true,
    allowedToolNames: ["read"],
  });

  expect(out).not.toBe(params.input);
  const firstAssistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
  const secondAssistant = out[2] as Extract<AgentMessage, { role: "assistant" }>;
  const firstToolCall = firstAssistant.content?.[params.firstToolCallIndex] as { id?: string };
  const secondToolCall = secondAssistant.content?.[1] as { id?: string };

  if (params.preservedTurn === "first") {
    expect(firstToolCall.id).toBe("call1");
    expect(secondToolCall.id).not.toBe("call1");
    expect((out[1] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe("call1");
    expect((out[3] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe(
      secondToolCall.id,
    );
  } else {
    expect(firstToolCall.id).not.toBe("call1");
    expect(secondToolCall.id).toBe("call1");
    expect((out[1] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe(
      firstToolCall.id,
    );
    expect((out[3] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe("call1");
  }

  expect(firstToolCall.id).not.toBe(secondToolCall.id);
}

describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
  describe("strict mode (default)", () => {
    it("is a no-op for already-valid non-colliding IDs", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "call1", name: "read", arguments: {} }],
        },
        {
          role: "toolResult",
          toolCallId: "call1",
          toolName: "read",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      expect(out).toBe(input);
    });

    it("strips non-alphanumeric characters from tool call IDs", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "call|item:123", name: "read", arguments: {} }],
        },
        {
          role: "toolResult",
          toolCallId: "call|item:123",
          toolName: "read",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      expect(out).not.toBe(input);
      // Strict mode strips all non-alphanumeric characters
      expectSingleToolCallRewrite(out, "callitem123", "strict");
    });

    it("avoids collisions when sanitization would produce duplicate IDs", () => {
      const input = buildDuplicateIdCollisionInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      expect(out).not.toBe(input);
      expectCollisionIdsRemainDistinct(out, "strict");
    });

    it("reuses one rewritten id when a tool result carries matching toolCallId and toolUseId", () => {
      const input = buildRepeatedSharedToolResultIdInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      expect(out).not.toBe(input);
      expectToolUseIdsFollowDistinctToolCallIds(out, "strict");
    });

    it("assigns distinct IDs when identical raw tool call ids repeat", () => {
      const input = buildRepeatedRawIdInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      expect(out).not.toBe(input);
      expectCollisionIdsRemainDistinct(out, "strict");
    });

    it("caps tool call IDs at 40 chars while preserving uniqueness", () => {
      const longA = `call_${"a".repeat(60)}`;
      const longB = `call_${"a".repeat(59)}b`;
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: longA, name: "read", arguments: {} },
            { type: "toolCall", id: longB, name: "read", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: longA,
          toolName: "read",
          content: [{ type: "text", text: "one" }],
        },
        {
          role: "toolResult",
          toolCallId: longB,
          toolName: "read",
          content: [{ type: "text", text: "two" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input);
      const { aId, bId } = expectCollisionIdsRemainDistinct(out, "strict");
      expect(aId.length).toBeLessThanOrEqual(40);
      expect(bId.length).toBeLessThanOrEqual(40);
    });
  });

  describe("strict mode (alphanumeric only)", () => {
    it("strips underscores and hyphens from tool call IDs", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            {
              type: "toolCall",
              id: "plugin_login_1768799841527_1",
              name: "login",
              arguments: {},
            },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "plugin_login_1768799841527_1",
          toolName: "login",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).not.toBe(input);
      // Strict mode strips all non-alphanumeric characters
      expectSingleToolCallRewrite(out, "pluginlogin17687998415271", "strict");
    });

    it("preserves native anthropic ids while sanitizing mixed-provider ids when requested", () => {
      const nativeId = "toolu_01ABCDEF1234567890";
      const nonNativeId = "call_123|fc_123";
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            { type: "toolUse", id: nativeId, name: "read", input: { path: "IDENTITY.md" } },
            { type: "toolUse", id: nonNativeId, name: "read", input: { path: "README.md" } },
          ],
        },
        {
          role: "toolResult",
          toolCallId: nativeId,
          toolUseId: nativeId,
          toolName: "read",
          content: [{ type: "text", text: "identity" }],
        },
        {
          role: "toolResult",
          toolCallId: nonNativeId,
          toolUseId: nonNativeId,
          toolName: "read",
          content: [{ type: "text", text: "readme" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict", {
        preserveNativeAnthropicToolUseIds: true,
      });

      expect(out).not.toBe(input);
      expect((out[0] as Extract<AgentMessage, { role: "assistant" }>).content).toEqual([
        { type: "toolUse", id: nativeId, name: "read", input: { path: "IDENTITY.md" } },
        { type: "toolUse", id: "call123fc123", name: "read", input: { path: "README.md" } },
      ]);
      expect(
        (out[1] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string })
          .toolCallId,
      ).toBe(nativeId);
      expect(
        (out[1] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string })
          .toolUseId,
      ).toBe(nativeId);
      expect(
        (out[2] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string })
          .toolCallId,
      ).toBe("call123fc123");
      expect(
        (out[2] as Extract<AgentMessage, { role: "toolResult" }> & { toolUseId?: string })
          .toolUseId,
      ).toBe("call123fc123");
    });

    it("preserves replay-safe signed-thinking tool ids when requested", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            { type: "thinking", thinking: "internal", thinkingSignature: "sig_1" },
            { type: "toolCall", id: "call_1", name: "read", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "call_1",
          toolName: "read",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict", {
        preserveReplaySafeThinkingToolCallIds: true,
        allowedToolNames: ["read"],
      });

      expect(out).toBe(input);
      expect(
        ((out[0] as Extract<AgentMessage, { role: "assistant" }>).content?.[1] as { id?: string })
          .id,
      ).toBe("call_1");
      expect((out[1] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe("call_1");
    });

    it("rewrites earlier mutable ids away from later preserved signed ids", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [readToolCall("call_1")],
        },
        buildToolResult({ toolCallId: "call_1", text: "first" }),
        signedReadAssistant("sig_1", "call1"),
        buildToolResult({ toolCallId: "call1", text: "second" }),
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict", {
        preserveReplaySafeThinkingToolCallIds: true,
        allowedToolNames: ["read"],
      });

      expect(out).not.toBe(input);
      const firstAssistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
      const firstToolCall = firstAssistant.content?.[0] as { id?: string };
      expect(firstToolCall.id).not.toBe("call1");

      expectReplaySafeSignedTurnOwnership({
        input,
        preservedTurn: "second",
        firstToolCallIndex: 0,
      });
    });

    it("rewrites later signed turns when an earlier signed turn already owns the raw id", () => {
      const input = castAgentMessages([
        signedReadAssistant("sig_1", "call1"),
        buildToolResult({ toolCallId: "call1", text: "first" }),
        signedReadAssistant("sig_2", "call1"),
        buildToolResult({ toolCallId: "call1", text: "second" }),
      ]);

      expectReplaySafeSignedTurnOwnership({
        input,
        preservedTurn: "first",
        firstToolCallIndex: 1,
      });
    });

    it("avoids collisions with alphanumeric-only suffixes", () => {
      const input = buildDuplicateIdCollisionInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).not.toBe(input);
      const { aId, bId } = expectCollisionIdsRemainDistinct(out, "strict");
      // Should not contain underscores or hyphens
      expect(aId).not.toMatch(/[_-]/);
      expect(bId).not.toMatch(/[_-]/);
    });

    it("assigns distinct strict IDs when identical raw tool call ids repeat", () => {
      const input = buildRepeatedRawIdInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).not.toBe(input);
      const { aId, bId } = expectCollisionIdsRemainDistinct(out, "strict");
      expect(aId).not.toMatch(/[_-]/);
      expect(bId).not.toMatch(/[_-]/);
    });

    it("preserves native Kimi function ids in direct strict sanitization", () => {
      expect(sanitizeToolCallId("functions.read:0", "strict")).toBe("functions.read:0");
      expect(sanitizeToolCallId("functions.bash_tool:12", "strict")).toBe("functions.bash_tool:12");
      expect(sanitizeToolCallId("functions.edit-file:3", "strict")).toBe("functions.edit-file:3");
      expect(isValidCloudCodeAssistToolId("functions.read:0", "strict")).toBe(true);
      expect(isValidCloudCodeAssistToolId("functions.read:0", "strict9")).toBe(false);
    });

    it("preserves native Kimi function ids across assistant/toolResult pairs", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "functions.read:0", name: "read", arguments: {} }],
        },
        {
          role: "toolResult",
          toolCallId: "functions.read:0",
          toolName: "read",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).toBe(input);
    });

    it("preserves native Kimi ids while sanitizing non-Kimi siblings", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: "functions.read:0", name: "read", arguments: {} },
            { type: "toolCall", id: "call_a|b", name: "read", arguments: {} },
          ],
        },
        buildToolResult({ toolCallId: "functions.read:0", text: "native" }),
        buildToolResult({ toolCallId: "call_a|b", text: "sanitized" }),
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).not.toBe(input);
      const assistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
      const native = assistant.content?.[0] as { id?: string };
      const sibling = assistant.content?.[1] as { id?: string };
      expect(native.id).toBe("functions.read:0");
      expect(sibling.id).toBe("callab");
      expect((out[1] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe(
        "functions.read:0",
      );
      expect((out[2] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe("callab");
    });

    it("disambiguates repeated native Kimi ids after preserving the first occurrence", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "functions.read:0", name: "read", arguments: {} }],
        },
        buildToolResult({ toolCallId: "functions.read:0", text: "one" }),
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "functions.read:0", name: "read", arguments: {} }],
        },
        buildToolResult({ toolCallId: "functions.read:0", text: "two" }),
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
      expect(out).not.toBe(input);
      const first = (out[0] as Extract<AgentMessage, { role: "assistant" }>).content?.[0] as {
        id?: string;
      };
      const second = (out[2] as Extract<AgentMessage, { role: "assistant" }>).content?.[0] as {
        id?: string;
      };
      expect(first.id).toBe("functions.read:0");
      expect(second.id).not.toBe("functions.read:0");
      expect(isValidCloudCodeAssistToolId(second.id as string, "strict")).toBe(true);
      expect((out[1] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe(
        "functions.read:0",
      );
      expect((out[3] as Extract<AgentMessage, { role: "toolResult" }>).toolCallId).toBe(second.id);
    });

    it("does not preserve malformed Kimi-like ids", () => {
      for (const bad of [
        "functions.read",
        "functions.:0",
        "functions.read:",
        "functions.read:x",
        "functions.read:0:extra",
        "xfunctions.read:0",
      ]) {
        expect(sanitizeToolCallId(bad, "strict")).not.toBe(bad);
        expect(isValidCloudCodeAssistToolId(bad, "strict")).toBe(false);
      }
    });
  });

  describe("strict9 mode (Mistral tool call IDs)", () => {
    it("is a no-op for already-valid 9-char alphanumeric IDs", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "abc123XYZ", name: "read", arguments: {} }],
        },
        {
          role: "toolResult",
          toolCallId: "abc123XYZ",
          toolName: "read",
          content: [{ type: "text", text: "ok" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
      expect(out).toBe(input);
    });

    it("enforces alphanumeric IDs with length 9", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: "call_abc|item:123", name: "read", arguments: {} },
            { type: "toolCall", id: "call_abc|item:456", name: "read", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "call_abc|item:123",
          toolName: "read",
          content: [{ type: "text", text: "one" }],
        },
        {
          role: "toolResult",
          toolCallId: "call_abc|item:456",
          toolName: "read",
          content: [{ type: "text", text: "two" }],
        },
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
      expectDistinctStrict9Ids(out, input);
    });

    it("assigns distinct strict9 IDs when identical raw tool call ids repeat", () => {
      const input = buildRepeatedRawIdInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
      expectDistinctStrict9Ids(out, input);
    });

    it("reuses one rewritten strict9 id when a tool result carries matching toolCallId and toolUseId", () => {
      const input = buildRepeatedSharedToolResultIdInput();

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
      expect(out).not.toBe(input);
      expectStrict9IdLengths(expectToolUseIdsFollowDistinctToolCallIds(out, "strict9"));
    });

    it("rewrites native Kimi function ids in strict9 mode", () => {
      const input = castAgentMessages([
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "functions.read:0", name: "read", arguments: {} }],
        },
        buildToolResult({ toolCallId: "functions.read:0", text: "ok" }),
      ]);

      const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
      expect(out).not.toBe(input);
      expectSingleToolCallRewrite(out, "functions", "strict9");
    });
  });
});