fix(kiro): retry once when a malformed tool_call wrapper precedes any output

.i18n-state.json

@@ -74498,7 +74498,7 @@
       }
     },
     "docs/reference/FEATURE_FLAGS.md": {
-      "source_hash": "9afe2c360388a5d6ff2f73f56652187c2cf217e9646ad26babb8a1558288e8ca",
+      "source_hash": "478b7577694ce8cd1dcc0f7624a6fd25dcba1476048bd56971caa4f81cc43fb5",
       "locales": {
         "ar": {
           "source_hash": "9afe2c360388a5d6ff2f73f56652187c2cf217e9646ad26babb8a1558288e8ca",

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

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": {

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)
 

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<string, unknown>
       )?.content as string) || "";
     const thinkingExpected = userContent.includes("<thinking_mode>enabled</thinking_mode>");
-    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();
     };
 

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<string, unknown>;
+
+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
+ * `<thinking_mode>` 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<Response>;
+  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<string | null> => {
+    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<Uint8Array>({
+    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",
+    },
+  });
+}