From 31221417718a6f2f5171eaee1e852027248ec47c Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:23:55 +0200 Subject: [PATCH 01/10] feat(endpoint): route Pro accounts to runtime.kiro.dev Accounts with a profileArn (Kiro Pro / Q Developer Pro) use the runtime.kiro.dev endpoint, which serves additional models (glm-5, minimax, etc.) not available on q.amazonaws.com. Free Builder ID accounts (no profileArn) keep using q.amazonaws.com, since runtime.kiro.dev returns 400 without a profileArn. Adds resolveKiroEndpoint() and tears down stale cached SDK clients on token rotation or endpoint change. (cherry picked from commit 0c0320cef896d8ef4007040761bc0a9cb8e9679f) --- src/constants.ts | 5 ++++ src/plugin/sdk-client.ts | 62 ++++++++++++++++++++++++++++++++++++---- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/src/constants.ts b/src/constants.ts index 43bb22f..61cb16a 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -38,7 +38,12 @@ export function extractRegionFromArn(arn: string | undefined): KiroRegion | unde export const KIRO_CONSTANTS = { REFRESH_URL: 'https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken', REFRESH_IDC_URL: 'https://oidc.{{region}}.amazonaws.com/token', + // Standard endpoint — works for all accounts (Builder ID + Pro) BASE_URL: 'https://q.{{region}}.amazonaws.com/generateAssistantResponse', + // Pro/runtime endpoint — Pro accounts with a profileArn use this; it serves + // additional models (glm-5, minimax, etc.) not available on the standard endpoint. + // Free Builder ID accounts (no profileArn) must NOT use this — it returns 400. + RUNTIME_URL: 'https://runtime.{{region}}.kiro.dev/generateAssistantResponse', USAGE_LIMITS_URL: 'https://q.{{region}}.amazonaws.com/getUsageLimits', DEFAULT_REGION: 'us-east-1' as KiroRegion, AXIOS_TIMEOUT: 120000, diff --git a/src/plugin/sdk-client.ts b/src/plugin/sdk-client.ts index 9ec4d56..96f6bfc 100644 --- a/src/plugin/sdk-client.ts +++ b/src/plugin/sdk-client.ts @@ -1,7 +1,35 @@ import { CodeWhispererStreamingClient } from '@aws/codewhisperer-streaming-client' -import { KIRO_CONSTANTS } from '../constants.js' +import * as crypto from 'crypto' +import { KIRO_CONSTANTS, buildUrl } from '../constants.js' import type { Effort, KiroAuthDetails } from './types' +const KIRO_VERSION = '0.11.63' +const KIRO_CLI_MAX_ATTEMPTS = 3 + +function getMachineId(auth: KiroAuthDetails): string { + const key = auth.profileArn || auth.email || 'default' + return crypto.createHash('sha256').update(key).digest('hex') +} + +/** + * Resolve the correct chat endpoint for the given auth details. + * + * - Accounts with a profileArn (Kiro Pro / Q Developer Pro) → runtime.kiro.dev + * This endpoint serves all models including third-party ones (glm-5, minimax, …). + * It requires a profileArn on every request and returns 400 without one. + * + * - Accounts without a profileArn (free AWS Builder ID) → q.amazonaws.com + * This endpoint accepts the same token + request shape but only serves + * Claude-family models. Using runtime.kiro.dev here causes a 400. + */ +export function resolveKiroEndpoint(auth: KiroAuthDetails): string { + const region = auth.region || 'us-east-1' + if (auth.profileArn) { + return buildUrl(KIRO_CONSTANTS.RUNTIME_URL, region as any) + } + return buildUrl(KIRO_CONSTANTS.BASE_URL, region as any) +} + /** * Cache key includes effort to ensure separate clients for different effort levels, * since middleware is configured at client creation time. @@ -9,38 +37,60 @@ import type { Effort, KiroAuthDetails } from './types' interface ClientCacheEntry { client: CodeWhispererStreamingClient token: string + endpoint: string effort?: Effort } const clientCache = new Map() -const KIRO_CLI_MAX_ATTEMPTS = 3 export function createSdkClient( auth: KiroAuthDetails, region: string, effort?: Effort ): CodeWhispererStreamingClient { - const cacheKey = `${region}:${auth.email || 'default'}:${effort || 'none'}` + const endpoint = resolveKiroEndpoint(auth) + // Cache key includes endpoint so a token refresh that also changes endpoint + // (unlikely but possible) gets a fresh client, and effort so different effort + // levels get separate clients (middleware is bound at creation time). + const cacheKey = `${region}:${auth.email || 'default'}:${endpoint}:${effort || 'none'}` const cached = clientCache.get(cacheKey) if (cached && cached.token === auth.access && cached.effort === effort) { return cached.client } + // Token rotated (refresh) or endpoint changed — tear down the stale client + // so its sockets/agent don't leak before we replace the cache entry. + if (cached) { + try { + cached.client.destroy() + } catch {} + } + + const machineId = getMachineId(auth) const token = auth.access + + // Strip the path portion — the SDK constructs the full URL from region + endpoint. + const endpointBase = endpoint.replace(/\/generateAssistantResponse$/, '') + const client = new CodeWhispererStreamingClient({ region, - endpoint: `https://q.${region}.amazonaws.com`, + endpoint: endpointBase, token: () => Promise.resolve({ token }), maxAttempts: KIRO_CLI_MAX_ATTEMPTS, retryMode: 'standard', - customUserAgent: [[KIRO_CONSTANTS.USER_AGENT]] + customUserAgent: [[`${KIRO_CONSTANTS.USER_AGENT}-${KIRO_VERSION}-${machineId}`]], + requestHandler: { + connectionTimeout: 10000, + requestTimeout: 120000 + } }) // Add Kiro-specific headers client.middlewareStack.add( (next: any) => async (args: any) => { args.request.headers['x-amzn-kiro-agent-mode'] = 'vibe' + args.request.headers['x-amzn-codewhisperer-optout'] = 'true' return next(args) }, { step: 'build', name: 'addKiroHeaders' } @@ -71,7 +121,7 @@ export function createSdkClient( ) } - clientCache.set(cacheKey, { client, token, effort }) + clientCache.set(cacheKey, { client, token, endpoint, effort }) return client } From d6fbf89811f536b23b49d0704ab945d0a64f2b3d Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:25:26 +0200 Subject: [PATCH 02/10] feat(streaming): robust SDK tool-call parsing and thinking passthrough Handle inline content_block_start tool calls and bracket-format tool calls without double-emitting, track per-tool stop state, and map shortened Kiro tool names back to their original names on the way out. Shorten long tool names when building history to stay within Kiro's limits, and thread thinkingRequested through the response handler so thinking output is surfaced when requested. (cherry picked from commit da71f356b79c1a03f9fb26ea623a4eb78ee5e41d) --- src/core/request/response-handler.ts | 30 ++- .../transformers/history-builder.ts | 6 +- .../transformers/tool-transformer.ts | 114 ++++++++++- .../streaming/sdk-stream-transformer.ts | 183 +++++++++++++----- src/plugin/streaming/stream-transformer.ts | 18 +- src/plugin/streaming/types.ts | 8 + 6 files changed, 299 insertions(+), 60 deletions(-) diff --git a/src/core/request/response-handler.ts b/src/core/request/response-handler.ts index e95c78f..9a81655 100644 --- a/src/core/request/response-handler.ts +++ b/src/core/request/response-handler.ts @@ -19,10 +19,18 @@ export class ResponseHandler { sdkResponse: any, model: string, conversationId: string, - streaming: boolean + streaming: boolean, + toolNameMapper?: (name: string) => string, + thinkingRequested = false ): Promise { if (streaming) { - return this.handleSdkStreaming(sdkResponse, model, conversationId) + return this.handleSdkStreaming( + sdkResponse, + model, + conversationId, + toolNameMapper, + thinkingRequested + ) } return this.handleSdkNonStreaming(sdkResponse, model, conversationId) } @@ -33,12 +41,13 @@ export class ResponseHandler { conversationId: string ): Promise { const s = transformKiroStream(response, model, conversationId) + const enc = new TextEncoder() return new Response( new ReadableStream({ async start(c) { try { for await (const e of s) { - c.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(e)}\n\n`)) + c.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`)) } c.close() } catch (err) { @@ -53,15 +62,24 @@ export class ResponseHandler { private async handleSdkStreaming( sdkResponse: any, model: string, - conversationId: string + conversationId: string, + toolNameMapper?: (name: string) => string, + thinkingRequested = false ): Promise { - const s = transformSdkStream(sdkResponse, model, conversationId) + const s = transformSdkStream( + sdkResponse, + model, + conversationId, + toolNameMapper, + thinkingRequested + ) + const enc = new TextEncoder() return new Response( new ReadableStream({ async start(c) { try { for await (const e of s) { - c.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(e)}\n\n`)) + c.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`)) } c.close() } catch (err) { diff --git a/src/infrastructure/transformers/history-builder.ts b/src/infrastructure/transformers/history-builder.ts index 951b914..0456898 100644 --- a/src/infrastructure/transformers/history-builder.ts +++ b/src/infrastructure/transformers/history-builder.ts @@ -6,7 +6,7 @@ import { } from '../../plugin/image-handler.js' import type { CodeWhispererMessage } from '../../plugin/types' import { getContentText } from './message-transformer.js' -import { deduplicateToolResults } from './tool-transformer.js' +import { deduplicateToolResults, shortenToolName } from './tool-transformer.js' /** * Collapse agentic loop sequences in the built history. @@ -154,7 +154,7 @@ export function buildHistory(msgs: any[], resolved: string): CodeWhispererMessag if (p.type === 'text') arm.content += p.text || '' else if (p.type === 'thinking') th += p.thinking || p.text || '' else if (p.type === 'tool_use') - tus.push({ input: p.input, name: p.name, toolUseId: p.id }) + tus.push({ input: p.input, name: shortenToolName(p.name), toolUseId: p.id }) } } else arm.content = getContentText(m) if (m.tool_calls && Array.isArray(m.tool_calls)) { @@ -164,7 +164,7 @@ export function buildHistory(msgs: any[], resolved: string): CodeWhispererMessag typeof tc.function?.arguments === 'string' ? JSON.parse(tc.function.arguments) : tc.function?.arguments, - name: tc.function?.name, + name: shortenToolName(tc.function?.name), toolUseId: tc.id }) } diff --git a/src/infrastructure/transformers/tool-transformer.ts b/src/infrastructure/transformers/tool-transformer.ts index ef785d9..aa60b92 100644 --- a/src/infrastructure/transformers/tool-transformer.ts +++ b/src/infrastructure/transformers/tool-transformer.ts @@ -1,9 +1,103 @@ +import * as crypto from 'crypto' + +const MAX_TOOL_NAME_LENGTH = 64 + +// Slice without breaking surrogate pairs. +function safeSlice(s: string, end: number): string { + if (end <= 0) return '' + if (end >= s.length) return s + const code = s.charCodeAt(end - 1) + if (code >= 0xd800 && code <= 0xdbff) end -= 1 + return s.slice(0, end) +} + +export function shortenToolName(name: string): string { + if (!name || name.length <= MAX_TOOL_NAME_LENGTH) return name + const hash = crypto.createHash('sha256').update(name).digest('hex').slice(0, 12) + const prefix = safeSlice(name, MAX_TOOL_NAME_LENGTH - hash.length - 1) + return `${prefix}_${hash}` +} + +export function buildToolNameMaps(tools: any[]): { + toKiroName: (name: string) => string + fromKiroName: (name: string) => string +} { + const originalToAlias = new Map() + const aliasToOriginal = new Map() + + for (const t of tools) { + const name = t.name || t.function?.name + if (!name) continue + const alias = shortenToolName(name) + originalToAlias.set(name, alias) + if (alias !== name) aliasToOriginal.set(alias, name) + } + + return { + toKiroName: (name: string) => originalToAlias.get(name) || shortenToolName(name), + fromKiroName: (name: string) => aliasToOriginal.get(name) || name + } +} + +function sanitizeToolInput(input: any): any { + if (!input || typeof input !== 'object' || Array.isArray(input)) return input + const result: Record = {} + for (const [key, value] of Object.entries(input)) { + if (key === '') continue + result[key] = value + } + return result +} + +function sanitizeSchema(schema: any, seen: WeakSet = new WeakSet()): any { + if (!schema || typeof schema !== 'object' || Array.isArray(schema)) return schema + if (seen.has(schema)) return {} + seen.add(schema) + + const result: Record = {} + for (const [key, value] of Object.entries(schema)) { + if (key === 'additionalProperties') continue + if (key === 'required' && Array.isArray(value) && value.length === 0) continue + + if ( + (key === 'properties' || + key === 'patternProperties' || + key === '$defs' || + key === 'definitions') && + typeof value === 'object' && + value !== null && + !Array.isArray(value) + ) { + const props: Record = {} + for (const [pk, pv] of Object.entries(value)) { + props[pk] = sanitizeSchema(pv, seen) + } + result[key] = props + } else if ( + (key === 'anyOf' || key === 'oneOf' || key === 'allOf' || key === 'prefixItems') && + Array.isArray(value) + ) { + result[key] = value.map((v) => sanitizeSchema(v, seen)) + } else if ( + (key === 'items' || key === 'not' || key === 'contains') && + typeof value === 'object' + ) { + result[key] = sanitizeSchema(value, seen) + } else { + result[key] = value + } + } + return result +} + export function convertToolsToCodeWhisperer(tools: any[]): any[] { return tools.map((t) => ({ toolSpecification: { - name: t.name || t.function?.name, + name: shortenToolName(t.name || t.function?.name || ''), description: (t.description || t.function?.description || '').substring(0, 9216), - inputSchema: { json: t.input_schema || t.function?.parameters || {} } + inputSchema: { + json: sanitizeSchema(sanitizeToolInput(t.input_schema || t.function?.parameters || {})) + } } })) } @@ -19,3 +113,19 @@ export function deduplicateToolResults(trs: any[]): any[] { } return u } + +export function deduplicateToolCallsByContent(toolCalls: any[]): any[] { + const seen = new Set() + const unique: any[] = [] + for (const tc of toolCalls) { + // \x00 as separator (can't appear in a tool name) + const name = tc.name || tc.function?.name || '' + const input = tc.input || tc.function?.arguments || '' + const key = `${name}\x00${typeof input === 'string' ? input : JSON.stringify(input)}` + if (!seen.has(key)) { + seen.add(key) + unique.push(tc) + } + } + return unique +} diff --git a/src/plugin/streaming/sdk-stream-transformer.ts b/src/plugin/streaming/sdk-stream-transformer.ts index 6787c3e..3bb780e 100644 --- a/src/plugin/streaming/sdk-stream-transformer.ts +++ b/src/plugin/streaming/sdk-stream-transformer.ts @@ -1,18 +1,20 @@ import { parseBracketToolCalls } from '../../infrastructure/transformers/tool-call-parser.js' +import { deduplicateToolCallsByContent } from '../../infrastructure/transformers/tool-transformer.js' +import * as logger from '../logger.js' import { getContextWindowSize } from '../models.js' import { estimateTokens } from '../response.js' import { convertToOpenAI } from './openai-converter.js' import { findRealTag } from './stream-parser.js' -import { createTextDeltaEvents, createThinkingDeltaEvents, stopBlock } from './stream-state.js' +import { createTextDeltaEvents, createThinkingDeltaEvents } from './stream-state.js' import { StreamState, THINKING_END_TAG, THINKING_START_TAG, ToolCallState } from './types.js' export async function* transformSdkStream( sdkResponse: any, model: string, - conversationId: string + conversationId: string, + toolNameMapper?: (name: string) => string, + thinkingRequested = false ): AsyncGenerator { - const thinkingRequested = true - const streamState: StreamState = { thinkingRequested, buffer: '', @@ -29,8 +31,11 @@ export async function* transformSdkStream( let outputTokens = 0 let inputTokens = 0 let contextUsagePercentage: number | null = null + let realInputTokens: number | undefined + let realOutputTokens: number | undefined const toolCalls: ToolCallState[] = [] - let currentToolCall: ToolCallState | null = null + const activeToolCalls = new Map() + let lastToolUseId: string | null = null const eventStream = sdkResponse.generateAssistantResponseResponse if (!eventStream) { @@ -92,7 +97,6 @@ export async function* transformSdkStream( streamState.inThinking = false streamState.thinkingExtracted = true deltaEvents.push(...createThinkingDeltaEvents('', streamState)) - deltaEvents.push(...stopBlock(streamState.thinkingBlockIndex, streamState)) if (streamState.buffer.startsWith('\n\n')) { streamState.buffer = streamState.buffer.slice(2) } @@ -126,40 +130,121 @@ export async function* transformSdkStream( } } else if (event.toolUseEvent) { const tc = event.toolUseEvent + // Only accumulate the tool *name* into totalContent — the input JSON is + // never bracket-format and can be hundreds of KB; including it causes + // catastrophic backtracking in parseBracketToolCalls. if (tc.name) totalContent += tc.name - if (tc.input) totalContent += tc.input - if (tc.name && tc.toolUseId) { - if (currentToolCall && currentToolCall.toolUseId === tc.toolUseId) { - currentToolCall.input += tc.input || '' - } else { - if (currentToolCall) toolCalls.push(currentToolCall) - currentToolCall = { + if (tc.toolUseId) { + const existing = activeToolCalls.get(tc.toolUseId) + if (existing) { + existing.input += tc.input || '' + if (tc.input) { + const _c = convertToOpenAI( + { + type: 'content_block_delta', + index: existing.blockIndex, + delta: { type: 'input_json_delta', partial_json: tc.input } + }, + conversationId, + model + ) + if (_c !== null) yield _c + } + } else if (tc.name) { + const blockIndex = streamState.nextBlockIndex++ + const newToolCall: ToolCallState = { toolUseId: tc.toolUseId, - name: tc.name, - input: tc.input || '' + name: toolNameMapper ? toolNameMapper(tc.name) : tc.name, + input: tc.input || '', + stopped: false, + blockIndex + } + activeToolCalls.set(tc.toolUseId, newToolCall) + lastToolUseId = tc.toolUseId + { + const _c = convertToOpenAI( + { + type: 'content_block_start', + index: blockIndex, + content_block: { + type: 'tool_use', + id: tc.toolUseId, + name: newToolCall.name, + input: {} + } + }, + conversationId, + model + ) + if (_c !== null) yield _c + } + if (tc.input) { + const _c = convertToOpenAI( + { + type: 'content_block_delta', + index: blockIndex, + delta: { type: 'input_json_delta', partial_json: tc.input } + }, + conversationId, + model + ) + if (_c !== null) yield _c } } - if (tc.stop && currentToolCall) { - toolCalls.push(currentToolCall) - currentToolCall = null + } + if (tc.stop) { + const stopId: string | null = tc.toolUseId ?? lastToolUseId + const stoppingCall = stopId ? activeToolCalls.get(stopId) : null + if (stoppingCall) { + stoppingCall.stopped = true + toolCalls.push(stoppingCall) + activeToolCalls.delete(stopId as string) + if (lastToolUseId === stopId) { + let last: string | null = null + for (const k of activeToolCalls.keys()) last = k + lastToolUseId = last + } } } } else if (event.metadataEvent) { if (event.metadataEvent.contextUsagePercentage) { contextUsagePercentage = event.metadataEvent.contextUsagePercentage } + if (event.metadataEvent.tokenUsage) { + const tu = event.metadataEvent.tokenUsage + if (typeof tu.inputTokens === 'number') realInputTokens = tu.inputTokens + if (typeof tu.outputTokens === 'number') realOutputTokens = tu.outputTokens + } } else if ((event as any).contextUsageEvent) { const cue = (event as any).contextUsageEvent if (cue.contextUsagePercentage) { contextUsagePercentage = cue.contextUsagePercentage } + } else if ((event as any).meteringEvent) { + const me = (event as any).meteringEvent + logger.debug( + `[CREDITS] usage=${me.usage} ${me.unit || 'credit'}${me.usage !== 1 ? 's' : ''}` + ) } } - if (currentToolCall) { - toolCalls.push(currentToolCall) - currentToolCall = null + if (activeToolCalls.size > 0) { + // Stream cut off mid-tool-call(s). Tell the model to retry with smaller + // chunks so it can self-correct on the next turn. + for (const truncated of activeToolCalls.values()) { + logger.debug( + `[STREAM] Truncated tool call: name=${truncated.name} id=${truncated.toolUseId} inputLen=${truncated.input.length}` + ) + for (const ev of createTextDeltaEvents( + `\n\n[Kiro: "${truncated.name}" truncated mid-stream — 64K token output limit exceeded. Write in chunks of ≤500 lines.]`, + streamState + )) { + const _c = convertToOpenAI(ev, conversationId, model) + if (_c !== null) yield _c + } + } + activeToolCalls.clear() } if (thinkingRequested && streamState.buffer) { @@ -173,10 +258,6 @@ export async function* transformSdkStream( const _c = convertToOpenAI(ev, conversationId, model) if (_c !== null) yield _c } - for (const ev of stopBlock(streamState.thinkingBlockIndex, streamState)) { - const _c = convertToOpenAI(ev, conversationId, model) - if (_c !== null) yield _c - } } else { for (const ev of createTextDeltaEvents(streamState.buffer, streamState)) { const _c = convertToOpenAI(ev, conversationId, model) @@ -186,28 +267,32 @@ export async function* transformSdkStream( } } - for (const ev of stopBlock(streamState.textBlockIndex, streamState)) { - const _c = convertToOpenAI(ev, conversationId, model) - if (_c !== null) yield _c - } - - const bracketToolCalls = parseBracketToolCalls(totalContent) + const bracketToolCalls = totalContent.includes('[Called ') + ? parseBracketToolCalls(totalContent) + : [] if (bracketToolCalls.length > 0) { for (const btc of bracketToolCalls) { toolCalls.push({ toolUseId: btc.toolUseId, name: btc.name, - input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input) + input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input), + stopped: true + // no blockIndex — these are emitted post-stream below }) } } - if (toolCalls.length > 0) { - const baseIndex = streamState.nextBlockIndex - for (let i = 0; i < toolCalls.length; i++) { - const tc = toolCalls[i] + const dedupedToolCalls = deduplicateToolCallsByContent(toolCalls) + + // SDK tool calls were already emitted inline (content_block_start/delta/stop + // during streaming). Only bracket-format tool calls (blockIndex undefined) + // need to be emitted here. + const postStreamCalls = dedupedToolCalls.filter((tc) => tc.blockIndex === undefined) + if (postStreamCalls.length > 0) { + for (let i = 0; i < postStreamCalls.length; i++) { + const tc = postStreamCalls[i] if (!tc) continue - const blockIndex = baseIndex + i + const blockIndex = streamState.nextBlockIndex++ { const _c = convertToOpenAI( @@ -229,9 +314,8 @@ export async function* transformSdkStream( let inputJson: string try { - const parsed = JSON.parse(tc.input) - inputJson = JSON.stringify(parsed) - } catch (e) { + inputJson = JSON.stringify(JSON.parse(tc.input)) + } catch { inputJson = tc.input } @@ -240,10 +324,7 @@ export async function* transformSdkStream( { type: 'content_block_delta', index: blockIndex, - delta: { - type: 'input_json_delta', - partial_json: inputJson - } + delta: { type: 'input_json_delta', partial_json: inputJson } }, conversationId, model @@ -270,11 +351,15 @@ export async function* transformSdkStream( inputTokens = Math.max(0, totalTokens - outputTokens) } + // Real token counts from Kiro's metadata win over the context-% estimate. + if (realInputTokens !== undefined) inputTokens = realInputTokens + if (realOutputTokens !== undefined) outputTokens = realOutputTokens + { const _c = convertToOpenAI( { type: 'message_delta', - delta: { stop_reason: toolCalls.length > 0 ? 'tool_use' : 'end_turn' }, + delta: { stop_reason: dedupedToolCalls.length > 0 ? 'tool_use' : 'end_turn' }, usage: { input_tokens: inputTokens, output_tokens: outputTokens, @@ -293,6 +378,14 @@ export async function* transformSdkStream( if (_c !== null) yield _c } } catch (e) { + logger.debug( + `[STREAM] Error in transformSdkStream: ${e instanceof Error ? e.message : String(e)}` + ) + for (const tc of activeToolCalls.values()) { + logger.debug( + `[STREAM] Incomplete tool call: name=${tc.name} id=${tc.toolUseId} inputLen=${tc.input.length}` + ) + } throw e } } diff --git a/src/plugin/streaming/stream-transformer.ts b/src/plugin/streaming/stream-transformer.ts index 7036cb9..c6e237d 100644 --- a/src/plugin/streaming/stream-transformer.ts +++ b/src/plugin/streaming/stream-transformer.ts @@ -43,7 +43,12 @@ export async function* transformKiroStream( try { while (true) { const { done, value } = await reader.read() - if (done) break + if (done) { + // Flush remaining multi-byte UTF-8 from TextDecoder + const tail = decoder.decode() + if (tail) rawBuffer += tail + break + } const chunk = decoder.decode(value, { stream: true }) rawBuffer += chunk @@ -163,11 +168,13 @@ export async function* transformKiroStream( currentToolCall = { toolUseId: tc.toolUseId, name: tc.name, - input: tc.input || '' + input: tc.input || '', + stopped: false } } if (tc.stop && currentToolCall) { + currentToolCall.stopped = true toolCalls.push(currentToolCall) currentToolCall = null } @@ -222,13 +229,16 @@ export async function* transformKiroStream( if (_c !== null) yield _c } - const bracketToolCalls = parseBracketToolCalls(totalContent) + const bracketToolCalls = totalContent.includes('[Called ') + ? parseBracketToolCalls(totalContent) + : [] if (bracketToolCalls.length > 0) { for (const btc of bracketToolCalls) { toolCalls.push({ toolUseId: btc.toolUseId, name: btc.name, - input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input) + input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input), + stopped: true }) } } diff --git a/src/plugin/streaming/types.ts b/src/plugin/streaming/types.ts index 1736a2d..fcae5f3 100644 --- a/src/plugin/streaming/types.ts +++ b/src/plugin/streaming/types.ts @@ -22,6 +22,14 @@ export interface ToolCallState { toolUseId: string name: string input: string + /** true only when the SDK sent a stop event for this tool call */ + stopped: boolean + /** + * Block index assigned when content_block_start was emitted inline during + * streaming. Undefined for bracket-format tool calls which are emitted + * post-stream. Used to avoid re-emitting SDK tool calls at the end. + */ + blockIndex?: number } export const THINKING_START_TAG = '' From f6d93368449f875eaf85164211e89cbaf863d514 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:30:02 +0200 Subject: [PATCH 03/10] feat(resilience): exclude retry sleeps from timeout budget Track slept time (excludedMs) so long backoff waits don't count against the overall retry deadline, and thread it through the error handler. Harden locked-operations against contention and tidy account selection. (cherry picked from commit 9b5e23b207d2f554b55a39abcb43c617038bbf11) --- src/core/request/error-handler.ts | 45 +++++++++++++++++++------ src/core/request/retry-strategy.ts | 12 +++++-- src/plugin/health.ts | 2 +- src/plugin/storage/locked-operations.ts | 26 +++++++++++++- 4 files changed, 71 insertions(+), 14 deletions(-) diff --git a/src/core/request/error-handler.ts b/src/core/request/error-handler.ts index c9f968f..5446f0b 100644 --- a/src/core/request/error-handler.ts +++ b/src/core/request/error-handler.ts @@ -1,11 +1,15 @@ import type { AccountRepository } from '../../infrastructure/database/account-repository' import type { AccountManager } from '../../plugin/accounts' +import * as logger from '../../plugin/logger' import type { ManagedAccount } from '../../plugin/types' type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void interface RequestContext { retry: number + // Time spent in rate-limit sleeps, propagated up so the retry strategy can + // exclude it from the request timeout budget. + excludedMs?: number } interface ErrorHandlerConfig { @@ -38,12 +42,16 @@ export class ErrorHandler { if (response.status === 400) { const reason = await readBody() + logger.warn(`HTTP 400 on ${account.email}: ${reason || 'unknown'}`) showToast(`400: ${reason || 'unknown'}`, 'error') return { shouldRetry: false } } if (response.status === 401 && context.retry < this.config.rate_limit_max_retries) { const reason = await readBody() + logger.warn( + `HTTP 401 on ${account.email} (retry ${context.retry}): ${reason || 'Unauthorized'}` + ) showToast(`401: ${reason || 'Unauthorized'}. Retrying...`, 'warning') return { shouldRetry: true, @@ -64,15 +72,18 @@ export class ErrorHandler { } } catch (e) {} + logger.warn(`HTTP 500 on ${account.email} (failCount ${account.failCount}): ${errorMessage}`) if (account.failCount < 5) { const delay = 1000 * Math.pow(2, account.failCount - 1) showToast(`500: ${errorMessage}. Retrying in ${Math.ceil(delay / 1000)}s...`, 'warning') await this.sleep(delay) return { shouldRetry: true } } else { + account.failCount = 9 // markUnhealthy will increment to 10 and set isHealthy=false this.accountManager.markUnhealthy( account, - `Server Error (500) after 5 attempts: ${errorMessage}` + `Server error (500) after 5 attempts: ${errorMessage}`, + Date.now() + 1800000 // 30 min recovery ) await this.repository.batchSave(this.accountManager.getAccounts()) showToast(`500: ${errorMessage}. Marking account as unhealthy and switching...`, 'warning') @@ -82,6 +93,7 @@ export class ErrorHandler { if (response.status === 429) { const w = parseInt(response.headers.get('retry-after') || '60') * 1000 + logger.warn(`HTTP 429 on ${account.email}: rate limited, retry-after=${Math.ceil(w / 1000)}s`) this.accountManager.markRateLimited(account, w) await this.repository.batchSave(this.accountManager.getAccounts()) const count = this.accountManager.getAccountCount() @@ -90,7 +102,11 @@ export class ErrorHandler { } showToast(`429: Rate limited. Waiting ${Math.ceil(w / 1000)}s...`, 'warning') await this.sleep(w) - return { shouldRetry: true } + // The wait is not request runtime — exclude it from the timeout budget. + return { + shouldRetry: true, + newContext: { ...context, excludedMs: (context.excludedMs ?? 0) + w } + } } if (response.status === 402 || response.status === 403) { @@ -114,16 +130,18 @@ export class ErrorHandler { errorReason = 'Account Suspended' isPermanent = true } - if ( - errorReason.includes('bearer token included in the request is invalid') || - errorReason.includes('The bearer token included in the request is invalid') - ) { - isPermanent = true - } - if (isPermanent) { - account.failCount = 10 + if (errorReason.includes('bearer token included in the request is invalid')) { + // Force token refresh on next retry + account.expiresAt = 0 + return { shouldRetry: true } } + logger.warn(`HTTP ${response.status} on ${account.email}: ${errorReason}`, { + isPermanent, + retry: context.retry, + reason: errorData?.reason + }) + if (this.accountManager.getAccountCount() > 1) { showToast(`${response.status}: ${errorReason}. Switching account...`, 'warning') this.accountManager.markUnhealthy(account, errorReason) @@ -131,6 +149,12 @@ export class ErrorHandler { return { shouldRetry: true, switchAccount: true } } + if (isPermanent) { + this.accountManager.markUnhealthy(account, errorReason) + await this.repository.batchSave(this.accountManager.getAccounts()) + return { shouldRetry: false } + } + if ( response.status === 403 && !isPermanent && @@ -150,6 +174,7 @@ export class ErrorHandler { } const reason = await readBody() + logger.warn(`HTTP ${response.status} on ${account.email}: ${reason || response.statusText}`) showToast(`${response.status}: ${reason || response.statusText}`, 'error') return { shouldRetry: false } } diff --git a/src/core/request/retry-strategy.ts b/src/core/request/retry-strategy.ts index 4acf5db..b72616e 100644 --- a/src/core/request/retry-strategy.ts +++ b/src/core/request/retry-strategy.ts @@ -6,6 +6,8 @@ interface RetryConfig { interface RetryContext { iterations: number startTime: number + // Time spent in rate-limit waits, excluded from the request timeout budget. + excludedMs: number } export class RetryStrategy { @@ -21,7 +23,8 @@ export class RetryStrategy { } } - if (Date.now() - context.startTime > this.config.request_timeout_ms) { + const elapsed = Date.now() - context.startTime - context.excludedMs + if (elapsed > this.config.request_timeout_ms) { return { canContinue: false, error: 'Request timeout' @@ -31,10 +34,15 @@ export class RetryStrategy { return { canContinue: true } } + markSleep(context: RetryContext, ms: number): void { + context.excludedMs += Math.max(0, ms) + } + createContext(): RetryContext { return { iterations: 0, - startTime: Date.now() + startTime: Date.now(), + excludedMs: 0 } } } diff --git a/src/plugin/health.ts b/src/plugin/health.ts index 10eb464..6916b69 100644 --- a/src/plugin/health.ts +++ b/src/plugin/health.ts @@ -9,6 +9,6 @@ export function isPermanentError(reason?: string): boolean { reason.includes('ExpiredClientException') || reason.includes('Client is expired') || reason.includes('HTTP_401') || - reason.includes('HTTP_403') + reason.includes('Account Suspended') ) } diff --git a/src/plugin/storage/locked-operations.ts b/src/plugin/storage/locked-operations.ts index dcfd592..9686160 100644 --- a/src/plugin/storage/locked-operations.ts +++ b/src/plugin/storage/locked-operations.ts @@ -15,8 +15,16 @@ const LOCK_OPTIONS = { realpath: false } +// In-process serialisation: prevents concurrent writes within the same process +// from racing each other without paying the file-lock cost on every request. +let inProcessLockChain: Promise = Promise.resolve() + export async function withDatabaseLock(dbPath: string, fn: () => Promise): Promise { - const lockPath = `${dbPath}.lock` + // Serialise within this process first (cheap). + let resolveInProcess!: () => void + const prev = inProcessLockChain + inProcessLockChain = new Promise((r) => (resolveInProcess = r)) + await prev if (!existsSync(dbPath)) { const dir = dbPath.substring(0, dbPath.lastIndexOf('/')) @@ -36,6 +44,22 @@ export async function withDatabaseLock(dbPath: string, fn: () => Promise): console.warn('Failed to release lock:', e) } } + resolveInProcess() + } +} + +// Lightweight in-process-only lock — skips the file-level lockfile entirely. +// Use for single-process writes where cross-process races are impossible +// (e.g. normal per-request account state updates). +export async function withInProcessLock(fn: () => T): Promise { + let resolveInProcess!: () => void + const prev = inProcessLockChain + inProcessLockChain = new Promise((r) => (resolveInProcess = r)) + await prev + try { + return fn() + } finally { + resolveInProcess() } } From e7afde2947405822ca19982c8e63b6a68b5a6718 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:32:02 +0200 Subject: [PATCH 04/10] feat(images): carry forward images across agentic turns OpenCode strips image parts from conversation state on later turns. Cache converted Kiro images per (workspace, fingerprint) and re-attach them to currentMessage when a turn arrives without fresh images, so the model keeps seeing them. Never restores onto tool-result turns (Kiro 400s on images+toolResults). Gated by the image_carry_forward config flag (default on); Kiro bills per request, not per token, so re-sending has no billing impact. (cherry picked from commit 0416dcc317691f8fcacc4d8dd46912fa6d629fb8) --- src/plugin/config/schema.ts | 13 +- src/plugin/image-cache.ts | 335 ++++++++++++++++++++++++++++++++++++ src/plugin/image-handler.ts | 66 ++++--- 3 files changed, 377 insertions(+), 37 deletions(-) create mode 100644 src/plugin/image-cache.ts diff --git a/src/plugin/config/schema.ts b/src/plugin/config/schema.ts index 18cd3cf..5a4c89c 100644 --- a/src/plugin/config/schema.ts +++ b/src/plugin/config/schema.ts @@ -95,7 +95,15 @@ export const KiroConfigSchema = z.object({ * When true (default), maps budget ranges to effort levels. * When false, only uses explicit effort config or falls back to 'medium'. */ - auto_effort_mapping: z.boolean().default(true) + auto_effort_mapping: z.boolean().default(true), + + // OpenCode strips image parts from conversation state across agentic turns. + // When true, the plugin caches converted images per conversation and re-attaches + // them to currentMessage on later turns so the model keeps "seeing" them. + // Kiro bills per session (request), not per token, so re-sending the same + // images each turn has no billing impact. Only disable if you hit the + // per-request 3.75MB image-payload cap on conversations with many heavy images. + image_carry_forward: z.boolean().default(true) }) export type KiroConfig = z.infer @@ -114,5 +122,6 @@ export const DEFAULT_CONFIG: KiroConfig = { usage_tracking_enabled: true, auto_sync_kiro_cli: true, enable_log_api_request: false, - auto_effort_mapping: true + auto_effort_mapping: true, + image_carry_forward: true } diff --git a/src/plugin/image-cache.ts b/src/plugin/image-cache.ts new file mode 100644 index 0000000..fa303c5 --- /dev/null +++ b/src/plugin/image-cache.ts @@ -0,0 +1,335 @@ +/** + * In-memory + on-disk cache of converted Kiro images, keyed per conversation + * (workspace + fingerprint). See AGENTS.md §3 "Image carry-forward" for the + * design rationale and the OpenCode-strip behaviour this works around. + * + * Bounded by 24h TTL, 20-conversation memory LRU, and the per-entry + * MAX_KIRO_IMAGES / MAX_KIRO_IMAGE_BYTES caps. Disk persistence uses + * write-temp-rename so a partial write can't corrupt the cache file. + */ +import { Buffer } from 'node:buffer' +import { createHash } from 'node:crypto' +import { + existsSync, + mkdirSync, + readFileSync, + readdirSync, + renameSync, + statSync, + unlinkSync, + writeFileSync +} from 'node:fs' +import { homedir } from 'node:os' +import { join } from 'node:path' + +import { MAX_KIRO_IMAGES, MAX_KIRO_IMAGE_BYTES, type KiroImage } from './image-handler.js' + +// Content fingerprint for dedup — format + byte length + first/last 16 bytes +// is unique enough in practice; matching head AND tail for two distinct +// screenshots of identical size is vanishingly rare. No crypto needed. +function dedupKey(img: KiroImage): string { + const b = img.source.bytes + const n = b.byteLength + if (n === 0) return `${img.format}:empty` + const hex = (start: number, end: number): string => + Array.from(b.subarray(start, end), (v) => v.toString(16).padStart(2, '0')).join('') + return `${img.format}:${n}:${hex(0, Math.min(16, n))}:${hex(Math.max(0, n - 16), n)}` +} + +// Mirrors the layout of logger.ts and kiro.json so everything sits under one +// base folder. Internal — only the module-level singleton uses this. +function defaultCacheDir(): string { + const platform = process.platform + const base = + platform === 'win32' + ? join(process.env.APPDATA || join(homedir(), 'AppData', 'Roaming'), 'opencode') + : join(process.env.XDG_CONFIG_HOME || join(homedir(), '.config'), 'opencode') + return join(base, 'kiro-images') +} + +function diskFileFor(cacheDir: string, workspace: string, fingerprint: string): string { + // Hash the composite key so the filename has no path separators or other + // shell-unsafe characters from the workspace path. + const id = createHash('sha256') + .update(workspace + '\0' + fingerprint) + .digest('hex') + .slice(0, 32) + return join(cacheDir, `${id}.json`) +} + +interface CacheEntry { + images: KiroImage[] + totalBytes: number + lastAccess: number + // True once this conversation has ever contained an image-bearing turn. + // Lets callers skip the history scan entirely on long sessions where no + // images were ever attached. + everHadImages: boolean +} + +export interface ImageCacheOptions { + ttlMs?: number + maxEntries?: number + // Injectable clock for deterministic tests. + now?: () => number + // Filesystem persistence directory. Omit (or pass null) to keep the cache + // memory-only — tests rely on this default. + cacheDir?: string | null +} + +const DEFAULT_TTL_MS = 24 * 60 * 60 * 1000 +const DEFAULT_MAX_ENTRIES = 20 + +export class ImageCache { + private cache = new Map() + private ttlMs: number + private maxEntries: number + private now: () => number + private cacheDir: string | null + + constructor(opts: ImageCacheOptions = {}) { + this.ttlMs = opts.ttlMs ?? DEFAULT_TTL_MS + this.maxEntries = opts.maxEntries ?? DEFAULT_MAX_ENTRIES + this.now = opts.now ?? (() => Date.now()) + this.cacheDir = opts.cacheDir ?? null + if (this.cacheDir) this.sweepDiskExpired() + } + + private k(workspace: string, fingerprint: string): string { + return `${workspace}\0${fingerprint}` + } + + set(workspace: string, fingerprint: string, images: KiroImage[]): void { + if (images.length === 0) return + const totalBytes = images.reduce((n, im) => n + im.source.bytes.byteLength, 0) + const key = this.k(workspace, fingerprint) + const entry: CacheEntry = { + images, + totalBytes, + lastAccess: this.now(), + everHadImages: true + } + // delete-then-set so the entry moves to the tail of the Map's insertion + // order — that's what makes the eviction sweep pick the truly oldest. + this.cache.delete(key) + this.cache.set(key, entry) + this.writeEntryToDisk(workspace, fingerprint, entry) + this.evict() + } + + /** + * Merge new images with the existing cache for this conversation. + * New images go to the front (most recent), then existing fills in. + * Duplicates (same content fingerprint) collapse. + * Result is capped at MAX_KIRO_IMAGES and MAX_KIRO_IMAGE_BYTES; + * when over budget, the OLDEST entries drop first (FIFO). + * + * Returns the final image count for diagnostics. + */ + upsert(workspace: string, fingerprint: string, newImages: KiroImage[]): number { + const key = this.k(workspace, fingerprint) + // Try in-memory first; fall back to disk so cross-restart upserts merge + // with what we already persisted instead of overwriting it. + let existing = this.cache.get(key)?.images + if (!existing && this.cacheDir) { + existing = this.loadEntryFromDisk(workspace, fingerprint)?.images + } + existing ??= [] + + if (newImages.length === 0 && existing.length === 0) return 0 + + const seen = new Set() + const merged: KiroImage[] = [] + const add = (img: KiroImage): void => { + if (merged.length >= MAX_KIRO_IMAGES) return + const k = dedupKey(img) + if (seen.has(k)) return + seen.add(k) + merged.push(img) + } + for (const img of newImages) add(img) + for (const img of existing) add(img) + + // Enforce byte budget by dropping the oldest entries (tail) until in cap. + let total = merged.reduce((n, im) => n + im.source.bytes.byteLength, 0) + while (total > MAX_KIRO_IMAGE_BYTES && merged.length > 0) { + total -= merged.pop()!.source.bytes.byteLength + } + + if (merged.length === 0) { + this.cache.delete(key) + this.deleteFromDisk(workspace, fingerprint) + return 0 + } + + const entry: CacheEntry = { + images: merged, + totalBytes: total, + lastAccess: this.now(), + everHadImages: true + } + this.cache.delete(key) + this.cache.set(key, entry) + this.writeEntryToDisk(workspace, fingerprint, entry) + this.evict() + return merged.length + } + + get(workspace: string, fingerprint: string): KiroImage[] | null { + const key = this.k(workspace, fingerprint) + let entry = this.cache.get(key) + + // Lazy-load from disk if we don't have it in memory (e.g. fresh plugin + // process resuming a prior conversation). + if (!entry && this.cacheDir) { + const loaded = this.loadEntryFromDisk(workspace, fingerprint) + if (loaded) { + entry = loaded + this.cache.set(key, entry) + this.evict() + } + } + + if (!entry) return null + + if (this.now() - entry.lastAccess > this.ttlMs) { + this.cache.delete(key) + this.deleteFromDisk(workspace, fingerprint) + return null + } + // Refresh LRU position. + entry.lastAccess = this.now() + this.cache.delete(key) + this.cache.set(key, entry) + return entry.images + } + + delete(workspace: string, fingerprint: string): void { + this.cache.delete(this.k(workspace, fingerprint)) + this.deleteFromDisk(workspace, fingerprint) + } + + // True if this conversation has ever carried images. Lets callers skip + // the history scan entirely on long text-only sessions. + hasEverHadImages(workspace: string, fingerprint: string): boolean { + const entry = this.cache.get(this.k(workspace, fingerprint)) + if (entry) return entry.everHadImages + if (this.cacheDir) { + try { + return existsSync(diskFileFor(this.cacheDir, workspace, fingerprint)) + } catch { + return false + } + } + return false + } + + // Memory-only clear; disk files survive and may be re-loaded by later get(). + clear(): void { + this.cache.clear() + } + + size(): number { + return this.cache.size + } + + private evict(): void { + const t = this.now() + for (const [key, e] of this.cache) { + if (t - e.lastAccess > this.ttlMs) this.cache.delete(key) + } + while (this.cache.size > this.maxEntries) { + const first = this.cache.keys().next().value + if (first === undefined) break + this.cache.delete(first) + } + } + + // ── Disk persistence (no-ops when cacheDir is null) ──────────────────────── + + private writeEntryToDisk(workspace: string, fingerprint: string, entry: CacheEntry): void { + if (!this.cacheDir) return + try { + mkdirSync(this.cacheDir, { recursive: true }) + const file = diskFileFor(this.cacheDir, workspace, fingerprint) + const tmp = `${file}.tmp` + const payload = { + workspace, + fingerprint, + updatedAt: entry.lastAccess, + images: entry.images.map((im) => ({ + format: im.format, + b64: Buffer.from(im.source.bytes).toString('base64') + })) + } + writeFileSync(tmp, JSON.stringify(payload)) + renameSync(tmp, file) + } catch { + // Disk failures must never break the request — in-memory cache still works. + } + } + + private loadEntryFromDisk(workspace: string, fingerprint: string): CacheEntry | null { + if (!this.cacheDir) return null + try { + const file = diskFileFor(this.cacheDir, workspace, fingerprint) + if (!existsSync(file)) return null + const stat = statSync(file) + // mtime acts as "last access" — if older than TTL, treat as expired + // and clean up. + if (this.now() - stat.mtimeMs > this.ttlMs) { + try { + unlinkSync(file) + } catch {} + return null + } + const raw = readFileSync(file, 'utf-8') + const data = JSON.parse(raw) + if (!data || !Array.isArray(data.images)) return null + const images: KiroImage[] = [] + for (const im of data.images) { + if (!im || typeof im.format !== 'string' || typeof im.b64 !== 'string') continue + const bytes = new Uint8Array(Buffer.from(im.b64, 'base64')) + images.push({ format: im.format, source: { bytes } }) + } + if (images.length === 0) return null + const totalBytes = images.reduce((n, im) => n + im.source.bytes.byteLength, 0) + return { images, totalBytes, lastAccess: stat.mtimeMs, everHadImages: true } + } catch { + return null + } + } + + private deleteFromDisk(workspace: string, fingerprint: string): void { + if (!this.cacheDir) return + try { + const file = diskFileFor(this.cacheDir, workspace, fingerprint) + if (existsSync(file)) unlinkSync(file) + } catch {} + } + + /** + * Remove expired files from the cache dir. Runs once at construction so + * a long-lived dir doesn't grow unbounded with stale conversations. + */ + private sweepDiskExpired(): void { + if (!this.cacheDir) return + try { + if (!existsSync(this.cacheDir)) return + const cutoff = this.now() - this.ttlMs + for (const name of readdirSync(this.cacheDir)) { + if (!name.endsWith('.json') && !name.endsWith('.tmp')) continue + const file = join(this.cacheDir, name) + try { + const stat = statSync(file) + if (stat.mtimeMs < cutoff) unlinkSync(file) + } catch {} + } + } catch {} + } +} + +// Shared instance for the plugin runtime — persisted to ~/.config/opencode/kiro-images/ +// (or %APPDATA%/opencode/kiro-images/ on Windows). Tests construct their own +// ImageCache with a tmpdir or no cacheDir at all. +export const imageCache = new ImageCache({ cacheDir: defaultCacheDir() }) diff --git a/src/plugin/image-handler.ts b/src/plugin/image-handler.ts index 3c7fab9..d4cfd69 100644 --- a/src/plugin/image-handler.ts +++ b/src/plugin/image-handler.ts @@ -1,12 +1,14 @@ +import { Buffer } from 'node:buffer' + interface UnifiedImage { mediaType: string data: string } -const MAX_KIRO_IMAGES = 4 -const MAX_KIRO_IMAGE_BYTES = 3_750_000 +export const MAX_KIRO_IMAGES = 4 +export const MAX_KIRO_IMAGE_BYTES = 3_750_000 -interface KiroImage { +export interface KiroImage { format: string source: { bytes: Uint8Array @@ -18,15 +20,26 @@ interface ImageConversionResult { omitted: number } -/** Decode base64 to a plain Uint8Array (NOT Buffer) to avoid Buffer.toJSON() trap. */ +/** Decode base64 to a plain Uint8Array. Uses Node Buffer when available (native + * C++, ~10x faster than the atob + charCodeAt loop) and falls back to atob + * in non-Node environments. Returns a fresh Uint8Array — Buffer's underlying + * ArrayBuffer is shared with Node's pool, so we copy to detach. */ function base64ToUint8Array(b64: string): Uint8Array { + if (typeof Buffer !== 'undefined') { + const buf = Buffer.from(b64, 'base64') + const out = new Uint8Array(buf.byteLength) + out.set(buf) + return out + } const bin = atob(b64) const arr = new Uint8Array(bin.length) for (let i = 0; i < bin.length; i++) arr[i] = bin.charCodeAt(i) return arr } -function extractImagesFromAnthropicFormat(content: any[]): UnifiedImage[] { +export function extractAllImages(content: any): UnifiedImage[] { + if (!Array.isArray(content)) return [] + const images: UnifiedImage[] = [] for (const item of content) { @@ -35,46 +48,29 @@ function extractImagesFromAnthropicFormat(content: any[]): UnifiedImage[] { mediaType: item.source.media_type || 'image/jpeg', data: item.source.data }) - } - } + } else if (item.type === 'image_url' && item.image_url?.url) { + const url = item.image_url.url + if (!url.startsWith('data:')) continue - return images -} + const comma = url.indexOf(',') + if (comma < 0) continue -function extractImagesFromOpenAI(content: any[]): UnifiedImage[] { - const images: UnifiedImage[] = [] + const data = url.slice(comma + 1) + if (!data) continue - for (const item of content) { - if (item.type === 'image_url' && item.image_url?.url) { - const url = item.image_url.url + const headerEnd = url.indexOf(';', 5) + const mediaType = headerEnd > 0 ? url.slice(5, headerEnd) : url.slice(5, comma) - if (url.startsWith('data:')) { - try { - const [header, data] = url.split(',', 2) - if (!data) continue - - const mediaType = header.split(';')[0].replace('data:', '') - - images.push({ - mediaType: mediaType || 'image/jpeg', - data: data - }) - } catch (e) { - continue - } - } + images.push({ + mediaType: mediaType || 'image/jpeg', + data + }) } } return images } -export function extractAllImages(content: any): UnifiedImage[] { - if (!Array.isArray(content)) return [] - - return [...extractImagesFromAnthropicFormat(content), ...extractImagesFromOpenAI(content)] -} - export function convertImagesToKiroFormat(images: UnifiedImage[]): ImageConversionResult { const selected: UnifiedImage[] = [] let totalBase64Chars = 0 From 79e20ea98de37da75d6423c7539df3b0186e9e4e Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:32:22 +0200 Subject: [PATCH 05/10] feat(effort): map reasoning effort to thinking token budgets Add THINKING_BUDGETS mapping (low/medium/high/default) capped at Kiro's 200k max_thinking_length, used to translate OpenCode's reasoningEffort selection into an explicit thinking budget. (cherry picked from commit aad98f5f7a66e60e4d8edb983c7800e3856875a2) --- src/constants.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/constants.ts b/src/constants.ts index 61cb16a..7f39d95 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -54,6 +54,18 @@ export const KIRO_CONSTANTS = { ORIGIN_AI_EDITOR: 'AI_EDITOR' } +// Thinking budgets per effort level (tokens). +// Kiro's max_thinking_length cap is 200 000. +export const THINKING_BUDGETS = { + low: 10_000, + medium: 24_000, + high: 200_000, + // Default when thinking is on but no effort level is specified (-thinking suffix) + default: 16_000 +} as const + +export type ThinkingEffort = keyof typeof THINKING_BUDGETS + export const MODEL_MAPPING: Record = { // Claude Haiku 'claude-haiku-4-5': 'claude-haiku-4.5', From 3f3bb4f24f9bf1968982571b920bfe45a2b868ba Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:32:41 +0200 Subject: [PATCH 06/10] feat(conv): persist conversations and isolate parallel sessions Mint a stable conversationId + agentContinuationId per (workspace, fingerprint) and persist them in SQLite so agentic turns continue the same Kiro conversation instead of starting fresh each request. Isolate parallel sessions in one workspace via the x-session-id header, and reset a stale conversationId on 400 ValidationException. Add a cross-process reauth lock (pid + TTL) so concurrent instances don't re-authenticate simultaneously, with a bounded oauth timeout. Dispose SDK clients, image cache, and DB on plugin shutdown. (cherry picked from commit b25be3af526ab402270f56ecd24d5fab20a1b87d) --- src/core/request/request-handler.ts | 244 +++++++++++++++++++++------- src/plugin.ts | 19 ++- src/plugin/logger.ts | 2 +- src/plugin/request.ts | 242 ++++++++++++++++++++++++--- src/plugin/storage/migrations.ts | 39 +++++ src/plugin/storage/sqlite.ts | 116 +++++++++++++ src/plugin/types.ts | 7 + 7 files changed, 592 insertions(+), 77 deletions(-) diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index 7dcfdf4..7aa5a9e 100644 --- a/src/core/request/request-handler.ts +++ b/src/core/request/request-handler.ts @@ -1,11 +1,14 @@ import { GenerateAssistantResponseCommand } from '@aws/codewhisperer-streaming-client' +import { THINKING_BUDGETS } from '../../constants' import type { AccountRepository } from '../../infrastructure/database/account-repository' import type { AccountManager } from '../../plugin/accounts' import type { KiroConfig } from '../../plugin/config' import { isPermanentError } from '../../plugin/health' +import { imageCache } from '../../plugin/image-cache' import * as logger from '../../plugin/logger' import { transformToSdkRequest } from '../../plugin/request' import { createSdkClient } from '../../plugin/sdk-client' +import { kiroDb } from '../../plugin/storage/sqlite' import { syncFromKiroCli } from '../../plugin/sync/kiro-cli' import type { KiroAuthDetails, ManagedAccount, SdkPreparedRequest } from '../../plugin/types' import { AccountSelector } from '../account/account-selector' @@ -17,8 +20,17 @@ import { RetryStrategy } from './retry-strategy' type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void -const KIRO_API_PATTERN = /^(https?:\/\/)?q\.[a-z0-9-]+\.amazonaws\.com/ +// Matches both the standard q.amazonaws.com endpoint and the Pro runtime.kiro.dev endpoint +const KIRO_API_PATTERN = + /^(https?:\/\/)?(q\.[a-z0-9-]+\.amazonaws\.com|runtime\.[a-z0-9-]+\.kiro\.dev)/ const REAUTH_FAILURE_COOLDOWN_MS = 60000 +const REAUTH_TIMEOUT_MS = 90_000 + +function extractSessionId(headers: unknown): string | undefined { + if (!headers) return undefined + const h = headers as Record + return h['x-session-id'] ?? h['x-session-affinity'] +} export class RequestHandler { private accountSelector: AccountSelector @@ -29,13 +41,13 @@ export class RequestHandler { private retryStrategy: RetryStrategy private reauthInFlight: Promise | null = null private lastFailedReauthAt = 0 - private static kiroRequestQueue: Promise = Promise.resolve() constructor( private accountManager: AccountManager, private config: KiroConfig, private repository: AccountRepository, - private client?: any + private client?: any, + private workspace = '' ) { this.accountSelector = new AccountSelector(accountManager, config, syncFromKiroCli, repository) this.tokenRefresher = new TokenRefresher(config, accountManager, syncFromKiroCli, repository) @@ -52,43 +64,46 @@ export class RequestHandler { return fetch(input, init) } - return this.enqueueKiroRequest(() => this.handleKiroRequest(url, init, showToast)) - } - - private async enqueueKiroRequest(run: () => Promise): Promise { - const previous = RequestHandler.kiroRequestQueue - let release!: () => void - - RequestHandler.kiroRequestQueue = new Promise((resolve) => { - release = resolve - }) + const sessionId = extractSessionId(init?.headers) - await previous.catch(() => {}) - - try { - return await run() - } finally { - release() - } + return this.handleKiroRequest(url, init, showToast, sessionId) } private async handleKiroRequest( url: string, init: any, - showToast: ToastFunction + showToast: ToastFunction, + sessionId?: string ): Promise { const body = init?.body ? JSON.parse(init.body) : {} const model = this.extractModel(url) || body.model || 'claude-sonnet-4-5' - const think = - model.endsWith('-thinking') || !!body.providerOptions?.thinkingConfig || !!body.thinkingConfig - const budget = - body.providerOptions?.thinkingConfig?.thinkingBudget || - body.thinkingConfig?.thinkingBudget || - body.thinkingConfig?.budget_tokens || - 20000 + + // Resolve thinking mode + budget. + // + // Priority order: + // 1. Model ID ends with '-thinking' → adaptive thinking, default budget + // 2. providerOptions["kiro"].reasoningEffort → adaptive mode, effort-based budget + // (OpenCode sends this when the user picks low/medium/high in the UI) + // 3. providerOptions.thinkingConfig.thinkingBudget → explicit budget (legacy) + // + // Budget mapping (Kiro max = 200 000 tokens): + // low → 10 000 | medium → 24 000 | high → 200 000 | default → 16 000 + const provOpts = body.providerOptions?.['kiro'] ?? body.providerOptions ?? {} + const reasoningEffort: string | undefined = provOpts.reasoningEffort + const thinkingConfig = body.providerOptions?.thinkingConfig + + const think = model.endsWith('-thinking') || !!reasoningEffort || !!thinkingConfig + + let uiEffort: 'low' | 'medium' | 'high' | 'default' = 'default' + if (reasoningEffort === 'low') uiEffort = 'low' + else if (reasoningEffort === 'medium') uiEffort = 'medium' + else if (reasoningEffort === 'high') uiEffort = 'high' + + const budget: number = thinkingConfig?.thinkingBudget || THINKING_BUDGETS[uiEffort] let retry = 0 let consecutiveNullAccounts = 0 + let forceNewConversation = false const retryContext = this.retryStrategy.createContext() while (true) { @@ -131,12 +146,28 @@ export class RequestHandler { continue } - const sdkPrep = this.prepareSdkRequest(init?.body, model, auth, think, budget, showToast) + const sdkPrep = this.prepareSdkRequest( + body, + model, + auth, + think, + budget, + showToast, + uiEffort, + sessionId + ) + + const histLen = (sdkPrep.conversationState as any).history?.length || 0 + const agentContId = (sdkPrep.conversationState as any).agentContinuationId || 'none' + logger.debug( + `[REQ] convId=${sdkPrep.conversationId} history=${histLen} agentCont=${agentContId} model=${model}` + ) const apiTimestamp = this.config.enable_log_api_request ? logger.getTimestamp() : null if (apiTimestamp) { this.logSdkRequest(sdkPrep, acc, apiTimestamp) } + try { const client = createSdkClient(auth, sdkPrep.region, sdkPrep.effort) const command = new GenerateAssistantResponseCommand({ @@ -153,13 +184,20 @@ export class RequestHandler { this.handleSuccessfulRequest(acc) this.usageTracker.syncUsage(acc, auth) - return await this.responseHandler.handleSdkSuccess( + const result = await this.responseHandler.handleSdkSuccess( sdkResponse, model, sdkPrep.conversationId, - sdkPrep.streaming + sdkPrep.streaming, + sdkPrep.toolNameMapper, + think ) + logger.debug(`[REQ] done convId=${sdkPrep.conversationId}`) + return result } catch (e: any) { + logger.warn( + `[REQ] error convId=${sdkPrep.conversationId}: ${e?.name || ''} ${e?.message?.slice(0, 200) || String(e).slice(0, 200)}` + ) const httpStatus = e?.$metadata?.httpStatusCode if (httpStatus) { @@ -180,13 +218,15 @@ export class RequestHandler { e, mockResponse, acc, - { retry }, + { retry, excludedMs: retryContext.excludedMs }, showToast ) if (errorResult.shouldRetry) { if (errorResult.newContext) { retry = errorResult.newContext.retry + const sleptMs = (errorResult.newContext.excludedMs ?? 0) - retryContext.excludedMs + if (sleptMs > 0) this.retryStrategy.markSleep(retryContext, sleptMs) } if (errorResult.switchAccount) { continue @@ -194,6 +234,24 @@ export class RequestHandler { continue } + if (httpStatus === 400 && e?.name === 'ValidationException' && !forceNewConversation) { + const { workspace, fingerprint } = sdkPrep.conversationKey + kiroDb.deleteConversationId(workspace, fingerprint) + // The conversation is starting fresh — drop any carried-forward + // images too so the new convId doesn't inherit stale state. + imageCache.delete(workspace, fingerprint) + logger.warn( + `[REQ] stale conversationId reset, retrying convId=${sdkPrep.conversationId}` + ) + forceNewConversation = true + continue + } + + if (this.allAccountsPermanentlyUnhealthy()) { + const reauthed = await this.triggerReauth(showToast) + if (reauthed) continue + } + throw new Error(`Kiro Error: ${httpStatus}`) } @@ -221,30 +279,41 @@ export class RequestHandler { auth: KiroAuthDetails, think: boolean, budget: number, - showToast?: (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void + showToast?: (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void, + _uiEffort: 'low' | 'medium' | 'high' | 'default' = 'default', + sessionId?: string ): SdkPreparedRequest { - return transformToSdkRequest(body, model, auth, think, budget, showToast, { - effort: this.config.effort, - autoEffortMapping: this.config.auto_effort_mapping - }) + return transformToSdkRequest( + body, + model, + auth, + think, + budget, + showToast, + this.workspace, + this.config.image_carry_forward, + sessionId, + { effort: this.config.effort, autoEffortMapping: this.config.auto_effort_mapping } + ) } private handleSuccessfulRequest(acc: ManagedAccount): void { - if (acc.failCount && acc.failCount > 0) { - if (!isPermanentError(acc.unhealthyReason)) { - acc.failCount = 0 - acc.isHealthy = true - delete acc.unhealthyReason - delete acc.recoveryTime - this.repository.save(acc).catch(() => {}) - } + // Only write to DB if the account was actually degraded — avoids a + // withDatabaseLock + full merge/dedup round-trip on every healthy request. + if (acc.failCount && acc.failCount > 0 && !isPermanentError(acc.unhealthyReason)) { + acc.failCount = 0 + acc.isHealthy = true + delete acc.unhealthyReason + delete acc.recoveryTime + this.repository.save(acc).catch(() => {}) } } private logSdkRequest(prep: SdkPreparedRequest, acc: ManagedAccount, timestamp: string): void { + this.logImageDiagnostic(prep) logger.logApiRequest( { - url: `https://q.${prep.region}.amazonaws.com/generateAssistantResponse`, + url: `${prep.endpoint}/generateAssistantResponse`, method: 'POST', headers: { 'x-amzn-kiro-agent-mode': 'vibe' }, body: { @@ -264,6 +333,32 @@ export class RequestHandler { ) } + private logImageDiagnostic(prep: SdkPreparedRequest): void { + const kb = (bytes: number): number => Math.round(bytes / 1024) + const sumBytes = (imgs: { source?: { bytes?: { byteLength?: number } } }[]): number => + imgs.reduce((n, im) => n + (im.source?.bytes?.byteLength ?? 0), 0) + + const cmImgs = prep.conversationState.currentMessage?.userInputMessage?.images ?? [] + const history = (prep.conversationState as any).history ?? [] + const histDetail: string[] = [] + let histImgs = 0 + let histKb = 0 + for (let i = 0; i < history.length; i++) { + const imgs = history[i]?.userInputMessage?.images ?? [] + if (imgs.length === 0) continue + const entryKb = kb(sumBytes(imgs)) + histDetail.push(`i=${i}:user:${imgs.length}(${entryKb}KB)`) + histImgs += imgs.length + histKb += entryKb + } + + const detail = histDetail.length ? ` detail=[${histDetail.join(',')}]` : '' + logger.log( + `[IMG] convId=${prep.conversationId} cur=${cmImgs.length}(${kb(sumBytes(cmImgs))}KB)` + + ` hist=${histImgs}/${history.length}(${histKb}KB)${detail}` + ) + } + private logSdkResponse(prep: SdkPreparedRequest, timestamp: string): void { logger.logApiResponse( { @@ -295,7 +390,7 @@ export class RequestHandler { if (!this.config.enable_log_api_request) { logger.logApiError( { - url: `https://q.${prep.region}.amazonaws.com/generateAssistantResponse`, + url: `${prep.endpoint}/generateAssistantResponse`, method: 'POST', headers: {}, body: null, @@ -327,9 +422,26 @@ export class RequestHandler { return this.reauthInFlight } + if (!kiroDb.acquireReauthLock()) { + logger.warn('Reauth lock held by another instance — polling for completion') + showToast('Another session is re-authenticating. Please wait...', 'info') + const deadline = Date.now() + 10_000 + while (Date.now() < deadline) { + await this.sleep(1000) + if (kiroDb.isReauthLockHeld()) continue + this.repository.invalidateCache() + const accounts = await this.repository.findAll() + for (const acc of accounts) this.accountManager.addAccount(acc) + return this.hasUsableAccount(accounts) + } + showToast('Re-authentication timed out. Please try again.', 'error') + return false + } + this.reauthInFlight = this.performReauth(showToast) const success = await this.reauthInFlight.finally(() => { this.reauthInFlight = null + kiroDb.releaseReauthLock() }) if (!success) this.lastFailedReauthAt = Date.now() return success @@ -338,15 +450,31 @@ export class RequestHandler { private async performReauth(showToast: ToastFunction): Promise { try { showToast('Session expired. Re-authenticating...', 'warning') - await this.client.provider.oauth.authorize({ - path: { id: 'kiro' }, - body: { method: 0 } - }) + logger.warn('Reauth: starting oauth flow') + + const withTimeout = (promise: Promise, label: string): Promise => { + let timer: ReturnType | undefined + return Promise.race([ + promise.finally(() => clearTimeout(timer)), + new Promise( + (_, reject) => + (timer = setTimeout( + () => reject(new Error(`Reauth timed out waiting for ${label}`)), + REAUTH_TIMEOUT_MS + )) + ) + ]) + } - await this.client.provider.oauth.callback({ - path: { id: 'kiro' }, - body: { method: 0 } - }) + await withTimeout( + this.client.provider.oauth.authorize({ path: { id: 'kiro' }, body: { method: 0 } }), + 'oauth.authorize' + ) + + await withTimeout( + this.client.provider.oauth.callback({ path: { id: 'kiro' }, body: { method: 0 } }), + 'oauth.callback' + ) this.repository.invalidateCache() const accounts = await this.repository.findAll() @@ -364,6 +492,12 @@ export class RequestHandler { return true } catch (e) { logger.error('Re-auth failed', e instanceof Error ? e : new Error(String(e))) + showToast( + e instanceof Error && e.message.includes('timed out') + ? 'Re-authentication timed out. Please try again.' + : 'Re-authentication failed. Please try again.', + 'error' + ) return false } } diff --git a/src/plugin.ts b/src/plugin.ts index a4552e2..2ffba94 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -6,6 +6,10 @@ import { AccountRepository } from './infrastructure/database/account-repository. import { AccountManager } from './plugin/accounts.js' import { bootstrapAuthIfNeeded } from './plugin/auth-bootstrap.js' import { loadConfig } from './plugin/config/index.js' +import { imageCache } from './plugin/image-cache.js' +import * as logger from './plugin/logger.js' +import { clearSdkClientCache } from './plugin/sdk-client.js' +import { kiroDb } from './plugin/storage/sqlite.js' type ToastFunction = (message: string, variant: string) => void @@ -27,7 +31,7 @@ export const createKiroPlugin = const accountManager = await AccountManager.loadFromDisk(config.account_selection_strategy) authHandler.setAccountManager(accountManager) - const requestHandler = new RequestHandler(accountManager, config, repository, client) + const requestHandler = new RequestHandler(accountManager, config, repository, client, directory) // Compute the base URL once so both the config hook and auth loader use the same value const baseURL = KIRO_CONSTANTS.BASE_URL.replace('/generateAssistantResponse', '').replace( @@ -176,6 +180,19 @@ export const createKiroPlugin = return normalized } + }, + dispose: async () => { + logger.debug('[DISPOSE] Kiro plugin shutting down') + try { + clearSdkClientCache() + } catch {} + try { + imageCache.clear() + } catch {} + try { + kiroDb.close() + } catch {} + logger.debug('[DISPOSE] Kiro plugin shutdown complete') } } } diff --git a/src/plugin/logger.ts b/src/plugin/logger.ts index b0f85fe..1938668 100644 --- a/src/plugin/logger.ts +++ b/src/plugin/logger.ts @@ -72,7 +72,7 @@ export function warn(message: string, ...args: unknown[]): void { } export function debug(message: string, ...args: unknown[]): void { - if (process.env.DEBUG) { + if (process.env.DEBUG || process.env.OPENCODE_LOG_LEVEL === 'debug') { writeToFile('DEBUG', message, ...args) } } diff --git a/src/plugin/request.ts b/src/plugin/request.ts index a9459ef..79f09ae 100644 --- a/src/plugin/request.ts +++ b/src/plugin/request.ts @@ -13,16 +13,24 @@ import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' import { + buildToolNameMaps, convertToolsToCodeWhisperer, - deduplicateToolResults + deduplicateToolResults, + shortenToolName } from '../infrastructure/transformers/tool-transformer.js' import { getEffectiveEffort } from './effort.js' +import { imageCache } from './image-cache.js' import { + MAX_KIRO_IMAGES, convertImagesToKiroFormat, extractAllImages, - extractTextFromParts + extractTextFromParts, + type KiroImage } from './image-handler.js' +import * as logger from './logger.js' import { resolveKiroModel } from './models.js' +import { resolveKiroEndpoint } from './sdk-client.js' +import { kiroDb } from './storage/sqlite.js' import type { CodeWhispererRequest, Effort, @@ -31,15 +39,51 @@ import type { SdkPreparedRequest } from './types' +interface EffortConfig { + effort?: Effort + autoEffortMapping?: boolean +} + +// Look up or mint a stable convId per (workspace, fingerprint) pair. +// On a cache miss (new session or new prompt) we mint UUIDs; on a hit we +// reuse the stored ones so continuations land in the same conversation. +function deriveConversationIds( + workspace: string, + firstUserContent: string +): { convId: string; agentContinuationId: string; fingerprint: string } { + const fingerprint = crypto + .createHash('sha256') + .update(workspace + '\0' + (firstUserContent || '_empty_')) + .digest('hex') + .slice(0, 32) + + const existing = kiroDb.getConversationId(workspace, fingerprint) + if (existing) { + if (!existing.agentContinuationId) { + existing.agentContinuationId = crypto.randomUUID() + kiroDb.setConversationId( + workspace, + fingerprint, + existing.convId, + existing.agentContinuationId + ) + } + return { ...existing, fingerprint } + } + + const convId = crypto.randomUUID() + const agentContinuationId = crypto.randomUUID() + kiroDb.setConversationId(workspace, fingerprint, convId, agentContinuationId) + return { convId, agentContinuationId, fingerprint } +} + interface TransformResult { request: CodeWhispererRequest resolved: string convId: string -} - -interface EffortConfig { - effort?: Effort - autoEffortMapping?: boolean + agentContinuationId: string + fingerprint: string + toolNameMapper?: (name: string) => string } type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void @@ -50,13 +94,15 @@ function buildCodeWhispererRequest( auth: KiroAuthDetails, think = false, budget = 20000, - showToast?: ToastFunction + showToast?: ToastFunction, + workspace = '', + carryForward = true, + sessionId?: string ): TransformResult { const req = typeof body === 'string' ? JSON.parse(body) : body const { messages, tools, system } = req - const convId = crypto.randomUUID() if (!messages || messages.length === 0) throw new Error('No messages') - const resolved = resolveKiroModel(model) + const systemMsgs = messages.filter((m: any) => m.role === 'system') const otherMsgs = messages.filter((m: any) => m.role !== 'system') let sys = system || '' @@ -71,6 +117,25 @@ function buildCodeWhispererRequest( const msgs = mergeAdjacentMessages([...otherMsgs]) const lastMsg = msgs[msgs.length - 1] if (lastMsg && lastMsg.role === 'assistant' && getContentText(lastMsg) === '{') msgs.pop() + + const firstUserMsg = msgs.find((m: any) => m.role === 'user') + // Use text only — image bytes in the content array would change the + // fingerprint once OpenCode strips them on subsequent turns. + const firstUserContent = firstUserMsg + ? typeof firstUserMsg.content === 'string' + ? firstUserMsg.content + : extractTextFromParts(firstUserMsg.content) + : '' + const workspaceKey = sessionId ? `sess:${sessionId}` : workspace + if (process.env.DEBUG || process.env.OPENCODE_LOG_LEVEL === 'debug') { + logger.debug(`[CONV] ws=${workspaceKey} sessionId=${sessionId ?? 'none'} msgs=${msgs.length}`) + } + const { convId, agentContinuationId, fingerprint } = deriveConversationIds( + workspaceKey, + firstUserContent + ) + const resolved = resolveKiroModel(model) + const toolMaps = tools ? buildToolNameMaps(tools) : undefined const cwTools = tools ? convertToolsToCodeWhisperer(tools) : [] let history = buildHistory(msgs, resolved) @@ -114,7 +179,7 @@ function buildCodeWhispererRequest( else if (p.type === 'thinking') th += p.thinking || p.text || '' else if (p.type === 'tool_use') { if (!arm.toolUses) arm.toolUses = [] - arm.toolUses.push({ input: p.input, name: p.name, toolUseId: p.id }) + arm.toolUses.push({ input: p.input, name: shortenToolName(p.name), toolUseId: p.id }) } } } else arm.content = getContentText(curMsg) @@ -126,7 +191,7 @@ function buildCodeWhispererRequest( typeof tc.function?.arguments === 'string' ? JSON.parse(tc.function.arguments) : tc.function?.arguments, - name: tc.function?.name, + name: shortenToolName(tc.function?.name), toolUseId: tc.id }) } @@ -186,6 +251,8 @@ function buildCodeWhispererRequest( } const request: CodeWhispererRequest = { conversationState: { + agentContinuationId, + agentTaskType: 'vibe', chatTriggerType: KIRO_CONSTANTS.CHAT_TRIGGER_TYPE_MANUAL, conversationId: convId, currentMessage: { @@ -209,7 +276,7 @@ function buildCodeWhispererRequest( if (originalCall) { orphanedTrs.push({ call: { - name: originalCall.name || originalCall.function?.name || 'tool', + name: shortenToolName(originalCall.name || originalCall.function?.name || 'tool'), toolUseId: tr.toolUseId, input: originalCall.input || @@ -277,7 +344,114 @@ function buildCodeWhispererRequest( } } - return { request, resolved, convId } + // Strip empty toolUses arrays from history (Kiro quirk) + for (const h of history) { + if (h.assistantResponseMessage?.toolUses && h.assistantResponseMessage.toolUses.length === 0) { + delete h.assistantResponseMessage.toolUses + } + } + + // Trim history if payload exceeds Kiro's ~615KB limit. + // Compute per-entry sizes once and shrink incrementally to avoid O(N²) + // re-stringifying the full request on every iteration. + const MAX_PAYLOAD_BYTES = 600_000 + if (history.length > 2) { + const sizes = history.map((h) => JSON.stringify(h).length + 1) + const baseRequest: any = { ...request, conversationState: { ...request.conversationState } } + delete baseRequest.conversationState.history + let totalSize = JSON.stringify(baseRequest).length + 2 // for `"history":[]` + for (const s of sizes) totalSize += s + + while (history.length > 2 && totalSize > MAX_PAYLOAD_BYTES) { + // Drop the two oldest entries (typically user + assistant pair). + totalSize -= (sizes.shift() || 0) + (sizes.shift() || 0) + history.splice(0, 2) + + // Strip leading orphans: assistantResponseMessage can't start the history. + while (history.length > 0 && history[0]?.assistantResponseMessage) { + totalSize -= sizes.shift() || 0 + history.shift() + } + + // Strip leading toolResult-only userInputMessages whose toolUseIds are gone. + const toolUseIds = new Set( + history.flatMap( + (h) => h.assistantResponseMessage?.toolUses?.map((tu: any) => tu.toolUseId) ?? [] + ) + ) + while (history.length > 0) { + const trs = history[0]?.userInputMessage?.userInputMessageContext?.toolResults + if (!trs) break + const allMatched = trs.every((tr: any) => toolUseIds.has(tr.toolUseId)) + if (allMatched) break + totalSize -= sizes.shift() || 0 + history.shift() + } + } + } + + // After trimming, drop any current-message toolResults whose tool_use was + // removed from history — sending a toolResult without a matching toolUse + // in the surviving history causes Bedrock 400 ValidationException. + if (uim?.userInputMessageContext?.toolResults) { + const survivingToolUseIds = new Set( + history.flatMap( + (h) => h.assistantResponseMessage?.toolUses?.map((tu: any) => tu.toolUseId) ?? [] + ) + ) + const filtered = uim.userInputMessageContext.toolResults.filter((tr: any) => + survivingToolUseIds.has(tr.toolUseId) + ) + if (filtered.length === 0) { + delete uim.userInputMessageContext.toolResults + if (Object.keys(uim.userInputMessageContext).length === 0) { + delete uim.userInputMessageContext + } + } else { + uim.userInputMessageContext.toolResults = filtered + } + } + + // Image carry-forward: cache surviving images; on turns without fresh images, + // restore from cache — but never onto a tool-result turn (Kiro 400 on + // images+toolResults). Skip the history scan entirely if this conversation + // has never carried images — saves O(N) over 1000+ entry sessions. + if (carryForward && uim) { + const cmImgs = (uim.images as KiroImage[] | undefined) ?? [] + const wireImages: KiroImage[] = cmImgs.length > 0 ? [...cmImgs] : [] + if ( + wireImages.length < MAX_KIRO_IMAGES && + imageCache.hasEverHadImages(workspaceKey, fingerprint) + ) { + for (const h of history) { + const imgs = h.userInputMessage?.images as KiroImage[] | undefined + if (!imgs || imgs.length === 0) continue + wireImages.push(...imgs) + if (wireImages.length >= MAX_KIRO_IMAGES) break + } + } + + if (wireImages.length > 0) { + imageCache.upsert(workspaceKey, fingerprint, wireImages) + } else { + const hasToolResults = (uim.userInputMessageContext?.toolResults?.length ?? 0) > 0 + if (!hasToolResults) { + const cached = imageCache.get(workspaceKey, fingerprint) + if (cached && cached.length > 0) { + uim.images = cached.slice(0, MAX_KIRO_IMAGES) + } + } + } + } + + return { + request, + resolved, + convId, + agentContinuationId, + fingerprint, + toolNameMapper: toolMaps?.fromKiroName + } } export function transformToCodeWhisperer( @@ -286,9 +460,22 @@ export function transformToCodeWhisperer( model: string, auth: KiroAuthDetails, think = false, - budget = 20000 + budget = 20000, + workspace = '', + carryForward = true, + sessionId?: string ): PreparedRequest { - const { request, resolved, convId } = buildCodeWhispererRequest(body, model, auth, think, budget) + const { request, resolved, convId } = buildCodeWhispererRequest( + body, + model, + auth, + think, + budget, + undefined, + workspace, + carryForward, + sessionId + ) const osP = os.platform(), osR = os.release(), nodeV = process.version.replace('v', '') @@ -325,15 +512,21 @@ export function transformToSdkRequest( think = false, budget = 20000, showToast?: ToastFunction, + workspace = '', + carryForward = true, + sessionId?: string, effortConfig?: EffortConfig ): SdkPreparedRequest { - const { request, resolved, convId } = buildCodeWhispererRequest( + const { request, resolved, convId, fingerprint, toolNameMapper } = buildCodeWhispererRequest( body, model, auth, think, budget, - showToast + showToast, + workspace, + carryForward, + sessionId ) // Resolve effort level based on config and model capabilities @@ -345,13 +538,22 @@ export function transformToSdkRequest( effortConfig?.autoEffortMapping ?? true ) + const region = extractRegionFromArn(auth.profileArn) ?? auth.region + // resolveKiroEndpoint returns the full URL including /generateAssistantResponse. + // Strip the path so the endpoint field holds only the base (what the SDK uses). + const endpointFull = resolveKiroEndpoint(auth) + const endpoint = endpointFull.replace(/\/generateAssistantResponse$/, '') + const workspaceKey = sessionId ? `sess:${sessionId}` : workspace return { conversationState: request.conversationState, profileArn: request.profileArn, streaming: true, effectiveModel: resolved, conversationId: convId, - region: extractRegionFromArn(auth.profileArn) ?? auth.region, - effort + conversationKey: { workspace: workspaceKey, fingerprint }, + region, + endpoint, + effort, + toolNameMapper } } diff --git a/src/plugin/storage/migrations.ts b/src/plugin/storage/migrations.ts index aaf952b..af18835 100644 --- a/src/plugin/storage/migrations.ts +++ b/src/plugin/storage/migrations.ts @@ -9,6 +9,45 @@ export function runMigrations(db: Database): void { migrateStartUrlColumn(db) migrateOidcRegionColumn(db) migrateDropRefreshTokenUniqueIndex(db) + migrateConversationsTable(db) + migrateReauthLockTable(db) + migrateConversationsAgentContinuationId(db) +} + +function migrateConversationsTable(db: Database): void { + const hasTable = db + .prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='conversations'") + .get() + if (hasTable) return + + db.exec(` + CREATE TABLE conversations ( + workspace TEXT NOT NULL, + fingerprint TEXT NOT NULL, + conv_id TEXT NOT NULL, + last_used INTEGER NOT NULL, + PRIMARY KEY (workspace, fingerprint) + ) + `) + db.exec('CREATE INDEX idx_conversations_last_used ON conversations(last_used)') +} + +function migrateReauthLockTable(db: Database): void { + db.exec(` + CREATE TABLE IF NOT EXISTS reauth_lock ( + id INTEGER PRIMARY KEY CHECK (id = 1), + pid INTEGER NOT NULL, + acquired_at INTEGER NOT NULL + ) + `) +} + +function migrateConversationsAgentContinuationId(db: Database): void { + const columns = db.prepare('PRAGMA table_info(conversations)').all() as any[] + const names = new Set(columns.map((c: any) => c.name)) + if (!names.has('agent_continuation_id')) { + db.exec('ALTER TABLE conversations ADD COLUMN agent_continuation_id TEXT') + } } function migrateToUniqueRefreshToken(db: Database): void { diff --git a/src/plugin/storage/sqlite.ts b/src/plugin/storage/sqlite.ts index 53cedfc..bb30524 100644 --- a/src/plugin/storage/sqlite.ts +++ b/src/plugin/storage/sqlite.ts @@ -197,9 +197,125 @@ export class KiroDatabase { } } + private static readonly REAUTH_LOCK_TTL_MS = 120_000 + + acquireReauthLock(): boolean { + const now = Date.now() + try { + this.db.exec('BEGIN IMMEDIATE') + } catch { + // Another write transaction is active — treat as lock held + return false + } + try { + const existing = this.db + .prepare('SELECT pid, acquired_at FROM reauth_lock WHERE id = 1') + .get() as { pid: number; acquired_at: number } | undefined + + if (existing) { + const expired = now - existing.acquired_at >= KiroDatabase.REAUTH_LOCK_TTL_MS + const dead = (() => { + try { + process.kill(existing.pid, 0) + return false + } catch { + return true + } + })() + if (expired || dead) { + this.db.prepare('DELETE FROM reauth_lock WHERE id = 1').run() + } else { + this.db.exec('ROLLBACK') + return false + } + } + + // INSERT OR REPLACE handles a race where two instances both saw the + // same dead/expired lock and both reach this branch. + this.db + .prepare('INSERT OR REPLACE INTO reauth_lock (id, pid, acquired_at) VALUES (1, ?, ?)') + .run(process.pid, now) + this.db.exec('COMMIT') + return true + } catch { + try { + this.db.exec('ROLLBACK') + } catch { + // already rolled back + } + return false + } + } + + isReauthLockHeld(): boolean { + const row = this.db.prepare('SELECT pid FROM reauth_lock WHERE id = 1').get() as + { pid: number } | undefined + if (!row) return false + try { + process.kill(row.pid, 0) + return true + } catch { + return false + } + } + + releaseReauthLock(): void { + this.db.prepare('DELETE FROM reauth_lock WHERE id = 1 AND pid = ?').run(process.pid) + } + close() { this.db.close() } + + getConversationId( + workspace: string, + fingerprint: string + ): { convId: string; agentContinuationId: string } | undefined { + const row = this.db + .prepare( + 'SELECT conv_id, agent_continuation_id FROM conversations WHERE workspace = ? AND fingerprint = ?' + ) + .get(workspace, fingerprint) as + { conv_id: string; agent_continuation_id: string | null } | undefined + return row + ? { convId: row.conv_id, agentContinuationId: row.agent_continuation_id || '' } + : undefined + } + + deleteConversationId(workspace: string, fingerprint: string): void { + this.db + .prepare('DELETE FROM conversations WHERE workspace = ? AND fingerprint = ?') + .run(workspace, fingerprint) + } + + /** + * Persist a conversationId and agentContinuationId, clean up entries older than ttlDays (default 7). + */ + setConversationId( + workspace: string, + fingerprint: string, + convId: string, + agentContinuationId: string, + ttlDays = 7 + ): void { + const now = Date.now() + const cutoff = now - ttlDays * 24 * 60 * 60 * 1000 + this.db.exec('BEGIN TRANSACTION') + try { + this.db + .prepare( + `INSERT INTO conversations (workspace, fingerprint, conv_id, agent_continuation_id, last_used) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(workspace, fingerprint) DO UPDATE SET conv_id = excluded.conv_id, agent_continuation_id = excluded.agent_continuation_id, last_used = excluded.last_used` + ) + .run(workspace, fingerprint, convId, agentContinuationId, now) + this.db.prepare('DELETE FROM conversations WHERE last_used < ?').run(cutoff) + this.db.exec('COMMIT') + } catch (e) { + this.db.exec('ROLLBACK') + throw e + } + } } export function createDatabase(path?: string): KiroDatabase { diff --git a/src/plugin/types.ts b/src/plugin/types.ts index a17e0d8..5819b77 100644 --- a/src/plugin/types.ts +++ b/src/plugin/types.ts @@ -83,6 +83,8 @@ export interface CodeWhispererMessage { export interface CodeWhispererRequest { conversationState: { + agentContinuationId?: string + agentTaskType?: string chatTriggerType: string conversationId: string history?: CodeWhispererMessage[] @@ -119,9 +121,14 @@ export interface SdkPreparedRequest { streaming: boolean effectiveModel: string conversationId: string + conversationKey: { workspace: string; fingerprint: string } region: string /** Resolved effort level for thinking models */ effort?: Effort + // Resolved endpoint base URL (q.amazonaws.com or runtime.kiro.dev). + // Set by transformToSdkRequest so callers and logs can show the real target. + endpoint: string + toolNameMapper?: (name: string) => string } export type AccountSelectionStrategy = 'sticky' | 'round-robin' | 'lowest-usage' From f3bf7057aba7c42dd099fa6c160310ec60827e79 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 09:51:58 +0200 Subject: [PATCH 07/10] test: cover request pipeline, streaming, images, conversation store Add unit tests for the SDK client endpoint routing, stream/tool-call transformer, response handler, error/retry handling, event-stream parser, image cache, and the conversation/reauth-lock parts of the SQLite store. --- src/__tests__/edge-cases.test.ts | 904 ++++++++++++++++++++++ src/__tests__/error-handler.test.ts | 231 ++++++ src/__tests__/event-stream-parser.test.ts | 89 +++ src/__tests__/image-cache.test.ts | 356 +++++++++ src/__tests__/response.test.ts | 69 ++ src/__tests__/sqlite.test.ts | 246 ++++++ src/__tests__/stream-transformer.test.ts | 672 ++++++++++++++++ src/__tests__/tool-call-parser.test.ts | 115 +++ 8 files changed, 2682 insertions(+) create mode 100644 src/__tests__/edge-cases.test.ts create mode 100644 src/__tests__/error-handler.test.ts create mode 100644 src/__tests__/event-stream-parser.test.ts create mode 100644 src/__tests__/image-cache.test.ts create mode 100644 src/__tests__/response.test.ts create mode 100644 src/__tests__/sqlite.test.ts create mode 100644 src/__tests__/stream-transformer.test.ts create mode 100644 src/__tests__/tool-call-parser.test.ts diff --git a/src/__tests__/edge-cases.test.ts b/src/__tests__/edge-cases.test.ts new file mode 100644 index 0000000..3ba24e8 --- /dev/null +++ b/src/__tests__/edge-cases.test.ts @@ -0,0 +1,904 @@ +import { describe, expect, mock, test } from 'bun:test' + +// Stateful mock so consecutive calls can resolve prior setConversationId +// entries — mirrors the real DB enough to exercise deriveConversationIds. +const convStore = new Map() + +mock.module('../plugin/storage/sqlite.js', () => ({ + kiroDb: { + getConversationId: (_ws: string, fp: string) => convStore.get(fp), + setConversationId: (_ws: string, fp: string, convId: string, agentContinuationId: string) => { + convStore.set(fp, { convId, agentContinuationId }) + }, + deleteConversationId: (ws: string, fp: string) => { + convStore.delete(`${ws}\0${fp}`) + }, + getAccounts: () => [], + upsertAccount: () => Promise.resolve(), + deleteAccount: () => Promise.resolve(), + batchUpsertAccounts: () => Promise.resolve() + } +})) + +mock.module('../plugin/sync/kiro-cli.js', () => ({ + syncFromKiroCli: () => Promise.resolve(), + writeToKiroCli: () => Promise.resolve() +})) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken, + accessTokenExpired: () => false +})) + +import { + buildHistory, + collapseAgenticLoops +} from '../infrastructure/transformers/history-builder.js' +import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' +import { + buildToolNameMaps, + convertToolsToCodeWhisperer, + deduplicateToolCallsByContent, + shortenToolName +} from '../infrastructure/transformers/tool-transformer.js' +import { imageCache } from '../plugin/image-cache.js' +import { transformToSdkRequest } from '../plugin/request.js' +import type { KiroAuthDetails } from '../plugin/types.js' + +const auth: KiroAuthDetails = { + refresh: 'refresh', + access: 'token', + expires: Date.now() + 3600000, + authMethod: 'idc', + region: 'us-east-1', + email: 'test@test.com', + profileArn: 'arn:aws:codewhisperer:us-east-1:123:profile/ABC' +} + +describe('shortenToolName edge cases', () => { + test('null input', () => { + expect(shortenToolName(null as any)).toBeFalsy() + }) + + test('undefined input', () => { + expect(shortenToolName(undefined as any)).toBeFalsy() + }) + + test('empty string', () => { + expect(shortenToolName('')).toBe('') + }) + + test('exactly 64 chars', () => { + const name = 'a'.repeat(64) + expect(shortenToolName(name)).toBe(name) + }) + + test('65 chars', () => { + const name = 'a'.repeat(65) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('200 chars', () => { + const name = 'a'.repeat(200) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('special chars', () => { + const name = '🎉'.repeat(40) + '\n\t' + 'ñ'.repeat(30) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('does not split surrogate pair', () => { + // 🎉 is U+1F389 = 2 UTF-16 code units. 40 of them = 80 UTF-16 chars. + // The naive slice would land between high+low surrogate. + const name = '🎉'.repeat(40) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + // No unpaired surrogate left in the prefix + expect(/[\uD800-\uDBFF](?![\uDC00-\uDFFF])/.test(result)).toBe(false) + expect(/(? { + const a = 'tool_a_'.repeat(20) + const b = 'tool_b_'.repeat(20) + expect(shortenToolName(a)).not.toBe(shortenToolName(b)) + }) + + test('same long name is deterministic', () => { + const name = 'x'.repeat(100) + expect(shortenToolName(name)).toBe(shortenToolName(name)) + }) +}) + +describe('sanitizeSchema via convertToolsToCodeWhisperer', () => { + test('deeply nested schema strips additionalProperties', () => { + const tools = [ + { + name: 'deep', + description: 'test', + input_schema: { + type: 'object', + additionalProperties: false, + required: [], + properties: { + nested: { + type: 'object', + additionalProperties: true, + required: [], + properties: { + items: { + type: 'array', + items: { + type: 'object', + additionalProperties: false, + required: [], + properties: { + value: { + anyOf: [ + { type: 'string', additionalProperties: false, required: [] }, + { type: 'number', additionalProperties: true } + ] + } + } + } + } + } + } + } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + const json = JSON.stringify(result) + expect(json).not.toContain('additionalProperties') + expect(json).not.toContain('"required":[]') + }) + + test('schema with empty required at multiple levels', () => { + const tools = [ + { + name: 'req', + description: 'test', + input_schema: { + type: 'object', + required: [], + properties: { + a: { type: 'object', required: [], properties: { b: { type: 'string', required: [] } } } + } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('"required":[]') + }) + + test('schema with ONLY additionalProperties and required:[]', () => { + const tools = [ + { + name: 'minimal', + description: 'test', + input_schema: { additionalProperties: false, required: [] } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.inputSchema.json).toEqual({}) + }) + + test('null schema', () => { + const tools = [{ name: 'n', description: 'test', input_schema: null }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('undefined schema', () => { + const tools = [{ name: 'u', description: 'test' }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('strips additionalProperties from "not" subschema', () => { + const tools = [ + { + name: 'n', + description: 'test', + input_schema: { type: 'object', not: { additionalProperties: false, required: [] } } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + expect(JSON.stringify(result)).not.toContain('"required":[]') + }) + + test('strips additionalProperties from patternProperties', () => { + const tools = [ + { + name: 'p', + description: 'test', + input_schema: { + type: 'object', + patternProperties: { '^x': { additionalProperties: false, required: [] } } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('strips additionalProperties from $defs', () => { + const tools = [ + { + name: 'd', + description: 'test', + input_schema: { + type: 'object', + $defs: { foo: { additionalProperties: false, required: [] } } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('strips additionalProperties from prefixItems', () => { + const tools = [ + { + name: 'pi', + description: 'test', + input_schema: { + type: 'array', + prefixItems: [{ additionalProperties: false }, { additionalProperties: true }] + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('handles circular schema reference without stack overflow', () => { + const obj: any = { type: 'object', additionalProperties: false, properties: {} } + obj.properties.self = obj + const tools = [{ name: 'c', description: 'test', input_schema: obj }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('does not mutate input schema', () => { + const original = { + type: 'object', + additionalProperties: false, + required: [], + properties: { x: { type: 'string', additionalProperties: false } } + } + const before = JSON.stringify(original) + convertToolsToCodeWhisperer([{ name: 't', description: 't', input_schema: original }]) + expect(JSON.stringify(original)).toBe(before) + }) +}) + +describe('payload trim preserves valid structure', () => { + test('large payload trimmed to ≤600KB', () => { + const messages: any[] = [] + for (let i = 0; i < 80; i++) { + messages.push({ role: 'user', content: 'x'.repeat(10000) }) + messages.push({ + role: 'assistant', + content: [ + { type: 'text', text: 'y'.repeat(5000) }, + { type: 'tool_use', id: `tool-${i}`, name: 'myTool', input: { q: 'test' } } + ] + }) + } + messages.push({ + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tool-79', content: 'result' }] + }) + messages.push({ role: 'user', content: 'final question' }) + + const body = { + messages, + tools: [{ name: 'myTool', description: 'a tool', input_schema: { type: 'object' } }] + } + const result = transformToSdkRequest(body, 'auto', auth) + const payload = JSON.stringify(result.conversationState) + expect(payload.length).toBeLessThanOrEqual(600_000) + }) + + test('history starts with userInputMessage after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 80; i++) { + messages.push({ role: 'user', content: 'x'.repeat(10000) }) + messages.push({ role: 'assistant', content: 'y'.repeat(10000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history + if (history && history.length > 0) { + expect(history[0]!.userInputMessage).toBeDefined() + } + }) + + test('toolResults have matching toolUseIds in history after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 40; i++) { + messages.push({ role: 'user', content: 'x'.repeat(20000) }) + messages.push({ + role: 'assistant', + content: [ + { type: 'text', text: 'y'.repeat(5000) }, + { type: 'tool_use', id: `t-${i}`, name: 'myTool', input: { q: 'test' } } + ] + }) + messages.push({ + role: 'user', + content: [{ type: 'tool_result', tool_use_id: `t-${i}`, content: 'result '.repeat(1000) }] + }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { + messages, + tools: [{ name: 'myTool', description: 'd', input_schema: { type: 'object' } }] + } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history || [] + + const allToolUseIds = new Set() + for (const h of history) { + if (h.assistantResponseMessage?.toolUses) { + for (const tu of h.assistantResponseMessage.toolUses) allToolUseIds.add(tu.toolUseId) + } + } + for (const h of history) { + if (h.userInputMessage?.userInputMessageContext?.toolResults) { + for (const tr of h.userInputMessage.userInputMessageContext.toolResults) { + expect(allToolUseIds.has(tr.toolUseId)).toBe(true) + } + } + } + }) +}) + +describe('mergeAdjacentMessages edge cases', () => { + test('empty array', () => { + expect(mergeAdjacentMessages([])).toEqual([]) + }) + + test('assistant string + assistant array with tool_use', () => { + const msgs = [ + { role: 'assistant', content: 'hello' }, + { role: 'assistant', content: [{ type: 'tool_use', id: '1', name: 'x', input: {} }] } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].content).toContainEqual({ type: 'tool_use', id: '1', name: 'x', input: {} }) + }) + + test('assistant with tool_calls + assistant with tool_calls', () => { + const msgs = [ + { + role: 'assistant', + content: 'a', + tool_calls: [{ id: '1', function: { name: 'x', arguments: '{}' } }] + }, + { + role: 'assistant', + content: 'b', + tool_calls: [{ id: '2', function: { name: 'y', arguments: '{}' } }] + } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].tool_calls).toHaveLength(2) + }) + + test('user string + user array', () => { + const msgs = [ + { role: 'user', content: 'hello' }, + { role: 'user', content: [{ type: 'text', text: 'world' }] } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(Array.isArray(result[0].content)).toBe(true) + }) + + test('3 consecutive same-role messages', () => { + const msgs = [ + { role: 'assistant', content: 'a' }, + { role: 'assistant', content: 'b' }, + { role: 'assistant', content: 'c' } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].content).toBe('a\nb\nc') + }) +}) + +describe('buildHistory with null/undefined content', () => { + test('assistant with content: null', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: null }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('assistant with content: undefined', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: undefined }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('user with content: null', () => { + const msgs = [ + { role: 'user', content: null }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('assistant with empty tool_calls array', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: 'text', tool_calls: [] }, + { role: 'user', content: 'bye' } + ] + const history = buildHistory(msgs, 'model') + const asst = history.find((h) => h.assistantResponseMessage) + expect(asst?.assistantResponseMessage?.toolUses).toBeUndefined() + }) + + test('assistant with tool_use where name is undefined', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: [{ type: 'tool_use', id: '1', name: undefined, input: {} }] }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) +}) + +describe('collapseAgenticLoops edge cases', () => { + test('single assistant+user pair passes through unchanged', () => { + const history = [ + { + assistantResponseMessage: { + content: 'text', + toolUses: [{ name: 'x', toolUseId: '1', input: {} }] + } + }, + { + userInputMessage: { + content: 'result', + userInputMessageContext: { + toolResults: [{ toolUseId: '1', content: [{ text: 'r' }], status: 'success' }] + } + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result).toHaveLength(2) + expect(result[0]!.assistantResponseMessage?.content).toBe('text') + }) + + test('assistant with toolUses without following user passes through', () => { + const history = [ + { userInputMessage: { content: 'hi' } }, + { + assistantResponseMessage: { + content: 'text', + toolUses: [{ name: 'x', toolUseId: '1', input: {} }] + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result).toHaveLength(2) + }) + + test('3 pairs collapse correctly', () => { + const history = [ + { + assistantResponseMessage: { + content: 'first', + toolUses: [{ name: 'a', toolUseId: '1', input: {} }] + } + }, + { + userInputMessage: { + content: 'r1', + userInputMessageContext: { + toolResults: [{ toolUseId: '1', content: [{ text: 'r' }], status: 'success' }] + } + } + }, + { + assistantResponseMessage: { + content: 'second', + toolUses: [{ name: 'b', toolUseId: '2', input: {} }] + } + }, + { + userInputMessage: { + content: 'r2', + userInputMessageContext: { + toolResults: [{ toolUseId: '2', content: [{ text: 'r' }], status: 'success' }] + } + } + }, + { + assistantResponseMessage: { + content: 'third', + toolUses: [{ name: 'c', toolUseId: '3', input: {} }] + } + }, + { + userInputMessage: { + content: 'r3', + userInputMessageContext: { + toolResults: [{ toolUseId: '3', content: [{ text: 'r' }], status: 'success' }] + } + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result[0]!.assistantResponseMessage?.content).toBe('first') + expect(result[2]!.assistantResponseMessage?.content).toBe('[system: tool calling continues]') + expect(result[4]!.assistantResponseMessage?.content).toBe('[system: tool calling continues]') + }) +}) + +describe('deduplicateToolCallsByContent edge cases', () => { + test('same name different string inputs kept', () => { + const calls = [ + { name: 'tool', input: 'input1' }, + { name: 'tool', input: 'input2' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) + + test('same name same string input deduplicated', () => { + const calls = [ + { name: 'tool', input: 'same' }, + { name: 'tool', input: 'same' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(1) + }) + + test('empty array', () => { + expect(deduplicateToolCallsByContent([])).toEqual([]) + }) + + test('separator collision: name="a-b" input="c" vs name="a" input="b-c"', () => { + // Old impl used `${name}-${input}` so these would collapse to "a-b-c". + const calls = [ + { name: 'a-b', input: 'c' }, + { name: 'a', input: 'b-c' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) + + test('object input is stringified before compare', () => { + const calls = [ + { name: 'tool', input: { x: 1 } }, + { name: 'tool', input: { x: 1 } }, + { name: 'tool', input: { x: 2 } } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) +}) + +describe('convertToolsToCodeWhisperer edge cases', () => { + test('empty tools array', () => { + expect(convertToolsToCodeWhisperer([])).toEqual([]) + }) + + test('tool with undefined name', () => { + const tools = [{ description: 'test', input_schema: { type: 'object' } }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('tool with description >9216 chars truncated', () => { + const tools = [ + { name: 'big', description: 'x'.repeat(10000), input_schema: { type: 'object' } } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.description.length).toBe(9216) + }) + + test('tool with empty input_schema', () => { + const tools = [{ name: 'empty', description: 'test', input_schema: {} }] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.inputSchema.json).toEqual({}) + }) +}) + +describe('buildToolNameMaps edge cases', () => { + test('empty tools array', () => { + const maps = buildToolNameMaps([]) + expect(maps.toKiroName('anything')).toBeDefined() + expect(maps.fromKiroName('anything')).toBe('anything') + }) + + test('two tools with same name', () => { + const tools = [{ name: 'dup' }, { name: 'dup' }] + expect(() => buildToolNameMaps(tools)).not.toThrow() + const maps = buildToolNameMaps(tools) + expect(maps.toKiroName('dup')).toBe('dup') + }) + + test('tool with undefined name skipped', () => { + const tools = [{ name: undefined }, { name: 'valid' }] + const maps = buildToolNameMaps(tools) + expect(maps.toKiroName('valid')).toBe('valid') + }) +}) + +describe('history alternation after trim', () => { + test('after splice first entry is userInputMessage', () => { + const messages: any[] = [] + for (let i = 0; i < 100; i++) { + messages.push({ role: 'user', content: 'u'.repeat(8000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(8000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history + if (history && history.length > 0) { + expect(history[0]!.userInputMessage).toBeDefined() + expect(history[0]!.assistantResponseMessage).toBeUndefined() + } + }) + + test('remaining history alternates after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 100; i++) { + messages.push({ role: 'user', content: 'u'.repeat(8000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(8000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history || [] + for (let i = 0; i < history.length - 1; i++) { + const curr = history[i]! + const next = history[i + 1]! + if (curr.userInputMessage) { + expect(next.assistantResponseMessage).toBeDefined() + } + if (curr.assistantResponseMessage) { + expect(next.userInputMessage).toBeDefined() + } + } + }) +}) + +describe('fingerprint stability across image-strip', () => { + // Regression: OpenCode strips image parts from the conversation state across + // agentic turns. If the fingerprint depended on the full content (incl. + // base64 image bytes), the convId cache and the image-carry-forward cache + // would both miss after the first turn — which is exactly the bug we saw. + test('same first-user-message text produces the same fingerprint with or without image parts', () => { + const withImages = transformToSdkRequest( + { + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'Here is the design' }, + { + type: 'image_url', + image_url: { + url: 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==' + } + } + ] + } + ] + }, + 'auto', + auth, + false, + 0, + undefined, + '/ws-fp-stability', + false + ) + + const withoutImages = transformToSdkRequest( + { + messages: [ + { + role: 'user', + content: [{ type: 'text', text: 'Here is the design' }] + } + ] + }, + 'auto', + auth, + false, + 0, + undefined, + '/ws-fp-stability', + false + ) + + expect(withImages.conversationKey.fingerprint).toBe(withoutImages.conversationKey.fingerprint) + }) +}) + +describe('image carry-forward: no images on tool-result turns', () => { + function seedCache(workspace: string, firstUserText: string) { + const crypto = require('crypto') + const fingerprint = crypto + .createHash('sha256') + .update(workspace + '\0' + firstUserText) + .digest('hex') + .slice(0, 32) + imageCache.set(workspace, fingerprint, [ + { format: 'png', source: { bytes: new Uint8Array([137, 80, 78, 71]) } } + ]) + } + + test('images are NOT carried forward onto a tool-result turn', () => { + const workspace = '/ws-carry-fwd-toolresult' + seedCache(workspace, 'look at this') + + const result = transformToSdkRequest( + { + messages: [ + { role: 'user', content: [{ type: 'text', text: 'look at this' }] }, + { + role: 'assistant', + content: [ + { type: 'tool_use', id: 'tu_1', name: 'think', input: { thought: 'thinking' } } + ] + }, + { + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tu_1', content: 'done' }] + } + ] + }, + 'auto', + auth, + false, + 0, + undefined, + workspace, + true + ) + + const uim = result.conversationState.currentMessage?.userInputMessage as any + expect((uim?.userInputMessageContext?.toolResults ?? []).length).toBeGreaterThan(0) + expect((uim?.images ?? []).length).toBe(0) + }) + + test('images ARE carried forward on a normal (non-tool-result) turn', () => { + const workspace = '/ws-carry-fwd-normal' + seedCache(workspace, 'look at this') + + const result = transformToSdkRequest( + { + messages: [ + { role: 'user', content: [{ type: 'text', text: 'look at this' }] }, + { role: 'assistant', content: 'I see it.' }, + { role: 'user', content: 'What color is it?' } + ] + }, + 'auto', + auth, + false, + 0, + undefined, + workspace, + true + ) + + const uim = result.conversationState.currentMessage?.userInputMessage as any + expect((uim?.images ?? []).length).toBeGreaterThan(0) + }) +}) + +describe('payload trim performance', () => { + test('trims very large history in <50ms', () => { + // 200 user/asst pairs of 10KB each = ~4MB raw — would be O(N²) without + // the incremental size accounting. + const messages: any[] = [] + for (let i = 0; i < 200; i++) { + messages.push({ role: 'user', content: 'u'.repeat(10000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(10000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const start = performance.now() + const result = transformToSdkRequest({ messages }, 'auto', auth) + const elapsed = performance.now() - start + + // Result must still be under the limit. + const size = JSON.stringify(result.conversationState).length + expect(size).toBeLessThanOrEqual(600_000) + // Allow 100ms on slow CI; fail if we regress to seconds (O(N²)). + expect(elapsed).toBeLessThan(100) + }) +}) + +// ── Session ID isolation ────────────────────────────────────────────────────── + +describe('parallel sessions in same workspace get distinct convIds', () => { + test('different sessionId → different convId (cache key prefix)', () => { + const body = { + messages: [{ role: 'user', content: 'hello' }] + } + const r1 = transformToSdkRequest( + body, + 'auto', + auth, + false, + 16000, + undefined, + '', + true, + 'sess-A' + ) + const r2 = transformToSdkRequest( + body, + 'auto', + auth, + false, + 16000, + undefined, + '', + true, + 'sess-B' + ) + expect(r1.conversationId).not.toBe(r2.conversationId) + expect(r1.conversationKey.workspace).toBe('sess:sess-A') + expect(r2.conversationKey.workspace).toBe('sess:sess-B') + }) + + test('same sessionId + same prompt → same convId on second call', () => { + const body = { + messages: [{ role: 'user', content: 'hello' }] + } + const r1 = transformToSdkRequest( + body, + 'auto', + auth, + false, + 16000, + undefined, + '', + true, + 'sess-X' + ) + const r2 = transformToSdkRequest( + body, + 'auto', + auth, + false, + 16000, + undefined, + '', + true, + 'sess-X' + ) + expect(r1.conversationId).toBe(r2.conversationId) + }) + + test('no sessionId falls back to workspace as cache key', () => { + const body = { + messages: [{ role: 'user', content: 'hello' }] + } + const r1 = transformToSdkRequest(body, 'auto', auth, false, 16000, undefined, '/some/dir', true) + expect(r1.conversationKey.workspace).toBe('/some/dir') + }) +}) diff --git a/src/__tests__/error-handler.test.ts b/src/__tests__/error-handler.test.ts new file mode 100644 index 0000000..9743846 --- /dev/null +++ b/src/__tests__/error-handler.test.ts @@ -0,0 +1,231 @@ +import { describe, expect, mock, test } from 'bun:test' +import { ErrorHandler } from '../core/request/error-handler.js' +import { AccountManager } from '../plugin/accounts.js' +import type { ManagedAccount } from '../plugin/types.js' + +mock.module('../plugin/storage/sqlite.js', () => ({ + kiroDb: { + getAccounts: () => [], + upsertAccount: () => Promise.resolve(), + deleteAccount: () => Promise.resolve(), + batchUpsertAccounts: () => Promise.resolve() + } +})) +mock.module('../plugin/sync/kiro-cli.js', () => ({ + syncFromKiroCli: () => Promise.resolve(), + writeToKiroCli: () => Promise.resolve() +})) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken, + accessTokenExpired: () => false +})) + +const defaultConfig = { rate_limit_max_retries: 3, rate_limit_retry_delay_ms: 100 } +const noToast = () => {} + +function makeAccount(overrides: Partial = {}): ManagedAccount { + return { + id: 'acc-1', + email: 'test@example.com', + authMethod: 'idc', + region: 'eu-central-1', + refreshToken: 'r', + accessToken: 'a', + expiresAt: Date.now() + 3600000, + rateLimitResetTime: 0, + isHealthy: true, + failCount: 0, + lastUsed: 0, + usedCount: 0, + limitCount: 0, + ...overrides + } +} + +function makeRepo(accounts: ManagedAccount[]) { + return { + findAll: async () => accounts, + batchSave: async () => {}, + save: async () => {}, + invalidateCache: () => {} + } as any +} + +function makeResponse(status: number, body: any, headers: Record = {}): Response { + return new Response(JSON.stringify(body), { + status, + headers: { 'Content-Type': 'application/json', ...headers } + }) +} + +// ── 400 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 400', () => { + test('returns shouldRetry=false', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(400, { message: 'Bad Request' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(false) + }) +}) + +// ── 401 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 401', () => { + test('retries when under max retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(401, { message: 'Unauthorized' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.retry).toBe(1) + }) + + test('stops retrying at max retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(401, { message: 'Unauthorized' }) + const result = await handler.handle(null, res, acc, { retry: 3 }, noToast) + expect(result.shouldRetry).toBe(false) + }) +}) + +// ── 403 single account ──────────────────────────────────────────────────────── + +describe('ErrorHandler: 403 single account', () => { + test('bearer token invalid 403 forces token refresh (sets expiresAt=0) and retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { + message: 'The bearer token included in the request is invalid' + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + // Should retry so the token refresher can get a fresh token + expect(result.shouldRetry).toBe(true) + // Account stays healthy — refresh will handle it + expect(acc.isHealthy).toBe(true) + // expiresAt zeroed so refreshIfNeeded triggers on next iteration + expect(acc.expiresAt).toBe(0) + }) + + test('TEMPORARILY_SUSPENDED marks account unhealthy', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { reason: 'TEMPORARILY_SUSPENDED', message: 'Suspended' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(false) + expect(acc.isHealthy).toBe(false) + }) + + test('non-permanent 403 retries with backoff', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { message: 'Forbidden' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.retry).toBe(1) + }) + + test('INVALID_MODEL_ID throws immediately', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { reason: 'INVALID_MODEL_ID', message: 'bad model' }) + await expect(handler.handle(null, res, acc, { retry: 0 }, noToast)).rejects.toThrow( + 'Invalid model: bad model' + ) + }) +}) + +// ── 403 multi account ───────────────────────────────────────────────────────── + +describe('ErrorHandler: 403 multi account', () => { + test('switches account on any 403, increments failCount', async () => { + const a = makeAccount({ id: 'a' }) + const b = makeAccount({ id: 'b', email: 'b@example.com' }) + const mgr = new AccountManager([a, b]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([a, b])) + const res = makeResponse(403, { message: 'Forbidden' }) + const result = await handler.handle(null, res, a, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.switchAccount).toBe(true) + // non-permanent 403: failCount incremented but account still available + expect(a.failCount).toBe(1) + expect(a.isHealthy).toBe(true) + }) +}) + +// ── 429 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 429', () => { + test('marks account rate-limited', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = new Response('', { + status: 429, + headers: { 'retry-after': '1' } // 1s not 30s + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(acc.rateLimitResetTime).toBeGreaterThan(Date.now() - 100) + }) + + test('with single account, sleep is excluded from request timeout budget', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = new Response('', { + status: 429, + headers: { 'retry-after': '1' } // 1s sleep + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.excludedMs).toBeGreaterThanOrEqual(1000) + }) + + test('with multiple accounts, switches without sleeping', async () => { + const acc1 = makeAccount({ id: 'a' }) + const acc2 = makeAccount({ id: 'b' }) + const mgr = new AccountManager([acc1, acc2]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc1, acc2])) + const res = new Response('', { status: 429, headers: { 'retry-after': '60' } }) + const start = Date.now() + const result = await handler.handle(null, res, acc1, { retry: 0 }, noToast) + const elapsed = Date.now() - start + expect(result.switchAccount).toBe(true) + expect(elapsed).toBeLessThan(500) // didn't sleep + }) +}) + +// ── 500 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 500', () => { + test('retries with backoff on first failure', async () => { + const acc = makeAccount({ failCount: 0 }) + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(500, { message: 'Internal Server Error' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(acc.failCount).toBe(1) + }) + + test('marks unhealthy after 5 failures', async () => { + const acc = makeAccount({ failCount: 4 }) + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(500, { message: 'Internal Server Error' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.switchAccount).toBe(true) + expect(acc.isHealthy).toBe(false) + }) +}) diff --git a/src/__tests__/event-stream-parser.test.ts b/src/__tests__/event-stream-parser.test.ts new file mode 100644 index 0000000..eabc5ae --- /dev/null +++ b/src/__tests__/event-stream-parser.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from 'bun:test' +import { + parseAwsEventStreamBuffer, + parseEventLine +} from '../infrastructure/transformers/event-stream-parser.js' + +// ── parseEventLine ──────────────────────────────────────────────────────────── + +describe('parseEventLine', () => { + test('parses valid JSON', () => { + expect(parseEventLine('{"content":"hello"}')).toEqual({ content: 'hello' }) + }) + + test('returns null on invalid JSON', () => { + expect(parseEventLine('not json')).toBeNull() + expect(parseEventLine('{')).toBeNull() + }) +}) + +// ── parseAwsEventStreamBuffer ───────────────────────────────────────────────── + +describe('parseAwsEventStreamBuffer', () => { + test('empty buffer returns empty array', () => { + expect(parseAwsEventStreamBuffer('')).toEqual([]) + }) + + test('parses content event', () => { + const result = parseAwsEventStreamBuffer('{"content":"Hello"}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('content') + expect(result[0]!.data).toBe('Hello') + }) + + test('skips followupPrompt as not content', () => { + const result = parseAwsEventStreamBuffer('{"content":"x","followupPrompt":"y"}') + expect(result).toHaveLength(0) + }) + + test('parses toolUse event', () => { + const result = parseAwsEventStreamBuffer( + '{"name":"bash","toolUseId":"t-1","input":"ls","stop":false}' + ) + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUse') + expect(result[0]!.data.name).toBe('bash') + expect(result[0]!.data.toolUseId).toBe('t-1') + }) + + test('parses toolUseInput event', () => { + const result = parseAwsEventStreamBuffer('{"input":"-la"}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUseInput') + expect(result[0]!.data.input).toBe('-la') + }) + + test('parses toolUseStop event', () => { + const result = parseAwsEventStreamBuffer('{"stop":true}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUseStop') + expect(result[0]!.data.stop).toBe(true) + }) + + test('parses contextUsage event', () => { + const result = parseAwsEventStreamBuffer('{"contextUsagePercentage":42}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('contextUsage') + expect(result[0]!.data.contextUsagePercentage).toBe(42) + }) + + test('parses multiple events from a single buffer', () => { + const buffer = '{"content":"Hello"}\n{"content":" world"}\n{"contextUsagePercentage":25}' + const result = parseAwsEventStreamBuffer(buffer) + expect(result).toHaveLength(3) + expect(result[0]!.data).toBe('Hello') + expect(result[1]!.data).toBe(' world') + expect(result[2]!.type).toBe('contextUsage') + }) + + test('handles incomplete JSON at end (no jsonEnd)', () => { + const result = parseAwsEventStreamBuffer('{"content":"truncated') + expect(result).toHaveLength(0) + }) + + test('handles escaped strings inside JSON values', () => { + const result = parseAwsEventStreamBuffer('{"content":"say \\"hello\\""}') + expect(result).toHaveLength(1) + expect(result[0]!.data).toBe('say "hello"') + }) +}) diff --git a/src/__tests__/image-cache.test.ts b/src/__tests__/image-cache.test.ts new file mode 100644 index 0000000..ebc2012 --- /dev/null +++ b/src/__tests__/image-cache.test.ts @@ -0,0 +1,356 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { existsSync, mkdtempSync, readdirSync, rmSync, utimesSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { ImageCache } from '../plugin/image-cache.js' +import { + convertImagesToKiroFormat, + extractAllImages, + type KiroImage +} from '../plugin/image-handler.js' + +function img(byteLen = 100, format = 'png', seed = 0): KiroImage { + // Use the seed to vary head bytes so dedupKey distinguishes distinct images + // of the same size/format. Default seed=0 reproduces "identical" images. + const bytes = new Uint8Array(byteLen) + if (seed !== 0) { + for (let i = 0; i < Math.min(16, byteLen); i++) bytes[i] = (seed + i) & 0xff + } + return { format, source: { bytes } } +} + +describe('ImageCache', () => { + test('round-trips images by workspace+fingerprint', () => { + const c = new ImageCache() + const a = img(50) + const b = img(75) + c.set('/ws', 'fp1', [a, b]) + + const got = c.get('/ws', 'fp1') + expect(got).not.toBeNull() + expect(got).toHaveLength(2) + expect(got![0]!.source.bytes.byteLength).toBe(50) + }) + + test('returns null when conversation is unknown', () => { + const c = new ImageCache() + expect(c.get('/ws', 'fp-missing')).toBeNull() + }) + + test('does nothing on empty image array', () => { + const c = new ImageCache() + c.set('/ws', 'fp1', []) + expect(c.size()).toBe(0) + }) + + test('separate fingerprints are isolated', () => { + const c = new ImageCache() + c.set('/ws', 'fp1', [img(10)]) + c.set('/ws', 'fp2', [img(20)]) + expect(c.get('/ws', 'fp1')![0]!.source.bytes.byteLength).toBe(10) + expect(c.get('/ws', 'fp2')![0]!.source.bytes.byteLength).toBe(20) + }) + + test('delete removes a single entry', () => { + const c = new ImageCache() + c.set('/ws', 'fp1', [img()]) + c.set('/ws', 'fp2', [img()]) + c.delete('/ws', 'fp1') + expect(c.get('/ws', 'fp1')).toBeNull() + expect(c.get('/ws', 'fp2')).not.toBeNull() + }) + + test('TTL prunes stale entries on lookup', () => { + let t = 1000 + const c = new ImageCache({ ttlMs: 500, now: () => t }) + c.set('/ws', 'fp1', [img()]) + t = 1300 + expect(c.get('/ws', 'fp1')).not.toBeNull() // within TTL + t = 2000 + expect(c.get('/ws', 'fp1')).toBeNull() // expired + expect(c.size()).toBe(0) + }) + + test('LRU evicts the oldest when over capacity', () => { + let t = 0 + const c = new ImageCache({ maxEntries: 3, now: () => ++t }) + c.set('/ws', 'a', [img()]) + c.set('/ws', 'b', [img()]) + c.set('/ws', 'c', [img()]) + c.set('/ws', 'd', [img()]) // pushes 'a' out + expect(c.size()).toBe(3) + expect(c.get('/ws', 'a')).toBeNull() + expect(c.get('/ws', 'b')).not.toBeNull() + expect(c.get('/ws', 'd')).not.toBeNull() + }) + + test('get() refreshes LRU position', () => { + let t = 0 + const c = new ImageCache({ maxEntries: 3, now: () => ++t }) + c.set('/ws', 'a', [img()]) + c.set('/ws', 'b', [img()]) + c.set('/ws', 'c', [img()]) + c.get('/ws', 'a') // 'a' is now most-recently-used + c.set('/ws', 'd', [img()]) // should evict 'b', not 'a' + expect(c.get('/ws', 'a')).not.toBeNull() + expect(c.get('/ws', 'b')).toBeNull() + }) + + test('clear empties everything', () => { + const c = new ImageCache() + c.set('/ws', 'a', [img()]) + c.set('/ws', 'b', [img()]) + c.clear() + expect(c.size()).toBe(0) + }) +}) + +describe('ImageCache.upsert (merge + dedup + caps)', () => { + test('merges new images with existing — new ones first', () => { + const c = new ImageCache() + c.set('/ws', 'fp', [img(100, 'png', 1), img(100, 'png', 2)]) + + const newImg = img(100, 'png', 3) + const count = c.upsert('/ws', 'fp', [newImg]) + + expect(count).toBe(3) + const got = c.get('/ws', 'fp')! + // Newest first + expect(got[0]!.source.bytes[0]).toBe((3 + 0) & 0xff) + expect(got).toHaveLength(3) + }) + + test('deduplicates identical content (same format + size + head/tail bytes)', () => { + const c = new ImageCache() + const a = img(100, 'png', 7) + c.set('/ws', 'fp', [a]) + + // Same content fingerprint as a — should not be added twice. + const aDup = img(100, 'png', 7) + const count = c.upsert('/ws', 'fp', [aDup]) + + expect(count).toBe(1) + expect(c.get('/ws', 'fp')).toHaveLength(1) + }) + + test('caps merged result at MAX_KIRO_IMAGES, dropping oldest first', () => { + const c = new ImageCache() + // Existing: 3 images + c.set('/ws', 'fp', [img(100, 'png', 1), img(100, 'png', 2), img(100, 'png', 3)]) + // Upsert 2 new → would be 5, MAX is 4, so the oldest (seed=3) drops. + const count = c.upsert('/ws', 'fp', [img(100, 'png', 9), img(100, 'png', 8)]) + + expect(count).toBe(4) + const got = c.get('/ws', 'fp')! + // Newest first + expect(got[0]!.source.bytes[0]).toBe((9 + 0) & 0xff) + expect(got[1]!.source.bytes[0]).toBe((8 + 0) & 0xff) + // Seed=3 should have been dropped from the tail. + const seedsPresent = got.map((g) => g.source.bytes[0]) + expect(seedsPresent).not.toContain((3 + 0) & 0xff) + }) + + test('caps total bytes at MAX_KIRO_IMAGE_BYTES (3.75MB)', () => { + const c = new ImageCache() + // Two 2MB images already cached + const TWO_MB = 2 * 1024 * 1024 + c.set('/ws', 'fp', [img(TWO_MB, 'png', 1), img(TWO_MB, 'png', 2)]) + + // Add another 2MB. Total would be 6MB; budget is 3.75MB. Oldest must drop. + c.upsert('/ws', 'fp', [img(TWO_MB, 'png', 3)]) + + const got = c.get('/ws', 'fp')! + const total = got.reduce((n, g) => n + g.source.bytes.byteLength, 0) + expect(total).toBeLessThanOrEqual(3_750_000) + // The newest image must still be there (seed=3 is most recent) + expect(got[0]!.source.bytes[0]).toBe((3 + 0) & 0xff) + }) + + test('upsert with empty input + no existing returns 0', () => { + const c = new ImageCache() + expect(c.upsert('/ws', 'fp', [])).toBe(0) + expect(c.size()).toBe(0) + }) + + test('upsert with empty input keeps existing intact', () => { + const c = new ImageCache() + c.set('/ws', 'fp', [img(100, 'png', 5)]) + // Empty newImages should not wipe the existing entry; returns the unchanged count. + expect(c.upsert('/ws', 'fp', [])).toBe(1) + expect(c.get('/ws', 'fp')).toHaveLength(1) + }) +}) + +describe('ImageCache.hasEverHadImages', () => { + test('returns false when no entry exists', () => { + const c = new ImageCache() + expect(c.hasEverHadImages('/ws', 'fp')).toBe(false) + }) + + test('returns true after set', () => { + const c = new ImageCache() + c.set('/ws', 'fp', [img(100)]) + expect(c.hasEverHadImages('/ws', 'fp')).toBe(true) + }) + + test('returns true after upsert with images', () => { + const c = new ImageCache() + c.upsert('/ws', 'fp', [img(100)]) + expect(c.hasEverHadImages('/ws', 'fp')).toBe(true) + }) + + test('returns false after delete', () => { + const c = new ImageCache() + c.set('/ws', 'fp', [img(100)]) + c.delete('/ws', 'fp') + expect(c.hasEverHadImages('/ws', 'fp')).toBe(false) + }) + + test('checks disk when in-memory cache misses', () => { + const dir = mkdtempSync(join(tmpdir(), 'kiro-img-cache-')) + try { + const writer = new ImageCache({ cacheDir: dir }) + writer.set('/ws', 'fp', [img(100)]) + const reader = new ImageCache({ cacheDir: dir }) + // No in-memory entry yet, but hasEverHadImages should still find it on disk + expect(reader.hasEverHadImages('/ws', 'fp')).toBe(true) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + test('hasEverHadImages avoids disk hit when cacheDir is null', () => { + const c = new ImageCache() + // Should be O(1), no disk access attempted + const start = performance.now() + for (let i = 0; i < 1000; i++) c.hasEverHadImages('/ws', 'fp-' + i) + expect(performance.now() - start).toBeLessThan(10) + }) +}) + +describe('ImageCache filesystem persistence', () => { + let dir: string + + beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), 'kiro-img-cache-')) + }) + + afterEach(() => { + rmSync(dir, { recursive: true, force: true }) + }) + + test('persists across instances', () => { + const writer = new ImageCache({ cacheDir: dir }) + writer.set('/ws', 'fp', [img(100, 'png', 5)]) + + // Fresh instance pointing at the same dir simulates an OpenCode restart. + const reader = new ImageCache({ cacheDir: dir }) + expect(reader.size()).toBe(0) // memory map starts empty + const got = reader.get('/ws', 'fp') + expect(got).not.toBeNull() + expect(got).toHaveLength(1) + expect(got![0]!.source.bytes[0]).toBe(5) + expect(reader.size()).toBe(1) // lazy-loaded into memory after first get + }) + + test('delete removes the on-disk file too', () => { + const c = new ImageCache({ cacheDir: dir }) + c.set('/ws', 'fp', [img()]) + expect(readdirSync(dir).filter((n) => n.endsWith('.json'))).toHaveLength(1) + c.delete('/ws', 'fp') + expect(readdirSync(dir).filter((n) => n.endsWith('.json'))).toHaveLength(0) + }) + + test('upsert merges with on-disk state from a previous instance', () => { + const first = new ImageCache({ cacheDir: dir }) + first.set('/ws', 'fp', [img(100, 'png', 1), img(100, 'png', 2)]) + + // Restart: new instance, in-memory is empty but disk has the entry. + const second = new ImageCache({ cacheDir: dir }) + const count = second.upsert('/ws', 'fp', [img(100, 'png', 3)]) + + expect(count).toBe(3) + const got = second.get('/ws', 'fp')! + // Newest (seed=3) is first; existing seeds 1 and 2 are still present. + expect(got[0]!.source.bytes[0]).toBe(3) + expect(got).toHaveLength(3) + }) + + test('constructor sweeps files older than the TTL', () => { + const c1 = new ImageCache({ cacheDir: dir, ttlMs: 1000 }) + c1.set('/ws', 'fp', [img()]) + const files = readdirSync(dir).filter((n) => n.endsWith('.json')) + expect(files).toHaveLength(1) + + // Backdate the file so it's far older than the TTL. + const ancient = new Date(Date.now() - 24 * 60 * 60 * 1000) + utimesSync(join(dir, files[0]!), ancient, ancient) + + // New instance with the same TTL should sweep the stale file at init. + new ImageCache({ cacheDir: dir, ttlMs: 1000 }) + expect(readdirSync(dir).filter((n) => n.endsWith('.json'))).toHaveLength(0) + }) + + test('lazy-load drops a file that has expired since it was written', () => { + const c1 = new ImageCache({ cacheDir: dir, ttlMs: 1000 }) + c1.set('/ws', 'fp', [img()]) + const file = readdirSync(dir).filter((n) => n.endsWith('.json'))[0]! + const ancient = new Date(Date.now() - 24 * 60 * 60 * 1000) + utimesSync(join(dir, file), ancient, ancient) + + // Big TTL on the reader so the constructor sweep doesn't wipe it; but the + // lazy-load path uses the reader's TTL too — make it small so it expires. + const reader = new ImageCache({ cacheDir: dir, ttlMs: 1000 }) + expect(reader.get('/ws', 'fp')).toBeNull() + // File should be cleaned up. + expect(existsSync(join(dir, file))).toBe(false) + }) + + test('disk persistence is OFF when no cacheDir is given', () => { + const c = new ImageCache() + c.set('/ws', 'fp', [img()]) + // Nothing should land in our tmpdir because the cache wasn't told about it. + expect(readdirSync(dir)).toHaveLength(0) + }) +}) + +// ── Image extraction + conversion performance ──────────────────────────────── + +describe('image extraction performance', () => { + test('extractAllImages + convertImagesToKiroFormat: 4 large images under 50ms', () => { + // Simulates an agentic turn carrying the 3.75MB Kiro payload limit worth of + // PNG screenshots. The base64 char-loop was the bottleneck — verify Buffer + // decoding keeps us well under the budget. + const big = 'A'.repeat(900_000) // ~675KB raw -> ~900KB base64 + const content = [ + { type: 'image', source: { type: 'base64', media_type: 'image/png', data: big } }, + { type: 'image', source: { type: 'base64', media_type: 'image/png', data: big } }, + { type: 'image', source: { type: 'base64', media_type: 'image/png', data: big } }, + { type: 'image', source: { type: 'base64', media_type: 'image/png', data: big } } + ] + const start = performance.now() + const unified = extractAllImages(content) + const { images, omitted } = convertImagesToKiroFormat(unified) + const elapsed = performance.now() - start + expect(unified.length).toBe(4) + expect(images.length).toBe(4) + expect(omitted).toBe(0) + expect(elapsed).toBeLessThan(50) + }) + + test('extractAllImages handles OpenAI data URLs in one pass', () => { + const content = [ + { type: 'text', text: 'hello' }, + { + type: 'image_url', + image_url: { + url: 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=' + } + }, + { type: 'image_url', image_url: { url: 'https://example.com/img.png' } } // non-data URL: skipped + ] + const result = extractAllImages(content) + expect(result).toHaveLength(1) + expect(result[0]!.mediaType).toBe('image/png') + }) +}) diff --git a/src/__tests__/response.test.ts b/src/__tests__/response.test.ts new file mode 100644 index 0000000..baca78f --- /dev/null +++ b/src/__tests__/response.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, test } from 'bun:test' +import { estimateTokens, parseEventStream } from '../plugin/response.js' + +// ── estimateTokens ──────────────────────────────────────────────────────────── + +describe('estimateTokens', () => { + test('empty string returns 0', () => { + expect(estimateTokens('')).toBe(0) + }) + + test('4-char string returns 1 token', () => { + expect(estimateTokens('abcd')).toBe(1) + }) + + test('5-char string returns 2 tokens (ceil)', () => { + expect(estimateTokens('abcde')).toBe(2) + }) + + test('100-char string returns 25 tokens', () => { + expect(estimateTokens('a'.repeat(100))).toBe(25) + }) +}) + +// ── parseEventStream ────────────────────────────────────────────────────────── + +describe('parseEventStream', () => { + test('parses plain text content', () => { + const result = parseEventStream('{"content":"Hello world"}') + expect(result.content).toBe('Hello world') + expect(result.toolCalls).toHaveLength(0) + expect(result.stopReason).toBe('end_turn') + }) + + test('parses multiple content chunks', () => { + const result = parseEventStream('{"content":"Hello "}{"content":"world"}') + expect(result.content).toBe('Hello world') + }) + + test('parses tool use', () => { + const raw = [ + '{"name":"bash","toolUseId":"t-1","input":"","stop":false}', + '{"input":"{\\\"command\\\":\\\"ls\\\"}"}', + '{"stop":true}' + ].join('\n') + const result = parseEventStream(raw) + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0]!.name).toBe('bash') + expect(result.stopReason).toBe('tool_use') + }) + + test('parses context usage into token counts', () => { + const raw = '{"content":"hi"}{"contextUsagePercentage":50}' + const result = parseEventStream(raw, 'CLAUDE_SONNET_4_5') + // With a known context window, inputTokens should be set + expect(result.outputTokens).toBeDefined() + expect(result.inputTokens).toBeDefined() + }) + + test('returns end_turn when no tool calls', () => { + const result = parseEventStream('{"content":"answer"}') + expect(result.stopReason).toBe('end_turn') + }) + + test('empty input returns empty content', () => { + const result = parseEventStream('') + expect(result.content).toBe('') + expect(result.toolCalls).toHaveLength(0) + }) +}) diff --git a/src/__tests__/sqlite.test.ts b/src/__tests__/sqlite.test.ts new file mode 100644 index 0000000..ec06644 --- /dev/null +++ b/src/__tests__/sqlite.test.ts @@ -0,0 +1,246 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { KiroDatabase } from '../plugin/storage/sqlite.js' +import type { ManagedAccount } from '../plugin/types.js' + +function makeAccount(overrides: Partial = {}): ManagedAccount { + return { + id: 'acc-1', + email: 'test@example.com', + authMethod: 'idc', + region: 'eu-central-1', + refreshToken: 'r', + accessToken: 'a', + expiresAt: Date.now() + 3600000, + rateLimitResetTime: 0, + isHealthy: true, + failCount: 0, + lastUsed: 0, + usedCount: 0, + limitCount: 0, + ...overrides + } +} + +let dir: string +let dbPath: string +let db: KiroDatabase + +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), 'kiro-test-')) + dbPath = join(dir, 'test.db') + db = new KiroDatabase(dbPath) +}) + +afterEach(() => { + db.close() + rmSync(dir, { recursive: true, force: true }) +}) + +// ── accounts CRUD ───────────────────────────────────────────────────────────── + +describe('KiroDatabase: accounts', () => { + test('starts empty', () => { + expect(db.getAccounts()).toHaveLength(0) + }) + + test('upsertAccount stores and retrieves account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].email).toBe('test@example.com') + expect(rows[0].is_healthy).toBe(1) + }) + + test('upsertAccount updates token fields on existing account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.upsertAccount({ ...acc, accessToken: 'new-token' }) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].access_token).toBe('new-token') + }) + + test('upsertAccount with permanent error sets isHealthy=false', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.upsertAccount({ + ...acc, + isHealthy: false, + unhealthyReason: 'ExpiredTokenException' + }) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].is_healthy).toBe(0) + }) + + test('batchUpsertAccounts stores multiple accounts', async () => { + const a = makeAccount({ id: 'a', email: 'a@example.com' }) + const b = makeAccount({ id: 'b', email: 'b@example.com' }) + await db.batchUpsertAccounts([a, b]) + expect(db.getAccounts()).toHaveLength(2) + }) + + test('deleteAccount removes account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.deleteAccount('acc-1') + expect(db.getAccounts()).toHaveLength(0) + }) + + test('deleteAccount on non-existent id is a no-op', async () => { + await expect(db.deleteAccount('does-not-exist')).resolves.toBeUndefined() + }) +}) + +// ── reauth lock ─────────────────────────────────────────────────────────────── + +describe('KiroDatabase: reauth lock', () => { + test('acquireReauthLock returns true when no lock held', () => { + expect(db.acquireReauthLock()).toBe(true) + }) + + test('acquireReauthLock returns false when lock already held by this process', () => { + db.acquireReauthLock() + // Same process — process.kill(pid, 0) succeeds, so it's not dead + expect(db.acquireReauthLock()).toBe(false) + }) + + test('isReauthLockHeld returns false when no lock', () => { + expect(db.isReauthLockHeld()).toBe(false) + }) + + test('isReauthLockHeld returns true after acquire', () => { + db.acquireReauthLock() + expect(db.isReauthLockHeld()).toBe(true) + }) + + test('releaseReauthLock clears the lock', () => { + db.acquireReauthLock() + db.releaseReauthLock() + expect(db.isReauthLockHeld()).toBe(false) + }) + + test('after release, lock can be acquired again', () => { + db.acquireReauthLock() + db.releaseReauthLock() + expect(db.acquireReauthLock()).toBe(true) + }) + + test('stale lock (dead pid) is evicted on acquire', () => { + // Insert a lock row with a pid that no process uses (high number) + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb.prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, 9999999, ?)').run( + Date.now() - 1000 // recent but dead pid + ) + rawDb.close() + // Re-open our db instance + db.close() + db = new KiroDatabase(dbPath) + // Should evict dead-pid lock and acquire + expect(db.acquireReauthLock()).toBe(true) + }) + + test('expired lock is evicted on acquire', () => { + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb.prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, ?, ?)').run( + process.pid, + Date.now() - 200_000 // 200s ago, well past 120s TTL + ) + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + expect(db.acquireReauthLock()).toBe(true) + }) + + test('race-safe: row-replacement when prior dead-pid row exists', () => { + // Simulate a row that the SELECT picked up as "expired" but is actually + // the same one INSERT will try to write. INSERT OR REPLACE handles this. + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb + .prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, 9999998, ?)') + .run(Date.now() - 200_000) + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + // Should not throw on PRIMARY KEY conflict + expect(db.acquireReauthLock()).toBe(true) + // The row now belongs to this process + expect(db.isReauthLockHeld()).toBe(true) + }) +}) + +// ── conversations ───────────────────────────────────────────────────────────── + +describe('KiroDatabase: conversations', () => { + test('getConversationId returns undefined when not set', () => { + expect(db.getConversationId('ws', 'fp')).toBeUndefined() + }) + + test('setConversationId and getConversationId round-trip', () => { + db.setConversationId('ws', 'fp', 'conv-123', 'agent-123') + expect(db.getConversationId('ws', 'fp')).toEqual({ + convId: 'conv-123', + agentContinuationId: 'agent-123' + }) + }) + + test('setConversationId updates existing entry', () => { + db.setConversationId('ws', 'fp', 'conv-1', 'agent-1') + db.setConversationId('ws', 'fp', 'conv-2', 'agent-2') + expect(db.getConversationId('ws', 'fp')).toEqual({ + convId: 'conv-2', + agentContinuationId: 'agent-2' + }) + }) + + test('different fingerprints are independent', () => { + db.setConversationId('ws', 'fp1', 'conv-A', 'agent-A') + db.setConversationId('ws', 'fp2', 'conv-B', 'agent-B') + expect(db.getConversationId('ws', 'fp1')).toEqual({ + convId: 'conv-A', + agentContinuationId: 'agent-A' + }) + expect(db.getConversationId('ws', 'fp2')).toEqual({ + convId: 'conv-B', + agentContinuationId: 'agent-B' + }) + }) + + test('TTL cleanup removes old entries', () => { + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb + .prepare( + 'INSERT INTO conversations (workspace, fingerprint, conv_id, agent_continuation_id, last_used) VALUES (?, ?, ?, ?, ?)' + ) + .run('ws', 'old', 'conv-old', 'agent-old', Date.now() - 10 * 24 * 3600000) // 10 days old + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + // Trigger cleanup by setting a new one (ttlDays=7) + db.setConversationId('ws', 'new', 'conv-new', 'agent-new', 7) + expect(db.getConversationId('ws', 'old')).toBeUndefined() + expect(db.getConversationId('ws', 'new')).toEqual({ + convId: 'conv-new', + agentContinuationId: 'agent-new' + }) + }) + + test('deleteConversationId removes entry so next lookup returns undefined', () => { + db.setConversationId('ws', 'fp', 'conv-del', 'agent-del') + expect(db.getConversationId('ws', 'fp')).toBeDefined() + db.deleteConversationId('ws', 'fp') + expect(db.getConversationId('ws', 'fp')).toBeUndefined() + }) + + test('deleteConversationId is a no-op for non-existent entry', () => { + expect(() => db.deleteConversationId('ws', 'missing')).not.toThrow() + }) +}) diff --git a/src/__tests__/stream-transformer.test.ts b/src/__tests__/stream-transformer.test.ts new file mode 100644 index 0000000..0a629ad --- /dev/null +++ b/src/__tests__/stream-transformer.test.ts @@ -0,0 +1,672 @@ +import { describe, expect, test } from 'bun:test' +import { transformKiroStream } from '../plugin/streaming/stream-transformer.js' + +// Helper: create a Response with a ReadableStream from text chunks +function makeStreamResponse(chunks: string[]): Response { + const stream = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(new TextEncoder().encode(chunk)) + } + controller.close() + } + }) + return new Response(stream) +} + +async function collect(gen: AsyncGenerator): Promise { + const result: any[] = [] + for await (const item of gen) result.push(item) + return result +} + +const MODEL = 'CLAUDE_SONNET_4_5' +const CONV = 'conv-test' + +// ── basic text content ──────────────────────────────────────────────────────── + +describe('transformKiroStream: text', () => { + test('plain text content produces text delta events', async () => { + const response = makeStreamResponse(['{"content":"Hello world"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Should contain at least a content_block_delta and message_stop + const deltas = events.filter( + (e) => e.choices?.[0]?.delta?.content !== undefined || e.choices?.[0]?.delta?.type === 'text' + ) + expect(deltas.length).toBeGreaterThan(0) + const stop = events.find((e) => e.object === 'chat.completion.chunk') + expect(stop).toBeDefined() + }) + + test('multiple content chunks are concatenated', async () => { + const response = makeStreamResponse(['{"content":"Hello "}', '{"content":"world"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const allText = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('Hello') + expect(allText).toContain('world') + }) + + test('empty body throws', async () => { + const response = new Response(null) + await expect(collect(transformKiroStream(response, MODEL, CONV))).rejects.toThrow( + 'Response body is null' + ) + }) +}) + +// ── thinking tags ───────────────────────────────────────────────────────────── + +describe('transformKiroStream: thinking tags', () => { + test('thinking content produces thinking delta events', async () => { + const response = makeStreamResponse(['{"content":"I think...Answer"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Thinking events have type 'thinking' in the delta + const thinkingEvents = events.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + }) + + test('text after thinking tags is emitted as normal text', async () => { + const response = makeStreamResponse(['{"content":"ThinkFinal answer"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const textContent = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(textContent).toContain('Final answer') + }) + + test('thinking tag inside code block is not treated as thinking', async () => { + const response = makeStreamResponse([ + '{"content":"```\\nnot thinking\\n```"}' + ]) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const thinkingEvents = events.filter((e) => e.choices?.[0]?.delta?.thinking) + expect(thinkingEvents.length).toBe(0) + }) +}) + +// ── tool use ────────────────────────────────────────────────────────────────── + +describe('transformKiroStream: tool use', () => { + test('toolUse events produce tool_call chunks', async () => { + const chunks = [ + '{"name":"bash","toolUseId":"t-1","input":"","stop":false}', + '{"input":"{\\\"command\\\":\\\"ls\\\"}"}', + '{"stop":true}' + ] + const response = makeStreamResponse(chunks) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Should have tool_calls in at least one event + const toolEvents = events.filter( + (e) => e.choices?.[0]?.delta?.tool_calls || e.choices?.[0]?.delta?.type === 'tool_use' + ) + expect(toolEvents.length).toBeGreaterThan(0) + }) +}) + +// ── context usage ───────────────────────────────────────────────────────────── + +describe('transformKiroStream: context usage', () => { + test('contextUsagePercentage is reflected in final usage', async () => { + const response = makeStreamResponse(['{"content":"hi"}', '{"contextUsagePercentage":50}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const delta = events.find((e) => e.choices?.[0]?.delta?.stop_reason || e.usage) + expect(delta).toBeDefined() + }) +}) + +// ── TextDecoder flush ───────────────────────────────────────────────────────── + +describe('transformKiroStream: TextDecoder flush', () => { + test('multi-byte UTF-8 split across chunks is decoded correctly', async () => { + // € = E2 82 AC (3 bytes), split: [E2 82] + [AC ...] + const euro = new TextEncoder().encode('{"content":"€"}') + const part1 = euro.slice(0, euro.length - 5) // cut before '€' completes + const part2 = euro.slice(euro.length - 5) + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(part1) + controller.enqueue(part2) + controller.close() + } + }) + const response = new Response(stream) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const allText = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('€') + }) +}) + +// ── SDK stream: tool call chunking ─────────────────────────────────────────── + +import { transformSdkStream } from '../plugin/streaming/sdk-stream-transformer.js' + +function makeSdkResponse(events: any[]) { + return { + generateAssistantResponseResponse: (async function* () { + for (const e of events) yield e + })() + } +} + +async function collectSdk(gen: AsyncGenerator): Promise { + const result: any[] = [] + for await (const item of gen) result.push(item) + return result +} + +describe('transformSdkStream: tool call streaming', () => { + test('tool call input chunks without name are concatenated', async () => { + const events = [ + { toolUseEvent: { toolUseId: 't-1', name: 'write', input: '{"file' } }, + { toolUseEvent: { toolUseId: 't-1', input: 'Path":"/tmp/x.txt","content":"hello"}' } }, + { toolUseEvent: { toolUseId: 't-1', stop: true } } + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + + // With inline streaming, args arrive as multiple partial_json deltas. + // Collect all argument chunks for this tool and concatenate. + const argChunks = result + .filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.arguments !== undefined) + .map((e: any) => e.choices[0].delta.tool_calls[0].function.arguments) + .join('') + const args = JSON.parse(argChunks) + expect(args.filePath).toBe('/tmp/x.txt') + expect(args.content).toBe('hello') + }) + + test('multiple tool calls in one response are parsed correctly', async () => { + const events = [ + { assistantResponseEvent: { content: 'I will run two tools.' } }, + { toolUseEvent: { toolUseId: 't-1', name: 'read', input: '{"path":"/a"}', stop: true } }, + { toolUseEvent: { toolUseId: 't-2', name: 'read', input: '{"path":"/b"}', stop: true } } + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBe(2) + }) + + test('stop event without toolUseId does not falsely trigger truncation warning', async () => { + // Regression: SDK sometimes sends { toolUseEvent: { stop: true } } without + // repeating the toolUseId. The old code had stop handling inside the + // `if (tc.toolUseId)` guard, so it was skipped — leaving currentToolCall + // set after the loop, which triggered the false truncation warning. + const events = [ + { toolUseEvent: { toolUseId: 't-1', name: 'bash', input: '{"command":"ls"}' } }, + { toolUseEvent: { stop: true } } // no toolUseId on stop + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + + // Must NOT emit any truncation warning text + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).not.toContain('truncated') + + // Tool call must still be emitted correctly + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBeGreaterThan(0) + }) + + test('truncated tool call (no stop event) closes the block and emits error text', async () => { + // Simulates Kiro cutting the stream mid-tool-call (large .frm JSON that + // exceeds the response size limit). + // With inline streaming: content_block_start IS emitted (tool name visible), + // then on truncation the block is closed and an error text is appended. + const events = [ + { assistantResponseEvent: { content: 'I will edit the file.' } }, + { + toolUseEvent: { + toolUseId: 't-trunc', + name: 'replaceFileContent', + input: '{"path":"/app.frm","content":{"forms":[{"name":"form1"' + } + } + // note: no stop event — stream cut off mid-JSON + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + + // content_block_start IS emitted (inline streaming — tool name is immediately visible) + const toolNameBlocks = result.filter( + (e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name + ) + expect(toolNameBlocks.length).toBeGreaterThan(0) + expect(toolNameBlocks[0]!.choices[0].delta.tool_calls[0].function.name).toBe( + 'replaceFileContent' + ) + + // Must emit a text delta explaining the truncation + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('truncated') + expect(allText).toContain('replaceFileContent') + }) + + test('tool input is NOT accumulated into totalContent — bracket parser never sees large JSON', async () => { + // Regression: previously `totalContent += tc.input` caused catastrophic + // regex backtracking when the tool input was hundreds of KB of JSON. + const largeInput = JSON.stringify( + Array.from({ length: 100 }, (_, i) => ({ id: i, value: 'x'.repeat(500) })) + ) + const events = [ + { + toolUseEvent: { + toolUseId: 't-big', + name: 'write', + input: largeInput.slice(0, largeInput.length / 2) + } + }, + { toolUseEvent: { toolUseId: 't-big', input: largeInput.slice(largeInput.length / 2) } }, + { toolUseEvent: { toolUseId: 't-big', stop: true } } + ] + const sdk = makeSdkResponse(events) + const start = performance.now() + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + const elapsed = performance.now() - start + + // Must complete in well under 1s — would hang for seconds with catastrophic backtracking + expect(elapsed).toBeLessThan(500) + + // Tool call must still be emitted correctly + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBeGreaterThan(0) + }) +}) + +// ── SDK stream: thinking tags ──────────────────────────────────────────────── + +describe('transformSdkStream: thinking tags', () => { + test('thinking content is extracted and emitted as reasoning_content', async () => { + const events = [ + { assistantResponseEvent: { content: 'deep thoughtAnswer' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const thinkingEvents = result.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('Answer') + }) + + test('text before thinking tag is emitted as normal text', async () => { + const events = [ + { assistantResponseEvent: { content: 'beforethoughtafter' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('before') + expect(allText).toContain('after') + }) + + test('partial thinking start tag at chunk boundary is buffered safely', async () => { + // Send text that ends mid-way through '' so safe-length flush is triggered + const events = [ + { assistantResponseEvent: { content: 'hello thoughtdone' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('hello') + expect(allText).toContain('done') + const thinkingEvents = result.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + }) + + test('thinking content split across chunks is assembled correctly', async () => { + const events = [ + { assistantResponseEvent: { content: 'part one ' } }, + { assistantResponseEvent: { content: 'part twoFinal' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const thinkingEvents = result.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + const allThinking = thinkingEvents + .map((e: any) => e.choices[0].delta.reasoning_content) + .join('') + expect(allThinking).toContain('part one') + expect(allThinking).toContain('part two') + }) + + test('thinking tag followed by \\n\\n strips the newlines', async () => { + const events = [ + { assistantResponseEvent: { content: 'thought\n\nAnswer' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('Answer') + expect(allText).not.toMatch(/^\n\n/) + }) + + test('unclosed thinking tag at stream end flushes remaining buffer', async () => { + const events = [{ assistantResponseEvent: { content: 'incomplete thought' } }] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const thinkingEvents = result.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + }) + + test('leftover buffer after thinking extracted is flushed as text', async () => { + // thinkingExtracted=true path: text arrives after in a later chunk + const events = [ + { assistantResponseEvent: { content: 'think' } }, + { assistantResponseEvent: { content: 'trailing text' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('trailing text') + }) + + test('leftover non-thinking buffer at stream end is flushed', async () => { + // After , thinkingExtracted=true. A subsequent chunk that is + // shorter than THINKING_END_TAG stays in the safe-length buffer. At stream + // end the else-branch (lines 288-292) flushes it. + const events = [ + { assistantResponseEvent: { content: 't' } }, + // 'hi' is 2 chars — much shorter than THINKING_END_TAG (10 chars), + // so safeLen=0 and it stays in buffer until stream end flush. + { assistantResponseEvent: { content: 'hi' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('hi') + }) + + test('thinkingExtracted path: text after in a new chunk is flushed via while loop', async () => { + // After thinkingExtracted=true, a chunk arrives that triggers the + // thinkingExtracted branch inside the while loop (line 118-125). + // Send a chunk that is long enough to pass safeLen > 0 check. + const events = [ + { assistantResponseEvent: { content: 'thought' } }, + { assistantResponseEvent: { content: 'this is a longer answer text' } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', undefined, true) + ) + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('longer answer') + }) + + test('toolNameMapper renames tool calls', async () => { + const events = [ + { toolUseEvent: { toolUseId: 't-1', name: 'original_name', input: '{"x":1}', stop: true } } + ] + const mapper = (name: string) => name.replace('original', 'mapped') + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1', mapper) + ) + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks[0]!.choices[0].delta.tool_calls[0].function.name).toBe('mapped_name') + }) +}) + +// ── SDK stream: contextUsageEvent / meteringEvent ──────────────────────────── + +describe('transformSdkStream: contextUsageEvent and meteringEvent', () => { + test('contextUsageEvent updates token counts', async () => { + const events = [ + { assistantResponseEvent: { content: 'hi' } }, + { contextUsageEvent: { contextUsagePercentage: 40 } } + ] + const result = await collectSdk( + transformSdkStream(makeSdkResponse(events), 'CLAUDE_SONNET_4_5', 'conv-1') + ) + const usageEvent = result.find((e) => e.usage) + expect(usageEvent).toBeDefined() + }) + + test('meteringEvent is silently consumed without error', async () => { + const events = [ + { assistantResponseEvent: { content: 'hi' } }, + { meteringEvent: { usage: 2, unit: 'credit' } } + ] + const result = await collectSdk(transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1')) + expect(result.length).toBeGreaterThan(0) + }) +}) + +// ── SDK stream: bracket tool calls ─────────────────────────────────────────── + +describe('transformSdkStream: bracket tool calls', () => { + test('bracket-format tool call in content is emitted post-stream', async () => { + const events = [ + { assistantResponseEvent: { content: '[Called bash with args: {"command":"ls"}]' } } + ] + const result = await collectSdk(transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1')) + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBeGreaterThan(0) + expect(toolBlocks[0]!.choices[0].delta.tool_calls[0].function.name).toBe('bash') + }) + + test('bracket tool call with invalid JSON falls back to raw string', async () => { + // parseBracketToolCalls skips invalid JSON, so we drive the fallback via + // the post-stream JSON.parse catch in the bracket tool emit path. + // Inject a pre-parsed bracket call with bad JSON directly via totalContent trick: + // use a valid bracket call format but with a nested object that the regex captures + // as valid but JSON.parse fails on — not easily achievable via the regex path. + // Instead test via a tool call whose input arrives as invalid JSON from the SDK. + // The catch path in postStreamCalls (line 351-354) handles bracket calls with bad input. + // We verify the tool is still emitted even with malformed JSON. + const events = [ + { assistantResponseEvent: { content: '[Called bash with args: {"command":"ls"}]' } } + ] + const result = await collectSdk(transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1')) + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBeGreaterThan(0) + }) + + test('stream error with active tool call logs and rethrows', async () => { + async function* failingStream() { + yield { toolUseEvent: { toolUseId: 't-1', name: 'bash', input: '{"command":"ls"}' } } + throw new Error('network failure') + } + const sdk = { generateAssistantResponseResponse: failingStream() } + await expect(collectSdk(transformSdkStream(sdk, 'auto', 'conv-1'))).rejects.toThrow( + 'network failure' + ) + }) + + test('stream error is re-thrown after logging', async () => { + async function* failingStream() { + yield { assistantResponseEvent: { content: 'hello' } } + throw new Error('network failure') + } + const sdk = { generateAssistantResponseResponse: failingStream() } + await expect(collectSdk(transformSdkStream(sdk, 'auto', 'conv-1'))).rejects.toThrow( + 'network failure' + ) + }) + + test('missing generateAssistantResponseResponse throws', async () => { + await expect(collectSdk(transformSdkStream({}, 'auto', 'conv-1'))).rejects.toThrow( + 'SDK response has no event stream' + ) + }) +}) + +// ── SDK stream: real token usage ───────────────────────────────────────────── + +describe('transformSdkStream: token usage', () => { + test('real tokenUsage from metadata wins over the context% estimate', async () => { + const events = [ + { assistantResponseEvent: { content: 'hello' } }, + { + metadataEvent: { + contextUsagePercentage: 80, + tokenUsage: { inputTokens: 1234, outputTokens: 56 } + } + } + ] + const result = await collectSdk(transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1')) + const usageEvent = result.find((e) => e.usage) + expect(usageEvent.usage.prompt_tokens).toBe(1234) + expect(usageEvent.usage.completion_tokens).toBe(56) + }) +}) + +// ── isNewThread after merge ────────────────────────────────────────────────── + +import { buildHistory } from '../infrastructure/transformers/history-builder.js' +import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' + +describe('isNewThread detection after merge', () => { + test('consecutive user messages merge to one — treated as new thread', () => { + const msgs = [ + { role: 'user', content: 'msg1' }, + { role: 'user', content: 'msg2' }, + { role: 'user', content: 'msg3' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(1) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(0) + }) + + test('user+assistant+user produces history', () => { + const msgs = [ + { role: 'user', content: 'hello' }, + { role: 'assistant', content: 'hi' }, + { role: 'user', content: 'next' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(3) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(2) + }) + + test('single user message produces empty history', () => { + const msgs = [{ role: 'user', content: 'only one' }] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(1) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(0) + }) + + test('consecutive assistant messages are merged', () => { + const msgs = [ + { role: 'user', content: 'start' }, + { role: 'assistant', content: 'part1' }, + { role: 'assistant', content: 'part2' }, + { role: 'user', content: 'next' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(3) + expect(merged[1].content).toContain('part1') + expect(merged[1].content).toContain('part2') + }) +}) + +// ── Tool name consistency (>64 chars must be shortened everywhere) ─────────── + +import { shortenToolName } from '../infrastructure/transformers/tool-transformer.js' + +describe('tool name >64 chars is shortened consistently in history', () => { + const longName = + 'a_very_long_tool_name_that_definitely_exceeds_the_sixty_four_character_limit_imposed_by_kiro' + + test('buildHistory shortens tool_use names', () => { + const msgs = [ + { role: 'user', content: 'do it' }, + { + role: 'assistant', + content: [ + { type: 'text', text: 'calling tool' }, + { type: 'tool_use', id: 'tu-1', name: longName, input: { x: 1 } } + ] + }, + { + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tu-1', content: 'done' }] + } + ] + const merged = mergeAdjacentMessages([...msgs]) + const history = buildHistory(merged, 'auto') + const toolUses = history.flatMap((h) => h.assistantResponseMessage?.toolUses || []) + expect(toolUses.length).toBeGreaterThan(0) + for (const tu of toolUses) { + expect(tu.name.length).toBeLessThanOrEqual(64) + expect(tu.name).toBe(shortenToolName(longName)) + } + }) + + test('buildHistory shortens tool_calls names', () => { + const msgs = [ + { role: 'user', content: 'do it' }, + { + role: 'assistant', + content: 'ok', + tool_calls: [{ id: 'tc-1', function: { name: longName, arguments: '{"a":1}' } }] + }, + { + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tc-1', content: 'result' }] + } + ] + const merged = mergeAdjacentMessages([...msgs]) + const history = buildHistory(merged, 'auto') + const toolUses = history.flatMap((h) => h.assistantResponseMessage?.toolUses || []) + expect(toolUses.length).toBeGreaterThan(0) + for (const tu of toolUses) { + expect(tu.name.length).toBeLessThanOrEqual(64) + expect(tu.name).toBe(shortenToolName(longName)) + } + }) +}) diff --git a/src/__tests__/tool-call-parser.test.ts b/src/__tests__/tool-call-parser.test.ts new file mode 100644 index 0000000..289ba5a --- /dev/null +++ b/src/__tests__/tool-call-parser.test.ts @@ -0,0 +1,115 @@ +import { describe, expect, test } from 'bun:test' +import { + cleanToolCallsFromText, + deduplicateToolCalls, + parseBracketToolCalls +} from '../infrastructure/transformers/tool-call-parser.js' + +describe('parseBracketToolCalls', () => { + test('returns empty array for plain text', () => { + expect(parseBracketToolCalls('Hello world')).toHaveLength(0) + }) + + test('parses single bracket tool call', () => { + const text = '[Called bash with args: {"command":"ls"}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(1) + expect(result[0]!.name).toBe('bash') + expect(result[0]!.input).toEqual({ command: 'ls' }) + }) + + test('parses multiple bracket tool calls', () => { + const text = + '[Called bash with args: {"command":"ls"}] [Called read with args: {"path":"/tmp"}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(2) + expect(result[0]!.name).toBe('bash') + expect(result[1]!.name).toBe('read') + }) + + test('skips malformed JSON args', () => { + const text = '[Called bash with args: {not valid json}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(0) + }) + + test('assigns unique toolUseIds', () => { + const text = + '[Called bash with args: {"command":"ls"}][Called bash with args: {"command":"pwd"}]' + const result = parseBracketToolCalls(text) + expect(result[0]!.toolUseId).not.toBe(result[1]!.toolUseId) + }) +}) + +describe('deduplicateToolCalls', () => { + test('returns empty array for empty input', () => { + expect(deduplicateToolCalls([])).toHaveLength(0) + }) + + test('keeps all unique tool calls', () => { + const calls = [ + { toolUseId: 'a', name: 'bash', input: {} }, + { toolUseId: 'b', name: 'read', input: {} } + ] + expect(deduplicateToolCalls(calls)).toHaveLength(2) + }) + + test('removes duplicates by toolUseId', () => { + const calls = [ + { toolUseId: 'a', name: 'bash', input: {} }, + { toolUseId: 'a', name: 'bash', input: { x: 1 } } + ] + const result = deduplicateToolCalls(calls) + expect(result).toHaveLength(1) + expect(result[0]!.input).toEqual({}) // first one kept + }) +}) + +describe('parseBracketToolCalls: large JSON content (catastrophic backtracking guard)', () => { + test('returns empty array fast when "[Called " is absent — even with huge JSON', () => { + // A 500 KB JSON array with many nested braces — previously caused catastrophic + // backtracking in the regex. The short-circuit guard must make this instant. + const largeJson = JSON.stringify( + Array.from({ length: 200 }, (_, i) => ({ + id: i, + data: 'x'.repeat(1000), + nested: { a: 1, b: [1, 2, 3] } + })) + ) + const start = performance.now() + const result = parseBracketToolCalls(largeJson) + const elapsed = performance.now() - start + expect(result).toHaveLength(0) + expect(elapsed).toBeLessThan(50) // must be near-instant, not seconds + }) + + test('still finds bracket tool calls when "[Called " is present in large text', () => { + const prefix = 'x'.repeat(10000) + const text = `${prefix}[Called bash with args: {"command":"ls"}]${prefix}` + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(1) + expect(result[0]!.name).toBe('bash') + }) +}) + +describe('cleanToolCallsFromText', () => { + test('removes bracket tool call from text', () => { + const text = 'Here is the result [Called bash with args: {"command":"ls"}] done.' + const toolCalls = [{ toolUseId: 'x', name: 'bash', input: {} }] + const result = cleanToolCallsFromText(text, toolCalls) + expect(result).not.toContain('[Called bash') + expect(result).toContain('done.') + }) + + test('leaves text unchanged when no tool calls match', () => { + const text = 'Hello world' + const result = cleanToolCallsFromText(text, []) + expect(result).toBe('Hello world') + }) + + test('trims and collapses whitespace', () => { + const text = ' hello world ' + const result = cleanToolCallsFromText(text, []) + expect(result).toBe('hello world') + }) +}) From fe67d865dbc3b57057fc2d9c846e0bd2419b9ec1 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 1 Jul 2026 14:32:41 +0200 Subject: [PATCH 08/10] fix(conv): configurable history trim limit to prevent context loss We trimmed conversation history at 600KB, over 6x too aggressive, which discarded older turns prematurely and caused context loss in longer agentic sessions. Kiro rejects oversized payloads with CONTENT_LENGTH_EXCEEDS_THRESHOLD; the hard limit is structure-dependent (verified against the live API): a single message survives up to ~7.6MB, but many-entry histories are rejected as low as ~5.9MB. Raise the trim threshold to a configurable max_payload_bytes (default 4MB), which stays safely below the lowest observed failure regardless of history shape while allowing far more context. Add [PAYLOAD] and [TRIM] debug logging so payload growth and trim events are visible. (cherry picked from commit ef532698c2d081638cf2455e9c05b07674256567) --- src/__tests__/edge-cases.test.ts | 43 +++++++++++++++++++++++------ src/core/request/request-handler.ts | 3 +- src/plugin/config/schema.ts | 15 ++++++++-- src/plugin/request.ts | 33 ++++++++++++++++++---- 4 files changed, 78 insertions(+), 16 deletions(-) diff --git a/src/__tests__/edge-cases.test.ts b/src/__tests__/edge-cases.test.ts index 3ba24e8..c2a1d07 100644 --- a/src/__tests__/edge-cases.test.ts +++ b/src/__tests__/edge-cases.test.ts @@ -277,9 +277,10 @@ describe('sanitizeSchema via convertToolsToCodeWhisperer', () => { }) describe('payload trim preserves valid structure', () => { - test('large payload trimmed to ≤600KB', () => { + test('oversized payload trimmed below Kiro threshold', () => { const messages: any[] = [] - for (let i = 0; i < 80; i++) { + // ~9MB of raw history so trimming must kick in (default cap is 4MB). + for (let i = 0; i < 300; i++) { messages.push({ role: 'user', content: 'x'.repeat(10000) }) messages.push({ role: 'assistant', @@ -291,7 +292,7 @@ describe('payload trim preserves valid structure', () => { } messages.push({ role: 'user', - content: [{ type: 'tool_result', tool_use_id: 'tool-79', content: 'result' }] + content: [{ type: 'tool_result', tool_use_id: 'tool-299', content: 'result' }] }) messages.push({ role: 'user', content: 'final question' }) @@ -301,7 +302,33 @@ describe('payload trim preserves valid structure', () => { } const result = transformToSdkRequest(body, 'auto', auth) const payload = JSON.stringify(result.conversationState) - expect(payload.length).toBeLessThanOrEqual(600_000) + expect(payload.length).toBeLessThanOrEqual(4_000_000) + }) + + test('respects a custom max_payload_bytes override', () => { + const messages: any[] = [] + for (let i = 0; i < 80; i++) { + messages.push({ role: 'user', content: 'x'.repeat(10000) }) + messages.push({ role: 'assistant', content: 'y'.repeat(10000) }) + } + messages.push({ role: 'user', content: 'final question' }) + + // Force a tight 500KB cap via the maxPayloadBytes argument. + const result = transformToSdkRequest( + { messages }, + 'auto', + auth, + false, + 20000, + undefined, + '', + true, + undefined, + undefined, + 500_000 + ) + const payload = JSON.stringify(result.conversationState) + expect(payload.length).toBeLessThanOrEqual(500_000) }) test('history starts with userInputMessage after trim', () => { @@ -810,10 +837,10 @@ describe('image carry-forward: no images on tool-result turns', () => { describe('payload trim performance', () => { test('trims very large history in <50ms', () => { - // 200 user/asst pairs of 10KB each = ~4MB raw — would be O(N²) without - // the incremental size accounting. + // 500 user/asst pairs of 10KB each = ~10MB raw — over the 4MB default cap, + // so trimming runs and would be O(N²) without the incremental accounting. const messages: any[] = [] - for (let i = 0; i < 200; i++) { + for (let i = 0; i < 500; i++) { messages.push({ role: 'user', content: 'u'.repeat(10000) }) messages.push({ role: 'assistant', content: 'a'.repeat(10000) }) } @@ -825,7 +852,7 @@ describe('payload trim performance', () => { // Result must still be under the limit. const size = JSON.stringify(result.conversationState).length - expect(size).toBeLessThanOrEqual(600_000) + expect(size).toBeLessThanOrEqual(4_000_000) // Allow 100ms on slow CI; fail if we regress to seconds (O(N²)). expect(elapsed).toBeLessThan(100) }) diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index 7aa5a9e..34531d6 100644 --- a/src/core/request/request-handler.ts +++ b/src/core/request/request-handler.ts @@ -293,7 +293,8 @@ export class RequestHandler { this.workspace, this.config.image_carry_forward, sessionId, - { effort: this.config.effort, autoEffortMapping: this.config.auto_effort_mapping } + { effort: this.config.effort, autoEffortMapping: this.config.auto_effort_mapping }, + this.config.max_payload_bytes ) } diff --git a/src/plugin/config/schema.ts b/src/plugin/config/schema.ts index 5a4c89c..1e51ddb 100644 --- a/src/plugin/config/schema.ts +++ b/src/plugin/config/schema.ts @@ -103,7 +103,17 @@ export const KiroConfigSchema = z.object({ // Kiro bills per session (request), not per token, so re-sending the same // images each turn has no billing impact. Only disable if you hit the // per-request 3.75MB image-payload cap on conversations with many heavy images. - image_carry_forward: z.boolean().default(true) + image_carry_forward: z.boolean().default(true), + + // Maximum conversation-state payload size (bytes) before the plugin trims the + // oldest history entries. Kiro's runtime endpoint rejects oversized payloads + // with CONTENT_LENGTH_EXCEEDS_THRESHOLD. The hard limit is structure-dependent + // (verified against the live API): a single message is accepted up to ~7.6MB, + // but conversations with many history entries are rejected as low as ~5.9MB. + // The 4MB default stays safely below the lowest observed failure regardless of + // structure, while allowing far more context than a conservative cap. Raising + // it risks 400s on long many-turn sessions; lowering it trims context sooner. + max_payload_bytes: z.number().min(100_000).max(5_500_000).default(4_000_000) }) export type KiroConfig = z.infer @@ -123,5 +133,6 @@ export const DEFAULT_CONFIG: KiroConfig = { auto_sync_kiro_cli: true, enable_log_api_request: false, auto_effort_mapping: true, - image_carry_forward: true + image_carry_forward: true, + max_payload_bytes: 4_000_000 } diff --git a/src/plugin/request.ts b/src/plugin/request.ts index 79f09ae..698f85a 100644 --- a/src/plugin/request.ts +++ b/src/plugin/request.ts @@ -97,7 +97,8 @@ function buildCodeWhispererRequest( showToast?: ToastFunction, workspace = '', carryForward = true, - sessionId?: string + sessionId?: string, + maxPayloadBytes = 4_000_000 ): TransformResult { const req = typeof body === 'string' ? JSON.parse(body) : body const { messages, tools, system } = req @@ -351,16 +352,24 @@ function buildCodeWhispererRequest( } } - // Trim history if payload exceeds Kiro's ~615KB limit. + // Trim history if the payload approaches Kiro's request-size limit. + // Kiro rejects oversized payloads with CONTENT_LENGTH_EXCEEDS_THRESHOLD. The + // hard limit is structure-dependent (verified against the live API): a single + // message survives up to ~7.6MB, but many-entry histories are rejected as low + // as ~5.9MB. The configurable maxPayloadBytes (default 4MB) stays safely below + // the lowest observed failure regardless of history shape. // Compute per-entry sizes once and shrink incrementally to avoid O(N²) // re-stringifying the full request on every iteration. - const MAX_PAYLOAD_BYTES = 600_000 + const MAX_PAYLOAD_BYTES = maxPayloadBytes + const trimStartLen = history.length + let trimSizeBefore = 0 if (history.length > 2) { const sizes = history.map((h) => JSON.stringify(h).length + 1) const baseRequest: any = { ...request, conversationState: { ...request.conversationState } } delete baseRequest.conversationState.history let totalSize = JSON.stringify(baseRequest).length + 2 // for `"history":[]` for (const s of sizes) totalSize += s + trimSizeBefore = totalSize while (history.length > 2 && totalSize > MAX_PAYLOAD_BYTES) { // Drop the two oldest entries (typically user + assistant pair). @@ -390,6 +399,12 @@ function buildCodeWhispererRequest( } } + if (trimStartLen !== history.length) { + logger.debug( + `[TRIM] history ${trimStartLen}→${history.length} entries, ~${Math.round(trimSizeBefore / 1024)}KB exceeded ${Math.round(MAX_PAYLOAD_BYTES / 1024)}KB cap` + ) + } + // After trimming, drop any current-message toolResults whose tool_use was // removed from history — sending a toolResult without a matching toolUse // in the surviving history causes Bedrock 400 ValidationException. @@ -515,7 +530,8 @@ export function transformToSdkRequest( workspace = '', carryForward = true, sessionId?: string, - effortConfig?: EffortConfig + effortConfig?: EffortConfig, + maxPayloadBytes = 4_000_000 ): SdkPreparedRequest { const { request, resolved, convId, fingerprint, toolNameMapper } = buildCodeWhispererRequest( body, @@ -526,7 +542,8 @@ export function transformToSdkRequest( showToast, workspace, carryForward, - sessionId + sessionId, + maxPayloadBytes ) // Resolve effort level based on config and model capabilities @@ -544,6 +561,12 @@ export function transformToSdkRequest( const endpointFull = resolveKiroEndpoint(auth) const endpoint = endpointFull.replace(/\/generateAssistantResponse$/, '') const workspaceKey = sessionId ? `sess:${sessionId}` : workspace + if (process.env.DEBUG || process.env.OPENCODE_LOG_LEVEL === 'debug') { + const finalBytes = JSON.stringify(request.conversationState).length + logger.debug( + `[PAYLOAD] convId=${convId} ~${Math.round(finalBytes / 1024)}KB history=${request.conversationState.history?.length ?? 0} entries` + ) + } return { conversationState: request.conversationState, profileArn: request.profileArn, From 93192aec107fd48741f383bcfd67db2bf4d9d4d4 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Wed, 15 Jul 2026 10:50:13 +0200 Subject: [PATCH 09/10] fix(stream): treat no-param tool calls without stop event as complete Kiro SDK does not send a stop event for tool calls with no parameters (e.g. time_currentTime). The stream ends with the tool call still in activeToolCalls, which was incorrectly treated as a 64K truncation. Now: tool calls that received no input at stream end are treated as complete. Only calls with partial input are flagged as truncated. --- src/__tests__/stream-transformer.test.ts | 22 +++++++++++++++++++ .../streaming/sdk-stream-transformer.ts | 16 ++++++++++---- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/__tests__/stream-transformer.test.ts b/src/__tests__/stream-transformer.test.ts index 0a629ad..d41783f 100644 --- a/src/__tests__/stream-transformer.test.ts +++ b/src/__tests__/stream-transformer.test.ts @@ -221,6 +221,28 @@ describe('transformSdkStream: tool call streaming', () => { expect(toolBlocks.length).toBeGreaterThan(0) }) + test('no-param tool call without stop event is treated as complete (not truncated)', async () => { + const events = [ + { toolUseEvent: { toolUseId: 't-noarg', name: 'time_currentTime' } }, + { toolUseEvent: { toolUseId: 't-noarg', name: 'time_currentTime', input: '' } } + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + + const allText = result + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e: any) => e.choices[0].delta.content) + .join('') + expect(allText).not.toContain('truncated') + + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBe(1) + expect(toolBlocks[0]!.choices[0].delta.tool_calls[0].function.name).toBe('time_currentTime') + + const stopReason = result.find((e) => e.choices?.[0]?.finish_reason) + expect(stopReason.choices[0].finish_reason).toBe('tool_calls') + }) + test('truncated tool call (no stop event) closes the block and emits error text', async () => { // Simulates Kiro cutting the stream mid-tool-call (large .frm JSON that // exceeds the response size limit). diff --git a/src/plugin/streaming/sdk-stream-transformer.ts b/src/plugin/streaming/sdk-stream-transformer.ts index 3bb780e..1b6d8cd 100644 --- a/src/plugin/streaming/sdk-stream-transformer.ts +++ b/src/plugin/streaming/sdk-stream-transformer.ts @@ -230,9 +230,18 @@ export async function* transformSdkStream( } if (activeToolCalls.size > 0) { - // Stream cut off mid-tool-call(s). Tell the model to retry with smaller - // chunks so it can self-correct on the next turn. - for (const truncated of activeToolCalls.values()) { + const trulyTruncated: ToolCallState[] = [] + for (const pending of activeToolCalls.values()) { + if (pending.input.length === 0) { + pending.stopped = true + toolCalls.push(pending) + } else { + trulyTruncated.push(pending) + } + } + activeToolCalls.clear() + + for (const truncated of trulyTruncated) { logger.debug( `[STREAM] Truncated tool call: name=${truncated.name} id=${truncated.toolUseId} inputLen=${truncated.input.length}` ) @@ -244,7 +253,6 @@ export async function* transformSdkStream( if (_c !== null) yield _c } } - activeToolCalls.clear() } if (thinkingRequested && streamState.buffer) { From 1e71f363b0a3f8a0137de2df33e68462377c4e12 Mon Sep 17 00:00:00 2001 From: Rene van Veen Date: Thu, 16 Jul 2026 16:10:20 +0200 Subject: [PATCH 10/10] fix(endpoint): keep signing region and endpoint host consistent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit resolveKiroEndpoint used auth.region (SSO/IDC home region) to pick the endpoint host, while transformToSdkRequest signs using the profile ARN's region when available. For IAM Identity Center accounts these can differ, causing a host/signing-region mismatch that surfaces as a spurious 400 that looks like model/region unavailability. Also reword the bare "Invalid model ID" 400 into an actionable message — Kiro rolls out model availability per region gradually, so this explains that instead of leaving a cryptic server string. --- src/__tests__/sdk-client.test.ts | 33 ++++++++++++++++++++++++++++- src/core/request/error-handler.ts | 23 +++++++++++++++++--- src/core/request/request-handler.ts | 3 ++- src/plugin/sdk-client.ts | 7 ++++-- 4 files changed, 59 insertions(+), 7 deletions(-) diff --git a/src/__tests__/sdk-client.test.ts b/src/__tests__/sdk-client.test.ts index 4a7c825..624427e 100644 --- a/src/__tests__/sdk-client.test.ts +++ b/src/__tests__/sdk-client.test.ts @@ -1,6 +1,6 @@ import { GenerateAssistantResponseCommand } from '@aws/codewhisperer-streaming-client' import { describe, expect, test } from 'bun:test' -import { clearSdkClientCache, createSdkClient } from '../plugin/sdk-client' +import { clearSdkClientCache, createSdkClient, resolveKiroEndpoint } from '../plugin/sdk-client' import type { KiroAuthDetails } from '../plugin/types' function auth(): KiroAuthDetails { @@ -71,3 +71,34 @@ describe('SDK client', () => { clearSdkClientCache() }) }) + +describe('resolveKiroEndpoint region consistency', () => { + test('Pro account: uses profileArn region even when auth.region (IDC/SSO home region) differs', () => { + // IAM Identity Center SSO portal region can differ from the region embedded + // in the CodeWhisperer profile ARN. The endpoint host must follow the + // profile ARN's region, not the SSO home region, or Kiro's backend rejects + // the request (host/signing-region mismatch looks like a spurious 400). + const a: KiroAuthDetails = { + ...auth(), + region: 'eu-west-1', // SSO/IDC home region + profileArn: 'arn:aws:codewhisperer:eu-central-1:123456789012:profile/ABCDEF' + } + expect(resolveKiroEndpoint(a)).toBe( + 'https://runtime.eu-central-1.kiro.dev/generateAssistantResponse' + ) + }) + + test('Builder ID account (no profileArn): falls back to auth.region', () => { + const a: KiroAuthDetails = { ...auth(), region: 'eu-central-1' } + expect(resolveKiroEndpoint(a)).toBe( + 'https://q.eu-central-1.amazonaws.com/generateAssistantResponse' + ) + }) + + test('missing region and profileArn falls back to us-east-1', () => { + const a: KiroAuthDetails = { ...auth(), region: undefined as any } + expect(resolveKiroEndpoint(a)).toBe( + 'https://q.us-east-1.amazonaws.com/generateAssistantResponse' + ) + }) +}) diff --git a/src/core/request/error-handler.ts b/src/core/request/error-handler.ts index 5446f0b..9ca7b01 100644 --- a/src/core/request/error-handler.ts +++ b/src/core/request/error-handler.ts @@ -29,7 +29,8 @@ export class ErrorHandler { response: Response, account: ManagedAccount, context: RequestContext, - showToast: ToastFunction + showToast: ToastFunction, + model?: string ): Promise<{ shouldRetry: boolean; newContext?: RequestContext; switchAccount?: boolean }> { const readBody = async (): Promise => { try { @@ -42,8 +43,9 @@ export class ErrorHandler { if (response.status === 400) { const reason = await readBody() - logger.warn(`HTTP 400 on ${account.email}: ${reason || 'unknown'}`) - showToast(`400: ${reason || 'unknown'}`, 'error') + const message = this.enrichIfInvalidModel(reason, account, model) + logger.warn(`HTTP 400 on ${account.email}: ${message || 'unknown'}`) + showToast(`400: ${message || 'unknown'}`, 'error') return { shouldRetry: false } } @@ -196,6 +198,21 @@ export class ErrorHandler { return { shouldRetry: false } } + // Kiro rolls out model availability per region gradually; reword its bare + // "Invalid model ID" into something actionable without guessing at regions. + private enrichIfInvalidModel(reason: string, account: ManagedAccount, model?: string): string { + if (!reason || !/invalid model id/i.test(reason)) { + return reason + } + const region = account.region || 'your account\u2019s region' + const modelPart = model ? `Model "${model}"` : 'This model' + return ( + `${modelPart} is not available via region "${region}". Kiro rolls out model ` + + `availability gradually per region — check https://kiro.dev/docs/models/ for current ` + + `availability. (${reason})` + ) + } + private isNetworkError(e: any): boolean { return ( e instanceof Error && /econnreset|etimedout|enotfound|network|fetch failed/i.test(e.message) diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index 34531d6..9bd86d2 100644 --- a/src/core/request/request-handler.ts +++ b/src/core/request/request-handler.ts @@ -219,7 +219,8 @@ export class RequestHandler { mockResponse, acc, { retry, excludedMs: retryContext.excludedMs }, - showToast + showToast, + model ) if (errorResult.shouldRetry) { diff --git a/src/plugin/sdk-client.ts b/src/plugin/sdk-client.ts index 96f6bfc..20547b8 100644 --- a/src/plugin/sdk-client.ts +++ b/src/plugin/sdk-client.ts @@ -1,6 +1,6 @@ import { CodeWhispererStreamingClient } from '@aws/codewhisperer-streaming-client' import * as crypto from 'crypto' -import { KIRO_CONSTANTS, buildUrl } from '../constants.js' +import { KIRO_CONSTANTS, buildUrl, extractRegionFromArn } from '../constants.js' import type { Effort, KiroAuthDetails } from './types' const KIRO_VERSION = '0.11.63' @@ -21,9 +21,12 @@ function getMachineId(auth: KiroAuthDetails): string { * - Accounts without a profileArn (free AWS Builder ID) → q.amazonaws.com * This endpoint accepts the same token + request shape but only serves * Claude-family models. Using runtime.kiro.dev here causes a 400. + * + * Region must match transformToSdkRequest's signing region (same fallback + * order) — SSO home region and profile ARN region can differ for IDC accounts. */ export function resolveKiroEndpoint(auth: KiroAuthDetails): string { - const region = auth.region || 'us-east-1' + const region = extractRegionFromArn(auth.profileArn) ?? auth.region ?? 'us-east-1' if (auth.profileArn) { return buildUrl(KIRO_CONSTANTS.RUNTIME_URL, region as any) }