diff --git a/.i18n-state.json b/.i18n-state.json index 4b3d320256a6..3e37d6937131 100644 --- a/.i18n-state.json +++ b/.i18n-state.json @@ -74498,7 +74498,7 @@ } }, "docs/reference/FEATURE_FLAGS.md": { - "source_hash": "9afe2c360388a5d6ff2f73f56652187c2cf217e9646ad26babb8a1558288e8ca", + "source_hash": "478b7577694ce8cd1dcc0f7624a6fd25dcba1476048bd56971caa4f81cc43fb5", "locales": { "ar": { "source_hash": "9afe2c360388a5d6ff2f73f56652187c2cf217e9646ad26babb8a1558288e8ca", diff --git a/changelog.d/fixes/16495-kiro-tool-call-wrapper-repair.md b/changelog.d/fixes/16495-kiro-tool-call-wrapper-repair.md new file mode 100644 index 000000000000..c28668d9be5f --- /dev/null +++ b/changelog.d/fixes/16495-kiro-tool-call-wrapper-repair.md @@ -0,0 +1 @@ +- **fix(kiro):** retry once with a corrective instruction when Kiro emits a malformed `tool_call` wrapper before any output, instead of failing the turn; gated by the `KIRO_TOOL_CALL_REPAIR_ENABLED` flag (on by default) ([#16495](https://github.com/diegosouzapw/OmniRoute/pull/16495)) — thanks @SemonCat diff --git a/config/quality/eslint-suppressions.json b/config/quality/eslint-suppressions.json index 800987162fc8..c1cd8b72104f 100644 --- a/config/quality/eslint-suppressions.json +++ b/config/quality/eslint-suppressions.json @@ -131,7 +131,7 @@ }, "open-sse/executors/kiro.ts": { "@typescript-eslint/no-unused-vars": { - "count": 3 + "count": 2 } }, "open-sse/executors/kiro/eventstream.ts": { diff --git a/docs/reference/FEATURE_FLAGS.md b/docs/reference/FEATURE_FLAGS.md index 725257598bb1..069cfc2e1ab3 100644 --- a/docs/reference/FEATURE_FLAGS.md +++ b/docs/reference/FEATURE_FLAGS.md @@ -46,7 +46,7 @@ A boolean flag is considered **enabled** when its effective value is `"true"`, ## Flag Catalog -85 flags across 6 categories. **Default** is the definition default — the value +86 flags across 6 categories. **Default** is the definition default — the value used when neither a DB override nor an environment variable is present. ### Security (10) @@ -103,7 +103,7 @@ used when neither a DB override nor an environment variable is present. | `USAGE_LIMIT_IGNORE_UNPRICED` | boolean | `false` | Count usage of models that have no price as $0 in per-key USD usage quotas instead of treating the quota as exceeded. Off by default: an unpriced model or routing alias can hide real spend, so the quota fails closed. | | `RADAR_ENABLED` | boolean | `false` | Enable the OmniRoute Radar module (catalog feed screens and sync). Off by default; enabling only unlocks the UI — data sync remains a separate opt-in. | -### Runtime (36) +### Runtime (37) | Key | Type | Default | Restart | Description | | ------------------------------------------- | ------- | ------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | @@ -143,6 +143,7 @@ used when neither a DB override nor an environment variable is present. | `BATCH_AND_FILE_AUTO_CLEANUP_ENABLED` | boolean | `false` | | Let the automatic cleanup sweep delete terminal (completed/failed/cancelled/expired) Batch API jobs older than `OMNIROUTE_BATCH_RETENTION_DAYS`, along with their per-line checkpoints, and clear the BLOB content of uploaded files past their own `expires_at`. Off by default: every existing install keeps this data exactly as before until an operator opts in. The operator-triggered `DELETE /api/v1/batches/delete-completed` route is unaffected either way — it is a separate, unconditional public API contract. | | `ANTIGRAVITY_ACCOUNT_LEASE_ENABLED` | boolean | `false` | | Reserve the selected Antigravity account for the streaming lifecycle of the request that picked it, so a concurrent retry or the credential handoff cannot re-pick an account already committed to an in-flight stream. The reservation is scoped to (connection, callable upstream model), so one account can still serve two different models at once. When every eligible account is already leased for that model, the request returns a structured 503 `antigravity_pool_busy` with a bounded `Retry-After` instead of piling onto a busy account. Off by default: account selection stays exactly as before, and no reservation is taken. | | `COMBO_AUTO_PRUNE_STALE_STEPS` | boolean | `false` | | After a successful model sync against an authoritative live catalog, remove combo steps pinned to models the catalog no longer lists, with one audit entry per removed step. Never prunes on a failed, degraded or free-only sync, and never empties a combo. Off by default: stale steps are only flagged. | +| `KIRO_TOOL_CALL_REPAIR_ENABLED` | boolean | `true` | | When a streaming Kiro response emits a malformed `tool_call` wrapper (missing `input.name` or `input.arguments`) before any visible output, suppress it and re-send the request once with a corrective instruction instead of ending the turn with `invalid_kiro_tool_call`. The retry costs one extra Kiro request; a second failure ends the turn with `kiro_tool_call_repair_retry_failed`. On by default. Set to `false` to fail the turn on the first malformed wrapper. | ### CLI (5) diff --git a/open-sse/executors/kiro.ts b/open-sse/executors/kiro.ts index 7f7c68f8cdb5..6c56b14b8bea 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -32,6 +32,10 @@ import { validateKiroToolUse, type PendingKiroWrapperToolCall, } from "./kiroToolCallValidation.ts"; +import { + createKiroToolCallRepairStream, + isKiroToolCallRepairEnabled, +} from "./kiro/toolCallRepair.ts"; export { validateKiroToolUse } from "./kiroToolCallValidation.ts"; @@ -415,25 +419,30 @@ export class KiroExecutor extends BaseExecutor { const headers = this.buildHeaders(credentials, stream); mergeUpstreamExtraHeaders(headers, upstreamExtraHeaders); const transformedBody = await this.transformRequest(model, body, stream, credentials); - const requestBody = JSON.stringify(transformedBody); - - let response!: Response; - let url = candidateUrls[0]; - for (let i = 0; i < candidateUrls.length; i++) { - url = candidateUrls[i]; - response = await fetch(url, { - method: "POST", - headers, - body: requestBody, - signal, - }); - const hasFallback = i + 1 < candidateUrls.length; - if (response.ok || !hasFallback) break; - const shouldFallback = - KIRO_ENDPOINT_FALLBACK_STATUSES.has(response.status) || - (await isBrandedGatewayProfileArnRejection(response)); - if (!shouldFallback) break; - } + + // Branded gateway first, raw CodeWhisperer host on an auth/profile-shaped failure. + // Shared with the one-shot tool_call repair retry so it takes the same path. + const send = async (payload: unknown, sendSignal: AbortSignal | null | undefined) => { + let res!: Response; + let target = candidateUrls[0]; + for (let i = 0; i < candidateUrls.length; i++) { + target = candidateUrls[i]; + res = await fetch(target, { + method: "POST", + headers, + body: JSON.stringify(payload), + signal: sendSignal, + }); + const hasFallback = i + 1 < candidateUrls.length; + if (res.ok || !hasFallback) break; + const shouldFallback = + KIRO_ENDPOINT_FALLBACK_STATUSES.has(res.status) || + (await isBrandedGatewayProfileArnRejection(res)); + if (!shouldFallback) break; + } + return { res, target }; + }; + const { res: response, target: url } = await send(transformedBody, signal); if (!response.ok) { return { response, url, headers, transformedBody }; @@ -459,9 +468,18 @@ export class KiroExecutor extends BaseExecutor { )?.userInputMessage as Record )?.content as string) || ""; const thinkingExpected = userContent.includes("enabled"); - const transformedResponse = this.transformEventStreamToSSE(response, model, { - thinkingExpected, - }); + const transformedResponse = + stream && isKiroToolCallRepairEnabled() + ? createKiroToolCallRepairStream({ + upstream: response, + transform: (raw, opts) => + this.transformEventStreamToSSE(raw, model, { thinkingExpected, ...opts }), + retry: async (payload, retrySignal) => (await send(payload, retrySignal)).res, + transformedBody, + signal, + log, + }) + : this.transformEventStreamToSSE(response, model, { thinkingExpected }); return { response: transformedResponse, url, headers, transformedBody }; } @@ -483,12 +501,17 @@ export class KiroExecutor extends BaseExecutor { transformEventStreamToSSE( response: Response, model: string, - opts: { thinkingExpected?: boolean } = {} + opts: { + thinkingExpected?: boolean; + /** Return true to suppress the error frame (the one-shot repair will retry). */ + onInvalidToolCall?: (message: string) => boolean; + responseId?: string; + } = {} ) { const thinkingExpected = !!opts.thinkingExpected; const buffer = new ByteQueue(); let chunkIndex = 0; - const responseId = `chatcmpl-${Date.now()}`; + const responseId = opts.responseId || `chatcmpl-${Date.now()}`; const created = Math.floor(Date.now() / 1000); const state: KiroStreamState = { endDetected: false, @@ -582,8 +605,10 @@ export class KiroExecutor extends BaseExecutor { }; state.invalidToolCall = true; state.finishEmitted = true; - controller.enqueue(encodeSse(`data: ${JSON.stringify(error)}\n\n`)); - controller.enqueue(encodeSse("data: [DONE]\n\n")); + if (!opts.onInvalidToolCall?.(message)) { + controller.enqueue(encodeSse(`data: ${JSON.stringify(error)}\n\n`)); + controller.enqueue(encodeSse("data: [DONE]\n\n")); + } controller.terminate(); }; diff --git a/open-sse/executors/kiro/toolCallRepair.ts b/open-sse/executors/kiro/toolCallRepair.ts new file mode 100644 index 000000000000..13f41f99fbfa --- /dev/null +++ b/open-sse/executors/kiro/toolCallRepair.ts @@ -0,0 +1,220 @@ +/** + * One-shot repair for malformed Kiro `tool_call` wrappers. + * + * Kiro sometimes emits the MCP wrapper tool `tool_call` without `input.name` or + * `input.arguments`. The SSE transform detects that once the wrapper input is + * complete and, without repair, ends the turn with `invalid_kiro_tool_call`. + * + * The repair runs inside the response stream, so the executor still hands the + * stream off on the first role-only chunk (#15260) and TTFT is unchanged: + * + * - Chunks are forwarded as they arrive. Role-only start chunks do not count as + * output; text, reasoning, tool calls and finish chunks do. + * - A malformed wrapper seen before any output is suppressed, the first upstream + * is cancelled, and the request is re-sent once with a corrective instruction + * appended to the current user message. The retry's role-only chunk is dropped + * because the client already received one. + * - A malformed wrapper after visible output (text or reasoning) cannot be + * replayed invisibly, so it surfaces as `invalid_kiro_tool_call` as before. + * - A retry that fails (non-OK status, network error, or a second malformed + * wrapper) ends the stream with `kiro_tool_call_repair_retry_failed`. Headers + * are already sent at that point, so a non-OK retry cannot reach account + * fallback; the turn fails exactly as it did before the repair existed. + * + * Gated by the KIRO_TOOL_CALL_REPAIR_ENABLED feature flag (on by default). + */ +import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags"; +import type { ExecutorLog } from "../base.ts"; +import { encodeSse } from "../kiroToolCallValidation.ts"; + +export const KIRO_TOOL_CALL_REPAIR_FLAG = "KIRO_TOOL_CALL_REPAIR_ENABLED"; +export const KIRO_TOOL_CALL_REPAIR_RETRY_FAILED_CODE = "kiro_tool_call_repair_retry_failed"; +const MAX_REASON_CHARS = 300; +const KIRO_TOOL_CALL_REPAIR_INSTRUCTION = [ + "Retry the previous response because its Kiro tool_call wrapper was malformed.", + "If you use the wrapper tool named tool_call, its input must be a JSON object with a non-empty string name and an arguments field.", + "Do not emit a tool_call wrapper without input.name and input.arguments.", +].join(" "); + +export function isKiroToolCallRepairEnabled(): boolean { + try { + return isFeatureFlagEnabled(KIRO_TOOL_CALL_REPAIR_FLAG); + } catch { + // An unreadable flag store keeps the definition default (on). + return true; + } +} + +type JsonRecord = Record; + +const asRecord = (value: unknown): JsonRecord | null => + value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : null; + +/** + * Copy the translated Kiro payload and append the corrective instruction to + * `conversationState.currentMessage.userInputMessage.content`, where the + * `` hint also lives. Only the validation message is echoed + * (capped at 300 chars); tool arguments never are. + */ +export function buildKiroToolCallRepairPayload(transformedBody: unknown, reason: string): unknown { + const repaired = structuredClone(transformedBody) as unknown; + const userInput = asRecord( + asRecord(asRecord(asRecord(repaired)?.conversationState)?.currentMessage)?.userInputMessage + ); + if (!userInput) return repaired; + const capped = String(reason || "invalid tool_call payload").slice(0, MAX_REASON_CHARS); + const instruction = `${KIRO_TOOL_CALL_REPAIR_INSTRUCTION} Previous validation error: ${capped}`; + const content = typeof userInput.content === "string" ? userInput.content : ""; + userInput.content = content ? `${content}\n\n${instruction}` : instruction; + return repaired; +} + +const decoder = new TextDecoder(); + +function parseSseData(chunk: Uint8Array): JsonRecord[] { + const out: JsonRecord[] = []; + for (const line of decoder.decode(chunk).split("\n")) { + if (!line.startsWith("data: ")) continue; + const data = line.slice(6).trim(); + if (!data || data === "[DONE]") continue; + try { + const parsed = asRecord(JSON.parse(data)); + if (parsed) out.push(parsed); + } catch { + // Not JSON: treat as output so it is never swallowed. + out.push({ raw: data }); + } + } + return out; +} + +/** True unless every data payload is a role-only `chat.completion.chunk`. */ +export function isMeaningfulKiroSseChunk(chunk: Uint8Array): boolean { + return parseSseData(chunk).some((payload) => { + const choices = Array.isArray(payload.choices) ? payload.choices : null; + if (!choices) return true; + return choices.some((choice) => { + const c = asRecord(choice); + if (!c) return true; + if (c.finish_reason !== null && c.finish_reason !== undefined) return true; + const delta = asRecord(c.delta) ?? {}; + return Object.keys(delta).some((key) => key !== "role"); + }); + }); +} + +function readChunkId(chunk: Uint8Array): string | undefined { + const id = parseSseData(chunk).find((payload) => typeof payload.id === "string")?.id; + return typeof id === "string" ? id : undefined; +} + +function repairErrorFrame(message: string): Uint8Array { + const error = { + error: { + message, + type: "invalid_request_error", + code: KIRO_TOOL_CALL_REPAIR_RETRY_FAILED_CODE, + }, + }; + return encodeSse(`data: ${JSON.stringify(error)}\n\ndata: [DONE]\n\n`); +} + +export type KiroSseTransform = ( + upstream: Response, + opts: { onInvalidToolCall: (message: string) => boolean; responseId?: string } +) => Response; + +export function createKiroToolCallRepairStream(params: { + upstream: Response; + transform: KiroSseTransform; + retry: (payload: unknown, signal: AbortSignal) => Promise; + transformedBody: unknown; + signal?: AbortSignal | null; + log?: ExecutorLog | null; +}): Response { + const { upstream, transform, retry, transformedBody, signal, log } = params; + const retryAbort = new AbortController(); + let phase: "first" | "retry" | "done" = "first"; + let cancelled = false; + let sawOutput = false; + let invalidReason: string | null = null; + let responseId: string | undefined; + + const openAttempt = (response: Response) => + transform(response, { + responseId, + onInvalidToolCall: (message) => { + if (sawOutput) return false; + invalidReason = message; + return true; + }, + }).body!.getReader(); + + let reader = openAttempt(upstream); + + const startRetry = async (): Promise => { + const reason = invalidReason ?? ""; + log?.warn?.("KIRO", "malformed tool_call wrapper before any output; retrying once"); + invalidReason = null; + phase = "retry"; + let response: Response; + try { + const retrySignal = signal ? AbortSignal.any([signal, retryAbort.signal]) : retryAbort.signal; + response = await retry(buildKiroToolCallRepairPayload(transformedBody, reason), retrySignal); + } catch (error) { + if (signal?.aborted || retryAbort.signal.aborted) throw error; + return "Kiro tool_call repair retry failed"; + } + if (cancelled || !response.ok) { + await response.body?.cancel().catch(() => {}); + return `Kiro tool_call repair retry failed with HTTP ${response.status}`; + } + reader = openAttempt(response); + return null; + }; + + const body = new ReadableStream({ + async pull(controller) { + while (phase !== "done") { + const { done, value } = await reader.read(); + if (cancelled) return; + if (!done) { + const meaningful = isMeaningfulKiroSseChunk(value); + if (!sawOutput && !meaningful && phase === "retry") continue; + if (phase === "first") responseId ??= readChunkId(value); + if (meaningful) sawOutput = true; + controller.enqueue(value); + return; + } + if (invalidReason === null) break; + if (phase === "first") { + const failure = await startRetry(); + if (cancelled) return; + if (failure === null) continue; + controller.enqueue(repairErrorFrame(failure)); + } else { + controller.enqueue(repairErrorFrame(invalidReason)); + } + break; + } + phase = "done"; + controller.close(); + }, + async cancel(reason) { + cancelled = true; + phase = "done"; + retryAbort.abort(reason); + await reader.cancel(reason).catch(() => {}); + }, + }); + + return new Response(body, { + status: upstream.status, + statusText: upstream.statusText, + headers: { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }, + }); +} diff --git a/open-sse/utils/error.ts b/open-sse/utils/error.ts index a953867f89a5..c914ffe5a7be 100644 --- a/open-sse/utils/error.ts +++ b/open-sse/utils/error.ts @@ -166,6 +166,7 @@ const SAFE_PUBLIC_ERROR_IDENTIFIERS = new Set([ "invalid_trailer", "key_allows_all_combos", "key_allows_all_models", + "kiro_tool_call_repair_retry_failed", "lease_action_invalid", "lease_api_key_invalid", "lease_authentication_required", @@ -535,8 +536,7 @@ export function sanitizeRecoveryHint( if (!r || typeof r !== "object") return undefined; const action = typeof r.action === "string" ? (r.action as ComboRecoveryAction) : null; if (!action || !RECOVERY_ACTIONS.has(action)) return undefined; - // Reject empty OR whitespace-only next_step — the value must render usefully as a - // header and as a body field. A whitespace-only string would print as a blank hint. + // Reject empty/whitespace-only next_step: it must render as a non-blank header + body hint. const next_step = clampRecoveryStr(r.next_step, 200).trim(); if (!next_step) return undefined; const hint: ComboRecoveryHint = { action, next_step }; diff --git a/src/shared/constants/featureFlagDefinitions.ts b/src/shared/constants/featureFlagDefinitions.ts index 2c967e9a5ace..2e648fd09022 100644 --- a/src/shared/constants/featureFlagDefinitions.ts +++ b/src/shared/constants/featureFlagDefinitions.ts @@ -886,6 +886,18 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [ requiresRestart: false, warningLevel: "caution", }, + { + key: "KIRO_TOOL_CALL_REPAIR_ENABLED", + label: "Kiro Tool Call Repair", + description: + "When a streaming Kiro response emits a malformed tool_call wrapper (missing input.name or input.arguments) before any visible output, suppress it and re-send the request once with a corrective instruction instead of ending the turn with invalid_kiro_tool_call. The retry costs one extra Kiro request; a second failure ends the turn with kiro_tool_call_repair_retry_failed. On by default. Set to false to fail the turn on the first malformed wrapper.", + descriptionI18nKey: "featureFlagKiroToolCallRepairEnabledDescription", + category: "runtime", + defaultValue: "true", + type: "boolean", + requiresRestart: false, + warningLevel: "info", + }, // ──────────────── CLI (5) ──────────────── { diff --git a/tests/unit/feature-flags-settings.test.ts b/tests/unit/feature-flags-settings.test.ts index 86e5aed55f95..5bbc9c562532 100644 --- a/tests/unit/feature-flags-settings.test.ts +++ b/tests/unit/feature-flags-settings.test.ts @@ -54,8 +54,9 @@ const { // pool member, default off) takes it to 82; USAGE_LIMIT_IGNORE_UNPRICED (#14799, // opt-in flag to count unpriced usage as $0 in USD quotas, default off) to 83; // REASONING_REPLAY_ENABLED (#12486, default on) to 84; COMBO_AUTO_PRUNE_STALE_STEPS -// (#13505, default off) to 85. -const EXPECTED_FEATURE_FLAG_COUNT = 85; +// (#13505, default off) to 85; KIRO_TOOL_CALL_REPAIR_ENABLED (one-shot Kiro +// tool_call wrapper repair, default on) to 86. +const EXPECTED_FEATURE_FLAG_COUNT = 86; // ────────────────────────────────────────────────────── // Test group 1 — Flag definitions registry diff --git a/tests/unit/fixtures/error-public-boundaries-hardening.fixture.ts b/tests/unit/fixtures/error-public-boundaries-hardening.fixture.ts index b775a585fa23..bcc604c8d257 100644 --- a/tests/unit/fixtures/error-public-boundaries-hardening.fixture.ts +++ b/tests/unit/fixtures/error-public-boundaries-hardening.fixture.ts @@ -269,6 +269,7 @@ test("public identifier vocabulary preserves current internal machine-readable c "orphan_tool_result", "bedrock_stream_error", "invalid_kiro_tool_call", + "kiro_tool_call_repair_retry_failed", "devin_cli_error", "upstream_websocket_error", "upstream_websocket_connect_failed", diff --git a/tests/unit/kiro-tool-call-repair.test.ts b/tests/unit/kiro-tool-call-repair.test.ts new file mode 100644 index 000000000000..dea8f7ad6dce --- /dev/null +++ b/tests/unit/kiro-tool-call-repair.test.ts @@ -0,0 +1,291 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { kiroEventFrame } from "../_helpers/kiroEventStreamFixture.ts"; + +const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-kiro-repair-")); +const ORIGINAL_DATA_DIR = process.env.DATA_DIR; +const ORIGINAL_FLAG = process.env.KIRO_TOOL_CALL_REPAIR_ENABLED; +process.env.DATA_DIR = dataDir; +delete process.env.KIRO_TOOL_CALL_REPAIR_ENABLED; + +const { KiroExecutor } = await import("../../open-sse/executors/kiro.ts"); +const { buildKiroToolCallRepairPayload, KIRO_TOOL_CALL_REPAIR_FLAG } = + await import("../../open-sse/executors/kiro/toolCallRepair.ts"); +const { FEATURE_FLAG_DEFINITIONS } = + await import("../../src/shared/constants/featureFlagDefinitions.ts"); +const { resetDbInstance } = await import("../../src/lib/db/core.ts"); + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; + delete process.env.KIRO_TOOL_CALL_REPAIR_ENABLED; +}); + +test.after(() => { + globalThis.fetch = originalFetch; + if (ORIGINAL_FLAG === undefined) delete process.env.KIRO_TOOL_CALL_REPAIR_ENABLED; + else process.env.KIRO_TOOL_CALL_REPAIR_ENABLED = ORIGINAL_FLAG; + resetDbInstance(); + if (ORIGINAL_DATA_DIR === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = ORIGINAL_DATA_DIR; + fs.rmSync(dataDir, { recursive: true, force: true }); +}); + +const MALFORMED = [ + kiroEventFrame("toolUseEvent", { + toolUseId: "wrapper_bad", + name: "tool_call", + input: { arguments: { query: "secret-arg-router" } }, + }), + kiroEventFrame("messageStopEvent", {}), +]; + +const VALID = [ + kiroEventFrame("toolUseEvent", { + toolUseId: "wrapper_ok", + name: "tool_call", + input: { name: "mcp_search", arguments: { query: "router" } }, + }), + kiroEventFrame("messageStopEvent", {}), +]; + +type Upstream = Uint8Array[] | { status: number }; + +function eventStreamResponse(frames: Uint8Array[], onCancel?: () => void): Response { + return new Response( + new ReadableStream({ + start(controller) { + for (const frame of frames) controller.enqueue(frame); + controller.close(); + }, + cancel() { + onCancel?.(); + }, + }), + { status: 200, headers: { "Content-Type": "application/vnd.amazon.eventstream" } } + ); +} + +function installFetch(sequence: Upstream[]) { + const bodies: string[] = []; + globalThis.fetch = async (_input, init) => { + bodies.push(String(init?.body ?? "")); + const next = sequence[bodies.length - 1]; + if (!next) throw new Error(`unexpected upstream fetch #${bodies.length}`); + if (!Array.isArray(next)) { + return new Response(JSON.stringify({ message: "Too many requests" }), { + status: next.status, + }); + } + return eventStreamResponse(next); + }; + return bodies; +} + +function runExecutor(stream = true, signal?: AbortSignal) { + return new KiroExecutor().execute({ + model: "claude-sonnet-4.6", + body: { + conversationState: { + currentMessage: { userInputMessage: { content: "hello" } }, + }, + }, + stream, + credentials: { + apiKey: "kiro-repair-fixture-only", + providerSpecificData: { authMethod: "api_key", region: "us-east-1" }, + }, + signal, + } as Parameters["execute"]>[0]); +} + +function sseChunks(text: string): Array> { + return text + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => line.slice(6).trim()) + .filter((line) => line && line !== "[DONE]") + .map((line) => JSON.parse(line) as Record); +} + +function roleChunkCount(text: string): number { + return sseChunks(text).filter((chunk) => { + const choices = Array.isArray(chunk.choices) ? chunk.choices : []; + const delta = (choices[0] as Record | undefined)?.delta as + Record | undefined; + return delta?.role === "assistant"; + }).length; +} + +function userContentOf(body: string): string { + const parsed = JSON.parse(body) as { + conversationState: { currentMessage: { userInputMessage: { content: string } } }; + }; + return parsed.conversationState.currentMessage.userInputMessage.content; +} + +test("the repair flag is defined, runtime-togglable and on by default", () => { + const def = FEATURE_FLAG_DEFINITIONS.find((d) => d.key === KIRO_TOOL_CALL_REPAIR_FLAG); + assert.ok(def, "KIRO_TOOL_CALL_REPAIR_ENABLED must be defined"); + assert.equal(def.defaultValue, "true"); + assert.equal(def.type, "boolean"); + assert.equal(def.requiresRestart, false); +}); + +test("flag on (default): a malformed wrapper is retried once and only the valid call reaches the client", async () => { + const bodies = installFetch([MALFORMED, VALID]); + const result = await runExecutor(); + const text = await result.response.text(); + + assert.equal(bodies.length, 2, "exactly one repair retry"); + assert.doesNotMatch(text, /invalid_kiro_tool_call/); + assert.doesNotMatch(text, /kiro_tool_call_repair_retry_failed/); + assert.match(text, /mcp_search/); + assert.equal(roleChunkCount(text), 1, "the retry must not emit a second role chunk"); + assert.equal(text.match(/data: \[DONE\]/g)?.length, 1); + assert.equal(result.response.headers.get("Content-Type"), "text/event-stream"); + + const retryContent = userContentOf(bodies[1]); + assert.match(retryContent, /^hello/); + assert.match(retryContent, /tool_call wrapper was malformed/); + assert.match(retryContent, /missing nested MCP tool name/); + assert.doesNotMatch(bodies[1], /secret-arg-router/, "tool arguments are never echoed"); + assert.equal(userContentOf(bodies[0]), "hello", "the first attempt is untouched"); +}); + +test("flag on: a second malformed wrapper ends with kiro_tool_call_repair_retry_failed", async () => { + const bodies = installFetch([MALFORMED, MALFORMED]); + const text = await (await runExecutor()).response.text(); + + assert.equal(bodies.length, 2, "the repair is bounded to one retry"); + assert.match(text, /kiro_tool_call_repair_retry_failed/); + assert.match(text, /missing nested MCP tool name/); + assert.doesNotMatch(text, /invalid_kiro_tool_call/); + assert.doesNotMatch(text, /secret-arg-router/); + assert.match(text, /data: \[DONE\]/); +}); + +test("flag on: a non-OK retry surfaces a bounded repair error without leaking the upstream body", async () => { + const bodies = installFetch([MALFORMED, { status: 429 }]); + const text = await (await runExecutor()).response.text(); + + assert.equal(bodies.length, 2); + assert.match(text, /kiro_tool_call_repair_retry_failed/); + assert.match(text, /HTTP 429/); + assert.doesNotMatch(text, /Too many requests/); + assert.match(text, /data: \[DONE\]/); +}); + +test("flag on: a malformed wrapper after visible text is not retried", async () => { + const bodies = installFetch([ + [kiroEventFrame("assistantResponseEvent", { content: "Let me search." }), ...MALFORMED], + ]); + const text = await (await runExecutor()).response.text(); + + assert.equal(bodies.length, 1); + assert.match(text, /Let me search\./); + assert.match(text, /invalid_kiro_tool_call/); +}); + +test("flag on: text is released before upstream completes and no retry happens", async () => { + let upstream!: ReadableStreamDefaultController; + let fetches = 0; + let upstreamCancelled = false; + globalThis.fetch = async () => { + fetches++; + return new Response( + new ReadableStream({ + start(controller) { + upstream = controller; + controller.enqueue(kiroEventFrame("assistantResponseEvent", { content: "EARLY_TEXT" })); + }, + cancel() { + upstreamCancelled = true; + }, + }), + { status: 200 } + ); + }; + + const result = await runExecutor(); + const reader = result.response.body!.getReader(); + const decoder = new TextDecoder(); + let seen = ""; + while (!seen.includes("EARLY_TEXT")) { + const { done, value } = await reader.read(); + assert.equal(done, false, "the stream must not end before upstream does"); + seen += decoder.decode(value, { stream: true }); + } + assert.match(seen, /EARLY_TEXT/); + + upstream.enqueue(kiroEventFrame("messageStopEvent", {})); + upstream.close(); + for (;;) { + const { done } = await reader.read(); + if (done) break; + } + assert.equal(fetches, 1); + assert.equal(upstreamCancelled, false); +}); + +test("flag on: cancelling the client stream cancels the upstream reader", async () => { + let cancelled = false; + globalThis.fetch = async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(kiroEventFrame("contextUsageEvent", { contextUsagePercentage: 5 })); + }, + cancel() { + cancelled = true; + }, + }), + { status: 200 } + ); + + const result = await runExecutor(); + const reader = result.response.body!.getReader(); + const first = await reader.read(); + assert.equal(first.done, false, "the role-only start chunk is handed off immediately"); + await reader.cancel("client went away"); + assert.equal(cancelled, true); +}); + +test("flag off: a malformed wrapper keeps the pre-repair behavior (single attempt, invalid_kiro_tool_call)", async () => { + process.env.KIRO_TOOL_CALL_REPAIR_ENABLED = "false"; + const bodies = installFetch([MALFORMED, VALID]); + const text = await (await runExecutor()).response.text(); + + assert.equal(bodies.length, 1, "no retry when the flag is off"); + assert.match(text, /invalid_kiro_tool_call/); + assert.doesNotMatch(text, /mcp_search/); +}); + +test("non-streaming requests are not retried", async () => { + const bodies = installFetch([MALFORMED, VALID]); + const text = await (await runExecutor(false)).response.text(); + + assert.equal(bodies.length, 1); + assert.match(text, /invalid_kiro_tool_call/); +}); + +test("buildKiroToolCallRepairPayload appends a capped reason to the current user message only", () => { + const original = { + conversationState: { + currentMessage: { userInputMessage: { content: "enabledq" } }, + history: [{ userInputMessage: { content: "earlier" } }], + }, + }; + const repaired = buildKiroToolCallRepairPayload(original, "x".repeat(1000)) as typeof original; + + assert.equal(original.conversationState.currentMessage.userInputMessage.content.length, 39); + const content = repaired.conversationState.currentMessage.userInputMessage.content; + assert.ok(content.startsWith("enabledq\n\n")); + assert.match(content, /Previous validation error: x{300}$/); + assert.deepEqual(repaired.conversationState.history, original.conversationState.history); +}); diff --git a/tests/unit/server-owned-tool-loop-flag.test.ts b/tests/unit/server-owned-tool-loop-flag.test.ts index 62877d15efeb..bbd3b42f348a 100644 --- a/tests/unit/server-owned-tool-loop-flag.test.ts +++ b/tests/unit/server-owned-tool-loop-flag.test.ts @@ -75,8 +75,8 @@ describe("feature-flags-settings count update", () => { // (#14657, default off) to 80; PROXY_OPERATOR_EGRESS_ENABLED (#15314, // default off) to 82; USAGE_LIMIT_IGNORE_UNPRICED (#14799, default off) to 83; // REASONING_REPLAY_ENABLED (#12486, default on) to 84; COMBO_AUTO_PRUNE_STALE_STEPS - // (#13505, default off) to 85. - assert.equal(FEATURE_FLAG_DEFINITIONS.length, 85); + // (#13505, default off) to 85; KIRO_TOOL_CALL_REPAIR_ENABLED (default on) to 86. + assert.equal(FEATURE_FLAG_DEFINITIONS.length, 86); }); });