diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 2ba164d6..3fa4409d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,7 @@ - Fixed `.pi` folder being created unnecessarily when only reading settings. The folder is now only created when writing project-specific settings. - Fixed extension-driven runtime theme changes to persist in settings so `/settings` reflects the active `currentTheme` after `ctx.ui.setTheme(...)` ([#1483](https://github.com/badlogic/pi-mono/pull/1483) by [@ferologics](https://github.com/ferologics)) +- Fixed interactive mode freezes during large streaming `write` tool calls by using incremental syntax highlighting while partial arguments stream, with a final full re-highlight after tool-call arguments complete. ## [0.54.1] - 2026-02-22 diff --git a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts index 6678e0d1..85de9fc2 100644 --- a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts +++ b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts @@ -25,6 +25,9 @@ import { truncateToVisualLines } from "./visual-truncate.js"; // Preview line limit for bash when not expanded const BASH_PREVIEW_LINES = 5; +// During partial write tool-call streaming, re-highlight the first N lines fully +// to keep multiline tokenization mostly correct without re-highlighting the full file. +const WRITE_PARTIAL_FULL_HIGHLIGHT_LINES = 50; /** * Convert absolute path to tilde notation if it's in home directory @@ -56,6 +59,14 @@ export interface ToolExecutionOptions { showImages?: boolean; // default: true (only used if terminal supports images) } +type WriteHighlightCache = { + rawPath: string | null; + lang: string; + rawContent: string; + normalizedLines: string[]; + highlightedLines: string[]; +}; + /** * Component that renders a tool call with its result (updateable) */ @@ -82,6 +93,8 @@ export class ToolExecutionComponent extends Container { private editDiffArgsKey?: string; // Track which args the preview is for // Cached converted images for Kitty protocol (which requires PNG), keyed by index private convertedImages: Map = new Map(); + // Incremental syntax highlighting cache for write tool call args + private writeHighlightCache?: WriteHighlightCache; constructor( toolName: string, @@ -129,14 +142,114 @@ export class ToolExecutionComponent extends Container { updateArgs(args: any): void { this.args = args; + if (this.toolName === "write" && this.isPartial) { + this.updateWriteHighlightCacheIncremental(); + } this.updateDisplay(); } + private highlightSingleLine(line: string, lang: string): string { + const highlighted = highlightCode(line, lang); + return highlighted[0] ?? ""; + } + + private refreshWriteHighlightPrefix(cache: WriteHighlightCache): void { + const prefixCount = Math.min(WRITE_PARTIAL_FULL_HIGHLIGHT_LINES, cache.normalizedLines.length); + if (prefixCount === 0) return; + + const prefixSource = cache.normalizedLines.slice(0, prefixCount).join("\n"); + const prefixHighlighted = highlightCode(prefixSource, cache.lang); + for (let i = 0; i < prefixCount; i++) { + cache.highlightedLines[i] = + prefixHighlighted[i] ?? this.highlightSingleLine(cache.normalizedLines[i] ?? "", cache.lang); + } + } + + private rebuildWriteHighlightCacheFull(rawPath: string | null, fileContent: string): void { + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + if (!lang) { + this.writeHighlightCache = undefined; + return; + } + + const normalized = replaceTabs(fileContent); + this.writeHighlightCache = { + rawPath, + lang, + rawContent: fileContent, + normalizedLines: normalized.split("\n"), + highlightedLines: highlightCode(normalized, lang), + }; + } + + private updateWriteHighlightCacheIncremental(): void { + const rawPath = str(this.args?.file_path ?? this.args?.path); + const fileContent = str(this.args?.content); + if (rawPath === null || fileContent === null) { + this.writeHighlightCache = undefined; + return; + } + + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + if (!lang) { + this.writeHighlightCache = undefined; + return; + } + + if (!this.writeHighlightCache) { + this.rebuildWriteHighlightCacheFull(rawPath, fileContent); + return; + } + + const cache = this.writeHighlightCache; + if (cache.lang !== lang || cache.rawPath !== rawPath) { + this.rebuildWriteHighlightCacheFull(rawPath, fileContent); + return; + } + + if (!fileContent.startsWith(cache.rawContent)) { + this.rebuildWriteHighlightCacheFull(rawPath, fileContent); + return; + } + + if (fileContent.length === cache.rawContent.length) { + return; + } + + const deltaRaw = fileContent.slice(cache.rawContent.length); + const deltaNormalized = replaceTabs(deltaRaw); + cache.rawContent = fileContent; + + if (cache.normalizedLines.length === 0) { + cache.normalizedLines.push(""); + cache.highlightedLines.push(""); + } + + const segments = deltaNormalized.split("\n"); + const lastIndex = cache.normalizedLines.length - 1; + cache.normalizedLines[lastIndex] += segments[0]; + cache.highlightedLines[lastIndex] = this.highlightSingleLine(cache.normalizedLines[lastIndex], cache.lang); + + for (let i = 1; i < segments.length; i++) { + cache.normalizedLines.push(segments[i]); + cache.highlightedLines.push(this.highlightSingleLine(segments[i], cache.lang)); + } + + this.refreshWriteHighlightPrefix(cache); + } + /** * Signal that args are complete (tool is about to execute). * This triggers diff computation for edit tool. */ setArgsComplete(): void { + if (this.toolName === "write") { + const rawPath = str(this.args?.file_path ?? this.args?.path); + const fileContent = str(this.args?.content); + if (rawPath !== null && fileContent !== null) { + this.rebuildWriteHighlightCacheFull(rawPath, fileContent); + } + } this.maybeComputeEditDiff(); } @@ -183,6 +296,13 @@ export class ToolExecutionComponent extends Container { ): void { this.result = result; this.isPartial = isPartial; + if (this.toolName === "write" && !isPartial) { + const rawPath = str(this.args?.file_path ?? this.args?.path); + const fileContent = str(this.args?.content); + if (rawPath !== null && fileContent !== null) { + this.rebuildWriteHighlightCacheFull(rawPath, fileContent); + } + } this.updateDisplay(); // Convert non-PNG images to PNG for Kitty protocol (async) this.maybeConvertImagesForKitty(); @@ -532,7 +652,28 @@ export class ToolExecutionComponent extends Container { text += `\n\n${theme.fg("error", "[invalid content arg - expected string]")}`; } else if (fileContent) { const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; - const lines = lang ? highlightCode(replaceTabs(fileContent), lang) : fileContent.split("\n"); + + let lines: string[]; + if (lang) { + const cache = this.writeHighlightCache; + if (cache && cache.lang === lang && cache.rawPath === rawPath && cache.rawContent === fileContent) { + lines = cache.highlightedLines; + } else { + const normalized = replaceTabs(fileContent); + lines = highlightCode(normalized, lang); + this.writeHighlightCache = { + rawPath, + lang, + rawContent: fileContent, + normalizedLines: normalized.split("\n"), + highlightedLines: lines, + }; + } + } else { + lines = fileContent.split("\n"); + this.writeHighlightCache = undefined; + } + const totalLines = lines.length; const maxLines = this.expanded ? lines.length : 10; const displayLines = lines.slice(0, maxLines); @@ -540,9 +681,7 @@ export class ToolExecutionComponent extends Container { text += "\n\n" + - displayLines - .map((line: string) => (lang ? replaceTabs(line) : theme.fg("toolOutput", replaceTabs(line)))) - .join("\n"); + displayLines.map((line: string) => (lang ? line : theme.fg("toolOutput", replaceTabs(line)))).join("\n"); if (remaining > 0) { text += theme.fg("muted", `\n... (${remaining} more lines, ${totalLines} total,`) +