From 65c65a0f56dfa567144a6fa2cd7d3023a4143d97 Mon Sep 17 00:00:00 2001 From: Edison42 Date: Fri, 10 Jul 2026 11:30:29 +0700 Subject: [PATCH] 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 --- open-sse/rtk/headroom.js | 133 ++++++++++++++++++++++++++++++++++ tests/unit/headroom.test.js | 139 ++++++++++++++++++++++++++++++++++++ 2 files changed, 272 insertions(+) diff --git a/open-sse/rtk/headroom.js b/open-sse/rtk/headroom.js index 8f3b1f33..880195ff 100644 --- 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" diff --git a/tests/unit/headroom.test.js b/tests/unit/headroom.test.js index 46468595..b7c414d0 100644 --- 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" }] };