diff --git a/extensions/user-input-fold/index.test.ts b/extensions/user-input-fold/index.test.ts new file mode 100644 index 00000000..48844996 --- /dev/null +++ b/extensions/user-input-fold/index.test.ts @@ -0,0 +1,250 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { foldUserMessage, transformUserMarkdown } from "./index.ts"; + +const MARKER = "… folded"; + +const numberedLines = (count: number) => + Array.from( + { length: count }, + (_, i) => `line ${String(i + 1).padStart(2, "0")}`, + ); + +const giantBlock = [ + "```js", + ...Array.from({ length: 300 }, (_, i) => `console.log(${i});`), + "```", +].join("\n"); + +test("messages below both thresholds round-trip byte-identical", () => { + const messages = [ + "hi", + "hello\n", + "a\r\nb", + " spaced \n\n", + "explanation\n```js\nfoo();\nbar();\n```\ndone", + "x".repeat(1200), + numberedLines(20).join("\n"), + ]; + for (const message of messages) { + assert.equal(foldUserMessage(message), message); + } +}); + +test("a message at exactly the line threshold is unchanged", () => { + const message = numberedLines(20).join("\n"); + assert.equal(foldUserMessage(message), message); +}); + +test("one line over the line threshold folds to a twelve-line preview", () => { + const message = numberedLines(21).join("\n"); + const out = foldUserMessage(message); + assert.ok(out.startsWith("line 01")); + assert.ok(out.includes("line 12")); + assert.ok(!out.includes("line 13")); + assert.equal(out.split("\n").length, 13); // 12 preview lines + marker + assert.ok( + out.endsWith("… folded 9 lines · full content was sent to the model"), + ); +}); + +test("a message at exactly the character threshold is unchanged", () => { + const message = "x".repeat(1200); + assert.equal(foldUserMessage(message), message); +}); + +test("one char over the char threshold with one line to hide is left alone", () => { + // Over the character threshold, but the fold would hide a single (mid-cut) + // line — far below the minimum-benefit gate, so it renders in full. + const message = "x".repeat(1201); + assert.equal(foldUserMessage(message), message); +}); + +test("long wrapped lines fold by character count, not just line count", () => { + // 20 lines of 61 chars: at the line threshold but 1239 chars overall. + const message = Array.from({ length: 20 }, () => "y".repeat(61)).join("\n"); + const out = foldUserMessage(message); + assert.equal(out.split("\n").length, 13); + assert.ok( + out.endsWith("… folded 8 lines · full content was sent to the model"), + ); +}); + +test("a char-heavy message with only a few foldable lines stays unchanged", () => { + // Like the /openpi-setup prompt: short config-recap lines plus one very + // long guidance line. Over the char threshold, but the fold would only + // hide a handful of lines the user is meant to read. + const shortLines = Array.from( + { length: 15 }, + (_, i) => `config line ${i}: ${"c".repeat(20)}`, + ); + const message = [ + "Configure the installed OpenPI package according to this request:", + ...shortLines, + "g".repeat(700), + ].join("\n"); + assert.ok(message.length > 1200); + assert.ok(message.split("\n").length <= 20); + assert.equal(foldUserMessage(message), message); +}); + +test("folding must hide at least 8 lines to earn its keep", () => { + // 19 lines of 70 chars: char-triggered (1348 chars), fold would hide 7. + const sevenHidden = Array.from({ length: 19 }, () => "z".repeat(70)).join( + "\n", + ); + assert.ok(sevenHidden.length > 1200); + assert.equal(foldUserMessage(sevenHidden), sevenHidden); + + // 20 lines of 70 chars: char-triggered (1419 chars), fold hides exactly 8. + const eightHidden = `${sevenHidden}\n${"z".repeat(70)}`; + const out = foldUserMessage(eightHidden); + assert.equal(out.split("\n").length, 13); + assert.ok( + out.endsWith("… folded 8 lines · full content was sent to the model"), + ); +}); + +test("a single giant fenced block folds to its first content lines", () => { + const out = foldUserMessage(giantBlock); + assert.ok(out.startsWith("```js")); + assert.ok(out.includes("console.log(3);")); + assert.ok(!out.includes("console.log(4);")); + assert.equal(out.split("\n").length, 8); // fence + 4 lines + … + fence + marker + assert.ok( + out.endsWith("… folded 296 lines · full content was sent to the model"), + ); +}); + +test("multiple fenced blocks fold per block while prose shares one budget", () => { + const message = [ + "intro prose", + "```js", + ...Array.from({ length: 10 }, (_, i) => `js-${i};`), + "```", + "middle prose", + "```py", + ...Array.from({ length: 10 }, (_, i) => `py-${i}`), + "```", + "tail prose", + ].join("\n"); + const out = foldUserMessage(message); + assert.ok(out.includes("intro prose")); + assert.ok(out.includes("js-3;")); + assert.ok(!out.includes("js-4;")); + assert.ok(out.includes("middle prose")); + assert.ok(out.includes("py-3")); + assert.ok(!out.includes("py-4")); + assert.ok(out.includes("tail prose")); + assert.ok( + out.endsWith("… folded 12 lines · full content was sent to the model"), + ); +}); + +test("CRLF messages fold without losing their line endings", () => { + const message = Array.from( + { length: 21 }, + (_, i) => `row ${String(i).padStart(2, "0")}`, + ).join("\r\n"); + const out = foldUserMessage(message); + assert.ok(out.includes("row 00\r")); + assert.ok(out.includes("row 11\r")); + assert.ok(!out.includes("row 12")); + assert.ok( + out.endsWith("… folded 9 lines · full content was sent to the model"), + ); +}); + +test("trailing newlines neither fold short messages nor pad the count", () => { + assert.equal(foldUserMessage("hello\n"), "hello\n"); + assert.equal(foldUserMessage("hello\n\n"), "hello\n\n"); + const withTrailing = `${numberedLines(21).join("\n")}\n`; + assert.ok( + foldUserMessage(withTrailing).endsWith( + "… folded 9 lines · full content was sent to the model", + ), + ); +}); + +test("a fenced message under the thresholds is not folded at all", () => { + const message = "explanation\n```js\nfoo();\nbar();\n```\ndone"; + assert.equal(foldUserMessage(message), message); +}); + +test("empty and whitespace-only messages are left alone", () => { + for (const message of ["", " ", "\n", "\t", " \n \n"]) { + assert.equal(foldUserMessage(message), message); + } +}); + +test("an unterminated fence folds conservatively as plain text", () => { + const message = [ + "```js", + ...Array.from({ length: 30 }, (_, i) => `stmt ${i};`), + ].join("\n"); + const out = foldUserMessage(message); + assert.ok(out.startsWith("```js")); + assert.ok(out.includes("stmt 10;")); + assert.ok(!out.includes("stmt 12;")); + assert.ok( + out.endsWith("… folded 19 lines · full content was sent to the model"), + ); +}); + +test("many individually short code blocks share one bounded preview budget", () => { + const block = ["```", "a", "b", "c", "d", "```"].join("\n"); + const message = Array.from({ length: 5 }, () => block).join("\n"); + const out = foldUserMessage(message); + assert.notEqual(out, message); + assert.ok(out.split("\n").length <= 21); + assert.equal((out.match(/```/g) ?? []).length % 2, 0); + assert.ok( + out.endsWith("… folded 12 lines · full content was sent to the model"), + ); +}); + +test("many truncated code blocks cannot make the folded output longer than the input", () => { + const block = ["```js", "a", "b", "c", "d", "e", "```"].join("\n"); + const message = Array.from({ length: 100 }, () => block).join("\n"); + const out = foldUserMessage(message); + assert.equal(message.split("\n").length, 700); + assert.ok(out.split("\n").length <= 21); + assert.ok(out.length < message.length); + assert.equal((out.match(/```/g) ?? []).length % 2, 0); + assert.ok( + out.endsWith("… folded 683 lines · full content was sent to the model"), + ); +}); + +test("folding is pure: the input is untouched and repeat calls agree", () => { + const snapshot = giantBlock; + const first = foldUserMessage(giantBlock); + const second = foldUserMessage(giantBlock); + assert.equal(giantBlock, snapshot); // no mutation of the input + assert.equal(first, second); // deterministic + assert.notEqual(first, giantBlock); // but something was folded + assert.ok(first.length < giantBlock.length); +}); + +test("the transformer only folds finalized user messages", () => { + const message = numberedLines(21).join("\n"); + assert.notEqual( + transformUserMarkdown(message, { messageType: "user", isStreaming: false }), + message, + ); + const untouched = [ + { messageType: "assistant", isStreaming: false }, + { messageType: "assistant-thinking", isStreaming: false }, + { messageType: "assistant", isStreaming: true }, + { messageType: "user", isStreaming: true }, + ]; + for (const context of untouched) { + assert.equal(transformUserMarkdown(message, context), message); + } +}); + +test("the marker always states that the model received the full content", () => { + const out = foldUserMessage(giantBlock); + assert.ok(out.includes(MARKER)); + assert.ok(out.includes("full content was sent to the model")); +}); diff --git a/extensions/user-input-fold/index.ts b/extensions/user-input-fold/index.ts new file mode 100644 index 00000000..3edbafa7 --- /dev/null +++ b/extensions/user-input-fold/index.ts @@ -0,0 +1,185 @@ +/** + * Fold very long user messages in the transcript display. + * + * A pasted log, stack trace, or whole file can wipe out several screens of + * chat. When a finalized user message exceeds the fold thresholds, the + * display keeps a short preview of the prose and of each fenced code block + * and closes with a marker line stating how much was folded. + * + * A fold that would hide nothing (or almost nothing) is skipped and the + * message rendered in full: hiding a handful of lines costs a marker row + * while saving almost no screen space, and char-heavy messages with few + * long lines do not shrink on screen when logical lines are removed — + * folding must hide at least MIN_FOLDED_LINES lines to earn its keep. + * + * This is display-only. The transformer runs through Pi's + * `registerMarkdownTransformer` hook, which changes only what the TUI + * renders: the session file and the model context keep the full message + * untouched, so the model always receives the complete paste. The pure + * `foldUserMessage` helper never mutates its input and has no side effects. + */ + +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; + +/** Fold a user message longer than this many lines. */ +const MAX_LINES = 20; +/** Fold a user message longer than this many characters. */ +const MAX_CHARS = 1_200; +/** Prose lines kept when a message is folded. */ +const PROSE_PREVIEW_LINES = 12; +/** Content lines kept per fenced code block when a message is folded. */ +const BLOCK_PREVIEW_LINES = 4; +/** Character budget for the prose part of a folded message. */ +const PROSE_PREVIEW_CHARS = 1_200; +/** Total rendered lines before the fold marker, across prose and code blocks. */ +const MAX_PREVIEW_LINES = 20; +/** Only fold when at least this many lines would be hidden. */ +const MIN_FOLDED_LINES = 8; + +type Segment = + | { kind: "prose"; lines: string[] } + | { kind: "code"; open: string; content: string[]; close: string }; + +const FENCE_OPEN = /^ {0,3}`{3,}/; +const FENCE_CLOSE = /^ {0,3}`{3,}[ \t]*$/; + +function countLines(markdown: string) { + const parts = markdown.split("\n"); + // A trailing newline ends the last line; it does not open a new one. + return parts.at(-1) === "" ? parts.length - 1 : parts.length; +} + +function splitLines(markdown: string) { + const lines = markdown.split("\n"); + if (lines.at(-1) === "") lines.pop(); + return lines; +} + +function parseSegments(lines: string[]): Segment[] { + const segments: Segment[] = []; + let prose: string[] = []; + let i = 0; + while (i < lines.length) { + if (!FENCE_OPEN.test(lines[i])) { + prose.push(lines[i]); + i += 1; + continue; + } + if (prose.length > 0) { + segments.push({ kind: "prose", lines: prose }); + prose = []; + } + const open = lines[i]; + const content: string[] = []; + let close: string | undefined; + let j = i + 1; + while (j < lines.length && close === undefined) { + if (FENCE_CLOSE.test(lines[j])) close = lines[j]; + else content.push(lines[j]); + j += 1; + } + if (close === undefined) { + // Unterminated fence: fold the whole message conservatively as text. + return [{ kind: "prose", lines }]; + } + segments.push({ kind: "code", open, content, close }); + i = j; + } + if (prose.length > 0) segments.push({ kind: "prose", lines: prose }); + return segments; +} + +/** + * Return the Markdown Pi should render instead of a long user message. + * Messages at or below both thresholds — and messages whose fold would + * hide fewer than MIN_FOLDED_LINES lines — are returned unchanged. Pure: + * the input string is never modified, and the model still sees the original. + */ +export function foldUserMessage(markdown: string): string { + const totalLines = countLines(markdown); + if (totalLines <= MAX_LINES && markdown.length <= MAX_CHARS) return markdown; + + const preview: string[] = []; + let proseLinesLeft = PROSE_PREVIEW_LINES; + let proseCharsLeft = PROSE_PREVIEW_CHARS; + let linesShown = 0; + + for (const segment of parseSegments(splitLines(markdown))) { + const remainingLines = MAX_PREVIEW_LINES - preview.length; + if (remainingLines <= 0) break; + if (segment.kind === "code") { + if (segment.content.length === 0) { + if (remainingLines < 2) break; + preview.push(segment.open, segment.close); + linesShown += 2; + continue; + } + // A partial block needs opening/closing fences plus an ellipsis. If that + // cannot fit, stop before the block instead of emitting broken Markdown. + if (remainingLines < 3) break; + let shown = Math.min( + segment.content.length, + BLOCK_PREVIEW_LINES, + remainingLines - 2, + ); + const truncated = () => shown < segment.content.length; + while (truncated() && shown + 3 > remainingLines) shown -= 1; + preview.push(segment.open, ...segment.content.slice(0, shown)); + if (truncated()) preview.push("…"); + preview.push(segment.close); + linesShown += 2 + shown; + continue; + } + for (const line of segment.lines) { + if (proseLinesLeft <= 0 || preview.length >= MAX_PREVIEW_LINES) break; + const cost = line.length + (preview.length > 0 ? 1 : 0); + if (preview.length > 0 && cost > proseCharsLeft) break; + if (preview.length === 0 && line.length > proseCharsLeft) { + // A single giant first line is the only case that cuts mid-line. + preview.push(`${line.slice(0, proseCharsLeft)}…`); + proseLinesLeft = 0; + proseCharsLeft = 0; + break; + } + preview.push(line); + linesShown += 1; + proseLinesLeft -= 1; + proseCharsLeft -= cost; + } + } + + const foldedLines = totalLines - linesShown; + if (foldedLines < MIN_FOLDED_LINES) { + // Folding must earn its keep. Hiding fewer lines than MIN_FOLDED_LINES + // costs a marker row, hides content, and saves almost nothing on screen + // (char-heavy messages with few long lines wrap regardless), so render + // the message in full instead. + return markdown; + } + const noun = foldedLines === 1 ? "line" : "lines"; + preview.push( + `… folded ${foldedLines} ${noun} · full content was sent to the model`, + ); + return preview.join("\n"); +} + +/** + * The registered transformer: fold only finalized user messages and leave + * assistant text, thinking blocks, and streaming updates untouched. + */ +export function transformUserMarkdown( + markdown: string, + context: { messageType: string; isStreaming: boolean }, +): string { + if (context.messageType !== "user" || context.isStreaming) return markdown; + try { + return foldUserMessage(markdown); + } catch { + // Display-only: a folding bug must never break rendering. + return markdown; + } +} + +export default function (pi: ExtensionAPI) { + pi.registerMarkdownTransformer(transformUserMarkdown); +}