Files
decolua--9router/open-sse/executors/kiro.js
T
wehub-resource-sync 05fcd08057
Deploy GitBook to 9router.github.io / build-deploy (push) Failing after 2s
chore: import upstream snapshot with attribution
2026-07-13 12:21:01 +08:00

593 lines
23 KiB
JavaScript

import { BaseExecutor } from "./base.js";
import { PROVIDERS } from "../config/providers.js";
import { 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";
import { getCapabilitiesForModel } from "../providers/capabilities.js";
/**
* KiroExecutor - Executor for Kiro AI (AWS CodeWhisperer)
* Uses AWS CodeWhisperer streaming API with AWS EventStream binary format
*/
export class KiroExecutor extends BaseExecutor {
constructor() {
super("kiro", PROVIDERS.kiro);
}
buildHeaders(credentials, stream = true) {
const headers = {
...this.config.headers,
"Amz-Sdk-Request": "attempt=1; max=3",
"Amz-Sdk-Invocation-Id": uuidv4()
};
// 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`
// header so CodeWhisperer treats it as a long-lived API key rather than an
// OIDC/social access token. Mirrors the Kiro IDE headless-auth behavior.
// Enterprise / Microsoft Entra (external_idp) tokens are OAuth access tokens,
// but CodeWhisperer requires TokenType=EXTERNAL_IDP to bind them to profiles.
const authMethod = credentials?.providerSpecificData?.authMethod;
const isApiKey = authMethod === "api_key";
const isExternalIdp = authMethod === "external_idp";
const apiKey = credentials?.apiKey || (isApiKey ? credentials?.accessToken : null);
if (isApiKey && apiKey) {
headers["Authorization"] = `Bearer ${apiKey}`;
headers["tokentype"] = "API_KEY";
} else if (credentials.accessToken) {
headers["Authorization"] = `Bearer ${credentials.accessToken}`;
if (isExternalIdp) {
headers["TokenType"] = "EXTERNAL_IDP";
}
}
return headers;
}
/**
* Auth-aware endpoint ordering.
*
* API-key Kiro connections store a raw CodeWhisperer credential (validated
* against codewhisperer.us-east-1.amazonaws.com via ListAvailableProfiles).
* 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
* 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.
*/
getOrderedBaseUrls(credentials) {
const baseUrls = this.getBaseUrls();
const authMethod = credentials?.providerSpecificData?.authMethod;
// IAM Identity Center (idc) tokens are AWS SSO access tokens — the same
// family as external_idp/api_key. The kiro.dev gateway rejects them with
// 403 "bearer token invalid", so they must hit the CodeWhisperer
// *.amazonaws.com surface, and in the region the token was minted in
// (the baseUrls are hardcoded us-east-1).
const isCodeWhispererSurface =
authMethod === "api_key" || authMethod === "external_idp" || authMethod === "idc";
if (!isCodeWhispererSurface) return baseUrls;
const region = (credentials?.providerSpecificData?.region || "us-east-1").trim();
const regionalize = (u) =>
region && region !== "us-east-1" && u.includes("amazonaws.com")
? u.replace(/([a-z]+)\.[a-z0-9-]+\.amazonaws\.com/, `$1.${region}.amazonaws.com`)
: u;
const amazon = baseUrls.filter((u) => u.includes("amazonaws.com")).map(regionalize);
const others = baseUrls.filter((u) => !u.includes("amazonaws.com"));
return amazon.length > 0 ? [...amazon, ...others] : baseUrls;
}
buildUrl(model, stream, urlIndex = 0, credentials = null) {
const baseUrls = this.getOrderedBaseUrls(credentials);
return baseUrls[urlIndex] || baseUrls[0] || this.config.baseUrl;
}
transformRequest(model, body, stream, credentials) {
return body;
}
/**
* Kiro execute — delegate to BaseExecutor for endpoint fallback + retry, then
* transform the binary AWS EventStream into OpenAI-shaped SSE on success.
*
* BaseExecutor.execute() walks config.baseUrls (runtime.us-east-1.kiro.dev →
* codewhisperer → q) advancing to the next host on 429 (shouldRetry) and on
* network/5xx errors, while tryRetry handles in-place retries per `retry: {429: 2}`.
* Note: api-key connections reorder these so the *.amazonaws.com hosts come
* first — see getOrderedBaseUrls/buildUrl above.
* Note: the baseUrls are alternate surfaces of one regional service, so rotation
* is edge-level failover — it does not grant fresh 429 quota. Per-account 429
* spreading is handled upstream by account rotation in sse/handlers/chat.js.
*
* Errors are returned untransformed so the upstream handler can read the body,
* classify the status, and trigger account fallback/cooldown.
*/
async execute(args) {
const result = await super.execute(args);
if (result?.response?.ok) {
result.response = this.transformEventStreamToSSE(result.response, args.model);
}
return result;
}
/**
* Transform AWS EventStream binary response to SSE text stream
* Using TransformStream instead of ReadableStream.pull() to avoid Workers timeout
*/
transformEventStreamToSSE(response, model) {
let buffer = new Uint8Array(0);
let chunkIndex = 0;
const responseId = `chatcmpl-${Date.now()}`;
const created = Math.floor(Date.now() / 1000);
const capabilityModel = resolveKiroModel(model).upstream;
const contextWindow = getCapabilitiesForModel("kiro", capabilityModel).contextWindow || 200000;
const state = {
endDetected: false,
finishEmitted: false,
hasToolCalls: false,
hasReasoningContent: false,
reasoningChunkCount: 0,
toolCallIndex: 0,
seenToolIds: new Map(),
inThinking: false
};
const transformStream = new TransformStream({
async transform(chunk, controller) {
// Track output so we can emit a keepalive if this frame yields no chunk.
const enqueueCountBefore = chunkIndex;
// Append to buffer
const newBuffer = new Uint8Array(buffer.length + chunk.length);
newBuffer.set(buffer);
newBuffer.set(chunk, buffer.length);
buffer = newBuffer;
// Parse events from buffer
let iterations = 0;
const maxIterations = 1000;
while (buffer.length >= 16 && iterations < maxIterations) {
iterations++;
const view = new DataView(buffer.buffer, buffer.byteOffset);
const totalLength = view.getUint32(0, false);
if (totalLength < 16 || totalLength > buffer.length || buffer.length < totalLength) break;
const eventData = buffer.slice(0, totalLength);
buffer = buffer.slice(totalLength);
const event = parseEventFrame(eventData);
if (!event) continue;
const eventType = event.headers[":event-type"] || "";
// Track total content length for token estimation
if (!state.totalContentLength) state.totalContentLength = 0;
if (!state.contextUsagePercentage) state.contextUsagePercentage = 0;
// Handle assistantResponseEvent
if (eventType === "assistantResponseEvent" && event.payload?.content) {
let content = event.payload.content;
// Kiro Claude models can leak <thinking> blocks into the content stream.
// We strip these literal tags to prevent duplication, as the reasoning
// is already routed correctly via reasoningContentEvent.
if (state.inThinking) {
if (content.includes("</thinking>")) {
state.inThinking = false;
const after = content.split("</thinking>").slice(1).join("</thinking>");
content = after.startsWith("\n") ? after.substring(1) : after;
} else {
content = ""; // Drop entirely while inside thinking block
}
} else if (content.includes("<thinking>")) {
state.inThinking = true;
if (content.includes("</thinking>")) {
state.inThinking = false;
const before = content.split("<thinking>")[0];
const after = content.split("</thinking>").slice(1).join("</thinking>");
content = before + (after.startsWith("\n") ? after.substring(1) : after);
} else {
content = content.split("<thinking>")[0];
}
}
if (!content && state.hasReasoningContent) {
// If we stripped everything, skip emitting an empty content chunk
continue;
}
state.totalContentLength += content.length;
const chunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: chunkIndex === 0
? { role: "assistant", content }
: { content },
finish_reason: null
}]
};
chunkIndex++;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
}
// Handle reasoningContentEvent (Kiro thinking / reasoning)
// Kiro returns reasoning as a separate event when the request system
// prompt contains <thinking_mode>enabled</thinking_mode>. Surface it
// as OpenAI delta.reasoning_content so downstream translators can map
// it back to Claude thinking blocks / Anthropic reasoning, etc.
if (eventType === "reasoningContentEvent") {
const reasoning = event.payload?.reasoningContentEvent || event.payload || {};
const reasoningText = (typeof reasoning === "string")
? reasoning
: (reasoning.text || reasoning.content || "");
if (reasoningText) {
state.hasReasoningContent = true;
state.totalContentLength += reasoningText.length;
const reasoningDelta = state.reasoningChunkCount === 0 && chunkIndex === 0
? { role: "assistant", reasoning_content: reasoningText }
: { reasoning_content: reasoningText };
const chunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: reasoningDelta,
finish_reason: null
}]
};
chunkIndex++;
state.reasoningChunkCount++;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
}
}
// Handle codeEvent
if (eventType === "codeEvent" && event.payload?.content) {
const chunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: { content: event.payload.content },
finish_reason: null
}]
};
chunkIndex++;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
}
// Handle toolUseEvent
if (eventType === "toolUseEvent" && event.payload) {
state.hasToolCalls = true;
const toolUse = event.payload;
const toolUses = Array.isArray(toolUse) ? toolUse : [toolUse];
for (const singleToolUse of toolUses) {
const toolCallId = singleToolUse.toolUseId || `call_${Date.now()}`;
const toolName = singleToolUse.name || "";
const toolInput = singleToolUse.input;
let toolIndex;
const isNewTool = !state.seenToolIds.has(toolCallId);
if (isNewTool) {
toolIndex = state.toolCallIndex++;
state.seenToolIds.set(toolCallId, toolIndex);
const startChunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: {
...(chunkIndex === 0 ? { role: "assistant" } : {}),
tool_calls: [{
index: toolIndex,
id: toolCallId,
type: "function",
function: {
name: toolName,
arguments: ""
}
}]
},
finish_reason: null
}]
};
chunkIndex++;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(startChunk)}\n\n`));
} else {
toolIndex = state.seenToolIds.get(toolCallId);
}
if (toolInput !== undefined) {
let argumentsStr;
if (typeof toolInput === 'string') {
argumentsStr = toolInput;
} else if (typeof toolInput === 'object') {
argumentsStr = JSON.stringify(toolInput);
} else {
continue;
}
const argsChunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: {
tool_calls: [{
index: toolIndex,
function: {
arguments: argumentsStr
}
}]
},
finish_reason: null
}]
};
chunkIndex++;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(argsChunk)}\n\n`));
}
}
}
// Handle messageStopEvent
if (eventType === "messageStopEvent") {
const chunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: {},
finish_reason: state.hasToolCalls ? "tool_calls" : "stop"
}]
};
state.finishEmitted = true;
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
}
// Handle contextUsageEvent to extract contextUsagePercentage
if (eventType === "contextUsageEvent" && event.payload?.contextUsagePercentage) {
state.contextUsagePercentage = event.payload.contextUsagePercentage;
// Mark that we received context usage event
state.hasContextUsage = true;
}
// Handle meteringEvent - mark that we received it
if (eventType === "meteringEvent") {
state.hasMeteringEvent = true;
}
// Handle metricsEvent for token usage
if (eventType === "metricsEvent") {
// Extract usage data from metricsEvent payload
const metrics = event.payload?.metricsEvent || event.payload;
if (metrics && typeof metrics === 'object') {
const inputTokens = metrics.inputTokens || 0;
const outputTokens = metrics.outputTokens || 0;
// ponytail: Amazon Q upstream does not expose cache fields today,
// but pick up cache_read_input_tokens / cache_creation_input_tokens
// if the event shape grows them so cost tracking stays accurate.
const cachedTokens = metrics.cacheReadInputTokens || metrics.cache_read_input_tokens || 0;
const cacheCreationInputTokens = metrics.cacheCreationInputTokens || metrics.cache_creation_input_tokens || 0;
if (inputTokens > 0 || outputTokens > 0) {
state.usage = {
prompt_tokens: inputTokens,
completion_tokens: outputTokens,
total_tokens: inputTokens + outputTokens
};
// Kiro is Claude-backed: inputTokens EXCLUDES cache (Claude convention),
// not inclusive like OpenAI's cached_tokens. Emit cache_read_input_tokens
// (not cached_tokens) so canonicalizeUsage takes the Claude fold path and
// correctly adds cache back into prompt_tokens instead of undercharging.
if (cachedTokens > 0) state.usage.cache_read_input_tokens = cachedTokens;
if (cacheCreationInputTokens > 0) state.usage.cache_creation_input_tokens = cacheCreationInputTokens;
}
}
}
// Emit final chunk only after receiving BOTH meteringEvent AND contextUsageEvent
if (state.hasMeteringEvent && state.hasContextUsage && !state.finishEmitted) {
state.finishEmitted = true;
// Estimate tokens if not available from events
if (!state.usage) {
// Estimate output tokens from content length
const estimatedOutputTokens = state.totalContentLength > 0
? Math.max(1, Math.floor(state.totalContentLength / 4))
: 0;
// Estimate input tokens from contextUsagePercentage
const estimatedInputTokens = state.contextUsagePercentage > 0
? Math.floor(state.contextUsagePercentage * contextWindow / 100)
: 0;
state.usage = {
prompt_tokens: estimatedInputTokens,
completion_tokens: estimatedOutputTokens,
total_tokens: estimatedInputTokens + estimatedOutputTokens
};
}
const finishChunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: {},
finish_reason: state.hasToolCalls ? "tool_calls" : "stop"
}]
};
// Include usage in final chunk if available
if (state.usage) {
finishChunk.usage = state.usage;
}
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(finishChunk)}\n\n`));
}
}
if (iterations >= maxIterations) {
console.warn("[Kiro] Max iterations reached in event parsing");
}
// No client chunk produced this frame — emit an SSE comment keepalive
// so the stall watchdog sees upstream activity (ignored by parser/client).
if (chunkIndex === enqueueCountBefore && !state.finishEmitted) {
controller.enqueue(new TextEncoder().encode(": ka\n\n"));
}
},
flush(controller) {
// Emit finish chunk if not already sent
if (!state.finishEmitted) {
state.finishEmitted = true;
const finishChunk = {
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{
index: 0,
delta: {},
finish_reason: state.hasToolCalls ? "tool_calls" : "stop"
}]
};
controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(finishChunk)}\n\n`));
}
// Send final done message
controller.enqueue(new TextEncoder().encode(SSE_DONE));
}
});
// Pipe response body through transform stream
if (!response.body) {
return new Response(SSE_DONE, { status: response.status, headers: { "Content-Type": "text/event-stream" } });
}
const transformedStream = response.body.pipeThrough(transformStream);
return new Response(transformedStream, {
status: response.status,
statusText: response.statusText,
headers: { ...SSE_HEADERS }
});
}
async refreshCredentials(credentials, log, proxyOptions = null) {
if (!credentials.refreshToken) return null;
try {
// Use centralized refreshKiroToken function (handles both AWS SSO OIDC and Social Auth)
const result = await refreshKiroToken(
credentials.refreshToken,
credentials.providerSpecificData,
log,
proxyOptions
);
return result;
} catch (error) {
log?.error?.("TOKEN", `Kiro refresh error: ${error.message}`);
return null;
}
}
}
/**
* Parse AWS EventStream frame
*/
function parseEventFrame(data) {
try {
const view = new DataView(data.buffer, data.byteOffset);
const headersLength = view.getUint32(4, false);
// Parse headers
const headers = {};
let offset = 12; // After prelude
const headerEnd = 12 + headersLength;
while (offset < headerEnd && offset < data.length) {
const nameLen = data[offset];
offset++;
if (offset + nameLen > data.length) break;
const name = new TextDecoder().decode(data.slice(offset, offset + nameLen));
offset += nameLen;
const headerType = data[offset];
offset++;
if (headerType === 7) { // String type
const valueLen = (data[offset] << 8) | data[offset + 1];
offset += 2;
if (offset + valueLen > data.length) break;
const value = new TextDecoder().decode(data.slice(offset, offset + valueLen));
offset += valueLen;
headers[name] = value;
} else {
break;
}
}
// Parse payload
const payloadStart = 12 + headersLength;
const payloadEnd = data.length - 4; // Exclude message CRC
let payload = null;
if (payloadEnd > payloadStart) {
const payloadStr = new TextDecoder().decode(data.slice(payloadStart, payloadEnd));
// Skip empty or whitespace-only payloads
if (!payloadStr || !payloadStr.trim()) {
return { headers, payload: null };
}
try {
payload = JSON.parse(payloadStr);
} catch (parseError) {
// Log parse error for debugging
console.warn(`[Kiro] Failed to parse payload: ${parseError.message} | payload: ${payloadStr.substring(0, 100)}`);
payload = { raw: payloadStr };
}
}
return { headers, payload };
} catch {
return null;
}
}
export default KiroExecutor;