diff --git a/src/responses/citation-markers.ts b/src/responses/citation-markers.ts index 5fe58142cf..b65477b66e 100644 --- a/src/responses/citation-markers.ts +++ b/src/responses/citation-markers.ts @@ -68,6 +68,15 @@ 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. * @@ -75,6 +84,9 @@ export interface CitationMarkerFilter { * 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. @@ -87,6 +99,11 @@ export function createCitationMarkerFilter(): CitationMarkerFilter { if (start === -1) return stripCitationMarkers(combined); const endAfterStart = combined.indexOf(CITATION_MARKER_END, start + 1); if (endAfterStart !== -1) return stripCitationMarkers(combined); + // Over the bound: this is not a citation span we will ever close. Emit it verbatim + // so neither the retained text nor the per-delta rescan grows without limit. + if (combined.length - start > MAX_STREAMING_MARKER_SPAN_LENGTH) { + return stripCitationMarkers(combined.slice(0, start)) + combined.slice(start); + } // The trailing span is still open: emit everything before it, hold the rest. held = combined.slice(start); return stripCitationMarkers(combined.slice(0, start)); diff --git a/tests/responses/citation-markers.test.ts b/tests/responses/citation-markers.test.ts index 0c1921750c..6145dbe688 100644 --- a/tests/responses/citation-markers.test.ts +++ b/tests/responses/citation-markers.test.ts @@ -87,4 +87,26 @@ 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(""); + }); });