From c2c81e9bc2fbf6e495484b8bfa08e8feaafb8166 Mon Sep 17 00:00:00 2001 From: quick Date: Fri, 7 Aug 2026 11:40:42 -0400 Subject: [PATCH 1/4] feat(coding-agent): add turn boundary renderer hook --- packages/coding-agent/docs/extensions.md | 22 +++++ .../extensions/turn-boundary-renderer.ts | 9 ++ .../coding-agent/src/core/agent-session.ts | 72 +++++++++++--- .../coding-agent/src/core/extensions/index.ts | 3 + .../src/core/extensions/loader.ts | 7 ++ .../src/core/extensions/runner.ts | 7 ++ .../coding-agent/src/core/extensions/types.ts | 18 ++++ packages/coding-agent/src/core/index.ts | 3 + packages/coding-agent/src/index.ts | 3 + .../components/turn-boundary-renderer.ts | 54 ++++++++++ .../src/modes/interactive/interactive-mode.ts | 28 +++++- .../test/extensions-runner.test.ts | 6 ++ .../interactive-mode-turn-boundary.test.ts | 99 +++++++++++++++++++ .../test/suite/agent-session-prompt.test.ts | 23 ++++- .../test/turn-boundary-renderer.test.ts | 96 ++++++++++++++++++ 15 files changed, 432 insertions(+), 18 deletions(-) create mode 100644 packages/coding-agent/examples/extensions/turn-boundary-renderer.ts create mode 100644 packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts create mode 100644 packages/coding-agent/test/interactive-mode-turn-boundary.test.ts create mode 100644 packages/coding-agent/test/turn-boundary-renderer.test.ts diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index a17fbe03be3..481baf0e2fe 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1667,6 +1667,25 @@ The hook is display-only. It does not modify persisted messages or model context See [built-in-message-renderer.ts](../examples/extensions/built-in-message-renderer.ts) for a complete wrapping example. +### pi.registerTurnBoundaryRenderer(transform) + +Wrap the visual boundary Pi inserts before each user turn after the first. Without an override, the boundary is Pi's existing blank one-line spacer. The transform receives the current renderer and returns another synchronous renderer; multiple extensions compose in load order, with later transforms wrapping earlier ones. + +The renderer receives `{ message, source, isReplay }` and the current theme. `message` is an isolated snapshot. `source` is `"interactive"`, `"rpc"`, or `"extension"` for live prompts that enter through the corresponding input path. It is undefined for restored messages and direct `steer()` or `followUp()` calls whose origin is not known. `isReplay` is true while Pi reconstructs a saved transcript. + +Pi does not invoke the hook for the first rendered user message because there is no preceding turn to separate. Normal and skill-invocation user turns share the same outer boundary. If a transform throws or returns an invalid component, Pi falls back to the previous renderer layer. + +```typescript +import { Text } from "@earendil-works/pi-tui"; + +pi.registerTurnBoundaryRenderer((_current) => (context, theme) => { + const source = context.isReplay ? "replay" : (context.source ?? "unknown"); + return new Text(theme.fg("dim", `── turn · ${source} ──`), 0, 0); +}); +``` + +See [turn-boundary-renderer.ts](../examples/extensions/turn-boundary-renderer.ts) for a complete example. + ### pi.registerMarkdownTransformer(transformer) Register a transformer for the Markdown in normal user text, assistant text, and thinking blocks. Transformers run in extension load order, and each transformer receives the Markdown returned by the previous transformer. After the chain finishes, Pi renders the transformed content with its built-in renderer. @@ -2937,6 +2956,8 @@ See [tui.md](tui.md) Pattern 7 for a complete example with mode indicator. Use `pi.registerBuiltInMessageRenderer()` to wrap the native user or assistant transcript component without changing the message stored in the session or sent to the model. The current renderer can be delegated to, padded differently, or placed inside a component-owned card. See [built-in-message-renderer.ts](../examples/extensions/built-in-message-renderer.ts). +Use `pi.registerTurnBoundaryRenderer()` to replace or wrap the blank spacer before each user turn after the first. See [turn-boundary-renderer.ts](../examples/extensions/turn-boundary-renderer.ts). + Register a custom renderer for messages with your `customType`. Use message renderers for content that should participate in LLM context: ```typescript @@ -3073,6 +3094,7 @@ All examples in [examples/extensions/](../examples/extensions/). | `status-line.ts` | Footer status indicator | `setStatus`, session events | | `working-indicator.ts` | Customize the streaming working indicator | `setWorkingIndicator`, `registerCommand` | | `built-in-message-renderer.ts` | Wrap native user and assistant transcript cards | `registerBuiltInMessageRenderer` | +| `turn-boundary-renderer.ts` | Replace the boundary before user turns | `registerTurnBoundaryRenderer` | | `github-issue-autocomplete.ts` | Add `#1234` issue completions on top of built-in autocomplete by preloading recent open issues from `gh issue list` | `addAutocompleteProvider`, `on("session_start")`, `exec` | | `custom-footer.ts` | Replace footer entirely | `registerCommand`, `setFooter` | | `custom-header.ts` | Replace startup header | `on("session_start")`, `setHeader` | diff --git a/packages/coding-agent/examples/extensions/turn-boundary-renderer.ts b/packages/coding-agent/examples/extensions/turn-boundary-renderer.ts new file mode 100644 index 00000000000..74f1250ad3a --- /dev/null +++ b/packages/coding-agent/examples/extensions/turn-boundary-renderer.ts @@ -0,0 +1,9 @@ +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { Text } from "@earendil-works/pi-tui"; + +export default function (pi: ExtensionAPI) { + pi.registerTurnBoundaryRenderer((_current) => (context, theme) => { + const source = context.isReplay ? "replay" : (context.source ?? "unknown"); + return new Text(theme.fg("dim", `── turn · ${source} ──`), 0, 0); + }); +} diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index df102097860..953d78ba796 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -145,7 +145,12 @@ export function parseSkillBlock(text: string): ParsedSkillBlock | null { /** Session-specific events that extend the core AgentEvent */ export type AgentSessionEvent = - | Exclude + | Exclude + | { + type: "message_start"; + message: AgentMessage; + source?: InputSource; + } | { type: "agent_end"; messages: AgentMessage[]; @@ -324,6 +329,7 @@ export class AgentSession { // Event subscription state private _unsubscribeAgent?: () => void; private _eventListeners: AgentSessionEventListener[] = []; + private _inputSources = new WeakMap(); private _isAgentRunActive = false; private _pendingExtensionMessageActions = 0; private _idleWaitPromise: Promise | undefined; @@ -637,8 +643,21 @@ export class AgentSession { // Emit to extensions first await this._emitExtensionEvent(event); + let sessionEvent: AgentSessionEvent; + if (event.type === "message_start") { + const source = event.message.role === "user" ? this._inputSources.get(event.message) : undefined; + if (event.message.role === "user") { + this._inputSources.delete(event.message); + } + sessionEvent = { ...event, source }; + } else if (event.type === "agent_end") { + sessionEvent = { ...event, willRetry: this._willRetryAfterAgentEnd(event) }; + } else { + sessionEvent = event; + } + // Notify all listeners - this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event); + this._emit(sessionEvent); // Handle session persistence if (event.type === "message_end") { @@ -1140,6 +1159,7 @@ export class AgentSession { async prompt(text: string, options?: PromptOptions): Promise { const expandPromptTemplates = options?.expandPromptTemplates ?? true; const preflightResult = options?.preflightResult; + const inputSource = options?.source ?? "interactive"; let messages: AgentMessage[] | undefined; try { @@ -1161,7 +1181,7 @@ export class AgentSession { const inputResult = await this._extensionRunner.emitInput( currentText, currentImages, - options?.source ?? "interactive", + inputSource, this.isStreaming ? options?.streamingBehavior : undefined, ); if (inputResult.action === "handled") { @@ -1189,9 +1209,9 @@ export class AgentSession { ); } if (options.streamingBehavior === "followUp") { - await this._queueFollowUp(expandedText, currentImages); + await this._queueFollowUp(expandedText, currentImages, inputSource); } else { - await this._queueSteer(expandedText, currentImages); + await this._queueSteer(expandedText, currentImages, inputSource); } preflightResult?.(true); return; @@ -1287,7 +1307,17 @@ export class AgentSession { } preflightResult?.(true); - await this._runAgentPrompt(messages); + const userMessage = messages.find((message) => message.role === "user"); + if (userMessage) { + this._inputSources.set(userMessage, inputSource); + } + try { + await this._runAgentPrompt(messages); + } finally { + if (userMessage) { + this._inputSources.delete(userMessage); + } + } } queueCommand(command: string, args = "", options?: QueueCommandOptions): void { @@ -1475,35 +1505,53 @@ export class AgentSession { /** * Internal: Queue a steering message (already expanded, no extension command check). */ - private async _queueSteer(text: string, images?: ImageContent[]): Promise { + private async _queueSteer(text: string, images?: ImageContent[], source?: InputSource): Promise { this._steeringMessages.push(text); this._emitQueueUpdate(); const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; if (images) { content.push(...images); } - this.agent.steer({ + const message: AgentMessage = { role: "user", content, timestamp: Date.now(), - }); + }; + if (source) { + this._inputSources.set(message, source); + } + try { + this.agent.steer(message); + } catch (error) { + this._inputSources.delete(message); + throw error; + } } /** * Internal: Queue a follow-up message (already expanded, no extension command check). */ - private async _queueFollowUp(text: string, images?: ImageContent[]): Promise { + private async _queueFollowUp(text: string, images?: ImageContent[], source?: InputSource): Promise { this._followUpMessages.push(text); this._emitQueueUpdate(); const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; if (images) { content.push(...images); } - this.agent.followUp({ + const message: AgentMessage = { role: "user", content, timestamp: Date.now(), - }); + }; + if (source) { + this._inputSources.set(message, source); + } + try { + this.agent.followUp(message); + } catch (error) { + this._inputSources.delete(message); + throw error; + } } /** diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 5c771d36a54..49032cfa80d 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -175,6 +175,9 @@ export type { TreeNavigationOptions, TreeNavigationSummary, TreePreparation, + TurnBoundaryContext, + TurnBoundaryRenderer, + TurnBoundaryRendererTransform, TurnEndEvent, TurnStartEvent, // Events - User Bash diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index c589f3ddee9..81c55fcc6f2 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -47,6 +47,7 @@ import type { RegisteredCommand, ToolDefinition, ToolTransform, + TurnBoundaryRendererTransform, } from "./types.ts"; /** Modules available to extensions via virtualModules (for compiled Bun binary) */ @@ -321,6 +322,12 @@ function createExtensionAPI( }); }, + registerTurnBoundaryRenderer(transform: TurnBoundaryRendererTransform): void { + runtime.assertActive(); + extension.turnBoundaryRendererTransforms ??= []; + extension.turnBoundaryRendererTransforms.push({ transform }); + }, + registerMarkdownTransformer(transformer: MarkdownTransformer): void { runtime.assertActive(); extension.markdownTransformer = transformer; diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 2a1fe57d16c..385017e4e6b 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -66,6 +66,7 @@ import type { ToolResultEvent, ToolResultEventResult, TreeNavigationOptions, + TurnBoundaryRendererTransform, UserBashEvent, UserBashEventResult, } from "./types.ts"; @@ -610,6 +611,12 @@ export class ExtensionRunner { ); } + getTurnBoundaryRendererTransforms(): TurnBoundaryRendererTransform[] { + return this.extensions.flatMap((extension) => + (extension.turnBoundaryRendererTransforms ?? []).map((registration) => registration.transform), + ); + } + getMarkdownTransformers(): MarkdownTransformer[] { return this.extensions.flatMap((ext) => (ext.markdownTransformer ? [ext.markdownTransformer] : [])); } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 224c4d66901..0182c4a7421 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -1257,6 +1257,20 @@ export interface RegisteredBuiltInMessageRendererTransform { transform: BuiltInMessageRendererTransform; } +export interface TurnBoundaryContext { + message: BuiltInMessageByRole["user"]; + source?: InputSource; + isReplay: boolean; +} + +export type TurnBoundaryRenderer = (context: TurnBoundaryContext, theme: Theme) => Component; + +export type TurnBoundaryRendererTransform = (renderer: TurnBoundaryRenderer) => TurnBoundaryRenderer; + +export interface RegisteredTurnBoundaryRendererTransform { + transform: TurnBoundaryRendererTransform; +} + export type EntryRenderer = ( entry: CustomEntry, options: EntryRenderOptions, @@ -1396,6 +1410,9 @@ export interface ExtensionAPI { transform: BuiltInMessageRendererTransform, ): void; + /** Wrap the visual boundary Pi inserts before each user turn after the first. */ + registerTurnBoundaryRenderer(transform: TurnBoundaryRendererTransform): void; + /** Register a transformer for user and assistant Markdown before Pi renders it in the interactive transcript. */ registerMarkdownTransformer(transformer: MarkdownTransformer): void; @@ -1816,6 +1833,7 @@ export interface Extension { toolTransforms?: RegisteredToolTransform[]; messageRenderers: Map; builtInMessageRendererTransforms?: RegisteredBuiltInMessageRendererTransform[]; + turnBoundaryRendererTransforms?: RegisteredTurnBoundaryRendererTransform[]; markdownTransformer?: MarkdownTransformer; entryRenderers?: Map; commands: Map; diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index 924932e6713..17d9c4371f3 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -73,6 +73,9 @@ export { type ToolDefinition, type ToolRenderResultOptions, type ToolResultEvent, + type TurnBoundaryContext, + type TurnBoundaryRenderer, + type TurnBoundaryRendererTransform, type TurnEndEvent, type TurnStartEvent, type WorkingIndicatorOptions, diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 5d38c54fe1c..b1875bfeda7 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -152,6 +152,9 @@ export type { ToolRenderResultOptions, ToolResultEvent, ToolTransform, + TurnBoundaryContext, + TurnBoundaryRenderer, + TurnBoundaryRendererTransform, TurnEndEvent, TurnStartEvent, UserBashEvent, diff --git a/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts b/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts new file mode 100644 index 00000000000..2acef7943b7 --- /dev/null +++ b/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts @@ -0,0 +1,54 @@ +import type { Component } from "@earendil-works/pi-tui"; +import type { + TurnBoundaryContext, + TurnBoundaryRenderer, + TurnBoundaryRendererTransform, +} from "../../../core/extensions/types.ts"; +import type { Theme } from "../theme/theme.ts"; + +function isComponent(value: unknown): value is Component { + if (typeof value !== "object" || value === null) return false; + const component = value as Partial; + return typeof component.render === "function" && typeof component.invalidate === "function"; +} + +function tryRender(renderer: TurnBoundaryRenderer, context: TurnBoundaryContext, theme: Theme): Component | undefined { + try { + const component = renderer(context, theme); + return isComponent(component) ? component : undefined; + } catch { + return undefined; + } +} + +export function composeTurnBoundaryRenderer( + fallback: TurnBoundaryRenderer, + transforms: readonly TurnBoundaryRendererTransform[], +): TurnBoundaryRenderer { + let renderer = fallback; + + for (const transform of transforms) { + const previous = renderer; + try { + const transformed = transform(previous); + if (typeof transformed !== "function") continue; + renderer = (context, theme) => tryRender(transformed, context, theme) ?? previous(context, theme); + } catch { + // Keep the previous renderer and continue with the next transform. + } + } + + const composed = renderer; + if (composed === fallback) return fallback; + + return (context, theme) => { + // Renderer transforms receive a snapshot so display code cannot mutate session or model state. + let isolatedContext: TurnBoundaryContext; + try { + isolatedContext = structuredClone(context); + } catch { + return fallback(context, theme); + } + return tryRender(composed, isolatedContext, theme) ?? fallback(context, theme); + }; +} diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 25b271f28cf..7b6bacdcbba 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -72,6 +72,7 @@ import type { ExtensionUIContext, ExtensionUIDialogOptions, ExtensionWidgetOptions, + InputSource, MarkdownTransformer, ProjectTrustContext, WorkingIndicatorOptions, @@ -145,6 +146,7 @@ import { import { ToolExecutionComponent } from "./components/tool-execution.ts"; import { TreeSelectorComponent } from "./components/tree-selector.ts"; import { TrustSelectorComponent } from "./components/trust-selector.ts"; +import { composeTurnBoundaryRenderer } from "./components/turn-boundary-renderer.ts"; import { UserMessageComponent } from "./components/user-message.ts"; import { UserMessageSelectorComponent } from "./components/user-message-selector.ts"; import { editInExternalEditor } from "./external-editor.ts"; @@ -2948,7 +2950,7 @@ export class InteractiveMode { this.addMessageToChat(event.message); this.ui.requestRender(); } else if (event.message.role === "user") { - this.addMessageToChat(event.message); + this.addMessageToChat(event.message, { source: event.source, isReplay: false }); this.updatePendingMessagesDisplay(); this.ui.requestRender(); } else if (event.message.role === "assistant") { @@ -3277,7 +3279,10 @@ export class InteractiveMode { this.chatContainer.addChild(component); } - private addMessageToChat(message: AgentMessage, options?: { populateHistory?: boolean }): void { + private addMessageToChat( + message: AgentMessage, + options?: { populateHistory?: boolean; source?: InputSource; isReplay?: boolean }, + ): void { switch (message.role) { case "bashExecution": { const component = new BashExecutionComponent(message.command, this.ui, message.excludeFromContext); @@ -3325,7 +3330,20 @@ export class InteractiveMode { const textContent = this.getUserMessageText(message); if (textContent) { if (this.chatContainer.children.length > 0) { - this.chatContainer.addChild(new Spacer(1)); + const renderBoundary = composeTurnBoundaryRenderer( + () => new Spacer(1), + this.session.extensionRunner.getTurnBoundaryRendererTransforms(), + ); + this.chatContainer.addChild( + renderBoundary( + { + message, + source: options?.source, + isReplay: options?.isReplay ?? false, + }, + theme, + ), + ); } const skillBlock = parseSkillBlock(textContent); if (skillBlock) { @@ -3394,7 +3412,7 @@ export class InteractiveMode { private renderSessionItems( items: readonly RenderSessionItem[], - options: { updateFooter?: boolean; populateHistory?: boolean } = {}, + options: { updateFooter?: boolean; populateHistory?: boolean; isReplay?: boolean } = {}, ): void { this.pendingTools.clear(); const renderedPendingTools = new Map(); @@ -3493,7 +3511,7 @@ export class InteractiveMode { } return sessionEntryToContextMessages(entry); }); - this.renderSessionItems(items, options); + this.renderSessionItems(items, { ...options, isReplay: true }); } /** diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 2b434e57833..e952018e7a4 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -672,12 +672,14 @@ describe("ExtensionRunner", () => { `export default function(pi) { pi.registerBuiltInMessageRenderer("user", (current) => current); pi.registerBuiltInMessageRenderer("assistant", (current) => current); + pi.registerTurnBoundaryRenderer((current) => current); }`, ); fs.writeFileSync( path.join(extensionsDir, "renderer-b.ts"), `export default function(pi) { pi.registerBuiltInMessageRenderer("user", (current) => current); + pi.registerTurnBoundaryRenderer((current) => current); }`, ); @@ -693,6 +695,10 @@ describe("ExtensionRunner", () => { ); expect(userTransforms).toEqual([firstUserRegistration?.transform, secondUserRegistration?.transform]); expect(runner.getBuiltInMessageRendererTransforms("assistant")).toHaveLength(1); + expect(runner.getTurnBoundaryRendererTransforms()).toEqual([ + result.extensions[0].turnBoundaryRendererTransforms?.[0]?.transform, + result.extensions[1].turnBoundaryRendererTransforms?.[0]?.transform, + ]); }); it("gets entry renderer by type", async () => { diff --git a/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts new file mode 100644 index 00000000000..da73cac8a31 --- /dev/null +++ b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts @@ -0,0 +1,99 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { Container, Spacer, Text } from "@earendil-works/pi-tui"; +import { describe, expect, test, vi } from "vitest"; +import type { InputSource, TurnBoundaryContext } from "../src/core/extensions/types.ts"; +import { InteractiveMode } from "../src/modes/interactive/interactive-mode.ts"; +import { initTheme } from "../src/modes/interactive/theme/theme.ts"; + +const addMessageToChat = Reflect.get(InteractiveMode.prototype, "addMessageToChat") as ( + this: ReturnType, + message: AgentMessage, + options?: { populateHistory?: boolean; source?: InputSource; isReplay?: boolean }, +) => void; + +function userMessage(content: string, timestamp: number): AgentMessage { + return { role: "user", content, timestamp }; +} + +function createFakeMode(contexts: TurnBoundaryContext[] = [], useBoundary = true) { + const transforms = useBoundary + ? [ + () => (context: TurnBoundaryContext) => { + contexts.push(context); + return new Text("boundary", 0, 0); + }, + ] + : []; + return { + chatContainer: new Container(), + getUserMessageText: (message: AgentMessage) => + message.role === "user" + ? typeof message.content === "string" + ? message.content + : (message.content.find((part) => part.type === "text")?.text ?? "") + : "", + session: { + extensionRunner: { + getTurnBoundaryRendererTransforms: () => transforms, + getBuiltInMessageRendererTransforms: () => [], + }, + }, + getMarkdownThemeWithSettings: () => undefined, + outputPad: 1, + getMarkdownTransformers: () => [], + toolOutputExpanded: false, + editor: { addToHistory: vi.fn() }, + }; +} + +describe("InteractiveMode turn boundaries", () => { + test("keeps the first message boundary-free and passes live context to later boundaries", () => { + initTheme("dark"); + const contexts: TurnBoundaryContext[] = []; + const fakeMode = createFakeMode(contexts); + + addMessageToChat.call(fakeMode, userMessage("first", 1), { source: "interactive", isReplay: false }); + addMessageToChat.call(fakeMode, userMessage("second", 2), { source: "rpc", isReplay: false }); + + expect(contexts).toEqual([ + expect.objectContaining({ + source: "rpc", + isReplay: false, + message: expect.objectContaining({ content: "second" }), + }), + ]); + expect(fakeMode.chatContainer.children).toHaveLength(3); + expect(fakeMode.chatContainer.children[1]?.render(40).map((line) => line.trimEnd())).toEqual(["boundary"]); + }); + + test("preserves the default spacer when no extension overrides it", () => { + initTheme("dark"); + const fakeMode = createFakeMode([], false); + + addMessageToChat.call(fakeMode, userMessage("first", 1)); + addMessageToChat.call(fakeMode, userMessage("second", 2)); + + expect(fakeMode.chatContainer.children).toHaveLength(3); + expect(fakeMode.chatContainer.children[1]).toBeInstanceOf(Spacer); + }); + + test("invokes the outer boundary for replayed skill user turns", () => { + initTheme("dark"); + const contexts: TurnBoundaryContext[] = []; + const fakeMode = createFakeMode(contexts); + const skill = '\nSkill body\n\n\nDo this'; + + addMessageToChat.call(fakeMode, userMessage("first", 1)); + addMessageToChat.call(fakeMode, userMessage(skill, 2), { isReplay: true }); + + expect(contexts).toEqual([ + expect.objectContaining({ + source: undefined, + isReplay: true, + message: expect.objectContaining({ content: skill }), + }), + ]); + expect(fakeMode.chatContainer.children[1]?.render(40).map((line) => line.trimEnd())).toEqual(["boundary"]); + expect(fakeMode.chatContainer.children.filter((child) => child instanceof Spacer)).toHaveLength(1); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts index 94e10a3e54a..55e4d2b799e 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -32,9 +32,14 @@ describe("AgentSession prompt characterization", () => { harnesses.push(harness); harness.setResponses([fauxAssistantMessage("hello")]); + let source: string | undefined; + harness.session.subscribe((event) => { + if (event.type === "message_start" && event.message.role === "user") source = event.source; + }); await harness.session.prompt("hi"); + expect(source).toBe("interactive"); expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(getMessageText(harness.session.messages[0]!)).toBe("hi"); expect(harness.getPendingResponseCount()).toBe(0); @@ -253,9 +258,14 @@ describe("AgentSession prompt characterization", () => { harnesses.push(harness); harness.setResponses([fauxAssistantMessage("response")]); + let source: string | undefined; + harness.session.subscribe((event) => { + if (event.type === "message_start" && event.message.role === "user") source = event.source; + }); await harness.session.sendUserMessage("from extension"); + expect(source).toBe("extension"); expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(getMessageText(harness.session.messages[0]!)).toBe("from extension"); }); @@ -286,6 +296,7 @@ describe("AgentSession prompt characterization", () => { releaseToolExecution = resolve; }); const inputEvents: InputEvent[] = []; + const userMessageSources: Array<{ text: string; source: string | undefined }> = []; const waitTool: AgentTool = { name: "wait", label: "Wait", @@ -324,14 +335,24 @@ describe("AgentSession prompt characterization", () => { }); }); + harness.session.subscribe((event) => { + if (event.type === "message_start" && event.message.role === "user") { + userMessageSources.push({ text: getMessageText(event.message), source: event.source }); + } + }); + const promptPromise = harness.session.prompt("start"); await sawToolStart; - await harness.session.prompt("queued", { streamingBehavior: "followUp" }); + await harness.session.prompt("queued", { streamingBehavior: "followUp", source: "rpc" }); expect(inputEvents.map((event) => event.streamingBehavior)).toEqual([undefined, "followUp"]); releaseToolExecution?.(); await promptPromise; + expect(userMessageSources).toEqual([ + { text: "start", source: "interactive" }, + { text: "queued", source: "rpc" }, + ]); }); it("throws when prompted during streaming without a streamingBehavior", async () => { diff --git a/packages/coding-agent/test/turn-boundary-renderer.test.ts b/packages/coding-agent/test/turn-boundary-renderer.test.ts new file mode 100644 index 00000000000..d1e90d74373 --- /dev/null +++ b/packages/coding-agent/test/turn-boundary-renderer.test.ts @@ -0,0 +1,96 @@ +import type { UserMessage } from "@earendil-works/pi-ai"; +import { Container, Text } from "@earendil-works/pi-tui"; +import { describe, expect, test } from "vitest"; +import type { TurnBoundaryRenderer, TurnBoundaryRendererTransform } from "../src/core/extensions/types.ts"; +import { composeTurnBoundaryRenderer } from "../src/modes/interactive/components/turn-boundary-renderer.ts"; +import { initTheme, theme } from "../src/modes/interactive/theme/theme.ts"; +import { stripAnsi } from "../src/utils/ansi.ts"; + +const message: UserMessage = { + role: "user", + content: "hello", + timestamp: 1, +}; + +const context = { message, source: "interactive" as const, isReplay: false }; + +function wrappingTransform(label: string, trace: string[]): TurnBoundaryRendererTransform { + return (current) => (currentContext, currentTheme) => { + trace.push(`${label}:before`); + const rendered = current(currentContext, currentTheme); + trace.push(`${label}:after`); + const component = new Container(); + component.addChild(new Text(label, 0, 0)); + component.addChild(rendered); + return component; + }; +} + +describe("composeTurnBoundaryRenderer", () => { + test("composes later transforms outside earlier transforms", () => { + initTheme("dark"); + const trace: string[] = []; + const fallback: TurnBoundaryRenderer = () => { + trace.push("fallback"); + return new Text("content", 0, 0); + }; + const renderer = composeTurnBoundaryRenderer(fallback, [ + wrappingTransform("inner", trace), + wrappingTransform("outer", trace), + ]); + + const result = renderer(context, theme); + + expect(trace).toEqual(["outer:before", "inner:before", "fallback", "inner:after", "outer:after"]); + expect(result.render(40).map((line) => stripAnsi(line).trimEnd())).toEqual(["outer", "inner", "content"]); + }); + + test("isolates the source context from transform mutations", () => { + initTheme("dark"); + const source = { + message: { ...message, content: [{ type: "text" as const, text: "original" }] }, + source: "rpc" as const, + isReplay: false, + }; + const mutate: TurnBoundaryRendererTransform = (current) => (currentContext, currentTheme) => { + if (typeof currentContext.message.content !== "string") { + const text = currentContext.message.content.find((part) => part.type === "text"); + if (text) text.text = "mutated"; + } + return current(currentContext, currentTheme); + }; + const fallback: TurnBoundaryRenderer = (currentContext) => { + const content = currentContext.message.content; + const text = + typeof content === "string" ? content : (content.find((part) => part.type === "text")?.text ?? ""); + return new Text(text, 0, 0); + }; + + const result = composeTurnBoundaryRenderer(fallback, [mutate])(source, theme); + + expect(result.render(40).map((line) => stripAnsi(line).trimEnd())).toEqual(["mutated"]); + expect(source.message.content).toEqual([{ type: "text", text: "original" }]); + }); + + test("keeps healthy transforms when registration or rendering fails", () => { + initTheme("dark"); + const trace: string[] = []; + const brokenRegistration: TurnBoundaryRendererTransform = () => { + throw new Error("registration failed"); + }; + const brokenRender = (() => () => { + throw new Error("render failed"); + }) as TurnBoundaryRendererTransform; + const fallback: TurnBoundaryRenderer = () => new Text("fallback", 0, 0); + const renderer = composeTurnBoundaryRenderer(fallback, [ + brokenRegistration, + brokenRender, + wrappingTransform("healthy", trace), + ]); + + const result = renderer(context, theme); + + expect(trace).toEqual(["healthy:before", "healthy:after"]); + expect(result.render(40).map((line) => stripAnsi(line).trimEnd())).toEqual(["healthy", "fallback"]); + }); +}); From 471f50211e5194afba12cb427d317ae0ab989bf4 Mon Sep 17 00:00:00 2001 From: quick Date: Fri, 7 Aug 2026 11:58:22 -0400 Subject: [PATCH 2/4] refactor(coding-agent): clarify turn source lifecycle --- .../coding-agent/src/core/agent-session.ts | 76 +++++++++---------- .../src/modes/interactive/interactive-mode.ts | 33 ++++---- .../interactive-mode-turn-boundary.test.ts | 2 + 3 files changed, 59 insertions(+), 52 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 953d78ba796..79897fc3b18 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -645,10 +645,7 @@ export class AgentSession { let sessionEvent: AgentSessionEvent; if (event.type === "message_start") { - const source = event.message.role === "user" ? this._inputSources.get(event.message) : undefined; - if (event.message.role === "user") { - this._inputSources.delete(event.message); - } + const source = this._consumeInputSource(event.message); sessionEvent = { ...event, source }; } else if (event.type === "agent_end") { sessionEvent = { ...event, willRetry: this._willRetryAfterAgentEnd(event) }; @@ -1097,7 +1094,12 @@ export class AgentSession { // Prompting // ========================================================================= - private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { + private async _runAgentPrompt(messages: AgentMessage | AgentMessage[], inputSource?: InputSource): Promise { + const userMessage = (Array.isArray(messages) ? messages : [messages]).find((message) => message.role === "user"); + if (userMessage && inputSource) { + this._inputSources.set(userMessage, inputSource); + } + this._isAgentRunActive = true; try { await this.agent.prompt(messages); @@ -1111,12 +1113,24 @@ export class AgentSession { await this.agent.continue(); } } finally { + if (userMessage) { + this._inputSources.delete(userMessage); + } this._systemPromptOverride = undefined; this._flushPendingBashMessages(); await this._emitAgentSettled(); } } + private _consumeInputSource(message: AgentMessage): InputSource | undefined { + if (message.role !== "user") { + return undefined; + } + const source = this._inputSources.get(message); + this._inputSources.delete(message); + return source; + } + private async _handlePostAgentRun(): Promise { const msg = this._lastAssistantMessage; this._lastAssistantMessage = undefined; @@ -1307,17 +1321,7 @@ export class AgentSession { } preflightResult?.(true); - const userMessage = messages.find((message) => message.role === "user"); - if (userMessage) { - this._inputSources.set(userMessage, inputSource); - } - try { - await this._runAgentPrompt(messages); - } finally { - if (userMessage) { - this._inputSources.delete(userMessage); - } - } + await this._runAgentPrompt(messages, inputSource); } queueCommand(command: string, args = "", options?: QueueCommandOptions): void { @@ -1506,34 +1510,26 @@ export class AgentSession { * Internal: Queue a steering message (already expanded, no extension command check). */ private async _queueSteer(text: string, images?: ImageContent[], source?: InputSource): Promise { - this._steeringMessages.push(text); - this._emitQueueUpdate(); - const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; - if (images) { - content.push(...images); - } - const message: AgentMessage = { - role: "user", - content, - timestamp: Date.now(), - }; - if (source) { - this._inputSources.set(message, source); - } - try { - this.agent.steer(message); - } catch (error) { - this._inputSources.delete(message); - throw error; - } + this._queueUserMessage("steer", text, images, source); } /** * Internal: Queue a follow-up message (already expanded, no extension command check). */ private async _queueFollowUp(text: string, images?: ImageContent[], source?: InputSource): Promise { - this._followUpMessages.push(text); + this._queueUserMessage("followUp", text, images, source); + } + + private _queueUserMessage( + delivery: "steer" | "followUp", + text: string, + images?: ImageContent[], + source?: InputSource, + ): void { + const queue = delivery === "steer" ? this._steeringMessages : this._followUpMessages; + queue.push(text); this._emitQueueUpdate(); + const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; if (images) { content.push(...images); @@ -1547,7 +1543,11 @@ export class AgentSession { this._inputSources.set(message, source); } try { - this.agent.followUp(message); + if (delivery === "steer") { + this.agent.steer(message); + } else { + this.agent.followUp(message); + } } catch (error) { this._inputSources.delete(message); throw error; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 7b6bacdcbba..f07162d0f7c 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -3279,6 +3279,24 @@ export class InteractiveMode { this.chatContainer.addChild(component); } + private createTurnBoundary( + message: Extract, + options?: { source?: InputSource; isReplay?: boolean }, + ): Component { + const renderBoundary = composeTurnBoundaryRenderer( + () => new Spacer(1), + this.session.extensionRunner.getTurnBoundaryRendererTransforms(), + ); + return renderBoundary( + { + message, + source: options?.source, + isReplay: options?.isReplay ?? false, + }, + theme, + ); + } + private addMessageToChat( message: AgentMessage, options?: { populateHistory?: boolean; source?: InputSource; isReplay?: boolean }, @@ -3330,20 +3348,7 @@ export class InteractiveMode { const textContent = this.getUserMessageText(message); if (textContent) { if (this.chatContainer.children.length > 0) { - const renderBoundary = composeTurnBoundaryRenderer( - () => new Spacer(1), - this.session.extensionRunner.getTurnBoundaryRendererTransforms(), - ); - this.chatContainer.addChild( - renderBoundary( - { - message, - source: options?.source, - isReplay: options?.isReplay ?? false, - }, - theme, - ), - ); + this.chatContainer.addChild(this.createTurnBoundary(message, options)); } const skillBlock = parseSkillBlock(textContent); if (skillBlock) { diff --git a/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts index da73cac8a31..71fd4605eb8 100644 --- a/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts +++ b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts @@ -5,6 +5,7 @@ import type { InputSource, TurnBoundaryContext } from "../src/core/extensions/ty import { InteractiveMode } from "../src/modes/interactive/interactive-mode.ts"; import { initTheme } from "../src/modes/interactive/theme/theme.ts"; +const createTurnBoundary = Reflect.get(InteractiveMode.prototype, "createTurnBoundary"); const addMessageToChat = Reflect.get(InteractiveMode.prototype, "addMessageToChat") as ( this: ReturnType, message: AgentMessage, @@ -26,6 +27,7 @@ function createFakeMode(contexts: TurnBoundaryContext[] = [], useBoundary = true : []; return { chatContainer: new Container(), + createTurnBoundary, getUserMessageText: (message: AgentMessage) => message.role === "user" ? typeof message.content === "string" From 3d7f317c8bf1b00bac7343758eb0d22097f0f467 Mon Sep 17 00:00:00 2001 From: baby-joel Date: Fri, 7 Aug 2026 16:08:49 +0000 Subject: [PATCH 3/4] fix(coding-agent): isolate turn boundary fallbacks --- .../components/turn-boundary-renderer.ts | 10 +++++++++- .../test/turn-boundary-renderer.test.ts | 13 +++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts b/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts index 2acef7943b7..fd28362a864 100644 --- a/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts +++ b/packages/coding-agent/src/modes/interactive/components/turn-boundary-renderer.ts @@ -32,7 +32,15 @@ export function composeTurnBoundaryRenderer( try { const transformed = transform(previous); if (typeof transformed !== "function") continue; - renderer = (context, theme) => tryRender(transformed, context, theme) ?? previous(context, theme); + renderer = (context, theme) => { + let isolatedContext: TurnBoundaryContext; + try { + isolatedContext = structuredClone(context); + } catch { + return previous(context, theme); + } + return tryRender(transformed, isolatedContext, theme) ?? previous(context, theme); + }; } catch { // Keep the previous renderer and continue with the next transform. } diff --git a/packages/coding-agent/test/turn-boundary-renderer.test.ts b/packages/coding-agent/test/turn-boundary-renderer.test.ts index d1e90d74373..10b5eaa7aba 100644 --- a/packages/coding-agent/test/turn-boundary-renderer.test.ts +++ b/packages/coding-agent/test/turn-boundary-renderer.test.ts @@ -93,4 +93,17 @@ describe("composeTurnBoundaryRenderer", () => { expect(trace).toEqual(["healthy:before", "healthy:after"]); expect(result.render(40).map((line) => stripAnsi(line).trimEnd())).toEqual(["healthy", "fallback"]); }); + + test("isolates a fallback layer from mutations made by a failed transform", () => { + initTheme("dark"); + const mutateThenThrow: TurnBoundaryRendererTransform = () => (currentContext) => { + currentContext.source = "rpc"; + throw new Error("render failed"); + }; + const fallback: TurnBoundaryRenderer = (currentContext) => new Text(currentContext.source ?? "unknown", 0, 0); + + const result = composeTurnBoundaryRenderer(fallback, [mutateThenThrow])(context, theme); + + expect(result.render(40).map((line) => stripAnsi(line).trimEnd())).toEqual(["interactive"]); + }); }); From 0f8475ed09c8c930eae0bf48604b5032a79d7005 Mon Sep 17 00:00:00 2001 From: baby-joel Date: Fri, 7 Aug 2026 16:09:57 +0000 Subject: [PATCH 4/4] fix(coding-agent): skip boundary hook for first user turn --- .../src/modes/interactive/interactive-mode.ts | 8 +++++++- .../test/interactive-mode-turn-boundary.test.ts | 13 +++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index f07162d0f7c..08e10a7aa8d 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -3348,7 +3348,13 @@ export class InteractiveMode { const textContent = this.getUserMessageText(message); if (textContent) { if (this.chatContainer.children.length > 0) { - this.chatContainer.addChild(this.createTurnBoundary(message, options)); + const hasRenderedUserMessage = this.chatContainer.children.some( + (child) => + child instanceof UserMessageComponent || child instanceof SkillInvocationMessageComponent, + ); + this.chatContainer.addChild( + hasRenderedUserMessage ? this.createTurnBoundary(message, options) : new Spacer(1), + ); } const skillBlock = parseSkillBlock(textContent); if (skillBlock) { diff --git a/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts index 71fd4605eb8..e4b63419620 100644 --- a/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts +++ b/packages/coding-agent/test/interactive-mode-turn-boundary.test.ts @@ -79,6 +79,19 @@ describe("InteractiveMode turn boundaries", () => { expect(fakeMode.chatContainer.children[1]).toBeInstanceOf(Spacer); }); + test("does not invoke the boundary hook for the first user message after non-user content", () => { + initTheme("dark"); + const contexts: TurnBoundaryContext[] = []; + const fakeMode = createFakeMode(contexts); + fakeMode.chatContainer.addChild(new Text("notice", 0, 0)); + + addMessageToChat.call(fakeMode, userMessage("first", 1), { source: "interactive", isReplay: false }); + + expect(contexts).toEqual([]); + expect(fakeMode.chatContainer.children).toHaveLength(3); + expect(fakeMode.chatContainer.children[1]).toBeInstanceOf(Spacer); + }); + test("invokes the outer boundary for replayed skill user turns", () => { initTheme("dark"); const contexts: TurnBoundaryContext[] = [];