fix(headroom): compress Kiro conversation state (#2488)

Project conversationState history/currentMessage into OpenAI-style messages for /v1/compress, then write compressed text back into the original Kiro fields while preserving provider payload shape. Fail open when the proxy returns malformed or reordered messages. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-10 11:30:29 +07:00
parent 7610f28f42
commit 65c65a0f56
2 changed files with 272 additions and 0 deletions
--- a/open-sse/rtk/headroom.js
+++ b/open-sse/rtk/headroom.js
@@ -18,6 +18,8 @@ function jsonBytes(value) {
 function messagePayload(body) {
  if (Array.isArray(body?.messages)) return body.messages;
  if (Array.isArray(body?.input)) return body.input;
+  const kiro = collectKiroHeadroomMessages(body);
+  if (kiro) return kiro.messages;
  return null;
 }

@@ -81,6 +83,121 @@ function hasUnsafeResponsesInputForCompression(body) {
  });
 }

+function collectKiroHeadroomMessages(body) {
+  const state = body?.conversationState;
+  if (!state || typeof state !== "object") return null;
+
+  const messages = [];
+  const targets = [];
+
+  const addTextTarget = (role, text, target, extra = {}) => {
+    if (typeof text !== "string") return;
+    messages.push({ role, content: text, ...extra });
+    targets.push(target);
+  };
+
+  const toToolCalls = (toolUses) => {
+    if (!Array.isArray(toolUses) || toolUses.length === 0) return undefined;
+    const calls = toolUses.map((toolUse) => ({
+      id: toolUse?.toolUseId,
+      type: "function",
+      function: {
+        name: toolUse?.name || "",
+        arguments: JSON.stringify(toolUse?.input || {}),
+      },
+    })).filter((call) => call.id || call.function.name);
+    return calls.length > 0 ? calls : undefined;
+  };
+
+  const visit = (item) => {
+    const user = item?.userInputMessage;
+    if (user) {
+      addTextTarget("system", user.systemInstruction, { object: user, key: "systemInstruction" });
+      addTextTarget("user", user.content, { object: user, key: "content" });
+
+      const toolResults = user.userInputMessageContext?.toolResults;
+      if (Array.isArray(toolResults)) {
+        for (const toolResult of toolResults) {
+          const content = toolResult?.content;
+          if (!Array.isArray(content)) continue;
+          for (const part of content) {
+            addTextTarget(
+              "tool",
+              part?.text,
+              { object: part, key: "text" },
+              toolResult?.toolUseId ? { tool_call_id: toolResult.toolUseId } : {}
+            );
+          }
+        }
+      }
+      return;
+    }
+
+    const assistant = item?.assistantResponseMessage;
+    if (assistant) {
+      const toolCalls = toToolCalls(assistant.toolUses);
+      addTextTarget(
+        "assistant",
+        assistant.content,
+        { object: assistant, key: "content" },
+        toolCalls ? { tool_calls: toolCalls } : {}
+      );
+    }
+  };
+
+  if (Array.isArray(state.history)) {
+    for (const item of state.history) visit(item);
+  }
+  if (state.currentMessage) visit(state.currentMessage);
+
+  return messages.length > 0 ? { messages, targets } : null;
+}
+
+function textFromHeadroomMessage(message) {
+  const content = message?.content;
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return null;
+
+  const parts = [];
+  for (const part of content) {
+    if (typeof part === "string") {
+      parts.push(part);
+    } else if (typeof part?.text === "string") {
+      parts.push(part.text);
+    }
+  }
+  return parts.length > 0 ? parts.join("\n") : null;
+}
+
+function applyKiroHeadroomMessages(projection, compressedMessages, diagnostics) {
+  if (!Array.isArray(compressedMessages) || compressedMessages.length !== projection.messages.length) {
+    setDiagnostic(diagnostics, "proxy response did not match Kiro message count");
+    return false;
+  }
+
+  const updates = [];
+  for (let i = 0; i < projection.messages.length; i++) {
+    const expected = projection.messages[i];
+    const actual = compressedMessages[i];
+    if (!actual || actual.role !== expected.role) {
+      setDiagnostic(diagnostics, "proxy response did not preserve Kiro message order");
+      return false;
+    }
+
+    const text = textFromHeadroomMessage(actual);
+    if (text === null) {
+      setDiagnostic(diagnostics, "proxy response missing Kiro text content");
+      return false;
+    }
+    updates.push({ target: projection.targets[i], text });
+  }
+
+  for (const update of updates) {
+    update.target.object[update.target.key] = update.text;
+  }
+  return true;
+}
+
 // POST messages to Headroom /v1/compress; returns compressed messages + stats or null.
 async function callCompress(url, messages, model, timeoutMs, compressUserMessages, diagnostics) {
  const endpoint = buildCompressEndpoint(url);
@@ -171,6 +288,22 @@ export async function compressWithHeadroom(body, { enabled, url, model, format,
      return data;
    }

+    // Kiro shape: conversationState.history/currentMessage are projected to
+    // OpenAI messages for the proxy, then copied back into the original Kiro
+    // fields. Keep the provider payload shape intact for Kiro's executor.
+    if (format === "kiro") {
+      const projection = collectKiroHeadroomMessages(body);
+      if (!projection) {
+        setDiagnostic(diagnostics, "Kiro request did not project to messages[]");
+        return null;
+      }
+      const data = await callCompress(url, projection.messages, model, timeoutMs, compressUserMessages, diagnostics || {});
+      if (!data) return null;
+      if (!applyKiroHeadroomMessages(projection, data.messages, diagnostics)) return null;
+      if (diagnostics) diagnostics.after = captureSizeSnapshot(body);
+      return data;
+    }
+
    // OpenAI shape: messages/input go straight to the proxy.
    const key = Array.isArray(body.messages) ? "messages"
      : Array.isArray(body.input) ? "input"
--- a/tests/unit/headroom.test.js
+++ b/tests/unit/headroom.test.js
@@ -31,6 +31,10 @@ describe("compressWithHeadroom", () => {
    expect(body.messages[0].content).toBe("short");
    expect(stats.tokens_saved).toBe(80);
    expect(global.fetch).toHaveBeenCalledWith("http://headroom:8787/v1/compress", expect.objectContaining({ method: "POST" }));
+    expect(JSON.parse(global.fetch.mock.calls[0][1].body)).toMatchObject({
+      model: "gpt-4o",
+      messages: [{ role: "user", content: "long" }],
+    });
  });

  it("compresses responses input in-place", async () => {
@@ -44,6 +48,141 @@ describe("compressWithHeadroom", () => {
    expect(body.input[0].content).toBe("short");
  });

+  it("compresses Kiro conversationState history/currentMessage in-place", async () => {
+    let requestPayload;
+    global.fetch = vi.fn(async (_url, init) => {
+      requestPayload = JSON.parse(init.body);
+      return new Response(JSON.stringify({
+        messages: [
+          { role: "user", content: "compressed earlier user" },
+          { role: "assistant", content: "compressed assistant", tool_calls: [{ id: "tool_1", type: "function", function: { name: "read_file", arguments: "{\"path\":\"a.js\"}" } }] },
+          { role: "system", content: "compressed system instruction" },
+          { role: "user", content: "compressed current user" },
+          { role: "tool", content: [{ type: "text", text: "compressed tool output" }], tool_call_id: "tool_1" },
+        ],
+        tokens_before: 100,
+        tokens_after: 40,
+        tokens_saved: 60,
+      }), { status: 200 });
+    });
+    const body = {
+      profileArn: "arn:test",
+      conversationState: {
+        chatTriggerType: "MANUAL",
+        conversationId: "conv-1",
+        history: [
+          {
+            userInputMessage: {
+              content: "earlier user",
+              modelId: "claude-sonnet-4.5",
+            },
+          },
+          {
+            assistantResponseMessage: {
+              content: "assistant response",
+              toolUses: [
+                {
+                  toolUseId: "tool_1",
+                  name: "read_file",
+                  input: { path: "a.js" },
+                },
+              ],
+            },
+          },
+        ],
+        currentMessage: {
+          userInputMessage: {
+            content: "current user",
+            modelId: "claude-sonnet-4.5",
+            systemInstruction: "native system instruction",
+            userInputMessageContext: {
+              tools: [{ toolSpecification: { name: "read_file" } }],
+              toolResults: [
+                {
+                  toolUseId: "tool_1",
+                  status: "success",
+                  content: [{ text: "long tool output" }],
+                },
+              ],
+            },
+          },
+        },
+      },
+    };
+
+    const stats = await compressWithHeadroom(body, {
+      enabled: true,
+      url: "http://localhost:8787",
+      model: "claude-sonnet-4.5",
+      format: "kiro",
+      compressUserMessages: true,
+    });
+
+    expect(stats.tokens_saved).toBe(60);
+    expect(requestPayload).toEqual({
+      model: "claude-sonnet-4.5",
+      config: { compress_user_messages: true },
+      messages: [
+        { role: "user", content: "earlier user" },
+        {
+          role: "assistant",
+          content: "assistant response",
+          tool_calls: [
+            {
+              id: "tool_1",
+              type: "function",
+              function: { name: "read_file", arguments: "{\"path\":\"a.js\"}" },
+            },
+          ],
+        },
+        { role: "system", content: "native system instruction" },
+        { role: "user", content: "current user" },
+        { role: "tool", content: "long tool output", tool_call_id: "tool_1" },
+      ],
+    });
+    expect(body.conversationState.history[0].userInputMessage.content).toBe("compressed earlier user");
+    expect(body.conversationState.history[1].assistantResponseMessage.content).toBe("compressed assistant");
+    expect(body.conversationState.currentMessage.userInputMessage.systemInstruction).toBe("compressed system instruction");
+    expect(body.conversationState.currentMessage.userInputMessage.content).toBe("compressed current user");
+    expect(body.conversationState.currentMessage.userInputMessage.userInputMessageContext.toolResults[0].content[0].text)
+      .toBe("compressed tool output");
+    expect(body.profileArn).toBe("arn:test");
+    expect(body.conversationState.currentMessage.userInputMessage.userInputMessageContext.tools)
+      .toEqual([{ toolSpecification: { name: "read_file" } }]);
+  });
+
+  it("fails open when Kiro Headroom output does not preserve message order", async () => {
+    global.fetch = vi.fn(async () => new Response(JSON.stringify({
+      messages: [{ role: "assistant", content: "wrong role" }],
+      tokens_saved: 10,
+    }), { status: 200 }));
+    const body = {
+      conversationState: {
+        currentMessage: {
+          userInputMessage: {
+            content: "original",
+            modelId: "claude-sonnet-4.5",
+          },
+        },
+        history: [],
+      },
+    };
+    const original = structuredClone(body);
+    const diagnostics = {};
+
+    const stats = await compressWithHeadroom(body, {
+      enabled: true,
+      url: "http://localhost:8787",
+      model: "claude-sonnet-4.5",
+      format: "kiro",
+      diagnostics,
+    });
+
+    expect(stats).toBeNull();
+    expect(body).toEqual(original);
+    expect(diagnostics.reason).toBe("proxy response did not preserve Kiro message order");
+  });
+
  it("fails open on bad response", async () => {
    global.fetch = vi.fn(async () => new Response(JSON.stringify({ error: "bad" }), { status: 500 }));
    const body = { messages: [{ role: "user", content: "long" }] };