Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
20 kB · 628 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629import type { AgentToolResult } from "@earendil-works/pi-agent-core";import { type Component, Container, Image, Text, type TUI } from "@earendil-works/pi-tui";import type { ToolDefinition, ToolRenderContext, ToolRenderResultOptions } from "../../../core/extensions/types.js";import type { KernelSentAgentMessage } from "../../../core/kernel/index.js";import { createBashToolDefinition } from "../../../core/tools/bash.js";import { createEditToolDefinition } from "../../../core/tools/edit.js";import { createAllToolDefinitions } from "../../../core/tools/index.js";import { formatToolDuration, getTextOutput as getRenderedTextOutput } from "../../../core/tools/render-utils.js";import type { AgentConnectionToolDefinition } from "../../agent-connection/index.js";import { type Theme, theme } from "../theme/theme.js";import { getWorkingPulseFrame, workingIconFrame } from "../theme/working-icon.js";import { getIpythonCodeFromArgs, IPythonCellComponent } from "./ipython-cell.js";import { expandCollapseHint } from "./keybinding-hints.js";import { type BackgroundShellHandle, readAssignedShellCommand, readBackgroundShellHandle, type ShellCompletion,} from "./shell-completion.js";import { ToolPanel } from "./tool-panel.js";
export interface ToolExecutionOptions { shouldAddLeadingSpace?: () => boolean; showImages?: boolean; /** Whether image metadata may parse dimensions from base64 data. */ includeImageDimensions?: boolean;}
export interface ToolExecutionRendererDefinition { renderShell?: "default" | "self"; renderCall?: (args: any, theme: Theme, context: ToolRenderContext<any, any>) => Component; renderResult?: ( result: AgentToolResult<any>, options: ToolRenderResultOptions, theme: Theme, context: ToolRenderContext<any, any>, ) => Component;}export type ToolExecutionDefinition = AgentConnectionToolDefinition & Partial<ToolExecutionRendererDefinition>;
function hasToolRenderer(toolDefinition: ToolExecutionDefinition | undefined): boolean { return toolDefinition?.renderCall !== undefined || toolDefinition?.renderResult !== undefined;}
function matchesBuiltInReplayMetadata(toolName: string, toolDefinition: ToolExecutionDefinition | undefined): boolean { if (!toolDefinition) { return true; } if (hasToolRenderer(toolDefinition)) { return false; } return toolDefinition.replayBuiltInToolName === toolName;}
function createReplayBuiltInToolDefinition( toolName: string, cwd: string, toolDefinition: ToolExecutionDefinition | undefined,): ToolDefinition<any, any> | undefined { if (toolName === "ipython") { return createAllToolDefinitions(cwd).ipython; } switch (toolName) { case "bash": { const builtInDefinition = createBashToolDefinition(cwd); return matchesBuiltInReplayMetadata(toolName, toolDefinition) ? builtInDefinition : undefined; } case "edit": { const builtInDefinition = createEditToolDefinition(cwd); return matchesBuiltInReplayMetadata(toolName, toolDefinition) ? builtInDefinition : undefined; } default: return undefined; }}
export class ToolExecutionComponent extends Container { private contentPanel: ToolPanel; private selfRenderContainer: Container; private callRendererComponent?: Component; private resultRendererComponent?: Component; private ipythonCellComponent?: IPythonCellComponent; private rendererState: any = {}; private imageComponents: Image[] = []; private toolName: string; private toolCallId: string; private args: any; private expanded = false; private agentMessagesExpanded = false; private editDiffsExpanded = false; private showExpandHint = true; private showImages: boolean; private includeImageDimensions: boolean; private readonly shouldAddLeadingSpace?: () => boolean; private isPartial = true; private toolDefinition?: ToolExecutionDefinition; private builtInToolDefinition?: ToolDefinition<any, any>; private ui: TUI; private cwd: string; private executionStarted = false; private argsComplete = false; private pendingSentAgentMessages: KernelSentAgentMessage[] = []; private result?: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; isError: boolean; details?: any; durationMs?: number; }; private startedAtMs?: number; private resultDurationMs?: number; private hideComponent = false; private shellCompletion?: ShellCompletion; private shellCompletionAmbiguous = false; private readonly resultListeners = new Set<() => void>();
constructor( toolName: string, toolCallId: string, args: any, options: ToolExecutionOptions = {}, toolDefinition: ToolExecutionDefinition | undefined, ui: TUI, cwd: string, ) { super(); this.toolName = toolName; this.toolCallId = toolCallId; this.args = args; this.toolDefinition = toolDefinition; this.builtInToolDefinition = createReplayBuiltInToolDefinition(toolName, cwd, toolDefinition); this.showImages = options.showImages ?? true; this.includeImageDimensions = options.includeImageDimensions ?? true; this.shouldAddLeadingSpace = options.shouldAddLeadingSpace; this.ui = ui; this.cwd = cwd;
// Always create both shell variants. contentPanel is the tool panel used // for default renderer-based composition (and the generic fallback when no // tool definition exists). selfRenderContainer is used when the tool // renders its own framing. this.contentPanel = new ToolPanel(); this.selfRenderContainer = new Container();
if (this.hasRendererDefinition() && this.getRenderShell() === "self") { this.addChild(this.selfRenderContainer); } else { this.addChild(this.contentPanel); }
this.updateDisplay(); }
private getCallRenderer(): ToolDefinition<any, any>["renderCall"] | undefined { if (!this.builtInToolDefinition) { return this.toolDefinition?.renderCall; } if (!this.toolDefinition) { return this.builtInToolDefinition.renderCall; } return this.toolDefinition.renderCall ?? this.builtInToolDefinition.renderCall; }
private getResultRenderer(): ToolDefinition<any, any>["renderResult"] | undefined { if (!this.builtInToolDefinition) { return this.toolDefinition?.renderResult; } if (!this.toolDefinition) { return this.builtInToolDefinition.renderResult; } return this.toolDefinition.renderResult ?? this.builtInToolDefinition.renderResult; }
private hasRendererDefinition(): boolean { return this.builtInToolDefinition !== undefined || this.toolDefinition !== undefined; }
private getRenderShell(): "default" | "self" { if (this.shouldUseIpythonRenderer()) { return "self"; } if (!this.builtInToolDefinition) { return this.toolDefinition?.renderShell ?? "default"; } if (!this.toolDefinition) { return this.builtInToolDefinition.renderShell ?? "default"; } return this.toolDefinition.renderShell ?? this.builtInToolDefinition.renderShell ?? "default"; }
private shouldUseIpythonRenderer(): boolean { return this.toolName === "ipython" && !this.toolDefinition?.renderCall && !this.toolDefinition?.renderResult; }
private isBuiltInEditTool(): boolean { return ( this.toolName === "edit" && (this.toolDefinition === undefined || this.toolDefinition.replayBuiltInToolName === "edit") ); }
private getRenderContext(lastComponent: Component | undefined): ToolRenderContext { return { args: this.args, toolCallId: this.toolCallId, invalidate: () => { this.invalidate(); this.ui.requestRender(); }, lastComponent, state: this.rendererState, cwd: this.cwd, executionStarted: this.executionStarted, argsComplete: this.argsComplete, isPartial: this.isPartial, expanded: this.isBuiltInEditTool() ? this.editDiffsExpanded : this.expanded, showExpandHint: this.showExpandHint, showImages: this.showImages, includeImageDimensions: this.includeImageDimensions, isError: this.result?.isError ?? false, }; }
private createCallFallback(): Component { return new Text(theme.fg("toolTitle", this.toolName), 0, 0); }
private createResultFallback(): Component | undefined { const output = this.getTextOutput(); if (!output) { return undefined; } return new Text(theme.fg("toolOutput", this.formatFallbackPreview(output)), 0, 0); }
updateArgs(args: any): void { this.args = args; this.updateDisplay(); }
markExecutionStarted(): void { this.executionStarted = true; this.startedAtMs ??= Date.now(); this.updateDisplay(); this.ui.requestRender(); }
setArgsComplete(): void { this.argsComplete = true; this.updateDisplay(); this.ui.requestRender(); }
updateResult( result: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; details?: any; isError: boolean; durationMs?: number; }, isPartial = false, ): void { if (typeof result.durationMs === "number") { this.resultDurationMs = result.durationMs; } const details = typeof result.details === "object" && result.details !== null ? (result.details as Record<string, unknown>) : {}; const sentAgentMessages = Array.isArray(details.sentAgentMessages) ? [...details.sentAgentMessages] : []; for (const message of this.pendingSentAgentMessages) { if ( !sentAgentMessages.some( (entry) => typeof entry === "object" && entry !== null && "id" in entry && entry.id === message.id, ) ) { sentAgentMessages.push(message); } } this.result = sentAgentMessages.length > 0 ? { ...result, details: { ...details, sentAgentMessages } } : result; this.isPartial = isPartial; this.updateDisplay(); for (const listener of this.resultListeners) listener(); }
getBackgroundShellHandle(): BackgroundShellHandle | undefined { return this.shouldUseIpythonRenderer() && !this.isPartial && !this.result?.isError ? readBackgroundShellHandle(getIpythonCodeFromArgs(this.args), this.result?.details) : undefined; }
getAssignedShellCommand(): string | undefined { return this.shouldUseIpythonRenderer() && !this.isPartial && !this.result?.isError ? readAssignedShellCommand(getIpythonCodeFromArgs(this.args), this.result?.details) : undefined; }
hasRunningBackgroundShell(): boolean { const handle = this.getBackgroundShellHandle(); return ( handle !== undefined && handle.exitCode === undefined && this.shellCompletion === undefined && !this.shellCompletionAmbiguous ); }
markShellCompletionAmbiguous(): void { if (!this.hasRunningBackgroundShell()) return; this.shellCompletionAmbiguous = true; this.updateDisplay(); this.ui.requestRender(); }
onResultUpdate(listener: () => void): () => void { this.resultListeners.add(listener); return () => this.resultListeners.delete(listener); }
isResultPending(): boolean { return this.isPartial; }
attachShellCompletion(completion: ShellCompletion): boolean { if (this.shellCompletion) return false; this.shellCompletion = completion; this.shellCompletionAmbiguous = false; this.updateDisplay(); this.ui.requestRender(); return true; }
appendSentAgentMessage(message: KernelSentAgentMessage): void { if (this.pendingSentAgentMessages.some((entry) => entry.id === message.id)) { return; } this.pendingSentAgentMessages.push(message); if (this.result) { this.updateResult(this.result, this.isPartial); } }
setExpanded(expanded: boolean): void { this.expanded = expanded; this.updateDisplay(); }
setAgentMessagesExpanded(expanded: boolean): void { if (this.agentMessagesExpanded === expanded) { return; } this.agentMessagesExpanded = expanded; this.updateDisplay(); }
setEditDiffsExpanded(expanded: boolean): void { if (this.editDiffsExpanded === expanded) { return; } this.editDiffsExpanded = expanded; this.updateDisplay(); }
setShowExpandHint(show: boolean): void { if (this.showExpandHint === show) { return; } this.showExpandHint = show; this.updateDisplay(); }
setShowImages(show: boolean): void { this.showImages = show; this.updateDisplay(); }
setIncludeImageDimensions(include: boolean): void { this.includeImageDimensions = include; this.updateDisplay(); }
override invalidate(): void { super.invalidate(); this.updateDisplay(); }
override render(width: number): string[] { if (this.hideComponent) { return []; } // Refresh the animated glyph without rebuilding the whole panel, for as long // as panelStatus() is still animating (including partial streaming results). if (this.isStatusAnimating() && !this.usesSelfRenderShell()) { this.contentPanel.setHeader(this.panelHeader()); } const lines = super.render(width); return this.expanded && this.shouldUseIpythonRenderer() && this.shouldAddLeadingSpace?.() ? ["", ...lines] : lines; }
private isStatusAnimating(): boolean { if (!this.executionStarted) { return false; } // Matches panelStatus(): animating until a non-partial or error result lands. if (this.result && !this.isPartial) { return false; } return !this.result?.isError; }
private usesSelfRenderShell(): boolean { return this.hasRendererDefinition() && this.getRenderShell() === "self"; }
private updateDisplay(): void { let hasContent = false; this.hideComponent = false; if (this.hasRendererDefinition() && this.getRenderShell() === "self") { this.selfRenderContainer.clear();
if (this.shouldUseIpythonRenderer()) { const state = { code: getIpythonCodeFromArgs(this.args), backgroundShell: this.getBackgroundShellHandle(), shellCompletion: this.shellCompletion, shellCompletionAmbiguous: this.shellCompletionAmbiguous, content: this.result?.content, details: this.result?.details, isPartial: this.isPartial, isError: this.result?.isError ?? false, expanded: this.expanded, agentMessagesExpanded: this.agentMessagesExpanded, editDiffsExpanded: this.editDiffsExpanded, executionStarted: this.executionStarted, argsComplete: this.argsComplete, showExpandHint: this.showExpandHint, showImages: this.showImages, cwd: this.cwd, }; if (!this.ipythonCellComponent) { this.ipythonCellComponent = new IPythonCellComponent(state); } else { this.ipythonCellComponent.update(state); } this.selfRenderContainer.addChild(this.ipythonCellComponent); hasContent = true; } else { hasContent = this.mountRenderers(this.selfRenderContainer, true); } } else { // Default shell: tool panel with a `label · status` header so the block // is self-identifying. The header replaces the bold-tool-name fallback. this.contentPanel.setHeader(this.panelHeader()); this.contentPanel.clear(); if (this.hasRendererDefinition()) { this.mountRenderers(this.contentPanel, false); } else { const fallbackText = this.formatToolExecution(); if (fallbackText) { this.contentPanel.addChild(new Text(fallbackText, 0, 0)); } } hasContent = true; }
for (const img of this.imageComponents) { this.removeChild(img); } this.imageComponents = [];
if (this.result) { const imageBlocks = this.result.content.filter((c) => c.type === "image"); for (let i = 0; i < imageBlocks.length; i++) { const img = imageBlocks[i]; if (!this.showImages || !img.data || !img.mimeType) continue;
const imageComponent = new Image( img.data, img.mimeType, { fallbackColor: (s: string) => theme.fg("toolOutput", s) }, { fallbackOnly: true, fallbackPrefix: " ╰─ ", }, ); this.imageComponents.push(imageComponent); this.addChild(imageComponent); } }
if (this.hasRendererDefinition() && !hasContent && this.imageComponents.length === 0) { this.hideComponent = true; } }
/** * Mount the call/result renderer components into the given shell container. * `useFallbacks` keeps the bold-tool-name call fallback for self-rendering * tools; the default panel shell already names the tool in its header. */ private mountRenderers(container: Container | ToolPanel, useFallbacks: boolean): boolean { let hasContent = false;
const callRenderer = this.getCallRenderer(); if (!callRenderer) { if (useFallbacks) { container.addChild(this.createCallFallback()); hasContent = true; } } else { try { const component = callRenderer(this.args, theme, this.getRenderContext(this.callRendererComponent)); this.callRendererComponent = component; container.addChild(component); hasContent = true; } catch { this.callRendererComponent = undefined; if (useFallbacks) { container.addChild(this.createCallFallback()); hasContent = true; } } }
if (this.result) { const resultRenderer = this.getResultRenderer(); if (!resultRenderer) { const component = this.createResultFallback(); if (component) { container.addChild(component); hasContent = true; } } else { try { const component = resultRenderer( { content: this.result.content as any, details: this.result.details }, { expanded: this.expanded, isPartial: this.isPartial }, theme, this.getRenderContext(this.resultRendererComponent), ); this.resultRendererComponent = component; container.addChild(component); hasContent = true; } catch { this.resultRendererComponent = undefined; const component = this.createResultFallback(); if (component) { container.addChild(component); hasContent = true; } } } }
return hasContent; }
private panelHeader(): string { const label = this.toolDefinition?.label ?? this.builtInToolDefinition?.label ?? this.toolName; return `${theme.fg("muted", label)}${theme.fg("dim", " · ")}${this.panelStatus()}`; }
private durationSuffix(durationMs: number | undefined): string { if (durationMs === undefined) { return ""; } return `${theme.fg("dim", " · ")}${theme.fg("muted", formatToolDuration(durationMs))}`; }
private panelStatus(): string { if (this.result && !this.isPartial) { // Only the runtime-measured duration is shown for finished calls; the // render-process clock is meaningless for replayed sessions. const status = this.result.isError ? theme.fg("error", "error") : theme.fg("success", "done"); return status + this.durationSuffix(this.resultDurationMs); } if (this.result?.isError) { return theme.fg("error", "error") + this.durationSuffix(this.resultDurationMs); } if (this.executionStarted) { const elapsedMs = this.startedAtMs === undefined ? undefined : Date.now() - this.startedAtMs; const running = theme.fg("bashMode", `${workingIconFrame(getWorkingPulseFrame())} running`); return running + this.durationSuffix(elapsedMs); } return theme.fg("muted", "queued"); }
private getTextOutput(): string { return getRenderedTextOutput(this.result, this.showImages, { includeImageDimensions: this.includeImageDimensions, }); }
private formatFallbackPreview(text: string): string { if (this.expanded) return text; const lines = text.split("\n"); if (lines.length <= 3) return text; return `${lines.slice(0, 3).join("\n")}\n${theme.fg("dim", `… ${lines.length - 3} more lines`)}${this.showExpandHint ? ` ${expandCollapseHint("app.tools.expand", false)}` : ""}`; }
private formatToolExecution(): string { const parts: string[] = []; const content = JSON.stringify(this.args, null, 2); if (content) { parts.push(this.formatFallbackPreview(content)); } const output = this.getTextOutput(); if (output) { parts.push(this.formatFallbackPreview(output)); } return parts.join("\n\n"); }}
export function selectLatestToolExpandHint( existingComponents: readonly Component[], latestComponent: ToolExecutionComponent,): void { for (let index = existingComponents.length - 1; index >= 0; index--) { const component = existingComponents[index]; if (component instanceof ToolExecutionComponent) { component.setShowExpandHint(false); break; } } latestComponent.setShowExpandHint(true);}