answer: Escape XML and guard abort race

Amolith created 1 week ago

Escape <, >, and & in question text, user answers, and notes before
embedding them in pseudo-XML tags. Without this, literal angle brackets
in content could confuse model parsing of the <qna> block.

Wrap the done() callback in a finished-flag guard so only the first
completion takes effect. The loader.onAbort handler and the in-flight
extraction promise could both call done() on user cancellation,
producing a double-completion race.

Change summary

packages/answer/src/index.ts | 34 +++++++++++++++++++++++++---------
1 file changed, 25 insertions(+), 9 deletions(-)

Detailed changes

packages/answer/src/index.ts

@@ -31,6 +31,11 @@ import {
 } from "@mariozechner/pi-tui";
 import { Type } from "@sinclair/typebox";
 
+/** Escape characters that would break pseudo-XML tag boundaries. */
+function escapeXml(s: string): string {
+	return s.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;");
+}
+
 // Structured output format for question extraction
 interface ExtractedQuestion {
 	question: string;
@@ -319,12 +324,12 @@ class QnAComponent implements Component {
 		for (let i = 0; i < this.questions.length; i++) {
 			const q = this.questions[i];
 			const a = this.answers[i]?.trim() || "(no answer)";
-			parts.push(`<q n="${i}">${q.question}</q>`);
-			parts.push(`<a n="${i}">${a}</a>`);
+			parts.push(`<q n="${i}">${escapeXml(q.question)}</q>`);
+			parts.push(`<a n="${i}">${escapeXml(a)}</a>`);
 		}
 		parts.push(`</qna>`);
 		if (this.notesText) {
-			parts.push(`\n<note>${this.notesText}</note>`);
+			parts.push(`\n<note>${escapeXml(this.notesText)}</note>`);
 		}
 
 		this.onDone(parts.join("\n").trim());
@@ -625,7 +630,18 @@ export default function (pi: ExtensionAPI) {
 		const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
 			const extractionModel = resolveExtractionModel(ctx, sessionModel);
 			const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);
-			loader.onAbort = () => done({ kind: "cancelled" });
+
+			// Guard against double-completion: loader.onAbort fires on user
+			// cancel, but the in-flight promise may also resolve/reject after
+			// the abort.  Only the first call to finish() takes effect.
+			let finished = false;
+			const finish = (result: ExtractionOutcome) => {
+				if (finished) return;
+				finished = true;
+				done(result);
+			};
+
+			loader.onAbort = () => finish({ kind: "cancelled" });
 
 			const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
 				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
@@ -687,21 +703,21 @@ export default function (pi: ExtensionAPI) {
 
 				switch (result.kind) {
 					case "ok":
-						return done({ kind: "ok", result: result.result });
+						return finish({ kind: "ok", result: result.result });
 					case "cancelled":
-						return done({ kind: "cancelled" });
+						return finish({ kind: "cancelled" });
 					case "model_error":
-						return done({
+						return finish({
 							kind: "error",
 							message: `${result.model.name} returned an error with no content`,
 						});
 					case "parse_error":
-						return done({ kind: "error", message: result.message });
+						return finish({ kind: "error", message: result.message });
 				}
 			};
 
 			doExtract().catch((err) =>
-				done({
+				finish({
 					kind: "error",
 					message: `${err?.message ?? err}`,
 				}),