diff --git a/apps/web/lib/source-annotations.test.ts b/apps/web/lib/source-annotations.test.ts index dce37e1a..d2aa10b2 100644 --- a/apps/web/lib/source-annotations.test.ts +++ b/apps/web/lib/source-annotations.test.ts @@ -84,6 +84,29 @@ describe("source annotation parsing", () => { ).toBe("Alpha Beta") }) + it("recovers a duplicated malformed closing tag from model output", () => { + const input = + 'Pinki got soup for dinner at the church todaytoday.' + + expect( + parseSourceAnnotatedMarkdown(input, new Set(["S1", "S2"])).markdown, + ).toBe( + "Pinki got soup [for dinner at the church](#sm-source:S1) [today](#sm-source:S2).", + ) + expect(stripSourceMarkup(input)).toBe( + "Pinki got soup for dinner at the church today.", + ) + }) + + it("accepts a quoted malformed closing tag without duplicated text", () => { + expect( + parseSourceAnnotatedMarkdown( + 'Alpha Beta Gamma', + new Set(["S1"]), + ).markdown, + ).toBe("Alpha [Beta](#sm-source:S1) Gamma") + }) + it("allows only source ids that are safe in internal fragments", () => { expect(isSafeSourceId("S1._:-")).toBe(true) expect(isSafeSourceId("bad/id")).toBe(false) diff --git a/apps/web/lib/source-annotations.ts b/apps/web/lib/source-annotations.ts index 7abe6859..eb64c422 100644 --- a/apps/web/lib/source-annotations.ts +++ b/apps/web/lib/source-annotations.ts @@ -4,9 +4,16 @@ export type ParsedSourceAnnotations = { const RESPONSE_OPEN_PREFIX = " 0 && SAFE_SOURCE_ID_RE.test(id) } @@ -56,6 +63,64 @@ function parseOpeningTag( return { end: tagEnd + 1, sourceId } } +function parseClosingTagAt( + text: string, + index: number, +): ResponseClosingTag | null { + if (!text.startsWith(RESPONSE_CLOSE_PREFIX, index)) return null + + const tagEnd = text.indexOf(">", index + RESPONSE_CLOSE_PREFIX.length) + if (tagEnd === -1) return null + + const suffix = text.slice(index + RESPONSE_CLOSE_PREFIX.length, tagEnd) + if (!/^[\s"']*$/.test(suffix)) return null + + return { + start: index, + end: tagEnd + 1, + malformed: text.slice(index, tagEnd + 1) !== RESPONSE_CLOSE_TAG, + } +} + +function findClosingTag( + text: string, + index: number, +): ResponseClosingTag | null { + let cursor = index + while (cursor < text.length) { + const candidate = text.indexOf(RESPONSE_CLOSE_PREFIX, cursor) + if (candidate === -1) return null + + const closing = parseClosingTagAt(text, candidate) + if (closing) return closing + cursor = candidate + RESPONSE_CLOSE_PREFIX.length + } + + return null +} + +function consumeDuplicatedMalformedClose( + text: string, + closing: ResponseClosingTag, + plainInner: string, +): number { + if (!closing.malformed) return closing.end + + const trailingClose = findClosingTag(text, closing.end) + if (!trailingClose || trailingClose.malformed) return closing.end + + const duplicate = text.slice(closing.end, trailingClose.start) + if ( + duplicate.includes("<") || + duplicate.trim() !== plainInner.trim() || + plainInner.trim().length === 0 + ) { + return closing.end + } + + return trailingClose.end +} + function advanceCodeState( text: string, index: number, @@ -111,27 +176,23 @@ export function parseSourceAnnotatedMarkdown( } if (opening) { - const closeIndex = text.indexOf(RESPONSE_CLOSE_TAG, opening.end) - if (closeIndex === -1) { + const closing = findClosingTag(text, opening.end) + if (!closing) { output.push(stripSourceMarkup(text.slice(opening.end))) break } - const inner = text.slice(opening.end, closeIndex) + const inner = text.slice(opening.end, closing.start) const hasNested = inner.includes(RESPONSE_OPEN_PREFIX) || - inner.includes(RESPONSE_CLOSE_TAG) + inner.includes(RESPONSE_CLOSE_PREFIX) const isAllowed = allowedSourceIds.has(opening.sourceId) if (hasNested) { - const outerCloseIndex = text.indexOf( - RESPONSE_CLOSE_TAG, - closeIndex + RESPONSE_CLOSE_TAG.length, - ) - const fallbackEnd = - outerCloseIndex === -1 ? closeIndex : outerCloseIndex + const outerClosing = findClosingTag(text, closing.end) + const fallbackEnd = outerClosing?.start ?? closing.start output.push(stripSourceMarkup(text.slice(opening.end, fallbackEnd))) - i = fallbackEnd + RESPONSE_CLOSE_TAG.length + i = outerClosing?.end ?? closing.end continue } @@ -144,19 +205,21 @@ export function parseSourceAnnotatedMarkdown( output.push(plainInner) } - i = closeIndex + RESPONSE_CLOSE_TAG.length + i = consumeDuplicatedMalformedClose(text, closing, plainInner) codeState.lineStart = output.length === 0 || output[output.length - 1]?.endsWith("\n") === true continue } - const nextClose = text.indexOf(RESPONSE_CLOSE_TAG, i) - if (nextClose !== -1) { + const nextClose = findClosingTag(text, i) + if (nextClose) { const tagEnd = text.indexOf(">", i) - if (tagEnd !== -1 && tagEnd < nextClose) { - output.push(stripSourceMarkup(text.slice(tagEnd + 1, nextClose))) - i = nextClose + RESPONSE_CLOSE_TAG.length + if (tagEnd !== -1 && tagEnd < nextClose.start) { + output.push( + stripSourceMarkup(text.slice(tagEnd + 1, nextClose.start)), + ) + i = nextClose.end continue } } @@ -174,12 +237,26 @@ export function stripSourceMarkup(text: string): string { let i = 0 while (i < text.length) { - if (text.startsWith(RESPONSE_CLOSE_TAG, i)) { - i += RESPONSE_CLOSE_TAG.length + const closing = parseClosingTagAt(text, i) + if (closing) { + i = closing.end continue } if (text.startsWith(RESPONSE_OPEN_PREFIX, i)) { + const opening = parseOpeningTag(text, i) + if (opening && opening !== "incomplete") { + const matchingClose = findClosingTag(text, opening.end) + if (matchingClose) { + const plainInner = stripSourceMarkup( + text.slice(opening.end, matchingClose.start), + ) + output += plainInner + i = consumeDuplicatedMalformedClose(text, matchingClose, plainInner) + continue + } + } + const tagEnd = text.indexOf(">", i + RESPONSE_OPEN_PREFIX.length) if (tagEnd === -1) break i = tagEnd + 1