From 2417adb46a6ddc7f15bd396a78b0ca9b83ac2d0d Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 21 Jun 2026 23:17:13 +0200 Subject: [PATCH] fix(coding-agent): preserve startup extension UI closes #5943 --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 3 +- .../src/modes/interactive/interactive-mode.ts | 52 ++- .../interactive-mode-clone-command.test.ts | 2 +- .../5943-session-start-notify.test.ts | 430 ++++++++++++++++++ 5 files changed, 463 insertions(+), 25 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/5943-session-start-notify.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 20d13deb..341921f2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed transient extension UI and session-start messages emitted during session replacement or reload so they remain visible, and kept reload input blocked until reload completes ([#5943](https://github.com/earendil-works/pi/issues/5943)). - Fixed the plan-mode example to preserve active custom tools, skip the action prompt when no plan is found, and queue refinement/execution follow-ups correctly from `agent_end` ([#5940](https://github.com/earendil-works/pi/issues/5940)). - Fixed `pi update` to install the exact version returned by the Pi update check, make `--force` reinstall that checked version, fail instead of falling back to an unversioned reinstall when no version is available, and report both the old and updated versions. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index f8af4320..a1033128 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2443,7 +2443,7 @@ export class AgentSession { }); } - async reload(): Promise { + async reload(options?: { beforeSessionStart?: () => void | Promise }): Promise { const previousFlagValues = this._extensionRunner.getFlagValues(); await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); await this.settingsManager.reload(); @@ -2462,6 +2462,7 @@ export class AgentSession { this._extensionShutdownHandler || this._extensionErrorListener; if (hasBindings) { + await options?.beforeSessionStart?.(); await this._extensionRunner.emit({ type: "session_start", reason: "reload" }); await this.extendResourcesFromExtensions("reload"); } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 613d588c..100feb6b 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -395,7 +395,7 @@ export class InteractiveMode { this.resetExtensionUI(); }); this.runtimeHost.setRebindSession(async () => { - await this.rebindCurrentSession(); + await this.rebindCurrentSession({ renderBeforeBind: true }); }); this.version = VERSION; this.ui = new TUI(new ProcessTerminal(), this.settingsManager.getShowHardwareCursor()); @@ -1543,12 +1543,7 @@ export class InteractiveMode { } this.statusContainer.clear(); try { - const result = await this.runtimeHost.newSession(options); - if (!result.cancelled) { - this.renderCurrentSessionState(); - this.ui.requestRender(); - } - return result; + return await this.runtimeHost.newSession(options); } catch (error: unknown) { return this.handleFatalRuntimeError("Failed to create session", error); } @@ -1557,7 +1552,6 @@ export class InteractiveMode { try { const result = await this.runtimeHost.fork(entryId, options); if (!result.cancelled) { - this.renderCurrentSessionState(); this.editor.setText(result.selectedText ?? ""); this.showStatus("Forked to new session"); } @@ -1631,12 +1625,18 @@ export class InteractiveMode { } } - private async rebindCurrentSession(): Promise { + private async rebindCurrentSession(options: { renderBeforeBind?: boolean } = {}): Promise { this.unsubscribe?.(); this.unsubscribe = undefined; this.applyRuntimeSettings(); - await this.bindCurrentSessionExtensions(); - this.subscribeToAgent(); + if (options.renderBeforeBind) { + this.renderCurrentSessionState(); + this.subscribeToAgent(); + await this.bindCurrentSessionExtensions(); + } else { + await this.bindCurrentSessionExtensions(); + this.subscribeToAgent(); + } await this.updateAvailableProviderCount(); this.updateEditorBorderColor(); this.updateTerminalTitle(); @@ -4376,7 +4376,6 @@ export class InteractiveMode { return; } - this.renderCurrentSessionState(); this.editor.setText(result.selectedText ?? ""); done(); this.showStatus("Forked to new session"); @@ -4409,7 +4408,6 @@ export class InteractiveMode { return; } - this.renderCurrentSessionState(); this.editor.setText(""); this.showStatus("Cloned to new session"); } catch (error: unknown) { @@ -4601,7 +4599,6 @@ export class InteractiveMode { if (result.cancelled) { return result; } - this.renderCurrentSessionState(); this.showStatus("Resumed session"); return result; } catch (error: unknown) { @@ -4619,7 +4616,6 @@ export class InteractiveMode { if (result.cancelled) { return result; } - this.renderCurrentSessionState(); this.showStatus("Resumed session in current cwd"); return result; } @@ -5072,8 +5068,20 @@ export class InteractiveMode { this.ui.requestRender(); }; + let chatRestoredBeforeSessionStart = false; + let reloadBoxDismissed = false; + const restoreChatBeforeSessionStart = () => { + if (chatRestoredBeforeSessionStart) { + return; + } + this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); + this.rebuildChatFromMessages(); + chatRestoredBeforeSessionStart = true; + }; + try { - await this.session.reload(); + await this.session.reload({ beforeSessionStart: restoreChatBeforeSessionStart }); + restoreChatBeforeSessionStart(); configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs()); this.keybindings.reload(); const activeHeader = this.customHeader ?? this.builtInHeader; @@ -5081,7 +5089,6 @@ export class InteractiveMode { activeHeader.setExpanded(this.toolOutputExpanded); } setRegisteredThemes(this.session.resourceLoader.getThemes().themes); - this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); await this.themeController.applyFromSettings(); const editorPaddingX = this.settingsManager.getEditorPaddingX(); const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible(); @@ -5096,8 +5103,6 @@ export class InteractiveMode { this.setupAutocompleteProvider(); const runner = this.session.extensionRunner; this.setupExtensionShortcuts(runner); - this.rebuildChatFromMessages(); - dismissReloadBox(this.editor as Component); this.showLoadedResources({ force: false, showDiagnosticsWhenQuiet: true, @@ -5112,8 +5117,12 @@ export class InteractiveMode { ? "Reloaded keybindings, extensions, skills, prompts, themes; saved project trust" : "Reloaded keybindings, extensions, skills, prompts, themes", ); + dismissReloadBox(this.editor as Component); + reloadBoxDismissed = true; } catch (error) { - dismissReloadBox(previousEditor as Component); + if (!reloadBoxDismissed) { + dismissReloadBox(previousEditor as Component); + } this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`); } } @@ -5187,7 +5196,6 @@ export class InteractiveMode { this.showStatus("Import cancelled"); return; } - this.renderCurrentSessionState(); this.showStatus(`Session imported from: ${inputPath}`); } catch (error: unknown) { if (error instanceof MissingSessionCwdError) { @@ -5201,7 +5209,6 @@ export class InteractiveMode { this.showStatus("Import cancelled"); return; } - this.renderCurrentSessionState(); this.showStatus(`Session imported from: ${inputPath}`); return; } @@ -5540,7 +5547,6 @@ export class InteractiveMode { if (result.cancelled) { return; } - this.renderCurrentSessionState(); this.chatContainer.addChild(new Spacer(1)); this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1)); this.ui.requestRender(); diff --git a/packages/coding-agent/test/interactive-mode-clone-command.test.ts b/packages/coding-agent/test/interactive-mode-clone-command.test.ts index fea69c5d..69864146 100644 --- a/packages/coding-agent/test/interactive-mode-clone-command.test.ts +++ b/packages/coding-agent/test/interactive-mode-clone-command.test.ts @@ -41,7 +41,7 @@ describe("InteractiveMode /clone", () => { await interactiveModePrototype.handleCloneCommand.call(context); expect(fork).toHaveBeenCalledWith("leaf-123", { position: "at" }); - expect(renderCurrentSessionState).toHaveBeenCalled(); + expect(renderCurrentSessionState).not.toHaveBeenCalled(); expect(setText).toHaveBeenCalledWith(""); expect(showStatus).toHaveBeenCalledWith("Cloned to new session"); expect(showError).not.toHaveBeenCalled(); diff --git a/packages/coding-agent/test/suite/regressions/5943-session-start-notify.test.ts b/packages/coding-agent/test/suite/regressions/5943-session-start-notify.test.ts new file mode 100644 index 00000000..db7ce26d --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5943-session-start-notify.test.ts @@ -0,0 +1,430 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it, vi } from "vitest"; +import type { AgentSessionEvent } from "../../../src/core/agent-session.ts"; +import type { ExtensionUIContext } from "../../../src/core/extensions/index.ts"; +import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts"; +import { initTheme, type Theme, theme } from "../../../src/modes/interactive/theme/theme.ts"; +import { createHarness } from "../harness.ts"; + +function createUiContext( + onNotify: (message: string, type: "info" | "warning" | "error" | undefined) => void, +): ExtensionUIContext { + return { + select: async () => undefined, + confirm: async () => false, + input: async () => undefined, + notify: onNotify, + onTerminalInput: () => () => {}, + setStatus: () => {}, + setWorkingMessage: () => {}, + setWorkingVisible: () => {}, + setWorkingIndicator: () => {}, + setHiddenThinkingLabel: () => {}, + setWidget: () => {}, + setFooter: () => {}, + setHeader: () => {}, + setTitle: () => {}, + custom: async () => undefined as T, + pasteToEditor: () => {}, + setEditorText: () => {}, + getEditorText: () => "", + editor: async () => undefined, + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + get theme() { + return theme; + }, + getAllThemes: () => [], + getTheme: () => undefined, + setTheme: (_theme: string | Theme) => ({ success: false, error: "Theme switching not available in tests" }), + getToolsExpanded: () => false, + setToolsExpanded: () => {}, + }; +} + +type RebindContext = { + unsubscribe?: () => void; + applyRuntimeSettings: () => void; + renderCurrentSessionState: () => void; + bindCurrentSessionExtensions: () => Promise; + subscribeToAgent: () => void; + updateAvailableProviderCount: () => Promise; + updateEditorBorderColor: () => void; + updateTerminalTitle: () => void; +}; + +type ReloadCommandContext = { + hideThinkingBlock: boolean; + session: { + isStreaming: boolean; + isCompacting: boolean; + reload: (options?: { beforeSessionStart?: () => void | Promise }) => Promise; + resourceLoader: { getThemes: () => { themes: [] } }; + extensionRunner: unknown; + modelRegistry: { getError: () => string | undefined }; + }; + settingsManager: { + getHttpIdleTimeoutMs: () => number; + getHideThinkingBlock: () => boolean; + getEditorPaddingX: () => number; + getAutocompleteMaxVisible: () => number; + getShowHardwareCursor: () => boolean; + getClearOnShrink: () => boolean; + }; + keybindings: { reload: () => void }; + customHeader?: unknown; + builtInHeader?: unknown; + editorContainer: { clear: () => void; addChild: (component: unknown) => void }; + ui: { + setFocus: (component: unknown) => void; + requestRender: (force?: boolean) => void; + setShowHardwareCursor: (enabled: boolean) => void; + setClearOnShrink: (enabled: boolean) => void; + }; + editor: unknown; + defaultEditor: { setPaddingX: (padding: number) => void; setAutocompleteMaxVisible: (maxVisible: number) => void }; + themeController: { applyFromSettings: () => Promise }; + resetExtensionUI: () => void; + rebuildChatFromMessages: () => void; + setupAutocompleteProvider: () => void; + setupExtensionShortcuts: (runner: unknown) => void; + showLoadedResources: (options: unknown) => void; + maybeSaveImplicitProjectTrustAfterReload: () => boolean; + showStatus: (message: string) => void; + showWarning: (message: string) => void; + showError: (message: string) => void; +}; + +type InteractiveModePrototype = { + rebindCurrentSession(this: RebindContext, options?: { renderBeforeBind?: boolean }): Promise; + handleReloadCommand(this: ReloadCommandContext): Promise; +}; + +const interactiveModePrototype = InteractiveMode.prototype as unknown as InteractiveModePrototype; + +type ReloadCommandContextOverrides = Omit< + Partial, + "session" | "settingsManager" | "keybindings" | "editorContainer" | "ui" | "defaultEditor" | "themeController" +> & { + session?: Partial; + settingsManager?: Partial; + keybindings?: Partial; + editorContainer?: Partial; + ui?: Partial; + defaultEditor?: Partial; + themeController?: Partial; +}; + +function createReloadCommandContext(overrides: ReloadCommandContextOverrides = {}): ReloadCommandContext { + const editor = overrides.editor ?? {}; + return { + hideThinkingBlock: overrides.hideThinkingBlock ?? false, + session: { + isStreaming: false, + isCompacting: false, + reload: async (options) => { + await options?.beforeSessionStart?.(); + }, + resourceLoader: { getThemes: () => ({ themes: [] }) }, + extensionRunner: {}, + modelRegistry: { getError: () => undefined }, + ...overrides.session, + }, + settingsManager: { + getHttpIdleTimeoutMs: () => 0, + getHideThinkingBlock: () => false, + getEditorPaddingX: () => 1, + getAutocompleteMaxVisible: () => 10, + getShowHardwareCursor: () => false, + getClearOnShrink: () => false, + ...overrides.settingsManager, + }, + keybindings: { reload: () => {}, ...overrides.keybindings }, + editorContainer: { clear: () => {}, addChild: () => {}, ...overrides.editorContainer }, + ui: { + setFocus: () => {}, + requestRender: () => {}, + setShowHardwareCursor: () => {}, + setClearOnShrink: () => {}, + ...overrides.ui, + }, + editor, + defaultEditor: { setPaddingX: () => {}, setAutocompleteMaxVisible: () => {}, ...overrides.defaultEditor }, + themeController: { applyFromSettings: async () => {}, ...overrides.themeController }, + customHeader: overrides.customHeader, + builtInHeader: overrides.builtInHeader, + resetExtensionUI: overrides.resetExtensionUI ?? (() => {}), + rebuildChatFromMessages: overrides.rebuildChatFromMessages ?? (() => {}), + setupAutocompleteProvider: overrides.setupAutocompleteProvider ?? (() => {}), + setupExtensionShortcuts: overrides.setupExtensionShortcuts ?? (() => {}), + showLoadedResources: overrides.showLoadedResources ?? (() => {}), + maybeSaveImplicitProjectTrustAfterReload: overrides.maybeSaveImplicitProjectTrustAfterReload ?? (() => false), + showStatus: overrides.showStatus ?? (() => {}), + showWarning: overrides.showWarning ?? (() => {}), + showError: overrides.showError ?? (() => {}), + }; +} + +type MessageEvent = Extract; + +function getMessageText(event: MessageEvent): string { + const message = event.message; + if (!("content" in message)) { + return ""; + } + const content = message.content; + if (typeof content === "string") { + return content; + } + return content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join(""); +} + +describe("regression #5943: session_start transient UI", () => { + it("renders replacement session state before session_start handlers can notify", async () => { + const events: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", (_event, ctx) => { + ctx.ui.notify("Hello Error", "error"); + }); + }, + ], + }); + + try { + const context: RebindContext = { + applyRuntimeSettings: () => events.push("apply"), + renderCurrentSessionState: () => events.push("render"), + bindCurrentSessionExtensions: async () => { + events.push("bind"); + await harness.session.bindExtensions({ + uiContext: createUiContext((message) => events.push(`notify:${message}`)), + mode: "tui", + }); + }, + subscribeToAgent: () => events.push("subscribe"), + updateAvailableProviderCount: async () => {}, + updateEditorBorderColor: () => {}, + updateTerminalTitle: () => {}, + }; + + await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true }); + + expect(events).toEqual(["apply", "render", "subscribe", "bind", "notify:Hello Error"]); + } finally { + harness.cleanup(); + } + }); + + it("subscribes before replacement session_start handlers send messages", async () => { + const events: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", () => { + pi.sendMessage({ + customType: "session-start", + content: "custom from start", + display: true, + }); + }); + }, + ], + }); + + try { + const context: RebindContext = { + applyRuntimeSettings: () => {}, + renderCurrentSessionState: () => events.push("render"), + bindCurrentSessionExtensions: async () => { + events.push("bind"); + await harness.session.bindExtensions({ + uiContext: createUiContext(() => {}), + mode: "tui", + }); + }, + subscribeToAgent: () => { + events.push("subscribe"); + harness.session.subscribe((event) => { + if (event.type !== "message_start" && event.type !== "message_end") { + return; + } + events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`); + }); + }, + updateAvailableProviderCount: async () => {}, + updateEditorBorderColor: () => {}, + updateTerminalTitle: () => {}, + }; + + await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true }); + + expect(events).toEqual([ + "render", + "subscribe", + "bind", + "message_start:custom:custom from start", + "message_end:custom:custom from start", + ]); + } finally { + harness.cleanup(); + } + }); + + it("subscribes before replacement session_start handlers send user messages", async () => { + const events: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", () => { + pi.sendUserMessage("user from start"); + }); + }, + ], + }); + harness.setResponses([fauxAssistantMessage("assistant from start")]); + + try { + const context: RebindContext = { + applyRuntimeSettings: () => {}, + renderCurrentSessionState: () => events.push("render"), + bindCurrentSessionExtensions: async () => { + events.push("bind"); + await harness.session.bindExtensions({ + uiContext: createUiContext(() => {}), + mode: "tui", + }); + }, + subscribeToAgent: () => { + events.push("subscribe"); + harness.session.subscribe((event) => { + if (event.type !== "message_start" && event.type !== "message_end") { + return; + } + events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`); + }); + }, + updateAvailableProviderCount: async () => {}, + updateEditorBorderColor: () => {}, + updateTerminalTitle: () => {}, + }; + + await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true }); + await harness.session.agent.waitForIdle(); + + expect(events.slice(0, 3)).toEqual(["render", "subscribe", "bind"]); + expect(events).toContain("message_start:user:user from start"); + expect(events).toContain("message_end:user:user from start"); + expect(events).toContain("message_end:assistant:assistant from start"); + } finally { + harness.cleanup(); + } + }); + + it("runs the reload render hook before reload session_start handlers can notify", async () => { + const events: string[] = []; + const beforeSessionStart = vi.fn(() => { + events.push("render"); + }); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", (event, ctx) => { + events.push(`start:${event.reason}`); + ctx.ui.notify(`notify:${event.reason}`, "error"); + }); + }, + ], + }); + + try { + await harness.session.bindExtensions({ + uiContext: createUiContext((message) => events.push(message)), + mode: "tui", + }); + expect(events).toEqual(["start:startup", "notify:startup"]); + + events.length = 0; + await harness.session.reload({ beforeSessionStart }); + + expect(beforeSessionStart).toHaveBeenCalledTimes(1); + expect(events).toEqual(["render", "start:reload", "notify:reload"]); + } finally { + harness.cleanup(); + } + }); + + it("refreshes hideThinkingBlock before rebuilding chat during reload", async () => { + initTheme("dark", false); + const events: string[] = []; + let context: ReloadCommandContext; + context = createReloadCommandContext({ + settingsManager: { getHideThinkingBlock: () => true }, + session: { + reload: async (options) => { + events.push("reload"); + await options?.beforeSessionStart?.(); + events.push(`start:${context.hideThinkingBlock}`); + }, + }, + rebuildChatFromMessages: () => { + events.push(`rebuild:${context.hideThinkingBlock}`); + }, + }); + + await interactiveModePrototype.handleReloadCommand.call(context); + + expect(context.hideThinkingBlock).toBe(true); + expect(events).toEqual(["reload", "rebuild:true", "start:true"]); + }); + + it("keeps the reload blocker focused until async reload completes", async () => { + initTheme("dark", false); + const editor = {}; + let focused: unknown; + let chatRestored = false; + let markReloadWaiting!: () => void; + let finishReload!: () => void; + const reloadWaiting = new Promise((resolve) => { + markReloadWaiting = resolve; + }); + const reloadFinished = new Promise((resolve) => { + finishReload = resolve; + }); + + const context = createReloadCommandContext({ + editor, + session: { + reload: async (options) => { + await options?.beforeSessionStart?.(); + markReloadWaiting(); + await reloadFinished; + }, + }, + ui: { + setFocus: (component) => { + focused = component; + }, + }, + rebuildChatFromMessages: () => { + chatRestored = true; + }, + }); + + const reloadPromise = interactiveModePrototype.handleReloadCommand.call(context); + await reloadWaiting; + + expect(chatRestored).toBe(true); + expect(focused).not.toBe(editor); + + finishReload(); + await reloadPromise; + + expect(focused).toBe(editor); + }); +});