fix(kiro): canonicalize tool history and route API keys correctly

Route API-key inference through Amazon Q first, enforce adjacent one-to-one tool use/result pairs after session replay, and treat payload-invalid HTTP 400 as terminal.
2026-09-22 13:38:31 +00:00 · 2026-07-29 19:27:41 +07:00
parent 44c7b34837
commit 16cb40fda1
14 changed files with 1050 additions and 465 deletions
@@ -0,0 +1,67 @@
+import { describe, expect, it } from "vitest";
+import { KiroExecutor } from "../../open-sse/executors/kiro.js";
+
+const RUNTIME = "https://runtime.us-east-1.kiro.dev/generateAssistantResponse";
+const CODEWHISPERER = "https://codewhisperer.us-east-1.amazonaws.com/generateAssistantResponse";
+const Q = "https://q.us-east-1.amazonaws.com/generateAssistantResponse";
+
+function credentials(authMethod, region = "us-east-1") {
+  return { providerSpecificData: { authMethod, region } };
+}
+
+describe("Kiro auth-aware endpoint routing", () => {
+  const executor = new KiroExecutor();
+
+  it("routes API-key inference through Amazon Q before other surfaces", () => {
+    expect(executor.getOrderedBaseUrls(credentials("api_key"))).toEqual([
+      Q,
+      CODEWHISPERER,
+      RUNTIME,
+    ]);
+  });
+
+  it("keeps Builder ID OAuth on the Kiro runtime surface", () => {
+    expect(executor.getOrderedBaseUrls(credentials("builder-id"))).toEqual([
+      RUNTIME,
+      CODEWHISPERER,
+      Q,
+    ]);
+  });
+
+  it("keeps external IdP on CodeWhisperer before Amazon Q", () => {
+    expect(executor.getOrderedBaseUrls(credentials("external_idp"))).toEqual([
+      CODEWHISPERER,
+      Q,
+      RUNTIME,
+    ]);
+  });
+
+  it("regionalizes AWS endpoints for IDC without changing Kiro runtime", () => {
+    expect(executor.getOrderedBaseUrls(credentials("idc", "eu-west-1"))).toEqual([
+      "https://codewhisperer.eu-west-1.amazonaws.com/generateAssistantResponse",
+      "https://q.eu-west-1.amazonaws.com/generateAssistantResponse",
+      RUNTIME,
+    ]);
+  });
+
+  it("retries only endpoint/auth-surface failures, not payload-invalid 400s", () => {
+    expect(executor.shouldRetry(400, 0)).toBe(false);
+    expect(executor.shouldRetry(401, 1)).toBe(true);
+    expect(executor.shouldRetry(403, 2)).toBe(false);
+    expect(executor.shouldRetry(422, 0)).toBe(false);
+  });
+
+  it("builds endpoint-specific headers", () => {
+    const auth = { accessToken: "test-key", providerSpecificData: { authMethod: "api_key" } };
+    const qHeaders = executor.buildHeaders(auth, true, Q);
+    const codeWhispererHeaders = executor.buildHeaders(auth, true, CODEWHISPERER);
+    const runtimeHeaders = executor.buildHeaders(auth, true, RUNTIME);
+
+    expect(qHeaders.TokenType).toBe("API_KEY");
+    expect(qHeaders["X-Amz-Target"]).toBeUndefined();
+    expect(codeWhispererHeaders["X-Amz-Target"]).toBe(
+      "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"
+    );
+    expect(runtimeHeaders["X-Amz-Target"]).toBeUndefined();
+  });
+});
@@ -0,0 +1,372 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import {
+  canonicalizeKiroConversation,
+  normalizeKiroToolSpecs,
+  validateKiroConversation,
+} from "../../open-sse/translator/concerns/kiroConversation.js";
+import { clearKiroSessionReplayStore } from "../../open-sse/utils/kiroSessionReplay.js";
+import { clearSessionStore } from "../../open-sse/utils/sessionManager.js";
+import { claudeToKiroRequest } from "../../open-sse/translator/request/claude-to-kiro.js";
+import { openaiToKiroRequest } from "../../open-sse/translator/request/openai-to-kiro.js";
+
+const modelId = "claude-opus-5";
+
+function tool(name, schema = { type: "object", properties: {} }) {
+  return { name, description: `Tool ${name}`, input_schema: schema };
+}
+
+function specState(names = ["first", "second"]) {
+  const source = names.map((name) => tool(name));
+  return normalizeKiroToolSpecs(source);
+}
+
+function user(content, toolResults = []) {
+  return {
+    userInputMessage: {
+      content,
+      modelId,
+      ...(toolResults.length > 0 && { userInputMessageContext: { toolResults } }),
+    },
+  };
+}
+
+function assistant(content, toolUses = []) {
+  return {
+    assistantResponseMessage: {
+      content,
+      ...(toolUses.length > 0 && { toolUses }),
+    },
+  };
+}
+
+function result(toolUseId, value, status = "success") {
+  return { toolUseId, status, content: [{ text: value }] };
+}
+
+describe("Kiro conversation canonicalizer", () => {
+  beforeEach(() => {
+    clearKiroSessionReplayStore();
+    clearSessionStore();
+  });
+
+  it("keeps complete parallel tool pairs structured", () => {
+    const { specs, nameMap } = specState();
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [
+          { toolUseId: "t1", name: "first", input: { n: 1 } },
+          { toolUseId: "t2", name: "second", input: { n: 2 } },
+        ]),
+      ],
+      currentMessage: user("continue", [result("t1", "one"), result("t2", "two")]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    const calls = canonical.history[1].assistantResponseMessage.toolUses;
+    const results = canonical.currentMessage.userInputMessage.userInputMessageContext.toolResults;
+    expect(calls.map((call) => call.toolUseId)).toEqual(["t1", "t2"]);
+    expect(results.map((item) => item.toolUseId)).toEqual(["t1", "t2"]);
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("keeps the answered parallel call and flattens only the missing one", () => {
+    const { specs, nameMap } = specState();
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [
+          { toolUseId: "t1", name: "first", input: {} },
+          { toolUseId: "t2", name: "second", input: {} },
+        ]),
+      ],
+      currentMessage: user("continue", [result("t1", "one")]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    const assistantMessage = canonical.history[1].assistantResponseMessage;
+    expect(assistantMessage.toolUses).toHaveLength(1);
+    expect(assistantMessage.toolUses[0].toolUseId).toBe("t1");
+    expect(assistantMessage.content).toContain("[Tool call: second(");
+    expect(canonical.repairs.missingResults).toBe(1);
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("flattens non-adjacent and orphaned tool results", () => {
+    const { specs, nameMap } = specState(["first"]);
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [{ toolUseId: "t1", name: "first", input: {} }]),
+        user("result missing here"),
+        assistant("later"),
+      ],
+      currentMessage: user("late result", [result("t1", "too late")]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    expect(JSON.stringify(canonical)).not.toContain('"toolUseId":"t1"');
+    expect(canonical.history[1].assistantResponseMessage.content).toContain("[Tool call:");
+    expect(canonical.currentMessage.userInputMessage.content).toContain("too late");
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("remaps duplicate tool IDs together with their adjacent results", () => {
+    const { specs, nameMap } = specState();
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [
+          { toolUseId: "duplicate", name: "first", input: {} },
+          { toolUseId: "duplicate", name: "second", input: {} },
+        ]),
+      ],
+      currentMessage: user("continue", [
+        result("duplicate", "one"),
+        result("duplicate", "two"),
+      ]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    const calls = canonical.history[1].assistantResponseMessage.toolUses;
+    const results = canonical.currentMessage.userInputMessage.userInputMessageContext.toolResults;
+    expect(new Set(calls.map((call) => call.toolUseId)).size).toBe(2);
+    expect(results.map((item) => item.toolUseId)).toEqual(calls.map((call) => call.toolUseId));
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("deduplicates extra results without losing their text", () => {
+    const { specs, nameMap } = specState(["first"]);
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [{ toolUseId: "t1", name: "first", input: {} }]),
+      ],
+      currentMessage: user("continue", [result("t1", "one"), result("t1", "duplicate")]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    const current = canonical.currentMessage.userInputMessage;
+    expect(current.userInputMessageContext.toolResults).toHaveLength(1);
+    expect(current.content).toContain("duplicate");
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("flattens a trailing unanswered assistant tool call and creates a current user turn", () => {
+    const { specs, nameMap } = specState(["first"]);
+    const canonical = canonicalizeKiroConversation({
+      history: [user("start")],
+      currentMessage: assistant("run", [{ toolUseId: "t1", name: "first", input: {} }]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    expect(canonical.currentMessage.userInputMessage.content).toBe("continue");
+    expect(canonical.history[1].assistantResponseMessage.toolUses).toBeUndefined();
+    expect(canonical.history[1].assistantResponseMessage.content).toContain("[Tool call:");
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("flattens malformed input and tool uses missing from the current specs", () => {
+    const { specs, nameMap } = specState(["first"]);
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [
+          { toolUseId: "t1", name: "first", input: "{bad json" },
+          { toolUseId: "t2", name: "removed_tool", input: {} },
+        ]),
+      ],
+      currentMessage: user("continue", [result("t1", "one"), result("t2", "two")]),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    expect(canonical.history[1].assistantResponseMessage.toolUses).toBeUndefined();
+    expect(canonical.currentMessage.userInputMessage.userInputMessageContext.toolResults).toBeUndefined();
+    expect(canonical.currentMessage.userInputMessage.content).toContain("one");
+    expect(canonical.currentMessage.userInputMessage.content).toContain("two");
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("repairs a 30-call parallel turn with one missing result", () => {
+    const names = Array.from({ length: 30 }, (_, index) => `tool_${index}`);
+    const { specs, nameMap } = specState(names);
+    const calls = names.map((name, index) => ({
+      toolUseId: `t${index}`,
+      name,
+      input: { index },
+    }));
+    const results = names.slice(0, -1).map((_, index) => result(`t${index}`, `r${index}`));
+    const canonical = canonicalizeKiroConversation({
+      history: [user("start"), assistant("run", calls)],
+      currentMessage: user("continue", results),
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    expect(canonical.history[1].assistantResponseMessage.toolUses).toHaveLength(29);
+    expect(canonical.currentMessage.userInputMessage.userInputMessageContext.toolResults).toHaveLength(29);
+    expect(canonical.repairs.missingResults).toBe(1);
+    expect(canonical.valid).toBe(true);
+  });
+
+  it("flattens structured history when the client sent no tool specs", () => {
+    const canonical = canonicalizeKiroConversation({
+      history: [
+        user("start"),
+        assistant("run", [{ toolUseId: "t1", name: "first", input: {} }]),
+      ],
+      currentMessage: user("continue", [result("t1", "one")]),
+      modelId,
+    });
+
+    expect(JSON.stringify(canonical)).not.toContain("toolUses");
+    expect(JSON.stringify(canonical)).not.toContain("toolResults");
+    expect(canonical.history[1].assistantResponseMessage.content).toContain("[Tool call:");
+    expect(canonical.currentMessage.userInputMessage.content).toContain("[Tool result:");
+  });
+
+  it("normalizes names and recursively removes unsupported schema fields", () => {
+    const longDescription = "x".repeat(11000);
+    const { specs, nameMap } = normalizeKiroToolSpecs([{
+      name: "bad tool/name",
+      description: longDescription,
+      input_schema: {
+        additionalProperties: false,
+        properties: {
+          nested: {
+            type: "object",
+            additionalProperties: true,
+            properties: {},
+            required: [],
+          },
+        },
+        required: [],
+      },
+    }]);
+
+    const specification = specs[0].toolSpecification;
+    expect(nameMap.get("bad tool/name")).toBe("bad_tool_name");
+    expect(specification.name.length).toBeLessThanOrEqual(64);
+    expect(specification.description.length).toBe(10237);
+    expect(JSON.stringify(specification.inputSchema.json)).not.toContain("additionalProperties");
+    expect(JSON.stringify(specification.inputSchema.json)).not.toContain('"required":[]');
+  });
+
+  it("does not mutate the source conversation or tool definitions", () => {
+    const sourceTools = [tool("first")];
+    const sourceHistory = [
+      user("start"),
+      assistant("run", [{ toolUseId: "t1", name: "first", input: {} }]),
+    ];
+    const sourceCurrent = user("continue", [result("t1", "one")]);
+    const before = JSON.stringify({ sourceTools, sourceHistory, sourceCurrent });
+    const { specs, nameMap } = normalizeKiroToolSpecs(sourceTools);
+
+    canonicalizeKiroConversation({
+      history: sourceHistory,
+      currentMessage: sourceCurrent,
+      modelId,
+      toolSpecs: specs,
+      nameMap,
+    });
+
+    expect(JSON.stringify({ sourceTools, sourceHistory, sourceCurrent })).toBe(before);
+  });
+
+  it("preserves Claude tool_result errors", () => {
+    const output = claudeToKiroRequest(modelId, {
+      tools: [tool("first")],
+      messages: [
+        { role: "user", content: "start" },
+        { role: "assistant", content: [{ type: "tool_use", id: "t1", name: "first", input: {} }] },
+        { role: "user", content: [{ type: "tool_result", tool_use_id: "t1", is_error: true, content: "failed" }] },
+      ],
+    }, true, {});
+
+    const item = output.conversationState.currentMessage.userInputMessage
+      .userInputMessageContext.toolResults[0];
+    expect(item.status).toBe("error");
+  });
+
+  it("repairs partial parallel results in both direct translators", () => {
+    const claude = claudeToKiroRequest(modelId, {
+      tools: [tool("first"), tool("second")],
+      messages: [
+        { role: "user", content: "start" },
+        { role: "assistant", content: [
+          { type: "tool_use", id: "t1", name: "first", input: {} },
+          { type: "tool_use", id: "t2", name: "second", input: {} },
+        ] },
+        { role: "user", content: [{ type: "tool_result", tool_use_id: "t1", content: "one" }] },
+      ],
+    }, true, {});
+    const openai = openaiToKiroRequest(modelId, {
+      tools: [
+        { type: "function", function: { name: "first", parameters: { type: "object", properties: {} } } },
+        { type: "function", function: { name: "second", parameters: { type: "object", properties: {} } } },
+      ],
+      messages: [
+        { role: "user", content: "start" },
+        { role: "assistant", content: "", tool_calls: [
+          { id: "t1", type: "function", function: { name: "first", arguments: "{}" } },
+          { id: "t2", type: "function", function: { name: "second", arguments: "{}" } },
+        ] },
+        { role: "tool", tool_call_id: "t1", content: "one" },
+      ],
+    }, true, {});
+
+    for (const payload of [claude, openai]) {
+      const state = payload.conversationState;
+      const validation = validateKiroConversation(
+        state.history,
+        state.currentMessage,
+        state.currentMessage.userInputMessage.userInputMessageContext.tools
+      );
+      expect(validation.valid).toBe(true);
+      expect(state.history[1].assistantResponseMessage.toolUses).toHaveLength(1);
+    }
+  });
+
+  it("does not let session replay replace a tool-result turn", () => {
+    const credentials = {
+      rawHeaders: { "x-session-id": "kiro-replay-tool-result-regression" },
+      connectionId: "kiro-account",
+    };
+    claudeToKiroRequest(modelId, {
+      messages: [{ role: "user", content: "frozen session start" }],
+    }, true, credentials);
+
+    const output = claudeToKiroRequest(modelId, {
+      tools: [tool("first")],
+      messages: [
+        { role: "assistant", content: [{ type: "tool_use", id: "t1", name: "first", input: {} }] },
+        { role: "user", content: [{ type: "tool_result", tool_use_id: "t1", content: "kept" }] },
+      ],
+    }, true, credentials);
+    const state = output.conversationState;
+    const allText = JSON.stringify(state);
+
+    expect(allText).toContain("frozen session start");
+    expect(allText).toContain("kept");
+    expect(validateKiroConversation(
+      state.history,
+      state.currentMessage,
+      state.currentMessage.userInputMessage.userInputMessageContext.tools
+    ).valid).toBe(true);
+  });
+});
@@ -4,10 +4,8 @@ import { KiroService } from "../../src/lib/oauth/services/kiro.js";
 /**
 * Regression tests for Kiro API-key auth.
 *
- * KiroService.validateApiKey resolves a profileArn with the key (via
- * CodeWhisperer ListAvailableProfiles) and returns a credential shaped for
- * persistence with authMethod="api_key". The response profile field name
- * varies (`arn` vs `profileArn`) — both are accepted by listAvailableProfiles.
+ * KiroService.validateApiKey validates against the Amazon Q model catalog and
+ * returns an account-bound credential without inventing a profileArn.
 *
 * Note: OAuth (Builder ID / IDC) profileArn resolution is handled upstream by
 * fetchKiroProfileArn in providers.js and is covered there — not here.
@@ -16,11 +14,10 @@ describe("kiro API-key auth (KiroService.validateApiKey)", () => {
  beforeEach(() => vi.restoreAllMocks());
  afterEach(() => vi.restoreAllMocks());

-  it("validates an API key and resolves a credential with profileArn", async () => {
-    const expectedArn = "arn:aws:codewhisperer:us-east-1:444:profile/KEY";
+  it("validates an API key against Amazon Q without inventing profileArn", async () => {
    const fetchMock = vi.spyOn(globalThis, "fetch").mockResolvedValue({
      ok: true,
-      json: async () => ({ profiles: [{ arn: expectedArn }] }),
+      json: async () => ({ models: [{ modelId: "claude-opus-5" }] }),
    });

    const svc = new KiroService();
@@ -29,17 +26,18 @@ describe("kiro API-key auth (KiroService.validateApiKey)", () => {
    expect(cred).toEqual({
      accessToken: "my-secret-key",
      refreshToken: null,
-      profileArn: expectedArn,
+      profileArn: null,
      region: "us-east-1",
      authMethod: "api_key",
    });

    const [url, init] = fetchMock.mock.calls[0];
-    expect(url).toBe("https://codewhisperer.us-east-1.amazonaws.com");
-    expect(init.headers.Authorization).toBe("Bearer my-secret-key");
-    expect(init.headers["x-amz-target"]).toBe(
-      "AmazonCodeWhispererService.ListAvailableProfiles"
+    expect(url).toBe(
+      "https://q.us-east-1.amazonaws.com/ListAvailableModels?origin=AI_EDITOR"
    );
+    expect(init.method).toBe("GET");
+    expect(init.headers.Authorization).toBe("Bearer my-secret-key");
+    expect(init.headers.TokenType).toBe("API_KEY");
  });

  it("rejects an empty API key without a network call", async () => {
@@ -60,4 +58,15 @@ describe("kiro API-key auth (KiroService.validateApiKey)", () => {
      /API key validation failed/
    );
  });
+
+  it("rejects a 200 response with an empty model catalog", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValue({
+      ok: true,
+      json: async () => ({ models: [] }),
+    });
+    const svc = new KiroService();
+    await expect(svc.validateApiKey("empty-key")).rejects.toThrow(
+      /returned no available models/
+    );
+  });
 });