diff --git a/open-sse/config/kiroConstants.js b/open-sse/config/kiroConstants.js index 2f97256d..e6408da1 100644 --- a/open-sse/config/kiroConstants.js +++ b/open-sse/config/kiroConstants.js @@ -20,6 +20,12 @@ import { effortToBudget } from "../translator/concerns/thinking.js"; export const KIRO_AGENTIC_SUFFIX = "-agentic"; export const KIRO_THINKING_SUFFIX = "-thinking"; +export const KIRO_TOOL_NAME_MAX_LENGTH = 64; +export const KIRO_TOOL_DESCRIPTION_MAX_LENGTH = 10237; +export const KIRO_TOOL_ID_MAX_LENGTH = 64; +export const KIRO_CODEWHISPERER_TARGET = + "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"; +export const KIRO_ENDPOINT_FALLBACK_STATUSES = new Set([401, 403, 404]); // Public default CodeWhisperer profile ARNs (us-east-1), keyed by auth method. // Used when an account cannot resolve its own profileArn. Builder ID and social diff --git a/open-sse/executors/base.js b/open-sse/executors/base.js index 71418deb..5a3c4467 100644 --- a/open-sse/executors/base.js +++ b/open-sse/executors/base.js @@ -126,7 +126,7 @@ export class BaseExecutor { for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) { const url = this.buildUrl(model, stream, urlIndex, credentials); const transformedBody = this.transformRequest(model, body, stream, credentials); - const headers = this.buildHeaders(credentials, stream); + const headers = this.buildHeaders(credentials, stream, url); if (!retryAttemptsByUrl[urlIndex]) retryAttemptsByUrl[urlIndex] = 0; diff --git a/open-sse/executors/kiro.js b/open-sse/executors/kiro.js index 90ec713e..1205522b 100644 --- a/open-sse/executors/kiro.js +++ b/open-sse/executors/kiro.js @@ -1,6 +1,10 @@ import { BaseExecutor } from "./base.js"; import { PROVIDERS } from "../config/providers.js"; -import { resolveKiroModel } from "../config/kiroConstants.js"; +import { + KIRO_CODEWHISPERER_TARGET, + KIRO_ENDPOINT_FALLBACK_STATUSES, + resolveKiroModel, +} from "../config/kiroConstants.js"; import { v4 as uuidv4 } from "uuid"; import { refreshKiroToken } from "../services/tokenRefresh.js"; import { SSE_DONE, SSE_HEADERS } from "../utils/sseConstants.js"; @@ -216,12 +220,17 @@ export class KiroExecutor extends BaseExecutor { super("kiro", PROVIDERS.kiro); } - buildHeaders(credentials, stream = true) { + buildHeaders(credentials, stream = true, url = "") { const headers = { ...this.config.headers, "Amz-Sdk-Request": "attempt=1; max=3", "Amz-Sdk-Invocation-Id": uuidv4() }; + if (url.includes("://codewhisperer.")) { + headers["X-Amz-Target"] = KIRO_CODEWHISPERER_TARGET; + } else { + delete headers["X-Amz-Target"]; + } // API-key auth: the key is stored as accessToken and sent as a bearer token // exactly like an OAuth access token, but with an extra `tokentype: API_KEY` @@ -236,8 +245,8 @@ export class KiroExecutor extends BaseExecutor { const apiKey = credentials?.apiKey || (isApiKey ? credentials?.accessToken : null); if (isApiKey && apiKey) { headers["Authorization"] = `Bearer ${apiKey}`; - headers["tokentype"] = "API_KEY"; - } else if (credentials.accessToken) { + headers["TokenType"] = "API_KEY"; + } else if (credentials?.accessToken) { headers["Authorization"] = `Bearer ${credentials.accessToken}`; if (isExternalIdp) { headers["TokenType"] = "EXTERNAL_IDP"; @@ -250,14 +259,14 @@ export class KiroExecutor extends BaseExecutor { /** * Auth-aware endpoint ordering. * - * API-key Kiro connections store a raw CodeWhisperer credential (validated - * against codewhisperer.us-east-1.amazonaws.com via ListAvailableProfiles). + * API-key Kiro connections use the Amazon Q surface. The legacy + * codewhisperer.* GenerateAssistantResponse endpoint can authenticate the key + * but rejects the same valid payload with REQUEST_BODY_INVALID. Since a 400 + * is terminal in BaseExecutor, putting CodeWhisperer first prevents the working + * q.* endpoint from ever being tried. Keep q.* first only for api_key accounts. + * * The Kiro IDE gateway (runtime.*.kiro.dev) expects Kiro OIDC/social tokens - * and rejects an `tokentype: API_KEY` token with 401/403 — which - * BaseExecutor.execute() returns immediately (only 429 / network errors fall - * through to the next host). So for api-key auth we must try the *.amazonaws.com - * CodeWhisperer hosts FIRST, mirroring the Kiro-Go reference fork which never - * routes api-key traffic through kiro.dev. External IdP enterprise tokens also + * and rejects TokenType=API_KEY. External IdP enterprise tokens instead * use the CodeWhisperer surface, with the `TokenType: EXTERNAL_IDP` header. * Other OAuth methods keep the default order (kiro.dev first) since their * tokens are what that gateway accepts. @@ -282,6 +291,14 @@ export class KiroExecutor extends BaseExecutor { const amazon = baseUrls.filter((u) => u.includes("amazonaws.com")).map(regionalize); const others = baseUrls.filter((u) => !u.includes("amazonaws.com")); + if (authMethod === "api_key") { + const q = amazon.filter((u) => u.includes("://q.")); + const remaining = amazon.filter((u) => !u.includes("://q.")); + return q.length > 0 + ? [...q, ...remaining, ...others] + : [...amazon, ...others]; + } + return amazon.length > 0 ? [...amazon, ...others] : baseUrls; } @@ -290,6 +307,14 @@ export class KiroExecutor extends BaseExecutor { return baseUrls[urlIndex] || baseUrls[0] || this.config.baseUrl; } + // Retry only endpoint/auth-surface failures. Payload-invalid HTTP 400 must be + // terminal: sending the same malformed body to every surface cannot repair it. + shouldRetry(status, urlIndex) { + const hasFallback = urlIndex + 1 < this.getFallbackCount(); + return super.shouldRetry(status, urlIndex) + || (hasFallback && KIRO_ENDPOINT_FALLBACK_STATUSES.has(status)); + } + transformRequest(model, body, stream, credentials) { return body; } diff --git a/open-sse/providers/registry/kiro.js b/open-sse/providers/registry/kiro.js index 1a47adaa..325745a4 100644 --- a/open-sse/providers/registry/kiro.js +++ b/open-sse/providers/registry/kiro.js @@ -29,7 +29,6 @@ export default { headers: { "Content-Type": "application/json", Accept: "application/vnd.amazon.eventstream", - "X-Amz-Target": "AmazonCodeWhispererStreamingService.GenerateAssistantResponse", "User-Agent": "AWS-SDK-JS/3.0.0 kiro-ide/1.0.0", "X-Amz-User-Agent": "aws-sdk-js/3.0.0 kiro-ide/1.0.0", }, diff --git a/open-sse/translator/concerns/kiroConversation.js b/open-sse/translator/concerns/kiroConversation.js new file mode 100644 index 00000000..11d49dc7 --- /dev/null +++ b/open-sse/translator/concerns/kiroConversation.js @@ -0,0 +1,435 @@ +import { + KIRO_TOOL_DESCRIPTION_MAX_LENGTH, + KIRO_TOOL_ID_MAX_LENGTH, + KIRO_TOOL_NAME_MAX_LENGTH, +} from "../../config/kiroConstants.js"; + +const TOOL_ID_PATTERN = /^[a-zA-Z0-9_-]+$/; +const TOOL_NAME_PATTERN = /[^a-zA-Z0-9_-]/g; + +function clone(value) { + return value == null ? value : JSON.parse(JSON.stringify(value)); +} + +function text(value) { + if (typeof value === "string") return value; + if (value == null) return ""; + try { + return JSON.stringify(value); + } catch { + return String(value); + } +} + +function appendText(target, extra) { + if (!extra) return; + target.content = target.content ? `${target.content}\n\n${extra}` : extra; +} + +function trimCodePoints(value, limit) { + return [...String(value || "")].slice(0, limit).join(""); +} + +function uniqueName(rawName, index, usedNames) { + const cleaned = String(rawName || "") + .trim() + .replace(TOOL_NAME_PATTERN, "_") + .replace(/_+/g, "_") + .replace(/^_+|_+$/g, ""); + const base = trimCodePoints(cleaned || `tool_${index + 1}`, KIRO_TOOL_NAME_MAX_LENGTH); + let candidate = base; + let suffix = 2; + while (usedNames.has(candidate)) { + const tail = `_${suffix++}`; + candidate = `${base.slice(0, KIRO_TOOL_NAME_MAX_LENGTH - tail.length)}${tail}`; + } + usedNames.add(candidate); + return candidate; +} + +function cleanSchemaValue(value) { + if (Array.isArray(value)) return value.map(cleanSchemaValue); + if (!value || typeof value !== "object") return value; + + const cleaned = {}; + for (const [key, child] of Object.entries(value)) { + if (key === "additionalProperties") continue; + if (key === "required" && Array.isArray(child) && child.length === 0) continue; + cleaned[key] = cleanSchemaValue(child); + } + return cleaned; +} + +function normalizeRootSchema(schema) { + const cleaned = cleanSchemaValue(schema && typeof schema === "object" ? clone(schema) : {}); + cleaned.type = "object"; + if (!cleaned.properties || typeof cleaned.properties !== "object" || Array.isArray(cleaned.properties)) { + cleaned.properties = {}; + } + if (Array.isArray(cleaned.required)) { + cleaned.required = [...new Set(cleaned.required.filter( + (name) => typeof name === "string" && Object.hasOwn(cleaned.properties, name) + ))]; + if (cleaned.required.length === 0) delete cleaned.required; + } + return cleaned; +} + +/** Normalize OpenAI- or Claude-shaped tool definitions into Kiro tool specs. */ +export function normalizeKiroToolSpecs(tools) { + const specs = []; + const nameMap = new Map(); + const usedNames = new Set(); + + for (const [index, tool] of (Array.isArray(tools) ? tools : []).entries()) { + if (!tool || typeof tool !== "object") continue; + const rawName = tool.function?.name ?? tool.name; + if (typeof rawName !== "string" || !rawName.trim()) continue; + + // A repeated definition with the same source name describes the same tool. + if (nameMap.has(rawName)) continue; + const name = uniqueName(rawName, index, usedNames); + nameMap.set(rawName, name); + + const rawDescription = tool.function?.description ?? tool.description ?? `Tool: ${rawName}`; + const description = trimCodePoints( + String(rawDescription || `Tool: ${rawName}`), + KIRO_TOOL_DESCRIPTION_MAX_LENGTH + ); + const schema = tool.function?.parameters ?? tool.parameters ?? tool.input_schema ?? {}; + specs.push({ + toolSpecification: { + name, + description, + inputSchema: { json: normalizeRootSchema(schema) }, + }, + }); + } + + return { specs, nameMap }; +} + +function toolCallText(toolUse) { + return `[Tool call: ${toolUse?.name || "unknown"}(${text(toolUse?.input || {})})]`; +} + +function toolResultText(toolResult) { + const content = Array.isArray(toolResult?.content) + ? toolResult.content.map((part) => text(part?.text ?? part)).filter(Boolean).join("\n") + : text(toolResult?.content); + return `[Tool result${toolResult?.status === "error" ? " (error)" : ""}: ${content}]`; +} + +function mergeUser(target, source) { + appendText(target, source.content); + if (Array.isArray(source.images) && source.images.length > 0) { + target.images = [...(target.images || []), ...source.images]; + } + const results = source.userInputMessageContext?.toolResults; + if (Array.isArray(results) && results.length > 0) { + target.userInputMessageContext ||= {}; + target.userInputMessageContext.toolResults = [ + ...(target.userInputMessageContext.toolResults || []), + ...results, + ]; + } +} + +function mergeAssistant(target, source) { + appendText(target, source.content); + if (Array.isArray(source.toolUses) && source.toolUses.length > 0) { + target.toolUses = [...(target.toolUses || []), ...source.toolUses]; + } +} + +function normalizeTurns(history, currentMessage, modelId) { + const rawTurns = [...(Array.isArray(history) ? history : [])]; + if (currentMessage) rawTurns.push(currentMessage); + const turns = []; + + for (const raw of rawTurns) { + const isUser = !!raw?.userInputMessage; + const isAssistant = !!raw?.assistantResponseMessage; + if (isUser === isAssistant) continue; + + const turn = isUser + ? { userInputMessage: clone(raw.userInputMessage) } + : { assistantResponseMessage: clone(raw.assistantResponseMessage) }; + const previous = turns[turns.length - 1]; + if (turn.userInputMessage && previous?.userInputMessage) { + mergeUser(previous.userInputMessage, turn.userInputMessage); + } else if (turn.assistantResponseMessage && previous?.assistantResponseMessage) { + mergeAssistant(previous.assistantResponseMessage, turn.assistantResponseMessage); + } else { + turns.push(turn); + } + } + + if (turns[0]?.assistantResponseMessage) { + turns.unshift({ userInputMessage: { content: "continue", modelId } }); + } + if (turns.length === 0 || turns[turns.length - 1]?.assistantResponseMessage) { + turns.push({ userInputMessage: { content: "continue", modelId } }); + } + + for (const turn of turns) { + if (turn.userInputMessage) { + turn.userInputMessage.content = text(turn.userInputMessage.content).trim() || "continue"; + turn.userInputMessage.modelId ||= modelId; + if (turn.userInputMessage.userInputMessageContext?.tools) { + delete turn.userInputMessage.userInputMessageContext.tools; + } + } else { + turn.assistantResponseMessage.content = + text(turn.assistantResponseMessage.content).trim() || "..."; + } + } + return turns; +} + +function rawId(value) { + return typeof value === "string" ? value : ""; +} + +function reserveToolId(value, turnIndex, callIndex, name, usedIds) { + const sanitized = rawId(value).replace(/[^a-zA-Z0-9_-]/g, ""); + const generated = `call_msg${turnIndex}_tc${callIndex}_${name || "tool"}`; + const base = trimCodePoints( + TOOL_ID_PATTERN.test(sanitized) && sanitized ? sanitized : generated, + KIRO_TOOL_ID_MAX_LENGTH + ); + let candidate = base; + let suffix = 2; + while (usedIds.has(candidate)) { + const tail = `_${suffix++}`; + candidate = `${base.slice(0, KIRO_TOOL_ID_MAX_LENGTH - tail.length)}${tail}`; + } + usedIds.add(candidate); + return candidate; +} + +function normalizeToolInput(input) { + if (input && typeof input === "object" && !Array.isArray(input)) return clone(input); + if (typeof input === "string") { + try { + const parsed = JSON.parse(input); + if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) return parsed; + } catch { + return null; + } + } + return input == null ? {} : null; +} + +function normalizeToolResult(result) { + const content = Array.isArray(result?.content) + ? result.content.map((part) => ({ text: text(part?.text ?? part) })) + : [{ text: text(result?.content) }]; + return { + toolUseId: rawId(result?.toolUseId), + status: result?.status === "error" ? "error" : "success", + content: content.length > 0 ? content : [{ text: "" }], + }; +} + +function flattenResults(userMessage, results) { + for (const result of results) appendText(userMessage, toolResultText(result)); +} + +function cleanUserContext(userMessage) { + const context = userMessage.userInputMessageContext; + if (!context) return; + if (!context.toolResults?.length) delete context.toolResults; + if (!context.tools?.length) delete context.tools; + if (Object.keys(context).length === 0) delete userMessage.userInputMessageContext; +} + +function reconcileToolPair(assistant, user, turnIndex, nameMap, specNames, usedIds, repairs) { + const calls = Array.isArray(assistant.toolUses) ? assistant.toolUses : []; + const results = Array.isArray(user.userInputMessageContext?.toolResults) + ? user.userInputMessageContext.toolResults.map(normalizeToolResult) + : []; + if (calls.length === 0) { + if (results.length > 0) { + flattenResults(user, results); + repairs.orphanResults += results.length; + } + if (user.userInputMessageContext) delete user.userInputMessageContext.toolResults; + cleanUserContext(user); + return; + } + + const callQueues = new Map(); + const callRecords = calls.map((call, callIndex) => { + const key = rawId(call?.toolUseId); + const mappedName = nameMap.get(call?.name) || call?.name; + const input = normalizeToolInput(call?.input); + const record = { call, callIndex, key, mappedName, input, result: null }; + const queue = callQueues.get(key) || []; + queue.push(record); + callQueues.set(key, queue); + return record; + }); + + const orphanResults = []; + for (const result of results) { + const queue = callQueues.get(rawId(result.toolUseId)); + const record = queue?.find((candidate) => !candidate.result); + if (record) record.result = result; + else orphanResults.push(result); + } + + const keptCalls = []; + const keptResults = []; + for (const record of callRecords) { + const hasSpec = typeof record.mappedName === "string" && specNames.has(record.mappedName); + const valid = !!record.result && hasSpec && record.input !== null; + if (!valid) { + appendText(assistant, toolCallText({ name: record.mappedName, input: record.call?.input })); + repairs.missingResults += record.result ? 0 : 1; + repairs.invalidToolUses += hasSpec && record.input !== null ? 0 : 1; + if (record.result) { + flattenResults(user, [record.result]); + repairs.orphanResults++; + } + continue; + } + + const toolUseId = reserveToolId( + record.key, + turnIndex, + record.callIndex, + record.mappedName, + usedIds + ); + keptCalls.push({ + toolUseId, + name: record.mappedName, + input: record.input, + }); + keptResults.push({ ...record.result, toolUseId }); + } + + if (orphanResults.length > 0) { + flattenResults(user, orphanResults); + repairs.orphanResults += orphanResults.length; + } + + if (keptCalls.length > 0) assistant.toolUses = keptCalls; + else delete assistant.toolUses; + user.userInputMessageContext ||= {}; + if (keptResults.length > 0) user.userInputMessageContext.toolResults = keptResults; + else delete user.userInputMessageContext.toolResults; + cleanUserContext(user); +} + +/** Validate the final Kiro wire conversation without mutating it. */ +export function validateKiroConversation(history, currentMessage, toolSpecs = []) { + const errors = []; + const turns = [...(history || []), currentMessage].filter(Boolean); + const specNames = new Set(toolSpecs.map((spec) => spec?.toolSpecification?.name).filter(Boolean)); + const usedIds = new Set(); + + for (let index = 0; index < turns.length; index++) { + const expectedUser = index % 2 === 0; + const isUser = !!turns[index]?.userInputMessage; + if (isUser !== expectedUser) errors.push(`role:${index}`); + if (!isUser) { + const calls = turns[index].assistantResponseMessage?.toolUses || []; + const results = turns[index + 1]?.userInputMessage?.userInputMessageContext?.toolResults || []; + const callIds = calls.map((call) => call.toolUseId); + const resultIds = results.map((result) => result.toolUseId); + if (calls.length !== results.length || callIds.some((id) => !resultIds.includes(id))) { + errors.push(`pair:${index}`); + } + for (const call of calls) { + if (!call.toolUseId || usedIds.has(call.toolUseId)) errors.push(`id:${index}`); + usedIds.add(call.toolUseId); + if (!specNames.has(call.name)) errors.push(`spec:${index}`); + } + } else if (index === 0) { + const results = turns[index].userInputMessage?.userInputMessageContext?.toolResults; + if (results?.length) errors.push("orphan:0"); + } + } + if (!currentMessage?.userInputMessage?.content) errors.push("current"); + return { valid: errors.length === 0, errors }; +} + +function flattenAllStructuredTools(turns, repairs) { + for (const turn of turns) { + if (turn.assistantResponseMessage?.toolUses?.length) { + for (const call of turn.assistantResponseMessage.toolUses) { + appendText(turn.assistantResponseMessage, toolCallText(call)); + } + repairs.invalidToolUses += turn.assistantResponseMessage.toolUses.length; + delete turn.assistantResponseMessage.toolUses; + } + const user = turn.userInputMessage; + const results = user?.userInputMessageContext?.toolResults; + if (results?.length) { + flattenResults(user, results); + repairs.orphanResults += results.length; + delete user.userInputMessageContext.toolResults; + cleanUserContext(user); + } + } +} + +/** + * Produce a strict Kiro conversation: alternating turns, current user message, + * adjacent one-to-one tool use/result pairs, and tool specs only on currentMessage. + */ +export function canonicalizeKiroConversation({ + history, + currentMessage, + modelId, + toolSpecs = [], + nameMap = new Map(), +} = {}) { + const turns = normalizeTurns(history, currentMessage, modelId); + const repairs = { missingResults: 0, orphanResults: 0, invalidToolUses: 0 }; + const specNames = new Set(toolSpecs.map((spec) => spec?.toolSpecification?.name).filter(Boolean)); + const usedIds = new Set(); + + for (let index = 0; index < turns.length; index += 2) { + const user = turns[index].userInputMessage; + if (index === 0) { + const leadingResults = user.userInputMessageContext?.toolResults || []; + if (leadingResults.length > 0) { + flattenResults(user, leadingResults); + repairs.orphanResults += leadingResults.length; + delete user.userInputMessageContext.toolResults; + cleanUserContext(user); + } + } + const assistant = turns[index + 1]?.assistantResponseMessage; + const nextUser = turns[index + 2]?.userInputMessage; + if (assistant && nextUser) { + reconcileToolPair(assistant, nextUser, index + 1, nameMap, specNames, usedIds, repairs); + } + } + + const finalCurrent = turns[turns.length - 1]; + finalCurrent.userInputMessage.userInputMessageContext ||= {}; + if (toolSpecs.length > 0) { + finalCurrent.userInputMessage.userInputMessageContext.tools = clone(toolSpecs); + } + cleanUserContext(finalCurrent.userInputMessage); + + let finalHistory = turns.slice(0, -1); + let validation = validateKiroConversation(finalHistory, finalCurrent, toolSpecs); + if (!validation.valid) { + flattenAllStructuredTools(turns, repairs); + finalHistory = turns.slice(0, -1); + validation = validateKiroConversation(finalHistory, finalCurrent, toolSpecs); + } + + return { + history: finalHistory, + currentMessage: finalCurrent, + repairs, + valid: validation.valid, + errors: validation.errors, + }; +} diff --git a/open-sse/translator/index.js b/open-sse/translator/index.js index 37e5bde6..48bd1530 100644 --- a/open-sse/translator/index.js +++ b/open-sse/translator/index.js @@ -62,8 +62,13 @@ export function translateRequest(sourceFormat, targetFormat, model, body, stream // Always ensure tool_calls have id (some providers require it) ensureToolCallIds(result); - // Fix missing tool responses (insert empty tool_result if needed) - fixMissingToolResponses(result); + // Kiro performs stricter source-aware reconciliation after session replay. + // The generic helper inserts OpenAI `role: tool` messages, which a direct + // Claude→Kiro translator cannot consume and which cannot repair partial + // parallel tool results. + if (targetFormat !== FORMATS.KIRO) { + fixMissingToolResponses(result); + } // Capture thinking intent from the original (pre-translation) body, before any // format conversion strips/renames the fields. Applied after translation. diff --git a/open-sse/translator/request/claude-to-kiro.js b/open-sse/translator/request/claude-to-kiro.js index 51b9c7f8..8651e819 100644 --- a/open-sse/translator/request/claude-to-kiro.js +++ b/open-sse/translator/request/claude-to-kiro.js @@ -6,17 +6,10 @@ * direct `claude:kiro` route in ../index.js uses; it is NOT reached through the * claude→openai→kiro pivot. * - * It reproduces the two 400-guards that live in openai-to-kiro.js so that a - * Claude client which omits the `tools` array on a follow-up turn (typical - * after client-side compaction) does not trip Kiro's schema validator and get - * "Improperly formed request" (HTTP 400): - * - * 1. flattenClaudeToolInteractions — when the client sent NO tools, collapse - * every tool_use / tool_result block to plain text so no structured tool - * reference survives to trigger the "tools required" rule. - * 2. reconcileOrphanedToolResults — when tools ARE present, fold any - * tool_result whose tool_use_id has no matching tool_use back into the - * user text instead of leaving a dangling structured reference. + * After session replay it delegates to the shared Kiro conversation + * canonicalizer. That layer enforces adjacent one-to-one tool use/results, + * repairs partial parallel calls, and flattens compacted structured references + * that can no longer be represented safely. * * It also handles the 9router-synthetic `-agentic` / `-thinking` suffixes and * the `enabled` reasoning trigger, matching @@ -38,82 +31,17 @@ import { } from "../../config/kiroConstants.js"; import { DEFAULT_IMAGE_MIME } from "../schema/index.js"; import { ROLE, CLAUDE_BLOCK } from "../schema/index.js"; - -/** Stringify a tool_use input as a readable line. */ -function toolUseToText(name, input) { - let argStr; - try { - argStr = typeof input === "string" ? input : JSON.stringify(input ?? {}); - } catch { - argStr = "{}"; - } - return `[Tool call: ${name || "unknown"}(${argStr})]`; -} - -/** Render a Claude tool_result block's content as a readable line. */ -function toolResultBlockToText(content) { - let text = ""; - if (typeof content === "string") { - text = content; - } else if (Array.isArray(content)) { - text = content - .map((c) => (typeof c === "string" ? c : c?.text || "")) - .filter(Boolean) - .join("\n"); - } else if (content) { - try { - text = JSON.stringify(content); - } catch { - text = ""; - } - } - return `[Tool result: ${text}]`; -} - -/** - * When the client sent no tools, rewrite every tool_use (assistant) and - * tool_result (user) content block into plain text. Keeps text + images. - * Returns a new messages array; never mutates the input. - */ -function flattenClaudeToolInteractions(messages) { - const out = []; - for (const msg of messages) { - if (!msg) continue; - - if (msg.role === ROLE.ASSISTANT && Array.isArray(msg.content)) { - const parts = []; - for (const block of msg.content) { - if (block.type === CLAUDE_BLOCK.TEXT && block.text) { - parts.push(block.text); - } else if (block.type === CLAUDE_BLOCK.TOOL_USE) { - parts.push(toolUseToText(block.name, block.input)); - } - } - out.push({ ...msg, content: parts.join("\n") }); - continue; - } - - if (msg.role === ROLE.USER && Array.isArray(msg.content)) { - const newContent = msg.content.map((block) => - block.type === CLAUDE_BLOCK.TOOL_RESULT - ? { type: CLAUDE_BLOCK.TEXT, text: toolResultBlockToText(block.content) } - : block - ); - out.push({ ...msg, content: newContent }); - continue; - } - - out.push(msg); - } - return out; -} +import { + canonicalizeKiroConversation, + normalizeKiroToolSpecs, +} from "../concerns/kiroConversation.js"; /** * Convert Claude messages to Kiro history + currentMessage. * Kiro requires alternating user/assistant turns; consecutive same-role * messages are merged. */ -function convertClaudeMessagesToKiro(messages, tools, model) { +function convertClaudeMessagesToKiro(messages, model) { const history = []; let currentMessage = null; @@ -122,27 +50,6 @@ function convertClaudeMessagesToKiro(messages, tools, model) { let pendingToolResults = []; let pendingImages = []; let currentRole = null; - let toolsInjected = false; - - const clientProvidedTools = Array.isArray(tools) && tools.length > 0; - - const buildToolSpecs = () => - tools.map((t) => { - const name = t.name; - const description = t.description || `Tool: ${name}`; - const schema = t.input_schema || {}; - const normalizedSchema = - Object.keys(schema).length === 0 - ? { type: "object", properties: {}, required: [] } - : { ...schema, required: schema.required ?? [] }; - return { - toolSpecification: { - name, - description, - inputSchema: { json: normalizedSchema }, - }, - }; - }); const flushPending = () => { if (currentRole === ROLE.USER) { @@ -157,15 +64,6 @@ function convertClaudeMessagesToKiro(messages, tools, model) { toolResults: pendingToolResults, }; } - // Attach tools to the first user turn only. - if (clientProvidedTools && !toolsInjected) { - if (!userMsg.userInputMessage.userInputMessageContext) { - userMsg.userInputMessage.userInputMessageContext = {}; - } - userMsg.userInputMessage.userInputMessageContext.tools = buildToolSpecs(); - toolsInjected = true; - } - history.push(userMsg); currentMessage = userMsg; pendingUserContent = []; @@ -209,7 +107,7 @@ function convertClaudeMessagesToKiro(messages, tools, model) { } pendingToolResults.push({ toolUseId: block.tool_use_id, - status: "success", + status: block.is_error ? "error" : "success", content: [{ text: resultContent }], }); } @@ -256,14 +154,7 @@ function convertClaudeMessagesToKiro(messages, tools, model) { } } - // Grab tools from the first history user turn before cleanup strips them. - const firstHistoryTools = - history[0]?.userInputMessage?.userInputMessageContext?.tools; - history.forEach((item) => { - if (item.userInputMessage?.userInputMessageContext?.tools) { - delete item.userInputMessage.userInputMessageContext.tools; - } if ( item.userInputMessage?.userInputMessageContext && Object.keys(item.userInputMessage.userInputMessageContext).length === 0 @@ -307,66 +198,9 @@ function convertClaudeMessagesToKiro(messages, tools, model) { currentMessage = { userInputMessage: { content: "", modelId: model } }; } - // Inject tools into currentMessage after cleanup if not already present. - if ( - firstHistoryTools?.length > 0 && - !currentMessage.userInputMessage.userInputMessageContext?.tools - ) { - if (!currentMessage.userInputMessage.userInputMessageContext) { - currentMessage.userInputMessage.userInputMessageContext = {}; - } - currentMessage.userInputMessage.userInputMessageContext.tools = - firstHistoryTools; - } - return { history: mergedHistory, currentMessage }; } -/** - * Fold orphaned toolResults (those whose toolUseId has no matching toolUse in - * any assistant turn) back into the user text, removing the dangling - * structured reference that makes Kiro 400. - */ -function reconcileOrphanedToolResults(history, currentMessage) { - const validIds = new Set(); - for (const h of history) { - const arm = h.assistantResponseMessage; - if (!arm) continue; - for (const tu of arm.toolUses || []) { - if (tu.toolUseId) validIds.add(tu.toolUseId); - } - } - - const carriers = currentMessage ? [...history, currentMessage] : history; - for (const item of carriers) { - const uim = item.userInputMessage; - const ctx = uim?.userInputMessageContext; - if (!ctx?.toolResults?.length) continue; - - const kept = []; - const salvaged = []; - for (const tr of ctx.toolResults) { - if (validIds.has(tr.toolUseId)) { - kept.push(tr); - } else { - const text = Array.isArray(tr.content) - ? tr.content.map((c) => c?.text || "").join("\n") - : ""; - salvaged.push(`[Tool result: ${text}]`); - } - } - - if (salvaged.length === 0) continue; - - const extra = salvaged.join("\n"); - uim.content = uim.content ? `${uim.content}\n\n${extra}` : extra; - ctx.toolResults = kept; - if (kept.length === 0 && !ctx.tools?.length) { - delete uim.userInputMessageContext; - } - } -} - function extractClaudeSystemText(system) { if (!system) return ""; if (typeof system === "string") return system; @@ -383,9 +217,8 @@ function extractClaudeSystemText(system) { * Build a Kiro payload directly from a Claude Messages API request body. */ export function claudeToKiroRequest(model, body, stream, credentials) { - let messages = Array.isArray(body.messages) ? body.messages : []; + const messages = Array.isArray(body.messages) ? body.messages : []; const tools = Array.isArray(body.tools) ? body.tools : []; - const clientProvidedTools = tools.length > 0; const maxTokens = body.max_tokens || 32000; const temperature = body.temperature; const topP = body.top_p; @@ -397,21 +230,8 @@ export function claudeToKiroRequest(model, body, stream, credentials) { const additionalModelRequestFields = buildKiroAdditionalModelRequestFieldsForModel(thinkingBody, upstreamModel); const usesNativeGptEffort = usesKiroNativeGptEffort(thinkingBody, upstreamModel); - // Guard 1: no client tools → flatten all tool interactions to text. - if (!clientProvidedTools) { - messages = flattenClaudeToolInteractions(messages); - } - - const { history, currentMessage } = convertClaudeMessagesToKiro( - messages, - tools, - upstreamModel - ); - - // Guard 2: tools present → reconcile dangling tool_results. - if (clientProvidedTools) { - reconcileOrphanedToolResults(history, currentMessage); - } + const { specs: toolSpecs, nameMap } = normalizeKiroToolSpecs(tools); + const { history, currentMessage } = convertClaudeMessagesToKiro(messages, upstreamModel); // api_key / idc / external_idp must never use the shared default ARN (belongs // to another account → 403 "bearer token invalid"); OAuth/social fall back to it. @@ -460,7 +280,14 @@ export function claudeToKiroRequest(model, body, stream, credentials) { history, currentMessage, }); - const replayCurrent = replay.currentMessage?.userInputMessage || {}; + const canonical = canonicalizeKiroConversation({ + history: replay.history, + currentMessage: replay.currentMessage, + modelId: upstreamModel, + toolSpecs, + nameMap, + }); + const replayCurrent = canonical.currentMessage.userInputMessage; const userInputMessage = { content: replayCurrent.content || "", modelId: upstreamModel, @@ -482,7 +309,7 @@ export function claudeToKiroRequest(model, body, stream, credentials) { currentMessage: { userInputMessage, }, - history: replay.history, + history: canonical.history, }, agentMode: "vibe", }; diff --git a/open-sse/translator/request/openai-to-kiro.js b/open-sse/translator/request/openai-to-kiro.js index 798c827f..aa776949 100644 --- a/open-sse/translator/request/openai-to-kiro.js +++ b/open-sse/translator/request/openai-to-kiro.js @@ -20,148 +20,10 @@ import { import { parseDataUri } from "../concerns/image.js"; import { DEFAULT_IMAGE_MIME } from "../schema/index.js"; import { ROLE, OPENAI_BLOCK, CLAUDE_BLOCK } from "../schema/index.js"; - -/** Render a single tool call as a readable text line. */ -function toolCallToText(name, input) { - let argStr; - try { - argStr = typeof input === "string" ? input : JSON.stringify(input ?? {}); - } catch { - argStr = "{}"; - } - return `[Tool call: ${name || "unknown"}(${argStr})]`; -} - -/** Render a tool result (string or content-block array) as a text line. */ -function toolResultToText(content) { - const text = Array.isArray(content) - ? content.map(c => (typeof c === "string" ? c : c.text || "")).join("\n") - : (typeof content === "string" ? content : ""); - return `[Tool result: ${text}]`; -} - -/** - * Flatten all tool calls/results in a conversation into plain text. - * - * Kiro's schema validator requires a non-empty - * currentMessage.userInputMessageContext.tools array whenever the history - * references any tool use; otherwise it returns "Improperly formed request" - * (HTTP 400). A client can hit this by omitting the `tools` array on a - * follow-up request — typically after client-side compaction (e.g. OpenCode). - * - * Rather than fabricate stub tool specs — which would advertise tool-calling - * capability the client never requested and may not handle, risking a phantom - * tool call on an otherwise plain turn — we collapse the tool interaction into - * text. The request stays honest, and since no structured tool content - * remains, the validator's "tools required" rule never fires. - * - * Only invoked when the client did NOT send tools; when tools are present the - * structured form is preserved. - */ -function flattenToolInteractions(messages) { - const out = []; - - for (const msg of messages) { - // OpenAI tool-result message → user text line - if (msg.role === ROLE.TOOL) { - out.push({ role: ROLE.USER, content: toolResultToText(msg.content) }); - continue; - } - - if (msg.role === ROLE.ASSISTANT) { - const parts = []; - if (Array.isArray(msg.content)) { - for (const c of msg.content) { - if (c.type === CLAUDE_BLOCK.TOOL_USE) { - parts.push(toolCallToText(c.name, c.input)); - } else if (c.type === OPENAI_BLOCK.TEXT || c.text) { - parts.push(c.text || ""); - } - } - } else if (typeof msg.content === "string") { - parts.push(msg.content); - } - for (const tc of msg.tool_calls || []) { - parts.push(toolCallToText(tc.function?.name, tc.function?.arguments)); - } - out.push({ role: ROLE.ASSISTANT, content: parts.filter(Boolean).join("\n") }); - continue; - } - - // User messages: replace tool_result blocks with text, keep text + images. - if (msg.role === ROLE.USER && Array.isArray(msg.content)) { - const newContent = msg.content.map(c => - c.type === CLAUDE_BLOCK.TOOL_RESULT - ? { type: OPENAI_BLOCK.TEXT, text: toolResultToText(c.content) } - : c - ); - out.push({ ...msg, content: newContent }); - continue; - } - - out.push(msg); - } - - return out; -} - -/** - * Reconcile orphaned toolResults — those whose toolUseId has no matching - * toolUse in any assistant message. This happens when client-side compaction - * truncates the conversation and removes the assistant message containing the - * tool_use, but keeps the user message with the corresponding tool_result. - * - * A dangling structured reference makes Kiro return 400, so it must be removed. - * But the client deliberately kept the result content through compaction, so - * rather than discard it we fold it back into the user message as text — the - * same shape flattenToolInteractions() produces. The 400 trigger (the - * structured reference) is gone; the content survives. - * - * `messages` is every carrier that can hold toolResults — both history items - * and the popped-out currentMessage (orphans can land on either). - */ -function reconcileOrphanedToolResults(history, currentMessage) { - // Phase 1: collect all valid toolUseIds from assistant messages in history. - // (currentMessage is always a user turn, so it carries no toolUses.) - const validIds = new Set(); - for (const h of history) { - const arm = h.assistantResponseMessage; - if (!arm) continue; - for (const tu of arm.toolUses || []) { - if (tu.toolUseId) validIds.add(tu.toolUseId); - } - } - - // Phase 2: across history + currentMessage, keep results with a matching - // toolUse and salvage the rest as text. - const carriers = currentMessage ? [...history, currentMessage] : history; - for (const item of carriers) { - const uim = item.userInputMessage; - const ctx = uim?.userInputMessageContext; - if (!ctx?.toolResults?.length) continue; - - const kept = []; - const salvaged = []; - for (const tr of ctx.toolResults) { - if (validIds.has(tr.toolUseId)) { - kept.push(tr); - } else { - salvaged.push(toolResultToText(tr.content)); - } - } - - if (salvaged.length === 0) continue; // no orphans — leave untouched - - // Fold orphaned result content into the user text so it is not lost - const extra = salvaged.join("\n"); - uim.content = uim.content ? `${uim.content}\n\n${extra}` : extra; - - ctx.toolResults = kept; - if (kept.length === 0 && !ctx.tools?.length) { - delete uim.userInputMessageContext; - } - } -} +import { + canonicalizeKiroConversation, + normalizeKiroToolSpecs, +} from "../concerns/kiroConversation.js"; /** * Safely parse JSON string, returning fallback on failure. @@ -177,26 +39,15 @@ function safeJSONParse(str, fallback) { * * Returns { history, currentMessage }. */ -function convertMessages(messages, tools, model) { +function convertMessages(messages, model) { let history = []; let currentMessage = null; - const clientProvidedTools = tools && tools.length > 0; - - // When the client did not send tools, flatten any tool calls/results in the - // history into plain text (see flattenToolInteractions). This keeps the - // request honest and sidesteps Kiro's "tools required" 400, since no - // structured tool content survives to trigger it. - if (!clientProvidedTools) { - messages = flattenToolInteractions(messages); - } - let pendingUserContent = []; let pendingAssistantContent = []; let pendingToolResults = []; let pendingImages = []; let currentRole = null; - let toolsInjectedToFirstUserMsg = false; const flushPending = () => { if (currentRole === "user") { @@ -219,39 +70,6 @@ function convertMessages(messages, tools, model) { }; } - // Add tools to the user message that has no preceding assistant messages, - // OR the first user message (whichever comes first after any opening - // assistant messages). We track whether any user message has already - // received tools via a flag on the history array. - if (clientProvidedTools && !toolsInjectedToFirstUserMsg) { - if (!userMsg.userInputMessage.userInputMessageContext) { - userMsg.userInputMessage.userInputMessageContext = {}; - } - userMsg.userInputMessage.userInputMessageContext.tools = tools.map(t => { - const name = t.function?.name || t.name; - let description = t.function?.description || t.description || ""; - - if (!description.trim()) { - description = `Tool: ${name}`; - } - - const schema = t.function?.parameters || t.parameters || t.input_schema || {}; - // Normalize schema: Kiro requires required[] and proper type/properties - const normalizedSchema = Object.keys(schema).length === 0 - ? { type: "object", properties: {}, required: [] } - : { ...schema, required: schema.required ?? [] }; - - return { - toolSpecification: { - name, - description, - inputSchema: { json: normalizedSchema } - } - }; - }); - toolsInjectedToFirstUserMsg = true; - } - history.push(userMsg); currentMessage = userMsg; pendingUserContent = []; @@ -327,7 +145,7 @@ function convertMessages(messages, tools, model) { pendingToolResults.push({ toolUseId: block.tool_use_id, - status: "success", + status: block.is_error ? "error" : "success", content: [{ text: text }] }); }); @@ -339,7 +157,7 @@ function convertMessages(messages, tools, model) { const toolContent = typeof msg.content === "string" ? msg.content : ""; pendingToolResults.push({ toolUseId: msg.tool_call_id, - status: "success", + status: msg.is_error || msg.status === "error" ? "error" : "success", content: [{ text: toolContent }] }); } else if (content) { @@ -413,14 +231,8 @@ function convertMessages(messages, tools, model) { } } - // Grab tools from first history item BEFORE cleanup removes them - const firstHistoryTools = history[0]?.userInputMessage?.userInputMessageContext?.tools; - // Clean up history for Kiro API compatibility history.forEach(item => { - if (item.userInputMessage?.userInputMessageContext?.tools) { - delete item.userInputMessage.userInputMessageContext.tools; - } if (item.userInputMessage?.userInputMessageContext && Object.keys(item.userInputMessage.userInputMessageContext).length === 0) { delete item.userInputMessage.userInputMessageContext; @@ -473,33 +285,6 @@ function convertMessages(messages, tools, model) { }; } - // Reconcile orphaned toolResults across history AND currentMessage — when - // client-side compaction removes assistant messages containing tool_use but - // keeps the tool_result, the dangling reference triggers a Kiro 400. Fold the - // content back into the user text instead of discarding it. Run after - // currentMessage is finalized (an orphan can be merged into it) and before - // tool injection (which may re-add userInputMessageContext). - // - // Only needed on the tools-present path: when the client sent no tools, - // flattenToolInteractions already collapsed every toolResult to text, so - // there is nothing structured left to orphan. - if (clientProvidedTools) { - reconcileOrphanedToolResults(mergedHistory, currentMessage); - } - - // Inject tools into currentMessage AFTER cleanup. Tools only exist here when - // the client explicitly sent them (otherwise flattenToolInteractions already - // collapsed all tool content to text upstream, so there is nothing to carry). - const resolvedTools = firstHistoryTools; - - if (resolvedTools?.length > 0 && - !currentMessage.userInputMessage.userInputMessageContext?.tools) { - if (!currentMessage.userInputMessage.userInputMessageContext) { - currentMessage.userInputMessage.userInputMessageContext = {}; - } - currentMessage.userInputMessage.userInputMessageContext.tools = resolvedTools; - } - return { history: mergedHistory, currentMessage }; } @@ -532,7 +317,8 @@ export function openaiToKiroRequest(model, body, stream, credentials) { const additionalModelRequestFields = buildKiroAdditionalModelRequestFieldsForModel(thinkingBody, upstreamModel); const usesNativeGptEffort = usesKiroNativeGptEffort(thinkingBody, upstreamModel); - const { history, currentMessage } = convertMessages(messages, tools, upstreamModel); + const { specs: toolSpecs, nameMap } = normalizeKiroToolSpecs(tools); + const { history, currentMessage } = convertMessages(messages, upstreamModel); // API-key (headless) auth uses a raw CodeWhisperer credential whose profile is // account-specific. Injecting the shared builder-id/social *default* placeholder @@ -586,7 +372,14 @@ export function openaiToKiroRequest(model, body, stream, credentials) { history, currentMessage, }); - const replayCurrent = replay.currentMessage?.userInputMessage || {}; + const canonical = canonicalizeKiroConversation({ + history: replay.history, + currentMessage: replay.currentMessage, + modelId: upstreamModel, + toolSpecs, + nameMap, + }); + const replayCurrent = canonical.currentMessage.userInputMessage; const payload = { conversationState: { @@ -607,7 +400,7 @@ export function openaiToKiroRequest(model, body, stream, credentials) { }) } }, - history: replay.history + history: canonical.history }, agentMode: "vibe", }; diff --git a/open-sse/utils/kiroSessionReplay.js b/open-sse/utils/kiroSessionReplay.js index 11cae9cd..d758eed6 100644 --- a/open-sse/utils/kiroSessionReplay.js +++ b/open-sse/utils/kiroSessionReplay.js @@ -42,6 +42,14 @@ function findFirstUserIndex(history) { return history.findIndex((item) => item?.userInputMessage); } +function hasToolResults(message) { + return !!message?.userInputMessage?.userInputMessageContext?.toolResults?.length; +} + +function canReplaceSessionStart(history, firstUserIndex) { + return firstUserIndex === 0 && !hasToolResults(history[firstUserIndex]); +} + function rememberSessionStart(key, entry) { if (sessionStartStore.size >= MAX_SESSION_STARTS) { sessionStartStore.delete(sessionStartStore.keys().next().value); @@ -73,10 +81,13 @@ export function applyKiroSessionReplay({ existing.lastUsed = Date.now(); const firstUserIndex = findFirstUserIndex(baseHistory); const sessionStart = ensureUserMessageModelId(clone(existing.sessionStart), modelId); - if (firstUserIndex >= 0) { + if (canReplaceSessionStart(baseHistory, firstUserIndex)) { baseHistory[firstUserIndex] = sessionStart; } else { baseHistory.unshift(sessionStart); + if (baseHistory.length === 1) { + baseHistory.push({ assistantResponseMessage: { content: "..." } }); + } } return { history: ensureHistoryModelIds(baseHistory, modelId), @@ -88,10 +99,18 @@ export function applyKiroSessionReplay({ const firstUserIndex = findFirstUserIndex(baseHistory); let sessionStart; let nextCurrent = ensureUserMessageModelId(baseCurrent, modelId); - if (firstUserIndex >= 0) { + if (canReplaceSessionStart(baseHistory, firstUserIndex)) { sessionStart = prefixUserMessage(baseHistory[firstUserIndex], contentPrefix, modelId); baseHistory[firstUserIndex] = clone(sessionStart); nextCurrent = prefixUserMessage(baseCurrent, currentContentPrefix, modelId); + } else if (firstUserIndex >= 0) { + sessionStart = prefixUserMessage( + { userInputMessage: { content: "", modelId } }, + contentPrefix, + modelId + ); + baseHistory.unshift(clone(sessionStart)); + nextCurrent = prefixUserMessage(baseCurrent, currentContentPrefix, modelId); } else { sessionStart = prefixUserMessage(baseCurrent, contentPrefix, modelId); nextCurrent = clone(sessionStart); diff --git a/src/app/api/oauth/kiro/api-key/route.js b/src/app/api/oauth/kiro/api-key/route.js index 139df9b5..bd2140b3 100644 --- a/src/app/api/oauth/kiro/api-key/route.js +++ b/src/app/api/oauth/kiro/api-key/route.js @@ -5,8 +5,8 @@ import { createProviderConnection } from "@/models"; /** * POST /api/oauth/kiro/api-key * Import a Kiro API key (headless auth). The key is a long-lived bearer - * credential — there is no refresh token. It is validated by listing - * CodeWhisperer profiles, then stored with authMethod="api_key". + * credential — there is no refresh token. It is validated against the Amazon + * Q model catalog, then stored with authMethod="api_key". */ export async function POST(request) { try { @@ -21,7 +21,7 @@ export async function POST(request) { const kiroService = new KiroService(); - // Validate the key and resolve its profileArn via ListAvailableProfiles + // Validate the key against the same Amazon Q surface used for inference. const credential = await kiroService.validateApiKey( apiKey, region || "us-east-1" @@ -40,7 +40,7 @@ export async function POST(request) { expiresAt: new Date(Date.now() + 365 * 24 * 60 * 60 * 1000).toISOString(), email: email || null, providerSpecificData: { - profileArn: credential.profileArn, + ...(credential.profileArn ? { profileArn: credential.profileArn } : {}), region: credential.region, authMethod: "api_key", provider: "API Key", diff --git a/src/lib/oauth/services/kiro.js b/src/lib/oauth/services/kiro.js index a739661e..f82089d8 100644 --- a/src/lib/oauth/services/kiro.js +++ b/src/lib/oauth/services/kiro.js @@ -260,11 +260,9 @@ export class KiroService { } /** - * List available CodeWhisperer profiles for a token (or API key) and return - * the best-matching profileArn. AWS SSO OIDC logins return no profileArn, so - * it must be fetched separately — the same call works for API-key auth. - * Accepts both `arn` and `profileArn` response field names (the API-key - * JSON-1.0 surface returns `arn`). + * List available CodeWhisperer profiles for OAuth/IDC tokens and return the + * best-matching profileArn. API keys use the Amazon Q model catalog instead; + * ListAvailableProfiles does not support TokenType=API_KEY. */ async listAvailableProfiles(accessToken, region = "us-east-1") { assertValidAwsRegion(region); @@ -294,10 +292,41 @@ export class KiroService { } /** - * Validate an API-key credential by listing profiles with it. API keys are - * long-lived bearer tokens (no refresh), so the only way to validate one is - * to make an authenticated CodeWhisperer call. Returns a credential object - * ready to persist as a "kiro" connection with authMethod="api_key". + * Validate an API key against the Amazon Q model catalog. A bearer-only call + * to ListAvailableProfiles can return HTTP 200 with an empty list for an + * arbitrary key, so it is not proof that the key can run inference. + */ + async listAvailableApiKeyModels(apiKey, region = "us-east-1") { + assertValidAwsRegion(region); + const params = new URLSearchParams({ origin: "AI_EDITOR" }); + const endpoint = `https://q.${region}.amazonaws.com/ListAvailableModels?${params}`; + const response = await fetch(endpoint, { + method: "GET", + headers: { + "Authorization": `Bearer ${apiKey}`, + "TokenType": "API_KEY", + "Accept": "application/json", + "User-Agent": "AWS-SDK-JS/3.0.0 kiro-ide/1.0.0", + "X-Amz-User-Agent": "aws-sdk-js/3.0.0 kiro-ide/1.0.0", + }, + }); + + if (!response.ok) { + const error = await response.text(); + throw new Error(`Failed to list API-key models: ${error}`); + } + + const data = await response.json(); + const models = Array.isArray(data?.models) ? data.models : []; + if (models.length === 0) { + throw new Error("API key returned no available models"); + } + return models; + } + + /** + * Validate an API-key credential through the same Amazon Q surface used for + * inference. API keys are account-bound but do not require a profileArn. */ async validateApiKey(apiKey, region = "us-east-1") { if (!apiKey || typeof apiKey !== "string" || !apiKey.trim()) { @@ -305,9 +334,8 @@ export class KiroService { } const trimmed = apiKey.trim(); - let profileArn = null; try { - profileArn = await this.listAvailableProfiles(trimmed, region); + await this.listAvailableApiKeyModels(trimmed, region); } catch (error) { throw new Error(`API key validation failed: ${error.message}`); } @@ -315,7 +343,7 @@ export class KiroService { return { accessToken: trimmed, refreshToken: null, - profileArn, + profileArn: null, region, authMethod: "api_key", }; diff --git a/tests/unit/kiro-api-key-endpoint-routing.test.js b/tests/unit/kiro-api-key-endpoint-routing.test.js new file mode 100644 index 00000000..a0750adc --- /dev/null +++ b/tests/unit/kiro-api-key-endpoint-routing.test.js @@ -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(); + }); +}); diff --git a/tests/unit/kiro-conversation-canonicalization.test.js b/tests/unit/kiro-conversation-canonicalization.test.js new file mode 100644 index 00000000..e19984d1 --- /dev/null +++ b/tests/unit/kiro-conversation-canonicalization.test.js @@ -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); + }); +}); diff --git a/tests/unit/kiro-profile-arn.test.js b/tests/unit/kiro-profile-arn.test.js index 1925582b..bfa3906c 100644 --- a/tests/unit/kiro-profile-arn.test.js +++ b/tests/unit/kiro-profile-arn.test.js @@ -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/ + ); + }); });