fix(responses): bound the streaming citation marker span (carry #3843)

src/adapters/kiro.ts

@@ -1,6 +1,7 @@
 import { decodeEventStream } from "../lib/eventstream-decoder";
 import { estimateTokens } from "../lib/token-estimate";
 import { debugProviderDiagnostic } from "../lib/debug";
+import { isDebugEnabled } from "../lib/debug-settings";
 import { resolveKiroApiRegion, resolveKiroRequestProfile } from "../oauth/kiro";
 import { KIRO_MODEL_CONTEXT_WINDOWS, normalizeKiroModelId } from "../providers/kiro-models";
 import { modelRecordValue } from "../reasoning-effort";
@@ -2120,17 +2121,21 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
     const rawContextInputEstimate = estimateKiroPayloadInputTokens(built.payload, parsed.modelId);
     const contextInputEstimate = calibrateKiroEstimate(built.conversationId, rawContextInputEstimate);
     const body = JSON.stringify(built.payload);
-    debugProviderDiagnostic("kiro", "request", {
-      region,
-      requestedModel: parsed.modelId,
-      completionMode: built.completionMode,
-      bodyBytes: new TextEncoder().encode(body).length,
-      messageCount: kiroPayloadMessages(parsed).length,
-      toolCount: parsed.context.tools?.length ?? 0,
-      hasProfileArn: Boolean(profileArn),
-      wireClient,
-      hasPreviousResponseId: Boolean(parsed.previousResponseId),
-    });
+    // Every field below is evaluated before the call, so an unguarded call re-encodes the
+    // whole request body on each request even when provider debug is off. Gate the details.
+    if (isDebugEnabled()) {
+      debugProviderDiagnostic("kiro", "request", {
+        region,
+        requestedModel: parsed.modelId,
+        completionMode: built.completionMode,
+        bodyBytes: new TextEncoder().encode(body).length,
+        messageCount: kiroPayloadMessages(parsed).length,
+        toolCount: parsed.context.tools?.length ?? 0,
+        hasProfileArn: Boolean(profileArn),
+        wireClient,
+        hasPreviousResponseId: Boolean(parsed.previousResponseId),
+      });
+    }
     return {
       request: {
         url: kiroRuntimeEndpoint(provider, region),

src/responses/citation-markers.ts

@@ -68,13 +68,25 @@ export interface CitationMarkerFilter {
   flush(): string;
 }
 
+/**
+ * Upper bound on the text withheld for one unterminated START.
+ *
+ * A real span is `cite` plus a few turn-scoped ids, so it is far under this. Without a
+ * bound, a backend that emits a START and never terminates it makes `held` grow for the
+ * whole response, and every later delta re-scans that accumulated prefix.
+ */
+const MAX_STREAMING_MARKER_SPAN_LENGTH = 4_096;
+
 /**
  * Streaming filter.
  *
  * A marker can straddle a delta boundary — `\uE200cite` in one chunk and the rest in the
  * next — so a stateless per-delta strip would emit the tail of a span it never recognized.
  * This holds back the text from an unterminated START and releases it once the END arrives
  * (removed) or the stream ends (verbatim, so nothing the model actually said is lost).
+ *
+ * A span that grows past `MAX_STREAMING_MARKER_SPAN_LENGTH` is malformed ordinary text, so
+ * it is released verbatim instead of withheld; a later START can still open a valid span.
  */
 export function createCitationMarkerFilter(): CitationMarkerFilter {
   // Text from an open START that has not been terminated yet.
@@ -83,13 +95,29 @@ export function createCitationMarkerFilter(): CitationMarkerFilter {
     push(delta: string): string {
       const combined = held + delta;
       held = "";
-      const start = combined.lastIndexOf(CITATION_MARKER_START);
-      if (start === -1) return stripCitationMarkers(combined);
-      const endAfterStart = combined.indexOf(CITATION_MARKER_END, start + 1);
-      if (endAfterStart !== -1) return stripCitationMarkers(combined);
-      // The trailing span is still open: emit everything before it, hold the rest.
-      held = combined.slice(start);
-      return stripCitationMarkers(combined.slice(0, start));
+      let start = combined.indexOf(CITATION_MARKER_START);
+      if (start === -1) return combined;
+      let out = combined.slice(0, start);
+      // Walk START-delimited segments independently so an earlier malformed START is never
+      // paired with a later span's END (the whole-string strip would do exactly that).
+      while (start !== -1) {
+        const nextStart = combined.indexOf(CITATION_MARKER_START, start + 1);
+        const segment = combined.slice(start, nextStart === -1 ? combined.length : nextStart);
+        const end = segment.indexOf(CITATION_MARKER_END, 1);
+        if (end !== -1) {
+          // A complete span: drop it, keep whatever trails it inside this segment.
+          out += segment.slice(end + 1);
+        } else if (nextStart === -1 && segment.length <= MAX_STREAMING_MARKER_SPAN_LENGTH) {
+          // Only a bounded trailing span can still be completed by a later delta.
+          held = segment;
+        } else {
+          // Superseded by a later START, or over the bound: ordinary text, emitted verbatim
+          // so neither the retained text nor the per-delta rescan grows without limit.
+          out += segment;
+        }
+        start = nextStart;
+      }
+      return out;
     },
     flush(): string {
       const rest = held;
@@ -98,4 +126,3 @@ export function createCitationMarkerFilter(): CitationMarkerFilter {
     },
   };
 }
-

tests/providers/kiro/kiro-stream.test.ts

@@ -16,6 +16,11 @@ import { parseKiroEvent } from "../../../src/adapters/kiro-events";
 import { resetKiroThrottleStateForTests } from "../../../src/adapters/kiro-retry";
 import { resetKiroCalibration } from "../../../src/adapters/kiro-calibration";
 import { buildResponseJSON } from "../../../src/bridge";
+import {
+  clearDebugSetting,
+  getDebugSettings,
+  setDebugSettings,
+} from "../../../src/lib/debug-settings";
 import { encodeMessage } from "../../../src/lib/eventstream-decoder";
 import { estimateTokens } from "../../../src/lib/token-estimate";
 import { createTranslatorBudget } from "../../../src/lib/translator-budget";
@@ -34,19 +39,26 @@ const origApiRegion = process.env.KIRO_API_REGION;
 const origArn = process.env.KIRO_PROFILE_ARN;
 const origCredsFile = process.env.KIRO_CREDS_FILE;
 const origCredentialsFile = process.env.KIRO_CREDENTIALS_FILE;
-const origDebugFrames = process.env.OCX_DEBUG_FRAMES;
+let origDebug: string | undefined;
+let origDebugFrames: string | undefined;
+let origDebugOverride: boolean | undefined;
 const realFetch = globalThis.fetch;
 let tmp: string;
 
 beforeEach(() => {
+  origDebug = process.env.OCX_DEBUG;
+  origDebugFrames = process.env.OCX_DEBUG_FRAMES;
+  origDebugOverride = getDebugSettings().runtimeOverride.debug;
   tmp = mkdtempSync(join(tmpdir(), "kiro-stream-"));
   process.env.HOME = tmp;
   process.env.KIRO_REGION = "us-east-1";
   delete process.env.KIRO_API_REGION;
   delete process.env.KIRO_PROFILE_ARN;
   delete process.env.KIRO_CREDS_FILE;
   delete process.env.KIRO_CREDENTIALS_FILE;
+  delete process.env.OCX_DEBUG;
   delete process.env.OCX_DEBUG_FRAMES;
+  clearDebugSetting("debug");
 });
 afterEach(() => {
   globalThis.fetch = realFetch;
@@ -57,7 +69,10 @@ afterEach(() => {
   if (origArn === undefined) delete process.env.KIRO_PROFILE_ARN; else process.env.KIRO_PROFILE_ARN = origArn;
   if (origCredsFile === undefined) delete process.env.KIRO_CREDS_FILE; else process.env.KIRO_CREDS_FILE = origCredsFile;
   if (origCredentialsFile === undefined) delete process.env.KIRO_CREDENTIALS_FILE; else process.env.KIRO_CREDENTIALS_FILE = origCredentialsFile;
+  if (origDebug === undefined) delete process.env.OCX_DEBUG; else process.env.OCX_DEBUG = origDebug;
   if (origDebugFrames === undefined) delete process.env.OCX_DEBUG_FRAMES; else process.env.OCX_DEBUG_FRAMES = origDebugFrames;
+  if (origDebugOverride === undefined) clearDebugSetting("debug");
+  else setDebugSettings({ debug: origDebugOverride });
   removeTreeWithRetry(tmp);
 });
 
@@ -196,6 +211,21 @@ describe("kiro adapter — parseStream", () => {
     expect(providerState).toEqual({ kiro: { conversationId: "returned-conversation-1" } });
   });
 
+  test("request diagnostics do not re-encode the body when provider debug is off", async () => {
+    const encodeSpy = spyOn(TextEncoder.prototype, "encode");
+    try {
+      const adapter = createKiroAdapter(provider);
+      const before = encodeSpy.mock.calls.length;
+      await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }]));
+      const during = encodeSpy.mock.calls.slice(before);
+      // The diagnostic argument list is evaluated eagerly, so an unguarded call encodes the
+      // full serialized request body on every request even with diagnostics disabled.
+      expect(during.some(([value]) => typeof value === "string" && value.includes("conversationState"))).toBe(false);
+    } finally {
+      encodeSpy.mockRestore();
+    }
+  });
+
   test("invalid returned message metadata cannot poison continuation state", async () => {
     const adapter = createKiroAdapter(provider);
     const request = await adapter.buildRequest(parsedWith([{ role: "user", content: "hi" }]));

tests/responses/citation-markers.test.ts

@@ -87,4 +87,34 @@ describe("streaming citation marker filter (#3150)", () => {
     const filter = createCitationMarkerFilter();
     expect(filter.push(`visible now ${S}cite`)).toBe("visible now ");
   });
+
+  test("an unterminated span past the bound is released instead of retained", () => {
+    // A backend that opens a span and never closes it must not make the filter accumulate
+    // the rest of the response, which every later delta would then re-scan.
+    const filter = createCitationMarkerFilter();
+    let out = filter.push(`kept ${S}cite`);
+    expect(out).toBe("kept ");
+    for (let i = 0; i < 5_000; i += 1) out += filter.push("x");
+
+    // Everything after the malformed START is emitted verbatim, so nothing is lost, and
+    // flush() has nothing left to release.
+    expect(out).toBe(`kept ${S}cite${"x".repeat(5_000)}`);
+    expect(filter.flush()).toBe("");
+  });
+
+  test("a later START still opens a valid span after a released malformed one", () => {
+    const filter = createCitationMarkerFilter();
+    let out = filter.push(`a${S}${"y".repeat(5_000)}`);
+    out += filter.push(`${S}cite${P}turn1view0${E} tail`);
+    expect(out).toBe(`a${S}${"y".repeat(5_000)} tail`);
+    expect(filter.flush()).toBe("");
+  });
+
+  test("an oversized malformed span survives a later valid marker in the same delta", () => {
+    const filter = createCitationMarkerFilter();
+    const malformed = `${S}${"y".repeat(5_000)}`;
+    expect(filter.push(`a${span}${malformed}${S}cite${P}turn1view0${E} tail`))
+      .toBe(`a${malformed} tail`);
+    expect(filter.flush()).toBe("");
+  });
 });