From feba7a71d388cc5b35abde2b1e3c19db60e15e27 Mon Sep 17 00:00:00 2001 From: Pi on Code Agent <313040514+pi-on-code-agent[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:54:11 +0000 Subject: [PATCH 1/7] feat: edit and fork historical user messages --- media/style.css | 101 ++++++++++++++-- src/extension.ts | 172 ++++++++++++++++++++++++++- src/test/user-message-branch.test.ts | 94 +++++++++++++++ src/user-message-branch.ts | 52 ++++++++ src/webview-panel.ts | 36 ++++++ src/webview/handlers/index.ts | 61 +++++++++- 6 files changed, 505 insertions(+), 11 deletions(-) create mode 100644 src/test/user-message-branch.test.ts create mode 100644 src/user-message-branch.ts diff --git a/media/style.css b/media/style.css index c527ac4..236544c 100644 --- a/media/style.css +++ b/media/style.css @@ -2954,13 +2954,23 @@ html.pi-viewport-recovering #pi-web-app { padding-right: 42px; } -.user-copy-btn, -.tool-copy-btn { +.user-actions { position: absolute; right: 8px; top: 50%; + display: flex; + gap: 4px; transform: translateY(-50%); opacity: 0; + transition: opacity 120ms ease; +} + +.message.user:hover .user-actions, +.message.user:focus-within .user-actions { + opacity: 1; +} + +.user-actions button { padding: 2px 7px; border: 1px solid var(--pi-line); border-radius: 3px; @@ -2968,14 +2978,69 @@ html.pi-viewport-recovering #pi-web-app { color: var(--pi-muted); font-size: 11px; cursor: pointer; + white-space: nowrap; transition: opacity 120ms ease; } -.message.user:hover .user-copy-btn, -.user-copy-btn:focus-visible, -.tool-block:hover .tool-copy-btn, -.bash-execution:hover .tool-copy-btn, -.tool-copy-btn:focus-visible { +.user-actions button:hover, +.user-actions button:focus-visible { + border-color: var(--pi-lavender); + color: var(--pi-text); +} + +.user-edit-overlay { + position: absolute; + z-index: 6; + inset: 4px; + display: flex; + flex-direction: column; + gap: 6px; + padding: 8px; + border: 1px solid var(--pi-lavender); + border-radius: 4px; + background: var(--pi-surface); +} + +.user-edit-input { + width: 100%; + min-height: 64px; + resize: vertical; + padding: 6px; + border: 1px solid var(--pi-line); + border-radius: 2px; + background: var(--vscode-input-background, var(--pi-surface)); + color: var(--vscode-input-foreground, var(--pi-text)); + font: inherit; +} + +.user-edit-bar { + display: flex; + justify-content: flex-end; + gap: 6px; +} + +.user-edit-bar button { + padding: 3px 10px; + border: 1px solid var(--pi-line); + border-radius: 3px; + background: var(--pi-surface-raised); + color: var(--pi-muted); + font-size: 11px; + cursor: pointer; +} + +.user-edit-save { + border-color: var(--pi-lavender) !important; + color: var(--pi-lavender) !important; +} + +.user-copy-btn, +.user-edit-btn, +.user-fork-btn { + position: static; + right: auto; + top: auto; + transform: none; opacity: 1; } @@ -2984,6 +3049,28 @@ html.pi-viewport-recovering #pi-web-app { position: relative; } +.tool-copy-btn { + position: absolute; + right: 8px; + top: 50%; + transform: translateY(-50%); + opacity: 0; + padding: 2px 7px; + border: 1px solid var(--pi-line); + border-radius: 3px; + background: var(--pi-surface-raised); + color: var(--pi-muted); + font-size: 11px; + cursor: pointer; + transition: opacity 120ms ease; +} + +.tool-block:hover .tool-copy-btn, +.bash-execution:hover .tool-copy-btn, +.tool-copy-btn:focus-visible { + opacity: 1; +} + /* Generic extension tools expose their arguments in the header. */ .tool-block .tool-header-arguments { min-width: 0; diff --git a/src/extension.ts b/src/extension.ts index 1f8c177..06d8a92 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -21,6 +21,10 @@ import { extractSessionId } from "./session-reference.js"; import { shouldRevealSessionPanel } from "./session-startup.js"; import { findReusableDraft, shouldPromoteDraft } from "./session-draft.js"; import { normalizeSessionRename } from "./session-rename.js"; +import { + forkSessionLabel, + resolveUserMessageEditTarget, +} from "./user-message-branch.js"; import { emptySessionListPreferences, isSessionArchived, @@ -419,6 +423,102 @@ async function deleteSidebarSession(target: PiSidebarDeleteTarget): Promise { + const sourceSm = sw.piService.sessionManagerInstance; + if (!sourceSm) { throw new Error("Session has no session manager."); } + const sourcePath = sw.piService.sessionFilePath ?? sw.restoringPath; + if (!sourcePath) { throw new Error("Session has no persisted file."); } + const title = cleanSessionTitle( + sw.piService.sessionName ?? sw.webviewPanel.summary ?? sw.label, + ); + + let forkedPath: string | null = null; + const tempPi = new PiService(context.secrets); + try { + const result = await tempPi.initialize({ openPath: sourcePath, cwd: sw.cwd }); + if (!result.success) { throw new Error(`Cannot open source session: ${result.error}`); } + const srcSm = tempPi.sessionManagerInstance; + if (!srcSm) { throw new Error("Source session has no session manager."); } + const entries = (srcSm.getEntries?.() ?? []) as Array<{ + id?: string; + type?: string; + message?: { role?: string }; + }>; + const target = resolveUserMessageEditTarget(entries, entryId); + if (!target) { throw new Error("Selected entry is not a persisted user message."); } + if (target.predecessorId) { + forkedPath = srcSm.createBranchedSession(target.predecessorId); + if (!forkedPath) { throw new Error("Failed to build the rewritten session."); } + } + } finally { + tempPi.dispose(); + } + + const replacement = forkedPath + ? createSessionWindow(context, { path: forkedPath, title }, false, sw.cwd) + : createSessionWindow(context, undefined, false, sw.cwd); + replacement.webviewPanel.initialWelcomeVisible = false; + setActiveSession(replacement); + void replacement.webviewPanel.show(); + sessionTreeProvider?.refresh(); + + const ok = await initSessionInBackground( + context, + replacement, + forkedPath ? { openPath: forkedPath } : { fresh: true }, + ); + if (!ok || !replacement.initialized) { + removeSession(replacement); + throw new Error("Failed to open the rewritten session."); + } + + // Start the AI reply from the edited text before retiring the old session so + // a prompt failure leaves the original conversation intact. + try { + await replacement.piService.sendPrompt(text); + } catch (error: unknown) { + replacement.webviewPanel.onDispose = null; + replacement.piService.dispose(); + removeSession(replacement); + replacement.webviewPanel.dispose(); + throw new Error( + `Rewrite ready but could not start the reply: ${error instanceof Error ? error.message : String(error)}`, + ); + } + replacement.piService.setSessionName(title); + replacement.label = title; + + // Retire the superseded window and delete its file: the edited history is + // now fully represented by the replacement session. + const wasActive = activeSessionWindow === sw; + sw.webviewPanel.onDispose = null; + sw.piService.dispose(); + removeSession(sw); + sw.webviewPanel.dispose(); + await deleteSessionFileIfPresent(sourcePath); + await refreshPastSessionsList(); + await saveOpenSessionPaths(); + sessionTreeProvider?.refresh(); + if (wasActive || activeSessionWindow !== replacement) { + await replacement.webviewPanel.show(); + } + vscode.window.showInformationMessage("Message rewritten; continuing from the new text."); +} + /** Create a new session window pair. Restore hints prevent transient duplicate rows. */ function createSessionWindow( context: vscode.ExtensionContext, @@ -460,6 +560,13 @@ function createSessionWindow( // If saved successfully, it will appear in Past Sessions on next refresh. webviewPanel.onDispose = handlePanelDispose(sw); + // History-message actions are routed through commands so the activate-scope + // session helpers can run them with the full session registry. + webviewPanel.onEditUserMessage = (entryId, text) => + void vscode.commands.executeCommand("pi-on-code.editHistoryMessage", sw.id, entryId, text); + webviewPanel.onForkUserMessage = (entryId) => + void vscode.commands.executeCommand("pi-on-code.forkHistoryMessage", sw.id, entryId); + sessions.push(sw); return sw; } @@ -805,7 +912,11 @@ export async function activate(context: vscode.ExtensionContext): Promise // ── Fork helpers ───────────────────────────────────── /** Fork at a specific entry within an already-open session. */ - async function doForkFromOpenEntry(sessionId: string, entryId: string): Promise { + async function doForkFromOpenEntry( + sessionId: string, + entryId: string, + forkLabel?: string, + ): Promise { const srcSw = sessions.find((s) => s.id === sessionId); if (!srcSw || !srcSw.piService.sessionManagerInstance) { throw new Error(`Source session not found (id=${sessionId}).`); @@ -850,7 +961,7 @@ export async function activate(context: vscode.ExtensionContext): Promise } piLog(`doForkFromOpenEntry: forked to ${forkedPath}`); - await openForkedSession(forkedPath, srcSw.cwd); + await openForkedSession(forkedPath, srcSw.cwd, forkLabel); } /** Fork a past session at its current leaf (opens the session, then forks). */ @@ -886,7 +997,11 @@ export async function activate(context: vscode.ExtensionContext): Promise } /** Create a new session window initialized from a forked session file. */ - async function openForkedSession(forkedPath: string, cwd = getWorkspaceCwd()): Promise { + async function openForkedSession( + forkedPath: string, + cwd = getWorkspaceCwd(), + label?: string, + ): Promise { const newSw = createSessionWindow(context, { path: forkedPath }, false, cwd); setActiveSession(newSw); void newSw.webviewPanel.show(); @@ -899,6 +1014,12 @@ export async function activate(context: vscode.ExtensionContext): Promise throw new Error("Failed to initialize forked session."); } + if (label) { + newSw.piService.setSessionName(label); + newSw.label = label; + sessionTreeProvider?.refresh(); + } + // sendInitialMessages() is already called during initialize() inside the // batch-start/batch-end wrapper — no need for a second call here. sessionTreeProvider?.refresh(); @@ -966,6 +1087,51 @@ export async function activate(context: vscode.ExtensionContext): Promise }), ); + // ── Edit / fork from a historical user message (transcript actions) ── + context.subscriptions.push( + vscode.commands.registerCommand( + "pi-on-code.editHistoryMessage", + async (sessionId: unknown, entryId: unknown, text: unknown) => { + if (typeof sessionId !== "string" || typeof entryId !== "string" || typeof text !== "string") { + return; + } + const sw = sessions.find((session) => session.id === sessionId); + if (!sw) { + vscode.window.showErrorMessage("Source session not found."); + return; + } + try { + await editHistoryMessageFromEntry(context, sw, entryId, text); +// eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (e: any) { + vscode.window.showErrorMessage(`Edit failed: ${e.message ?? e}`); + } + }, + ), + ); + context.subscriptions.push( + vscode.commands.registerCommand( + "pi-on-code.forkHistoryMessage", + async (sessionId: unknown, entryId: unknown) => { + if (typeof sessionId !== "string" || typeof entryId !== "string") { return; } + const sw = sessions.find((session) => session.id === sessionId); + if (!sw) { + vscode.window.showErrorMessage("Source session not found."); + return; + } + try { + const title = cleanSessionTitle( + sw.piService.sessionName ?? sw.webviewPanel.summary ?? sw.label, + ); + await doForkFromOpenEntry(sw.id, entryId, forkSessionLabel(title)); +// eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (e: any) { + vscode.window.showErrorMessage(`Fork failed: ${e.message ?? e}`); + } + }, + ), + ); + // ── Compact session context ──────────────────────── context.subscriptions.push( vscode.commands.registerCommand("pi-on-code.compact", async () => { diff --git a/src/test/user-message-branch.test.ts b/src/test/user-message-branch.test.ts new file mode 100644 index 0000000..4ab3346 --- /dev/null +++ b/src/test/user-message-branch.test.ts @@ -0,0 +1,94 @@ +import * as assert from "node:assert"; +import { readFileSync } from "node:fs"; +import { + forkSessionLabel, + isUserMessageEntry, + resolveUserMessageEditTarget, + type SessionEntryLike, +} from "../user-message-branch.js"; + +const userEntry = (id: string, role = "user"): SessionEntryLike => ({ + id, + type: "message", + message: { role }, +}); + +suite("User message branch helpers", () => { + test("recognizes only persisted user message entries", () => { + assert.ok(isUserMessageEntry(userEntry("u1"))); + assert.ok(!isUserMessageEntry({ id: "a1", type: "message", message: { role: "assistant" } })); + assert.ok(!isUserMessageEntry({ id: "c1", type: "custom_message" })); + assert.ok(!isUserMessageEntry({ id: "x" })); + }); + + test("branches at the preceding entry when the target is not first", () => { + const entries = [ + userEntry("u1"), + { id: "a1", type: "message", message: { role: "assistant" } }, + userEntry("u2"), + { id: "t1", type: "message", message: { role: "assistant" } }, + ]; + assert.deepStrictEqual(resolveUserMessageEditTarget(entries, "u2"), { + entryIndex: 2, + predecessorId: "a1", + }); + }); + + test("has no predecessor for the first user message", () => { + const entries = [userEntry("u1"), userEntry("u2")]; + assert.deepStrictEqual(resolveUserMessageEditTarget(entries, "u1"), { + entryIndex: 0, + predecessorId: null, + }); + }); + + test("rejects missing or non-user targets", () => { + const entries = [userEntry("u1"), { id: "a1", type: "message", message: { role: "assistant" } }]; + assert.strictEqual(resolveUserMessageEditTarget(entries, "missing"), null); + assert.strictEqual(resolveUserMessageEditTarget(entries, "a1"), null); + }); + + test("formats fork labels from the original title", () => { + assert.strictEqual(forkSessionLabel("Fix lint errors"), ""); + assert.strictEqual(forkSessionLabel(" "), ""); + assert.strictEqual(forkSessionLabel(""), ""); + }); + + test("caps fork labels by code points", () => { + const long = "title ".repeat(30); + const label = forkSessionLabel(long); + assert.ok(label.startsWith("")); + assert.ok(Array.from(label.slice(7, -1)).length <= 41, "label body stays near the cap"); + }); + + test("wires transcript actions through panel callbacks and commands", () => { + const handlers = readFileSync( + new URL("../../src/webview/handlers/index.ts", import.meta.url), + "utf8", + ); + assert.match(handlers, /user-message-edit/); + assert.match(handlers, /user-message-fork/); + assert.match(handlers, /className = "user-edit-btn"/); + assert.match(handlers, /className = "user-fork-btn"/); + assert.match(handlers, /\.user-edit-overlay/); + + const panel = readFileSync( + new URL("../../src/webview-panel.ts", import.meta.url), + "utf8", + ); + assert.match(panel, /case "user-message-edit"/); + assert.match(panel, /case "user-message-fork"/); + + const extension = readFileSync( + new URL("../../src/extension.ts", import.meta.url), + "utf8", + ); + assert.match(extension, /"pi-on-code\.editHistoryMessage"/); + assert.match(extension, /"pi-on-code\.forkHistoryMessage"/); + assert.match(extension, /resolveUserMessageEditTarget\(entries, entryId\)/); + assert.match(extension, /createBranchedSession\(target\.predecessorId\)/); + assert.match(extension, /sendPrompt\(text\)/); + assert.match(extension, /forkSessionLabel\(title\)/); + }); +}); diff --git a/src/user-message-branch.ts b/src/user-message-branch.ts new file mode 100644 index 0000000..5c296be --- /dev/null +++ b/src/user-message-branch.ts @@ -0,0 +1,52 @@ +/** + * Helpers for editing or forking a conversation at an arbitrary user message. + * The session file itself is only ever rewritten by composing SDK primitives + * (`createBranchedSession`) so the extension never hand-edits JSONL entries. + */ + +export interface SessionEntryLike { + id?: string; + type?: string; + message?: { role?: string }; +} + +export interface UserMessageActionTarget { + entryIndex: number; + /** Entry id the branch should stop after. Null when the message is first. */ + predecessorId: string | null; +} + +export function isUserMessageEntry(entry: SessionEntryLike): boolean { + return entry.type === "message" && entry.message?.role === "user"; +} + +/** + * Locate a persisted user message entry and compute its edit branch point. + * + * Editing is expressed as: branch everything *before* the target message, then + * append the rewritten text as a fresh user message. Entries before the target + * keep their ids (prefix cache reuse); the target and everything after it are + * excluded from the new session. + */ +export function resolveUserMessageEditTarget( + entries: readonly SessionEntryLike[], + targetId: string, +): UserMessageActionTarget | null { + const entryIndex = entries.findIndex((entry) => entry.id === targetId); + if (entryIndex < 0 || !isUserMessageEntry(entries[entryIndex])) { return null; } + const predecessor = entryIndex > 0 ? entries[entryIndex - 1] : undefined; + return { + entryIndex, + predecessorId: predecessor?.id ?? null, + }; +} + +/** Display label for a fork, e.g. ``. */ +export function forkSessionLabel(originalTitle: string): string { + const title = originalTitle.replace(/\s+/g, " ").trim() || "Untitled session"; + const codePoints = Array.from(title); + const capped = codePoints.length > 40 + ? `${codePoints.slice(0, 40).join("").replace(/[\s.…]+$/u, "")}…` + : title; + return ``; +} diff --git a/src/webview-panel.ts b/src/webview-panel.ts index 61d6b20..2313edd 100644 --- a/src/webview-panel.ts +++ b/src/webview-panel.ts @@ -67,6 +67,8 @@ export class PiWebviewPanel { /** Callback invoked when the panel is disposed (VS Code tab closed) */ private _onDispose: PanelDisposeCallback | null = null; private _onBeforePrompt: ((text: string) => void) | null = null; + private _onEditUserMessage: ((entryId: string, text: string) => void | Promise) | null = null; + private _onForkUserMessage: ((entryId: string) => void | Promise) | null = null; private _initialWelcomeVisible = false; constructor( @@ -83,6 +85,12 @@ export class PiWebviewPanel { /** Register a callback that fires when the panel/webview is closed. */ set onDispose(cb: PanelDisposeCallback | null) { this._onDispose = cb; } set onBeforePrompt(cb: ((text: string) => void) | null) { this._onBeforePrompt = cb; } + set onEditUserMessage(cb: ((entryId: string, text: string) => void | Promise) | null) { + this._onEditUserMessage = cb; + } + set onForkUserMessage(cb: ((entryId: string) => void | Promise) | null) { + this._onForkUserMessage = cb; + } set initialWelcomeVisible(value: boolean) { this._initialWelcomeVisible = value; } /** Register a callback that fires when this panel/view becomes active. */ @@ -679,6 +687,34 @@ export class PiWebviewPanel { void this.triggerContextBudgetPicker(); break; + // Rewrite the conversation from an edited user message + case "user-message-edit": + if (typeof message.entryId === "string" && typeof message.text === "string") { + try { + await this._onEditUserMessage?.(message.entryId, message.text); + } catch (error: unknown) { + this.postMessage({ + type: "error", + data: { message: `Edit failed: ${error instanceof Error ? error.message : String(error)}` }, + }); + } + } + break; + + // Fork an independent session at a user message + case "user-message-fork": + if (typeof message.entryId === "string") { + try { + await this._onForkUserMessage?.(message.entryId); + } catch (error: unknown) { + this.postMessage({ + type: "error", + data: { message: `Fork failed: ${error instanceof Error ? error.message : String(error)}` }, + }); + } + } + break; + // Request settings state (#2, #8) case "resendUserMessage": if (message.text) { diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 7a05c24..ddad1f5 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -610,13 +610,72 @@ export function handleChatMessage(data: any) { if (editorContext) { mc.appendChild(editorContext); } } if (data.role === "user" && typeof data.content === "string" && data.content.trim()) { + var actions = document.createElement("div"); + actions.className = "user-actions"; + var userCopyButton = document.createElement("button"); userCopyButton.type = "button"; userCopyButton.className = "user-copy-btn"; userCopyButton.textContent = "Copy"; userCopyButton.setAttribute("aria-label", "Copy user message"); (userCopyButton as HTMLButtonElement & { _copyText?: string })._copyText = data.content; - el.appendChild(userCopyButton); + actions.appendChild(userCopyButton); + + // Edit / fork operate on the persisted entry, so they need its entry id + // and are hidden while a run is in progress. + if (data.entryId && !state.isStreaming) { + var editButton = document.createElement("button"); + editButton.type = "button"; + editButton.className = "user-edit-btn"; + editButton.textContent = "Edit"; + editButton.title = "Edit this message and continue from here"; + editButton.setAttribute("aria-label", "Edit user message and continue from here"); + editButton.addEventListener("click", function () { + if (el.querySelector(".user-edit-overlay")) { return; } + var overlay = document.createElement("div"); + overlay.className = "user-edit-overlay"; + var area = document.createElement("textarea"); + area.className = "user-edit-input"; + area.value = data.content; + area.setAttribute("aria-label", "Edited user message"); + var bar = document.createElement("div"); + bar.className = "user-edit-bar"; + var saveButton = document.createElement("button"); + saveButton.type = "button"; + saveButton.className = "user-edit-save"; + saveButton.textContent = "Save and continue"; + var cancelButton = document.createElement("button"); + cancelButton.type = "button"; + cancelButton.className = "user-edit-cancel"; + cancelButton.textContent = "Cancel"; + bar.append(cancelButton, saveButton); + overlay.append(area, bar); + el.appendChild(overlay); + area.focus(); + area.select(); + cancelButton.addEventListener("click", function () { overlay.remove(); }); + saveButton.addEventListener("click", function () { + var next = area.value.trim(); + if (next && next !== data.content) { + window.__vscode.postMessage({ type: "user-message-edit", entryId: data.entryId, text: next }); + overlay.remove(); + } + }); + }); + actions.appendChild(editButton); + + var forkButton = document.createElement("button"); + forkButton.type = "button"; + forkButton.className = "user-fork-btn"; + forkButton.textContent = "Fork"; + forkButton.title = "Fork a new session at this message"; + forkButton.setAttribute("aria-label", "Fork session at this message"); + forkButton.addEventListener("click", function () { + window.__vscode.postMessage({ type: "user-message-fork", entryId: data.entryId }); + }); + actions.appendChild(forkButton); + } + el.appendChild(actions); } state.chatContainer.appendChild(el); requestLocalMarkdownImages(el); From f555be49cc8889474528e33f87cf44ee44bdc6cc Mon Sep 17 00:00:00 2001 From: Pi on Code Agent <313040514+pi-on-code-agent[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:06:06 +0000 Subject: [PATCH 2/7] fix: guard rewrites during runs and commit on reply start --- media/style.css | 6 ++ src/edit-flow.ts | 57 ++++++++++++++++++ src/extension.ts | 86 +++++++++++++++++++++------- src/test/edit-flow.test.ts | 54 +++++++++++++++++ src/test/user-message-branch.test.ts | 8 +++ src/webview/handlers/index.ts | 11 ++++ 6 files changed, 201 insertions(+), 21 deletions(-) create mode 100644 src/edit-flow.ts create mode 100644 src/test/edit-flow.test.ts diff --git a/media/style.css b/media/style.css index 236544c..5fda022 100644 --- a/media/style.css +++ b/media/style.css @@ -3034,6 +3034,12 @@ html.pi-viewport-recovering #pi-web-app { color: var(--pi-lavender) !important; } +.user-edit-btn:disabled, +.user-fork-btn:disabled { + opacity: 0.45; + pointer-events: none; +} + .user-copy-btn, .user-edit-btn, .user-fork-btn { diff --git a/src/edit-flow.ts b/src/edit-flow.ts new file mode 100644 index 0000000..35fde8b --- /dev/null +++ b/src/edit-flow.ts @@ -0,0 +1,57 @@ +/** + * Decision helpers for the destructive message-rewrite flow. Keeping these + * pure makes the rollback / delete-ordering semantics testable without an SDK. + */ + +export interface ReplyStartObservation { + /** True once the replacement agent run actually started streaming. */ + started: boolean; + /** True once the sendPrompt promise settled (success or error). */ + promptDone: boolean; +} + +export interface WaitForReplyStartOptions { + timeoutMs?: number; + delayMs?: number; +} + +/** + * Watch a freshly issued prompt until either the agent run starts, the prompt + * settles without starting, or the timeout elapses. Delays are real timers so + * tests may pass small values. + */ +export async function waitForReplyStart( + isStreaming: () => boolean, + isPromptDone: () => boolean, + options: WaitForReplyStartOptions = {}, +): Promise { + const timeoutMs = options.timeoutMs ?? 10_000; + const delayMs = options.delayMs ?? 40; + const deadline = Date.now() + timeoutMs; + for (;;) { + const started = isStreaming(); + const promptDone = isPromptDone(); + if (started || promptDone) { return { started, promptDone }; } + if (Date.now() >= deadline) { return { started: false, promptDone: false }; } + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } +} + +export type EditOutcome = "keep" | "rollback"; + +/** + * Commit the rewrite once the reply has started (content is being produced): + * later aborts or API errors keep the replacement session. Roll back (and + * delete the forked file) only when the prompt never produced a run — either + * it failed or it stalled without starting. + */ +export function resolveEditOutcome(observation: ReplyStartObservation): EditOutcome { + if (observation.started) { return "keep"; } + if (observation.promptDone) { return "rollback"; } + return "rollback"; +} + +/** A running session must never be branched from or superseded. */ +export function canEditSession(isStreaming: boolean): boolean { + return !isStreaming; +} diff --git a/src/extension.ts b/src/extension.ts index 06d8a92..45c7e4f 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -25,6 +25,11 @@ import { forkSessionLabel, resolveUserMessageEditTarget, } from "./user-message-branch.js"; +import { + canEditSession, + resolveEditOutcome, + waitForReplyStart, +} from "./edit-flow.js"; import { emptySessionListPreferences, isSessionArchived, @@ -440,6 +445,9 @@ async function editHistoryMessageFromEntry( ): Promise { const sourceSm = sw.piService.sessionManagerInstance; if (!sourceSm) { throw new Error("Session has no session manager."); } + if (!canEditSession(sw.isStreaming)) { + throw new Error("Stop the current run before editing a historical message."); + } const sourcePath = sw.piService.sessionFilePath ?? sw.restoringPath; if (!sourcePath) { throw new Error("Session has no persisted file."); } const title = cleanSessionTitle( @@ -482,40 +490,72 @@ async function editHistoryMessageFromEntry( forkedPath ? { openPath: forkedPath } : { fresh: true }, ); if (!ok || !replacement.initialized) { + // Nothing was produced yet: clean up the branch file and window so no + // orphan session leaks into Past Sessions. + replacement.webviewPanel.onDispose = null; + replacement.piService.dispose(); removeSession(replacement); + replacement.webviewPanel.dispose(); + if (forkedPath) { await deleteSessionFileIfPresent(forkedPath); } + await refreshPastSessionsList(); throw new Error("Failed to open the rewritten session."); } - // Start the AI reply from the edited text before retiring the old session so - // a prompt failure leaves the original conversation intact. - try { - await replacement.piService.sendPrompt(text); - } catch (error: unknown) { + replacement.piService.setSessionName(title); + replacement.label = title; + + // Start the reply from the edited text. The rewrite is committed as soon as + // the agent run actually starts (or the prompt settles cleanly without being + // observed): a later abort or API failure keeps the replacement session. + // Only when no run ever started do we roll back and delete the branch file. + let promptError: unknown = null; + const promptDone = { value: false }; + const promptPromise = replacement.piService.sendPrompt(text); + promptPromise.then( + () => { promptDone.value = true; }, + (error: unknown) => { promptDone.value = true; promptError = error; }, + ); + const observation = await waitForReplyStart( + () => replacement.isStreaming, + () => promptDone.value, + ); + if (resolveEditOutcome(observation) === "rollback") { replacement.webviewPanel.onDispose = null; replacement.piService.dispose(); removeSession(replacement); replacement.webviewPanel.dispose(); - throw new Error( - `Rewrite ready but could not start the reply: ${error instanceof Error ? error.message : String(error)}`, - ); + if (forkedPath) { await deleteSessionFileIfPresent(forkedPath); } + await refreshPastSessionsList(); + const detail = promptError instanceof Error ? promptError.message : String(promptError ?? "the reply did not start"); + throw new Error(`Rewrite rolled back (${detail}); the original session is unchanged.`); } - replacement.piService.setSessionName(title); - replacement.label = title; - // Retire the superseded window and delete its file: the edited history is - // now fully represented by the replacement session. - const wasActive = activeSessionWindow === sw; - sw.webviewPanel.onDispose = null; - sw.piService.dispose(); - removeSession(sw); - sw.webviewPanel.dispose(); - await deleteSessionFileIfPresent(sourcePath); + // Commit: retire and delete the superseded session now that the rewritten + // conversation is producing a reply. Re-check it is idle so a file that is + // still being appended is never deleted. + if (!canEditSession(sw.isStreaming)) { + vscode.window.showWarningMessage( + "The original session started running during the rewrite; its file was kept.", + ); + } else { + const wasActive = activeSessionWindow === sw; + sw.webviewPanel.onDispose = null; + sw.piService.dispose(); + removeSession(sw); + sw.webviewPanel.dispose(); + await deleteSessionFileIfPresent(sourcePath); + if (wasActive || activeSessionWindow !== replacement) { + await replacement.webviewPanel.show(); + } + } await refreshPastSessionsList(); await saveOpenSessionPaths(); sessionTreeProvider?.refresh(); - if (wasActive || activeSessionWindow !== replacement) { - await replacement.webviewPanel.show(); - } + void promptPromise.then(undefined, (error: unknown) => { + void vscode.window.showErrorMessage( + `Reply stopped after the rewrite: ${error instanceof Error ? error.message : String(error)}`, + ); + }); vscode.window.showInformationMessage("Message rewritten; continuing from the new text."); } @@ -1123,6 +1163,10 @@ export async function activate(context: vscode.ExtensionContext): Promise const title = cleanSessionTitle( sw.piService.sessionName ?? sw.webviewPanel.summary ?? sw.label, ); + if (!canEditSession(sw.isStreaming)) { + vscode.window.showWarningMessage("Stop the current run before forking a historical message."); + return; + } await doForkFromOpenEntry(sw.id, entryId, forkSessionLabel(title)); // eslint-disable-next-line @typescript-eslint/no-explicit-any } catch (e: any) { diff --git a/src/test/edit-flow.test.ts b/src/test/edit-flow.test.ts new file mode 100644 index 0000000..4d85b70 --- /dev/null +++ b/src/test/edit-flow.test.ts @@ -0,0 +1,54 @@ +import * as assert from "node:assert"; +import { + canEditSession, + resolveEditOutcome, + waitForReplyStart, +} from "../edit-flow.js"; + +suite("Message rewrite flow decisions", () => { + test("never allows editing or superseding a running session", () => { + assert.ok(canEditSession(false)); + assert.ok(!canEditSession(true)); + }); + + test("commits once the replacement reply actually starts", async () => { + let streaming = false; + const start = waitForReplyStart( + () => streaming, + () => false, + { timeoutMs: 1000, delayMs: 5 }, + ); + await new Promise((resolve) => setTimeout(resolve, 15)); + streaming = true; + const observation = await start; + assert.deepStrictEqual(observation, { started: true, promptDone: false }); + assert.strictEqual(resolveEditOutcome(observation), "keep"); + }); + + test("rolls back when the prompt settles without starting a run", async () => { + let done = false; + const waiting = waitForReplyStart( + () => false, + () => done, + { timeoutMs: 1000, delayMs: 5 }, + ); + done = true; + const observation = await waiting; + assert.deepStrictEqual(observation, { started: false, promptDone: true }); + assert.strictEqual(resolveEditOutcome(observation), "rollback"); + }); + + test("rolls back a stalled prompt that never starts or settles", async () => { + const observation = await waitForReplyStart( + () => false, + () => false, + { timeoutMs: 20, delayMs: 5 }, + ); + assert.deepStrictEqual(observation, { started: false, promptDone: false }); + assert.strictEqual(resolveEditOutcome(observation), "rollback"); + }); + + test("keeps the replacement even when a started run later fails", () => { + assert.strictEqual(resolveEditOutcome({ started: true, promptDone: true }), "keep"); + }); +}); diff --git a/src/test/user-message-branch.test.ts b/src/test/user-message-branch.test.ts index 4ab3346..cebeaa8 100644 --- a/src/test/user-message-branch.test.ts +++ b/src/test/user-message-branch.test.ts @@ -72,6 +72,9 @@ suite("User message branch helpers", () => { assert.match(handlers, /className = "user-edit-btn"/); assert.match(handlers, /className = "user-fork-btn"/); assert.match(handlers, /\.user-edit-overlay/); + assert.match(handlers, /function setUserEditActionsEnabled\(enabled: boolean\)/); + assert.match(handlers, /setUserEditActionsEnabled\(false\)/); + assert.match(handlers, /setUserEditActionsEnabled\(true\)/); const panel = readFileSync( new URL("../../src/webview-panel.ts", import.meta.url), @@ -90,5 +93,10 @@ suite("User message branch helpers", () => { assert.match(extension, /createBranchedSession\(target\.predecessorId\)/); assert.match(extension, /sendPrompt\(text\)/); assert.match(extension, /forkSessionLabel\(title\)/); + assert.match(extension, /canEditSession\(sw\.isStreaming\)/); + assert.match(extension, /waitForReplyStart\(/); + assert.match(extension, /resolveEditOutcome\(observation\)/); + assert.match(extension, /deleteSessionFileIfPresent\(forkedPath\)/); + assert.match(extension, /Stop the current run before editing a historical message\./); }); }); diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index ddad1f5..3b25381 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -252,9 +252,19 @@ function handleExtensionMessage(msg: any): void { // ═══ Agent Lifecycle ═══════════════════════════════════ // ═══ Agent Lifecycle ═══════════════════════════════════ +/** Keep Edit/Fork disabled while an agent run is active so history rewrites + * can never race a running session. */ +function setUserEditActionsEnabled(enabled: boolean): void { + document.querySelectorAll(".user-edit-btn, .user-fork-btn").forEach((button) => { + button.disabled = !enabled; + button.setAttribute("aria-disabled", String(!enabled)); + }); +} + export function handleAgentStart() { logEvent("agent-start", { bashBlocksN: Object.keys(state.bashBlocks).length, toolBlocksN: Object.keys(state.currentToolBlocks).length }); state.isStreaming = true; + setUserEditActionsEnabled(false); state.queueMode = "steer"; // reset to default on new stream state.assistantToolCallIds = {}; // Do NOT clear the live panel here — extension cards (like tldr summaries) @@ -333,6 +343,7 @@ export function handleAgentEnd() { toolKeys: Object.keys(state.currentToolBlocks), }); state.isStreaming = false; + setUserEditActionsEnabled(true); state.isRetrying = false; state.assistantToolCallIds = {}; removeWorkingIndicator(); From c6f73ac66cab273521c3af1e00e3c0de2a2c32df Mon Sep 17 00:00:00 2001 From: auchan Date: Sun, 6 Sep 2026 14:13:42 +0800 Subject: [PATCH 3/7] feat: refine edit and fork of user messages --- media/style.css | 117 +++++++------ src/extension.ts | 129 +++++++++++---- src/pi-service.ts | 41 +++++ src/shared/protocol.ts | 6 + src/test/user-message-branch.test.ts | 34 ++-- src/test/user-message-style.test.ts | 19 ++- src/test/webview-history-autofill.test.ts | 37 +++++ src/user-message-branch.ts | 17 +- src/webview-panel.ts | 19 ++- src/webview/handlers/index.ts | 193 +++++++++++++++++----- src/webview/main.ts | 2 + src/webview/render/history-autofill.ts | 31 ++++ 12 files changed, 482 insertions(+), 163 deletions(-) create mode 100644 src/test/webview-history-autofill.test.ts create mode 100644 src/webview/render/history-autofill.ts diff --git a/media/style.css b/media/style.css index 5fda022..8a32b5d 100644 --- a/media/style.css +++ b/media/style.css @@ -2107,12 +2107,6 @@ html.pi-viewport-recovering #pi-web-app { min-width: 0; align-self: stretch; margin: 4px 0; - padding: 12px 14px; - border: 1px solid var(--pi-line); - border-left: 3px solid var(--pi-lavender); - border-radius: 6px; - background: var(--vscode-input-background, var(--pi-surface)); - background: color-mix(in srgb, var(--pi-lavender) 10%, var(--pi-surface)); } .message.user::before { @@ -2785,7 +2779,7 @@ html.pi-viewport-recovering #pi-web-app { width: min(340px, calc(100vw - 36px)); } - .message.user { + .message.user .message-content { padding: 10px 11px; } @@ -2921,11 +2915,17 @@ html.pi-viewport-recovering #pi-web-app { position: relative; } -.assistant-copy-btn { +.assistant-turn-actions { position: absolute; bottom: 6px; left: 0; + display: flex; + gap: 4px; opacity: 0; + transition: opacity 120ms ease; +} + +.assistant-turn-actions button { padding: 2px 7px; border: 1px solid var(--pi-line); border-radius: 3px; @@ -2933,34 +2933,52 @@ html.pi-viewport-recovering #pi-web-app { color: var(--pi-muted); font-size: 11px; cursor: pointer; - transition: opacity 120ms ease; + white-space: nowrap; + transition: border-color 120ms ease, color 120ms ease; } -.message.assistant:hover .assistant-copy-btn, -.assistant-copy-btn:focus-visible { - opacity: 1; +.assistant-turn-actions button:hover, +.assistant-turn-actions button:focus-visible { + border-color: var(--pi-lavender); + color: var(--pi-text); +} +.message.assistant:hover .assistant-turn-actions, +.assistant-turn-actions:focus-within { + opacity: 1; } -.message.assistant:has(.assistant-copy-btn) { +.message.assistant:has(.assistant-turn-actions) { padding-bottom: 28px; } .message.user { position: relative; + display: flex; + flex-direction: column; + align-items: stretch; } -.message.user:has(.user-copy-btn) { - padding-right: 42px; +/* User card chrome on the content bubble (declared after the shared + .message-content base so the border/background are not reset). The action + row is appended outside the bubble, below it, right-aligned. */ +.message.user .message-content { + position: relative; + align-self: stretch; + max-width: none; + padding: 12px 14px; + border: 1px solid var(--pi-line); + border-left: 3px solid var(--pi-lavender); + border-radius: 6px; + background: var(--vscode-input-background, var(--pi-surface)); + background: color-mix(in srgb, var(--pi-lavender) 10%, var(--pi-surface)); } .user-actions { - position: absolute; - right: 8px; - top: 50%; + align-self: flex-end; display: flex; - gap: 4px; - transform: translateY(-50%); + gap: 6px; + margin-top: 8px; opacity: 0; transition: opacity 120ms ease; } @@ -2970,6 +2988,7 @@ html.pi-viewport-recovering #pi-web-app { opacity: 1; } + .user-actions button { padding: 2px 7px; border: 1px solid var(--pi-line); @@ -2979,7 +2998,6 @@ html.pi-viewport-recovering #pi-web-app { font-size: 11px; cursor: pointer; white-space: nowrap; - transition: opacity 120ms ease; } .user-actions button:hover, @@ -2988,38 +3006,36 @@ html.pi-viewport-recovering #pi-web-app { color: var(--pi-text); } -.user-edit-overlay { - position: absolute; - z-index: 6; - inset: 4px; - display: flex; - flex-direction: column; - gap: 6px; - padding: 8px; - border: 1px solid var(--pi-lavender); - border-radius: 4px; - background: var(--pi-surface); +.message.user:has(.user-edit-input) .user-actions { + display: none; } +/* Inline editor: replaces the bubble content, so it inherits the message + bubble appearance (no extra border/background) and keeps at least two rows. */ .user-edit-input { + display: block; width: 100%; - min-height: 64px; + min-height: 3.6em; resize: vertical; - padding: 6px; - border: 1px solid var(--pi-line); - border-radius: 2px; - background: var(--vscode-input-background, var(--pi-surface)); - color: var(--vscode-input-foreground, var(--pi-text)); + overflow-y: auto; + padding: 0; + border: 0; + border-radius: 0; + outline: 0; + background: transparent; + color: inherit; font: inherit; + line-height: 1.6; } -.user-edit-bar { +.user-edit-actions { + align-self: flex-end; display: flex; - justify-content: flex-end; gap: 6px; + margin-top: 8px; } -.user-edit-bar button { +.user-edit-actions button { padding: 3px 10px; border: 1px solid var(--pi-line); border-radius: 3px; @@ -3029,27 +3045,22 @@ html.pi-viewport-recovering #pi-web-app { cursor: pointer; } +.user-edit-actions button:hover, +.user-edit-actions button:focus-visible { + border-color: var(--pi-lavender); + color: var(--pi-text); +} + .user-edit-save { border-color: var(--pi-lavender) !important; color: var(--pi-lavender) !important; } -.user-edit-btn:disabled, -.user-fork-btn:disabled { +.user-edit-btn:disabled { opacity: 0.45; pointer-events: none; } -.user-copy-btn, -.user-edit-btn, -.user-fork-btn { - position: static; - right: auto; - top: auto; - transform: none; - opacity: 1; -} - .tool-block, .bash-execution { position: relative; diff --git a/src/extension.ts b/src/extension.ts index 45c7e4f..527ab85 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -22,8 +22,7 @@ import { shouldRevealSessionPanel } from "./session-startup.js"; import { findReusableDraft, shouldPromoteDraft } from "./session-draft.js"; import { normalizeSessionRename } from "./session-rename.js"; import { - forkSessionLabel, - resolveUserMessageEditTarget, + forkSessionTitle, } from "./user-message-branch.js"; import { canEditSession, @@ -428,6 +427,43 @@ async function deleteSidebarSession(target: PiSidebarDeleteTarget): Promise unknown; + getEntries?: () => Array<{ id?: string; type?: string; message?: { id?: string; role?: string; content?: unknown } }>; + }, + entryId: string, + content?: string, +): unknown | null { + const entries = sm.getEntries?.() ?? []; + const direct = sm.getEntry?.(entryId); + if (direct) { return direct; } + const byId = entries.find((entry) => entry?.id === entryId); + if (byId) { return byId; } + const byMessageId = entries.find( + (entry) => entry?.type === "message" && entry?.message?.id === entryId, + ); + if (byMessageId) { return byMessageId; } + if (content) { + const text = content.trim(); + for (let index = entries.length - 1; index >= 0; index--) { + const entry = entries[index]; + if (entry?.type !== "message" || entry?.message?.role !== "user") { continue; } + const raw = entry.message.content; + const body = typeof raw === "string" + ? raw.trim() + : Array.isArray(raw) + ? raw.filter((item: { type?: string; text?: string }) => item?.type === "text").map((item) => item.text ?? "").join("\n").trim() + : ""; + if (body === text) { return entry; } + } + } + return null; +} + /** * Rewrite a conversation from an edited historical user message. * @@ -442,6 +478,7 @@ async function editHistoryMessageFromEntry( sw: SessionWindow, entryId: string, text: string, + originalText?: string, ): Promise { const sourceSm = sw.piService.sessionManagerInstance; if (!sourceSm) { throw new Error("Session has no session manager."); } @@ -455,6 +492,8 @@ async function editHistoryMessageFromEntry( ); let forkedPath: string | null = null; + // Resolve and branch through an isolated temporary manager so the live + // source session is never mutated by the SDK's createBranchedSession. const tempPi = new PiService(context.secrets); try { const result = await tempPi.initialize({ openPath: sourcePath, cwd: sw.cwd }); @@ -464,12 +503,19 @@ async function editHistoryMessageFromEntry( const entries = (srcSm.getEntries?.() ?? []) as Array<{ id?: string; type?: string; - message?: { role?: string }; + message?: { id?: string; role?: string; content?: unknown }; }>; - const target = resolveUserMessageEditTarget(entries, entryId); - if (!target) { throw new Error("Selected entry is not a persisted user message."); } - if (target.predecessorId) { - forkedPath = srcSm.createBranchedSession(target.predecessorId); + const resolved = resolveMessageEntry(srcSm, entryId, originalText); + const targetIndex = entries.findIndex((entry) => entry === resolved); + const candidate = targetIndex >= 0 ? entries[targetIndex] : undefined; + const isUserMessage = + candidate?.type === "message" && candidate?.message?.role === "user"; + if (!candidate || !isUserMessage) { + throw new Error("Selected entry is not a persisted user message."); + } + const predecessorId = targetIndex > 0 ? entries[targetIndex - 1]?.id ?? null : null; + if (predecessorId) { + forkedPath = srcSm.createBranchedSession(predecessorId); if (!forkedPath) { throw new Error("Failed to build the rewritten session."); } } } finally { @@ -602,10 +648,21 @@ function createSessionWindow( // History-message actions are routed through commands so the activate-scope // session helpers can run them with the full session registry. - webviewPanel.onEditUserMessage = (entryId, text) => - void vscode.commands.executeCommand("pi-on-code.editHistoryMessage", sw.id, entryId, text); - webviewPanel.onForkUserMessage = (entryId) => - void vscode.commands.executeCommand("pi-on-code.forkHistoryMessage", sw.id, entryId); + webviewPanel.onEditUserMessage = (entryId, text, content) => + void vscode.commands.executeCommand( + "pi-on-code.editHistoryMessage", + sw.id, + entryId, + text, + content, + ); + webviewPanel.onForkUserMessage = (entryId, content) => + void vscode.commands.executeCommand( + "pi-on-code.forkHistoryMessage", + sw.id, + entryId, + content, + ); sessions.push(sw); return sw; @@ -956,31 +1013,32 @@ export async function activate(context: vscode.ExtensionContext): Promise sessionId: string, entryId: string, forkLabel?: string, + content?: string, ): Promise { const srcSw = sessions.find((s) => s.id === sessionId); if (!srcSw || !srcSw.piService.sessionManagerInstance) { throw new Error(`Source session not found (id=${sessionId}).`); } - - // Get the source file path — we open a fresh SessionManager to branch - // so the source session is not mutated. const sourcePath = srcSw.piService.sessionFilePath; - if (!sourcePath) { - throw new Error("Source session has no persisted file."); - } + if (!sourcePath) { throw new Error("Source session has no persisted file."); } - // Open a temporary PiService to get an isolated SessionManager for branching + // Branch through an isolated temporary manager so the live source session + // is never mutated by the SDK's createBranchedSession. Messages are + // persisted as they are written, and live-entry sync keeps the ids in DOM + // canonical by the time the reply completes. const tempPi = new PiService(context.secrets); let forkedPath: string; try { const result = await tempPi.initialize({ openPath: sourcePath, cwd: srcSw.cwd }); - if (!result.success) { - throw new Error(`Cannot open source session: ${result.error}`); - } + if (!result.success) { throw new Error(`Cannot open source session: ${result.error}`); } const srcSm = tempPi.sessionManagerInstance; if (!srcSm) { throw new Error("Source session has no session manager."); } - const entry = srcSm.getEntry(entryId); + const entry = resolveMessageEntry(srcSm, entryId, content) as { + type?: string; + message?: { role?: string }; + id?: string; + } | null; if (!entry) { throw new Error("Entry not found in source session."); } const isUserMsg = entry.type === "message" && entry.message?.role === "user"; @@ -990,12 +1048,8 @@ export async function activate(context: vscode.ExtensionContext): Promise throw new Error("Fork only works on user, assistant, or custom messages. Selected entry type: " + (entry.type ?? "unknown")); } - // Fork at the selected entry (include it in the branch) - const targetLeafId = entryId; - forkedPath = srcSm.createBranchedSession(targetLeafId); - if (!forkedPath) { - throw new Error("Failed to create forked session file."); - } + forkedPath = srcSm.createBranchedSession(entry.id ?? entryId); + if (!forkedPath) { throw new Error("Failed to create forked session file."); } } finally { tempPi.dispose(); } @@ -1131,7 +1185,7 @@ export async function activate(context: vscode.ExtensionContext): Promise context.subscriptions.push( vscode.commands.registerCommand( "pi-on-code.editHistoryMessage", - async (sessionId: unknown, entryId: unknown, text: unknown) => { + async (sessionId: unknown, entryId: unknown, text: unknown, content: unknown) => { if (typeof sessionId !== "string" || typeof entryId !== "string" || typeof text !== "string") { return; } @@ -1141,7 +1195,13 @@ export async function activate(context: vscode.ExtensionContext): Promise return; } try { - await editHistoryMessageFromEntry(context, sw, entryId, text); + await editHistoryMessageFromEntry( + context, + sw, + entryId, + text, + typeof content === "string" ? content : undefined, + ); // eslint-disable-next-line @typescript-eslint/no-explicit-any } catch (e: any) { vscode.window.showErrorMessage(`Edit failed: ${e.message ?? e}`); @@ -1152,7 +1212,7 @@ export async function activate(context: vscode.ExtensionContext): Promise context.subscriptions.push( vscode.commands.registerCommand( "pi-on-code.forkHistoryMessage", - async (sessionId: unknown, entryId: unknown) => { + async (sessionId: unknown, entryId: unknown, content: unknown) => { if (typeof sessionId !== "string" || typeof entryId !== "string") { return; } const sw = sessions.find((session) => session.id === sessionId); if (!sw) { @@ -1167,7 +1227,12 @@ export async function activate(context: vscode.ExtensionContext): Promise vscode.window.showWarningMessage("Stop the current run before forking a historical message."); return; } - await doForkFromOpenEntry(sw.id, entryId, forkSessionLabel(title)); + await doForkFromOpenEntry( + sw.id, + entryId, + forkSessionTitle(title), + typeof content === "string" ? content : undefined, + ); // eslint-disable-next-line @typescript-eslint/no-explicit-any } catch (e: any) { vscode.window.showErrorMessage(`Fork failed: ${e.message ?? e}`); diff --git a/src/pi-service.ts b/src/pi-service.ts index 6503c6d..3d711bd 100644 --- a/src/pi-service.ts +++ b/src/pi-service.ts @@ -428,6 +428,27 @@ export class PiService { private _userMessages: Array<{ id: string; text: string; timestamp?: number }> = []; private pendingLiveUserMessageIds: string[] = []; private liveUserEntryAliases = new Map(); + private liveEntryIdMapJson = ""; + + /** Emit canonical entry ids for messages once they exist, so the Webview can + * replace transient SDK message ids used for freshly typed messages. */ + private syncLiveEntryIds(): void { + const entries = (this.sessionManager?.getEntries?.() ?? []) as Array<{ + id?: string; + type?: string; + message?: { id?: string }; + }>; + const mapping: Record = {}; + for (const entry of entries) { + if (entry?.type === "message" && entry?.message?.id && entry?.id && entry.id !== entry.message.id) { + mapping[entry.message.id] = entry.id; + } + } + const json = JSON.stringify(mapping); + if (json === this.liveEntryIdMapJson) { return; } + this.liveEntryIdMapJson = json; + this.emit({ type: "live-entry-ids", data: { mapping } }); + } // Lazily replayed session history. The cursor points to the oldest entry // currently rendered in the Webview. @@ -1301,6 +1322,24 @@ export class PiService { } this.historyCursor = findHistoryPageStart(this.historyEntries, this.historyEntries.length); + // A tool-heavy tail can make the newest page contain only assistant/tool + // content. Pull the start back to the most recent user message so the + // conversation always opens with at least one user turn loaded. + if (this.historyCursor > 0) { + const newestSlice = this.historyEntries.slice(this.historyCursor); + const hasUser = newestSlice.some( + (entry) => entry?.type === "message" && entry?.message?.role === "user", + ); + if (!hasUser) { + for (let index = this.historyCursor - 1; index >= 0; index--) { + const entry = this.historyEntries[index]; + if (entry?.type === "message" && entry?.message?.role === "user") { + this.historyCursor = index; + break; + } + } + } + } await this.replayHistoryEntries( this.historyEntries.slice(this.historyCursor), this.historyToolResultsById, @@ -1625,6 +1664,7 @@ export class PiService { this._isStreaming = false; this.currentAssistantToolCalls.clear(); this.turnIndex = 0; + this.syncLiveEntryIds(); this.emit({ type: "agent-end", data: { messages: event.messages } }); this.emitConversationTurns(undefined, true); this.reportStatus(); @@ -1679,6 +1719,7 @@ export class PiService { const entry = byMessageId.get(event.message.id); this.emit({ type: "assistant-start", data: { messageId: event.message.id, entryId: entry?.id ?? event.message.id } }); } + this.syncLiveEntryIds(); break; } diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 21206d8..4038cc2 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -324,6 +324,12 @@ const ExtensionToWebviewSchema = z.discriminatedUnion("type", [ })), }), }), + z.object({ + type: z.literal("live-entry-ids"), + data: z.object({ + mapping: z.record(z.string(), z.string()), + }), + }), // Thinking z.object({ diff --git a/src/test/user-message-branch.test.ts b/src/test/user-message-branch.test.ts index cebeaa8..a7cbc09 100644 --- a/src/test/user-message-branch.test.ts +++ b/src/test/user-message-branch.test.ts @@ -1,7 +1,7 @@ import * as assert from "node:assert"; import { readFileSync } from "node:fs"; import { - forkSessionLabel, + forkSessionTitle, isUserMessageEntry, resolveUserMessageEditTarget, type SessionEntryLike, @@ -48,18 +48,12 @@ suite("User message branch helpers", () => { assert.strictEqual(resolveUserMessageEditTarget(entries, "a1"), null); }); - test("formats fork labels from the original title", () => { - assert.strictEqual(forkSessionLabel("Fix lint errors"), ""); - assert.strictEqual(forkSessionLabel(" "), ""); - assert.strictEqual(forkSessionLabel(""), ""); - }); - - test("caps fork labels by code points", () => { - const long = "title ".repeat(30); - const label = forkSessionLabel(long); - assert.ok(label.startsWith("")); - assert.ok(Array.from(label.slice(7, -1)).length <= 41, "label body stays near the cap"); + test("numbers fork titles from the base title", () => { + assert.strictEqual(forkSessionTitle("Fix lint errors"), "Fix lint errors (2)"); + assert.strictEqual(forkSessionTitle("Fix lint errors (2)"), "Fix lint errors (3)"); + assert.strictEqual(forkSessionTitle("Fix lint errors (12)"), "Fix lint errors (13)"); + assert.strictEqual(forkSessionTitle(" "), "Untitled session (2)"); + assert.strictEqual(forkSessionTitle(""), "Fix lint errors (2)"); }); test("wires transcript actions through panel callbacks and commands", () => { @@ -70,8 +64,11 @@ suite("User message branch helpers", () => { assert.match(handlers, /user-message-edit/); assert.match(handlers, /user-message-fork/); assert.match(handlers, /className = "user-edit-btn"/); - assert.match(handlers, /className = "user-fork-btn"/); - assert.match(handlers, /\.user-edit-overlay/); + assert.match(handlers, /className = "assistant-fork-btn"/); + assert.doesNotMatch(handlers, /className = "user-fork-btn"/); + assert.match(handlers, /className = "assistant-turn-actions"/); + assert.match(handlers, /className = "user-edit-input"/); + assert.match(handlers, /className = "user-edit-actions"/); assert.match(handlers, /function setUserEditActionsEnabled\(enabled: boolean\)/); assert.match(handlers, /setUserEditActionsEnabled\(false\)/); assert.match(handlers, /setUserEditActionsEnabled\(true\)/); @@ -89,10 +86,11 @@ suite("User message branch helpers", () => { ); assert.match(extension, /"pi-on-code\.editHistoryMessage"/); assert.match(extension, /"pi-on-code\.forkHistoryMessage"/); - assert.match(extension, /resolveUserMessageEditTarget\(entries, entryId\)/); - assert.match(extension, /createBranchedSession\(target\.predecessorId\)/); + assert.match(extension, /Selected entry is not a persisted user message\./); + assert.match(extension, /createBranchedSession\(predecessorId\)/); assert.match(extension, /sendPrompt\(text\)/); - assert.match(extension, /forkSessionLabel\(title\)/); + assert.match(extension, /forkSessionTitle\(title\)/); + assert.match(extension, /\.message\?\.id === entryId/); assert.match(extension, /canEditSession\(sw\.isStreaming\)/); assert.match(extension, /waitForReplyStart\(/); assert.match(extension, /resolveEditOutcome\(observation\)/); diff --git a/src/test/user-message-style.test.ts b/src/test/user-message-style.test.ts index 04967bb..022dcc6 100644 --- a/src/test/user-message-style.test.ts +++ b/src/test/user-message-style.test.ts @@ -8,14 +8,17 @@ const styles = readFileSync( suite("User message styling", () => { test("renders user messages as distinct full-width cards", () => { - const rule = styles.match(/\.message\.user \{\n min-width:[\s\S]*?\n\}/)?.[0]; - assert.ok(rule, "final user message rule not found"); - assert.match(rule, /align-self: stretch/); - assert.match(rule, /padding: 12px 14px/); - assert.match(rule, /border: 1px solid var\(--pi-line\)/); - assert.match(rule, /border-left: 3px solid var\(--pi-lavender\)/); - assert.match(rule, /border-radius: 6px/); - assert.match(rule, /background: color-mix\(in srgb, var\(--pi-lavender\) 10%, var\(--pi-surface\)\)/); + // The card chrome sits on the message content; the row stacks it and the + // right-aligned action row below so actions live outside the bubble. + const card = styles.match(/\.message\.user \.message-content \{\n position: relative;[\s\S]*?\n\}/)?.[0]; + assert.ok(card, "user message card rule not found"); + assert.match(card, /padding: 12px 14px/); + assert.match(card, /border: 1px solid var\(--pi-line\)/); + assert.match(card, /border-left: 3px solid var\(--pi-lavender\)/); + assert.match(card, /border-radius: 6px/); + assert.match(card, /background: color-mix\(in srgb, var\(--pi-lavender\) 10%, var\(--pi-surface\)\)/); + assert.match(styles, /\.message\.user \{\n position: relative;\n display: flex;\n flex-direction: column;/); + assert.match(styles, /\.user-actions \{\n align-self: flex-end;/); }); test("keeps a theme-safe background fallback", () => { diff --git a/src/test/webview-history-autofill.test.ts b/src/test/webview-history-autofill.test.ts new file mode 100644 index 0000000..16b506d --- /dev/null +++ b/src/test/webview-history-autofill.test.ts @@ -0,0 +1,37 @@ +import * as assert from "node:assert"; +import { decideAutoLoadOlder } from "../webview/render/history-autofill.js"; + +const idle = { + hasMore: true, + loading: false, + hasUserMessage: true, + scrollableRoom: 0, + scrollTop: 0, + autoFillCount: 0, +}; + +suite("History auto-fill decisions", () => { + test("loads when no user message has been loaded yet", () => { + assert.strictEqual(decideAutoLoadOlder({ ...idle, hasUserMessage: false }), true); + }); + + test("tops up a transcript that cannot be scrolled while at the top", () => { + assert.strictEqual(decideAutoLoadOlder({ ...idle, scrollableRoom: 40 }), true); + assert.strictEqual(decideAutoLoadOlder({ ...idle, scrollableRoom: 120 }), false); + }); + + test("never auto-loads while a page is already loading or history is empty", () => { + assert.strictEqual(decideAutoLoadOlder({ ...idle, loading: true }), false); + assert.strictEqual(decideAutoLoadOlder({ ...idle, hasMore: false }), false); + assert.strictEqual(decideAutoLoadOlder({ ...idle, hasUserMessage: false, hasMore: false }), false); + }); + + test("skips loading when the user is not at the top edge", () => { + assert.strictEqual(decideAutoLoadOlder({ ...idle, scrollTop: 60, scrollableRoom: 0 }), false); + }); + + test("bounds a single auto-fill burst", () => { + assert.strictEqual(decideAutoLoadOlder({ ...idle, autoFillCount: 19 }), true); + assert.strictEqual(decideAutoLoadOlder({ ...idle, autoFillCount: 20 }), false); + }); +}); diff --git a/src/user-message-branch.ts b/src/user-message-branch.ts index 5c296be..a8de43c 100644 --- a/src/user-message-branch.ts +++ b/src/user-message-branch.ts @@ -41,12 +41,13 @@ export function resolveUserMessageEditTarget( }; } -/** Display label for a fork, e.g. ``. */ -export function forkSessionLabel(originalTitle: string): string { - const title = originalTitle.replace(/\s+/g, " ").trim() || "Untitled session"; - const codePoints = Array.from(title); - const capped = codePoints.length > 40 - ? `${codePoints.slice(0, 40).join("").replace(/[\s.…]+$/u, "")}…` - : title; - return ``; +/** Next fork title: append an incrementing (N) to the base title. + * Forking a session already titled "… (2)" yields "… (3)", never "… (2) (2)". */ +export function forkSessionTitle(originalTitle: string): string { + const cleaned = originalTitle.replace(/\s+/g, " ").trim(); + const legacy = cleaned.replace(/^\s*$/, "").trim(); + const base = legacy.replace(/\s+\(\d+\)\s*$/u, "").trim() || "Untitled session"; + const match = legacy.match(/\((\d+)\)\s*$/u); + const next = match ? Number(match[1]) + 1 : 2; + return `${base} (${next})`; } diff --git a/src/webview-panel.ts b/src/webview-panel.ts index 2313edd..385a58c 100644 --- a/src/webview-panel.ts +++ b/src/webview-panel.ts @@ -67,8 +67,8 @@ export class PiWebviewPanel { /** Callback invoked when the panel is disposed (VS Code tab closed) */ private _onDispose: PanelDisposeCallback | null = null; private _onBeforePrompt: ((text: string) => void) | null = null; - private _onEditUserMessage: ((entryId: string, text: string) => void | Promise) | null = null; - private _onForkUserMessage: ((entryId: string) => void | Promise) | null = null; + private _onEditUserMessage: ((entryId: string, text: string, content?: string) => void | Promise) | null = null; + private _onForkUserMessage: ((entryId: string, content?: string) => void | Promise) | null = null; private _initialWelcomeVisible = false; constructor( @@ -85,10 +85,10 @@ export class PiWebviewPanel { /** Register a callback that fires when the panel/webview is closed. */ set onDispose(cb: PanelDisposeCallback | null) { this._onDispose = cb; } set onBeforePrompt(cb: ((text: string) => void) | null) { this._onBeforePrompt = cb; } - set onEditUserMessage(cb: ((entryId: string, text: string) => void | Promise) | null) { + set onEditUserMessage(cb: ((entryId: string, text: string, content?: string) => void | Promise) | null) { this._onEditUserMessage = cb; } - set onForkUserMessage(cb: ((entryId: string) => void | Promise) | null) { + set onForkUserMessage(cb: ((entryId: string, content?: string) => void | Promise) | null) { this._onForkUserMessage = cb; } set initialWelcomeVisible(value: boolean) { this._initialWelcomeVisible = value; } @@ -691,7 +691,11 @@ export class PiWebviewPanel { case "user-message-edit": if (typeof message.entryId === "string" && typeof message.text === "string") { try { - await this._onEditUserMessage?.(message.entryId, message.text); + await this._onEditUserMessage?.( + message.entryId, + message.text, + typeof message.content === "string" ? message.content : undefined, + ); } catch (error: unknown) { this.postMessage({ type: "error", @@ -705,7 +709,10 @@ export class PiWebviewPanel { case "user-message-fork": if (typeof message.entryId === "string") { try { - await this._onForkUserMessage?.(message.entryId); + await this._onForkUserMessage?.( + message.entryId, + typeof message.content === "string" ? message.content : undefined, + ); } catch (error: unknown) { this.postMessage({ type: "error", diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 3b25381..2b2e6a0 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -11,6 +11,7 @@ import { shortenPath, renderCodeBlockHTML, } from "../render/engine.js"; import { restoreScrollAfterPrepend } from "../render/history-pagination.js"; +import { decideAutoLoadOlder } from "../render/history-autofill.js"; import { collapseExecutionProcesses, openExecutionProcessForElement, @@ -196,6 +197,9 @@ function handleExtensionMessage(msg: any): void { case "conversation-turns-update": window.dispatchEvent(new CustomEvent(CONVERSATION_TURNS_EVENT, { detail: msg.data?.turns ?? [] })); break; + case "live-entry-ids": + handleLiveEntryIds(msg.data); + break; case "localImageResolved": applyResolvedLocalImage(msg.data); break; @@ -287,39 +291,89 @@ function assistantMessageText(message: HTMLElement): string { return (clone.textContent || "").replace(/\s+$/g, "").trim(); } +/** Replace transient SDK message ids with canonical persisted entry ids once + * the SDK exposes them, so reveal/edit/fork target the real session entries. */ +function handleLiveEntryIds(data: unknown): void { + const mapping = (data && typeof data === "object" && "mapping" in data + ? (data as { mapping?: Record }).mapping + : undefined); + if (!mapping) { return; } + for (const messageId of Object.keys(mapping)) { + const entryId = mapping[messageId]; + if (typeof entryId !== "string" || entryId === messageId) { continue; } + const escaped = CSS.escape(messageId); + const byDataId = state.chatContainer.querySelector(`.message[data-entry-id="${escaped}"]`); + if (byDataId) { byDataId.setAttribute("data-entry-id", entryId); } + const byElementId = document.getElementById("entry-" + messageId); + if (byElementId) { byElementId.id = "entry-" + entryId; } + } + // Per-turn action rows may have captured the transient id for their fork + // target; rebuild them with the canonical ids. + updateLastAssistantCopyButton(state.chatContainer); +} + function updateLastAssistantCopyButton(root: HTMLElement): void { - // Clear any stale buttons first. - root.querySelectorAll(".assistant-copy-btn").forEach((button) => button.remove()); + // Clear any stale controls first. + root.querySelectorAll(".assistant-turn-actions").forEach((row) => row.remove()); - // Every completed user turn keeps a copy control below its final assistant - // message, so each turn's answer can be copied independently. A turn group - // is the run of assistant messages that follows a user message (a leading - // run without a user message is its own group too). + // Every completed user turn keeps an action row (Copy + Fork) below its final + // assistant message. A turn group is the run of assistant messages that + // follows a user message (a leading run without a user message is its own + // group and has no fork target). const groups: HTMLElement[][] = []; let current: HTMLElement[] = []; + const flush = (): void => { + if (current.length > 0) { groups.push(current); } + current = []; + }; const children = Array.from(root.children as HTMLCollectionOf); for (const child of children) { if (!child.classList.contains("message")) { continue; } if (child.classList.contains("user")) { - groups.push(current); - current = []; + flush(); } else if (child.classList.contains("assistant")) { current.push(child); } } - groups.push(current); + flush(); - for (const turnAssistants of groups) { - if (turnAssistants.length === 0) { continue; } - const texts = turnAssistants.map(assistantMessageText).filter(Boolean); + for (const group of groups) { + if (group.length === 0) { continue; } + const texts = group.map(assistantMessageText).filter(Boolean); if (texts.length === 0) { continue; } + + const row = document.createElement("div"); + row.className = "assistant-turn-actions"; + const copyButton = document.createElement("button"); copyButton.type = "button"; copyButton.className = "assistant-copy-btn"; copyButton.textContent = "Copy"; copyButton.setAttribute("aria-label", "Copy assistant response"); (copyButton as HTMLButtonElement & { _copyText?: string })._copyText = texts.join("\n\n"); - turnAssistants[turnAssistants.length - 1]!.appendChild(copyButton); + row.appendChild(copyButton); + + // Fork at the assistant message this row is attached to, so the branch + // includes that answer. The entry id is the persisted one after the + // live-entry sync; otherwise the SDK message id is resolved by the host. + const host = group[group.length - 1]!; + const forkId = + host.getAttribute("data-entry-id") + ?? host.id.replace(/^entry-/, ""); + if (forkId) { + const forkButton = document.createElement("button"); + forkButton.type = "button"; + forkButton.className = "assistant-fork-btn"; + forkButton.textContent = "Fork"; + forkButton.title = "Fork a new session at this message"; + forkButton.setAttribute("aria-label", "Fork session at this message"); + forkButton.addEventListener("click", () => { + window.__vscode.postMessage({ type: "user-message-fork", entryId: forkId }); + }); + row.appendChild(forkButton); + } + + host.appendChild(row); } } @@ -600,6 +654,9 @@ export function handleChatMessage(data: any) { el.id = "entry-" + data.entryId; el.setAttribute("data-entry-id", data.entryId); } + if (data.role === "user" && typeof data.content === "string") { + el.setAttribute("data-user-text", data.content); + } var mc = el.querySelector(".message-content"); if (mc) { var imageGallery = createMessageImages(images); @@ -632,25 +689,45 @@ export function handleChatMessage(data: any) { (userCopyButton as HTMLButtonElement & { _copyText?: string })._copyText = data.content; actions.appendChild(userCopyButton); - // Edit / fork operate on the persisted entry, so they need its entry id - // and are hidden while a run is in progress. - if (data.entryId && !state.isStreaming) { + // Edit operates on the persisted entry, so it needs its entry id. It is + // always rendered (when possible) and merely disabled during a run so a + // reply that starts before a replayed page finishes never loses it. + if (data.entryId) { var editButton = document.createElement("button"); editButton.type = "button"; editButton.className = "user-edit-btn"; editButton.textContent = "Edit"; editButton.title = "Edit this message and continue from here"; editButton.setAttribute("aria-label", "Edit user message and continue from here"); + editButton.disabled = state.isStreaming; editButton.addEventListener("click", function () { - if (el.querySelector(".user-edit-overlay")) { return; } - var overlay = document.createElement("div"); - overlay.className = "user-edit-overlay"; + if (el.querySelector(".user-edit-input")) { return; } + var actionsRow = el.querySelector(".user-actions"); + if (actionsRow) { actionsRow.style.display = "none"; } + var content = el.querySelector(".message-content"); + var originalHtml = content ? content.innerHTML : ""; + var startHeight = content ? content.clientHeight || 0 : 0; var area = document.createElement("textarea"); area.className = "user-edit-input"; area.value = data.content; area.setAttribute("aria-label", "Edited user message"); + if (content) { + content.innerHTML = ""; + content.appendChild(area); + } else { + el.appendChild(area); + } + // Start at the original message height (min-height keeps it at least + // two rows) and grow with the edited text as the user types. + var autosize = function () { + area.style.height = "auto"; + var target = Math.max(startHeight, area.scrollHeight); + area.style.height = target + "px"; + }; + autosize(); + area.addEventListener("input", autosize); var bar = document.createElement("div"); - bar.className = "user-edit-bar"; + bar.className = "user-edit-actions"; var saveButton = document.createElement("button"); saveButton.type = "button"; saveButton.className = "user-edit-save"; @@ -659,32 +736,37 @@ export function handleChatMessage(data: any) { cancelButton.type = "button"; cancelButton.className = "user-edit-cancel"; cancelButton.textContent = "Cancel"; - bar.append(cancelButton, saveButton); - overlay.append(area, bar); - el.appendChild(overlay); + bar.append(saveButton, cancelButton); + el.appendChild(bar); area.focus(); area.select(); - cancelButton.addEventListener("click", function () { overlay.remove(); }); + var restore = function () { + if (content) { content.innerHTML = originalHtml; } + bar.remove(); + if (actionsRow) { actionsRow.style.display = ""; } + }; + cancelButton.addEventListener("click", restore); + area.addEventListener("keydown", function (event) { + if (event.key === "Escape") { event.preventDefault(); restore(); } + else if (event.key === "Enter" && !event.shiftKey) { + event.preventDefault(); + saveButton.click(); + } + }); saveButton.addEventListener("click", function () { var next = area.value.trim(); + restore(); if (next && next !== data.content) { - window.__vscode.postMessage({ type: "user-message-edit", entryId: data.entryId, text: next }); - overlay.remove(); + window.__vscode.postMessage({ + type: "user-message-edit", + entryId: data.entryId, + text: next, + content: data.content, + }); } }); }); actions.appendChild(editButton); - - var forkButton = document.createElement("button"); - forkButton.type = "button"; - forkButton.className = "user-fork-btn"; - forkButton.textContent = "Fork"; - forkButton.title = "Fork a new session at this message"; - forkButton.setAttribute("aria-label", "Fork session at this message"); - forkButton.addEventListener("click", function () { - window.__vscode.postMessage({ type: "user-message-fork", entryId: data.entryId }); - }); - actions.appendChild(forkButton); } el.appendChild(actions); } @@ -703,7 +785,10 @@ export function handleAssistantStart(data: any) { // Create the assistant container eagerly before any content arrives state.currentAssistantEl = createMessageEl("assistant"); // #9: Entry ID for scroll-to - if (data.entryId) {state.currentAssistantEl.id = "entry-" + data.entryId;} + if (data.entryId) { + state.currentAssistantEl.id = "entry-" + data.entryId; + state.currentAssistantEl.setAttribute("data-entry-id", data.entryId); + } state.currentThinkingEl = null; state._streamPrevTokens = []; // Reset token tracker for new message state.assistantToolCallIds = {}; @@ -1050,6 +1135,7 @@ export function handleStatus(data: any) { } export function handleBatchStart(data: any) { + historyAutoFillCount = 0; // Hide the welcome screen before enabling batch mode because hideWelcome // intentionally ignores ordinary replay events while a batch is active. if (data?.hasEntries) { hideWelcome(); } @@ -1058,6 +1144,35 @@ export function handleBatchStart(data: any) { document.body.classList.add("no-animate"); } +/** + * Auto-load earlier history when the transcript is too short to scroll (the + * execution-process collapsing can shrink a loaded page below the viewport, so + * reaching the top to trigger a manual load may be impossible). Bounded so a + * pathological session cannot flood the extension with requests. + */ +let historyAutoFillCount = 0; + +export function resetHistoryAutoFill(): void { + historyAutoFillCount = 0; +} + +function autoFillOlderHistory(): void { + const container = state.chatContainer; + const hasUserMessage = container.querySelector(".message.user") !== null; + const should = decideAutoLoadOlder({ + hasMore: state.historyHasMore, + loading: state.historyLoading, + hasUserMessage, + scrollableRoom: container.scrollHeight - container.scrollTop - container.clientHeight, + scrollTop: container.scrollTop, + autoFillCount: historyAutoFillCount, + }); + if (!should) { return; } + historyAutoFillCount++; + state.historyLoading = true; + window.__vscode.postMessage({ type: "loadOlderHistory" }); +} + function settleScrollToBottom(): void { // Async rendering (images, local Markdown data URIs, syntax highlighting) can // grow the transcript after the initial follow scroll. Re-apply the follow a @@ -1091,6 +1206,7 @@ export function handleBatchEnd(data: any) { }); }); settleScrollToBottom(); + autoFillOlderHistory(); } interface HistoryPrependContext { @@ -1230,6 +1346,7 @@ export function handleHistoryPageEnd(data: any) { context.previousScrollTop, ); state.historyLoading = false; + autoFillOlderHistory(); } function submitFollowUpQueue(messages: string[]): void { diff --git a/src/webview/main.ts b/src/webview/main.ts index 8fab434..0e43a49 100644 --- a/src/webview/main.ts +++ b/src/webview/main.ts @@ -25,6 +25,7 @@ import { ScrollToBottomButton } from "./components/scroll-to-bottom-button.js"; // Side-effect imports (self-register on load) import "./tools/index.js"; import "./handlers/index.js"; +import { resetHistoryAutoFill } from "./handlers/index.js"; // ── Initialize ────────────────────────────────────────────── @@ -279,6 +280,7 @@ state.chatContainer.addEventListener("scroll", () => { owner: state.scrollOwner, })) { state.historyLoading = true; + resetHistoryAutoFill(); vscode.postMessage({ type: "loadOlderHistory" }); } }); diff --git a/src/webview/render/history-autofill.ts b/src/webview/render/history-autofill.ts new file mode 100644 index 0000000..c1104c3 --- /dev/null +++ b/src/webview/render/history-autofill.ts @@ -0,0 +1,31 @@ +/** Pure decision logic for auto-loading earlier history pages. Keeping this + * separate makes the pagination behaviour testable without a DOM. */ + +export interface HistoryAutoFillInput { + hasMore: boolean; + loading: boolean; + /** True once at least one user message is present in the loaded transcript. */ + hasUserMessage: boolean; + /** Remaining scrollable space at the current position (px). */ + scrollableRoom: number; + scrollTop: number; + /** Pages already auto-loaded in the current burst. */ + autoFillCount: number; + autoFillCap?: number; +} + +export function decideAutoLoadOlder(input: HistoryAutoFillInput): boolean { + if (!input.hasMore || input.loading) { return false; } + const cap = input.autoFillCap ?? 20; + if (input.autoFillCount >= cap) { return false; } + + // A page may contain only collapsed execution content (tool runs, thinking) + // with no visible user turn. Keep loading until at least one user message + // has arrived so the conversation always has an anchor to continue from. + if (!input.hasUserMessage) { return true; } + + // Otherwise top up when the transcript cannot be scrolled (execution-process + // collapsing can shrink a loaded page below the viewport, making a manual + // upward scroll impossible) and the user is already at the top edge. + return input.scrollTop <= 8 && input.scrollableRoom <= 60; +} From fc349d2900f32854ed0763fa40b9e487f1ddda38 Mon Sep 17 00:00:00 2001 From: auchan Date: Sun, 6 Sep 2026 15:06:20 +0800 Subject: [PATCH 4/7] feat: show fork divider and anchor live-turn forks --- media/style.css | 56 ++++++++++++++++++++++ src/extension.ts | 87 +++++++++++++++++++++++++++++++++-- src/pi-service.ts | 13 ++++++ src/webview/handlers/index.ts | 59 ++++++++++++++++++++++-- 4 files changed, 208 insertions(+), 7 deletions(-) diff --git a/media/style.css b/media/style.css index 8a32b5d..c01c088 100644 --- a/media/style.css +++ b/media/style.css @@ -2943,6 +2943,62 @@ html.pi-viewport-recovering #pi-web-app { color: var(--pi-text); } +.fork-divider { + display: flex; + align-items: center; + gap: 10px; + margin: 14px 8px; + color: var(--pi-muted); +} + +.fork-divider-rule { + flex: 1; + height: 1px; + background: var(--pi-line); +} + +.fork-divider-icon { + display: grid; + width: 18px; + height: 18px; + place-items: center; + border-radius: 50%; + background: var(--pi-hover); + color: var(--pi-lavender); +} + +.fork-divider-icon svg { + display: block; + width: 13px; + height: 13px; + fill: none; + stroke: currentColor; + stroke-width: 1.3; + stroke-linecap: round; + stroke-linejoin: round; +} + +.fork-divider-text { + color: var(--pi-faint); + font-size: 11px; + white-space: nowrap; +} + +.fork-divider-link { + padding: 0; + border: 0; + background: transparent; + color: var(--pi-lavender); + font: inherit; + cursor: pointer; + text-decoration: underline; + text-underline-offset: 2px; +} + +.fork-divider-link:hover { + color: var(--pi-strong); +} + .message.assistant:hover .assistant-turn-actions, .assistant-turn-actions:focus-within { opacity: 1; diff --git a/src/extension.ts b/src/extension.ts index 527ab85..142e727 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,5 +1,7 @@ import * as vscode from "vscode"; import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; import { PiService } from "./pi-service.js"; import { PiWebviewPanel } from "./webview-panel.js"; import { PiPackageService, type ManagedCapability } from "./pi-package-service.js"; @@ -803,7 +805,14 @@ export async function activate(context: vscode.ExtensionContext): Promise vscode.window.showErrorMessage("Cannot fork: session has no entries."); return; } - await vscode.commands.executeCommand("pi-on-code.forkSession", session.id, leafId); + await vscode.commands.executeCommand( + "pi-on-code.forkSession", + session.id, + leafId, + forkSessionTitle(cleanSessionTitle( + session.piService.sessionName ?? session.webviewPanel.summary ?? session.label, + )), + ); } else if (target.kind === "past" && target.path) { await vscode.commands.executeCommand("pi-on-code.forkSession", target.path); } @@ -1014,6 +1023,7 @@ export async function activate(context: vscode.ExtensionContext): Promise entryId: string, forkLabel?: string, content?: string, + sourceName?: string, ): Promise { const srcSw = sessions.find((s) => s.id === sessionId); if (!srcSw || !srcSw.piService.sessionManagerInstance) { @@ -1044,12 +1054,29 @@ export async function activate(context: vscode.ExtensionContext): Promise const isUserMsg = entry.type === "message" && entry.message?.role === "user"; const isAssistantMsg = entry.type === "message" && entry.message?.role === "assistant"; const isCustomMsg = entry.type === "custom_message"; - if (!isUserMsg && !isAssistantMsg && !isCustomMsg) { + // session_info entries are SDK metadata leaves (renames / branch marks); + // forking at the leaf is still a fork of the whole conversation. + const isSessionInfo = entry.type === "session_info"; + if (!isUserMsg && !isAssistantMsg && !isCustomMsg && !isSessionInfo) { throw new Error("Fork only works on user, assistant, or custom messages. Selected entry type: " + (entry.type ?? "unknown")); } forkedPath = srcSm.createBranchedSession(entry.id ?? entryId); if (!forkedPath) { throw new Error("Failed to create forked session file."); } + + // Persist a fork divider at the fork leaf so reloads render it in the + // same file position (SDK appendCustomMessageEntry keeps the natural + // line order; session_info is reserved for renames). + if (sourceName) { + const appender = srcSm as { + appendCustomMessageEntry?: (type: string, content: unknown, display: boolean, details?: unknown) => string | undefined; + }; + if (typeof appender.appendCustomMessageEntry === "function") { + appender.appendCustomMessageEntry("fork-info", sourceName, true, { + sourceId: readSessionId(sourcePath) ?? null, + }); + } + } } finally { tempPi.dispose(); } @@ -1143,7 +1170,11 @@ export async function activate(context: vscode.ExtensionContext): Promise try { if (cmdArgs.length >= 2) { - await doForkFromOpenEntry(cmdArgs[0] as string, cmdArgs[1] as string); + await doForkFromOpenEntry( + cmdArgs[0] as string, + cmdArgs[1] as string, + typeof cmdArgs[2] === "string" ? cmdArgs[2] : undefined, + ); } else { await doForkFromPastSession(cmdArgs[0] as string); } @@ -1232,6 +1263,7 @@ export async function activate(context: vscode.ExtensionContext): Promise entryId, forkSessionTitle(title), typeof content === "string" ? content : undefined, + title, ); // eslint-disable-next-line @typescript-eslint/no-explicit-any } catch (e: any) { @@ -1350,6 +1382,55 @@ export async function activate(context: vscode.ExtensionContext): Promise }), ); + // ── Open a session by its file session id (used by the fork divider link) ── + const findSessionFilePathById = (sessionId: string): string | undefined => { + const openSession = sessions.find( + (candidate) => candidate.piService.sessionIdValue === sessionId, + ); + const openPath = openSession + ? openSession.piService.sessionFilePath ?? openSession.restoringPath + : undefined; + if (openPath) { return openPath; } + const configuredDir = vscode.workspace + .getConfiguration("pi-on-code") + .get("sessionDir") + ?.trim(); + const root = configuredDir + || path.join(os.homedir(), ".pi", "agent", "sessions"); + const matches: string[] = []; + const walk = (directory: string): void => { + let entries: fs.Dirent[]; + try { entries = fs.readdirSync(directory, { withFileTypes: true }); } catch { return; } + for (const entry of entries) { + const full = path.join(directory, entry.name); + if (entry.isDirectory()) { walk(full); continue; } + if (entry.isFile() && entry.name.endsWith(".jsonl") && entry.name.endsWith(`_${sessionId}.jsonl`)) { + matches.push(full); + } + } + }; + walk(root); + return matches[0]; + }; + context.subscriptions.push( + vscode.commands.registerCommand("pi-on-code.resumeSessionById", async (sessionId?: unknown) => { + if (typeof sessionId !== "string" || !sessionId) { return; } + const openSession = sessions.find( + (candidate) => candidate.piService.sessionIdValue === sessionId, + ); + if (openSession) { + await vscode.commands.executeCommand("pi-on-code.focusSession", openSession.id); + return; + } + const filePath = findSessionFilePathById(sessionId); + if (!filePath) { + vscode.window.showErrorMessage("Source session file not found."); + return; + } + await vscode.commands.executeCommand("pi-on-code.resumePastSession", filePath); + }), + ); + // ── Delete a past session from the tree view ────── context.subscriptions.push( vscode.commands.registerCommand("pi-on-code.deletePastSession", async (filePath?: SessionTreeItem | string) => { diff --git a/src/pi-service.ts b/src/pi-service.ts index 3d711bd..d653085 100644 --- a/src/pi-service.ts +++ b/src/pi-service.ts @@ -427,6 +427,7 @@ export class PiService { // User message history for the resend/reuse feature (#2) private _userMessages: Array<{ id: string; text: string; timestamp?: number }> = []; private pendingLiveUserMessageIds: string[] = []; + private liveUserEntryAliases = new Map(); private liveEntryIdMapJson = ""; @@ -1446,6 +1447,18 @@ export class PiService { this.emit({ type: "bash-start", data: { toolCallId: bashEntryId, command: msg.command ?? "", entryId: entry.id } }); this.emit({ type: "bash-end", data: { toolCallId: bashEntryId, command: msg.command ?? "", exitCode: msg.exitCode, cancelled: msg.cancelled, output: msg.output ?? "", isError: msg.exitCode !== 0 && msg.exitCode !== null, entryId: entry.id } }); } + } else if (entry.type === "custom_message") { + this.emit({ + type: "custom-message", + data: { + customType: entry.customType ?? "custom", + content: entry.content, + display: entry.display === true, + details: entry.details, + timestamp: this._toTimestamp(entry.timestamp), + entryId: entry.id, + }, + }); } else if (entry.type === "compaction") { this.emit({ type: "compaction-summary-message", diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 2b2e6a0..6c38c90 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -784,10 +784,16 @@ export function handleAssistantStart(data: any) { // Create the assistant container eagerly before any content arrives state.currentAssistantEl = createMessageEl("assistant"); - // #9: Entry ID for scroll-to - if (data.entryId) { - state.currentAssistantEl.id = "entry-" + data.entryId; - state.currentAssistantEl.setAttribute("data-entry-id", data.entryId); + // #9: Entry ID for scroll-to. Fresh live messages may not have their + // canonical entry id yet, so fall back to the SDK message id; the host + // replaces it after the entry sync. + const assistantEntryValue = + typeof data.entryId === "string" ? data.entryId + : typeof data.messageId === "string" ? data.messageId + : ""; + if (assistantEntryValue) { + state.currentAssistantEl.id = "entry-" + assistantEntryValue; + state.currentAssistantEl.setAttribute("data-entry-id", assistantEntryValue); } state.currentThinkingEl = null; state._streamPrevTokens = []; // Reset token tracker for new message @@ -2750,9 +2756,54 @@ export function renderInlineCustomMessage(data: any) { scrollToBottom(); } +function renderForkDivider(sourceName: string, details: { sourceId?: unknown } | null): void { + hideWelcome(); + const row = document.createElement("div"); + row.className = "fork-divider"; + const ruleLeft = document.createElement("span"); + ruleLeft.className = "fork-divider-rule"; + const icon = document.createElement("span"); + icon.className = "fork-divider-icon"; + icon.innerHTML = html``; + const text = document.createElement("span"); + text.className = "fork-divider-text"; + if (sourceName) { + const sourceId = details?.sourceId; + if (typeof sourceId === "string" && sourceId) { + const link = document.createElement("button"); + link.type = "button"; + link.className = "fork-divider-link"; + link.textContent = sourceName; + link.title = "Open the source session"; + link.addEventListener("click", () => { + window.__vscode.postMessage({ type: "open-session", sessionId: sourceId }); + }); + text.append("continued from ", link); + } else { + text.textContent = `continued from ${sourceName}`; + } + } else { + text.textContent = "continued from another session"; + } + const ruleRight = document.createElement("span"); + ruleRight.className = "fork-divider-rule"; + row.append(ruleLeft, icon, text, ruleRight); + state.chatContainer.appendChild(row); + scrollToBottom(); +} + export function handleCustomMessage(data: any) { var customType = data.customType || "custom"; + // Fork divider rendered in the transcript at the fork point. + if (customType === "fork-info") { + renderForkDivider( + String(data.content ?? ""), + data && typeof data === "object" ? (data as { details?: { sourceId?: unknown } | null }).details ?? null : null, + ); + return; + } + // ── display: true → inline in conversation stream ────── if (data.display === true) { hideWelcome(); From 6c9b6a6df57fbcd97a7e480adfa7bcfc02a2a733 Mon Sep 17 00:00:00 2001 From: auchan Date: Sun, 6 Sep 2026 15:06:36 +0800 Subject: [PATCH 5/7] feat: support resuming a session by id from the webview --- src/history-pagination.ts | 2 ++ src/webview-panel.ts | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/src/history-pagination.ts b/src/history-pagination.ts index f661719..2ca0270 100644 --- a/src/history-pagination.ts +++ b/src/history-pagination.ts @@ -2,12 +2,14 @@ export const HISTORY_PAGE_SIZE = 20; interface HistoryEntryLike { type?: unknown; + display?: unknown; message?: { role?: unknown; customType?: unknown; display?: unknown }; } /** Entries that produce visible conversation content during session replay. */ export function isVisibleHistoryEntry(entry: HistoryEntryLike): boolean { if (entry.type === "compaction") { return true; } + if (entry.type === "custom_message") { return entry.display === true; } if (entry.type !== "message") { return false; } const role = entry.message?.role; if (role === "custom") { diff --git a/src/webview-panel.ts b/src/webview-panel.ts index 385a58c..ea6d534 100644 --- a/src/webview-panel.ts +++ b/src/webview-panel.ts @@ -723,6 +723,11 @@ export class PiWebviewPanel { break; // Request settings state (#2, #8) + case "open-session": + if (typeof message.sessionId === "string") { + void vscode.commands.executeCommand("pi-on-code.resumeSessionById", message.sessionId); + } + break; case "resendUserMessage": if (message.text) { await this.piService.sendPrompt(message.text); From d844a6602f77415102ed4b223886dd505f03108e Mon Sep 17 00:00:00 2001 From: auchan Date: Sun, 6 Sep 2026 15:12:33 +0800 Subject: [PATCH 6/7] fix: allow repeated identical prompts across fork turns --- src/webview/handlers/index.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 6c38c90..0ed992f 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -398,6 +398,9 @@ export function handleAgentEnd() { }); state.isStreaming = false; setUserEditActionsEnabled(true); + // A new turn may reuse the same prompt text (e.g. after a fork); do not + // deduplicate it against a previous turn's identical message. + state.lastUserMessageContent = null; state.isRetrying = false; state.assistantToolCallIds = {}; removeWorkingIndicator(); @@ -2758,6 +2761,8 @@ export function renderInlineCustomMessage(data: any) { function renderForkDivider(sourceName: string, details: { sourceId?: unknown } | null): void { hideWelcome(); + // A forked continuation is a new turn even if it repeats the prompt text. + state.lastUserMessageContent = null; const row = document.createElement("div"); row.className = "fork-divider"; const ruleLeft = document.createElement("span"); From ee802ad27ce9a4c157745ebb48aeeed08de99ac1 Mon Sep 17 00:00:00 2001 From: auchan Date: Sun, 6 Sep 2026 15:26:37 +0800 Subject: [PATCH 7/7] fix: render repeated prompts and always offer fork with text fallback --- src/extension.ts | 6 ++++-- src/webview/handlers/index.ts | 40 ++++++++++++++++++++++------------- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 142e727..a41668b 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -453,8 +453,10 @@ function resolveMessageEntry( const text = content.trim(); for (let index = entries.length - 1; index >= 0; index--) { const entry = entries[index]; - if (entry?.type !== "message" || entry?.message?.role !== "user") { continue; } - const raw = entry.message.content; + const message = entry?.message; + const role = message?.role; + if (entry?.type !== "message" || (role !== "user" && role !== "assistant")) { continue; } + const raw = message?.content; const body = typeof raw === "string" ? raw.trim() : Array.isArray(raw) diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 0ed992f..5d8c221 100644 --- a/src/webview/handlers/index.ts +++ b/src/webview/handlers/index.ts @@ -355,23 +355,27 @@ function updateLastAssistantCopyButton(root: HTMLElement): void { // Fork at the assistant message this row is attached to, so the branch // includes that answer. The entry id is the persisted one after the - // live-entry sync; otherwise the SDK message id is resolved by the host. + // live-entry sync; otherwise the SDK message id is resolved by the host, + // and as a last resort the reply text is used to locate the entry. const host = group[group.length - 1]!; const forkId = host.getAttribute("data-entry-id") ?? host.id.replace(/^entry-/, ""); - if (forkId) { - const forkButton = document.createElement("button"); - forkButton.type = "button"; - forkButton.className = "assistant-fork-btn"; - forkButton.textContent = "Fork"; - forkButton.title = "Fork a new session at this message"; - forkButton.setAttribute("aria-label", "Fork session at this message"); - forkButton.addEventListener("click", () => { - window.__vscode.postMessage({ type: "user-message-fork", entryId: forkId }); + const forkText = texts.join("\n\n"); + const forkButton = document.createElement("button"); + forkButton.type = "button"; + forkButton.className = "assistant-fork-btn"; + forkButton.textContent = "Fork"; + forkButton.title = "Fork a new session at this message"; + forkButton.setAttribute("aria-label", "Fork session at this message"); + forkButton.addEventListener("click", () => { + window.__vscode.postMessage({ + type: "user-message-fork", + entryId: forkId || "", + content: forkText, }); - row.appendChild(forkButton); - } + }); + row.appendChild(forkButton); host.appendChild(row); } @@ -637,13 +641,16 @@ export function handleChatMessage(data: any) { var userMessageKey = String(data.content || "") + "\u0000" + imageKey + "\u0000" + contextKey; // Dedup repeated SDK/replay events while keeping image-only messages distinct. - if (data.role === "user" && userMessageKey === state.lastUserMessageContent) {return;} - if (data.role === "user") { + // Replay (batch/history) contains real historical entries that may repeat the + // same text (e.g. two identical “ping” prompts), so deduplication only applies + // to live streaming events. + if (data.role === "user" && !state._inBatch) { + if (userMessageKey === state.lastUserMessageContent) { return; } state.lastUserMessageContent = userMessageKey; // Populate state.userMessageHistory for up-arrow recall (#2). if (data.content) { state.userMessageHistory.unshift({ text: data.content }); - if (state.userMessageHistory.length > 50) {state.userMessageHistory.pop();} + if (state.userMessageHistory.length > 50) { state.userMessageHistory.pop(); } } } @@ -1215,6 +1222,9 @@ export function handleBatchEnd(data: any) { }); }); settleScrollToBottom(); + // Replaying a session records its last user text; a freshly typed prompt + // that matches it (e.g. “ping” twice) must not be deduplicated away. + state.lastUserMessageContent = null; autoFillOlderHistory(); }