diff --git a/media/style.css b/media/style.css index c527ac4..c01c088 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,28 +2933,195 @@ 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); +} + +.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:has(.assistant-copy-btn) { +.message.assistant:hover .assistant-turn-actions, +.assistant-turn-actions:focus-within { + opacity: 1; +} + +.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 { + align-self: flex-end; + display: flex; + gap: 6px; + margin-top: 8px; + 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; + background: var(--pi-surface-raised); + color: var(--pi-muted); + font-size: 11px; + cursor: pointer; + white-space: nowrap; +} + +.user-actions button:hover, +.user-actions button:focus-visible { + border-color: var(--pi-lavender); + color: var(--pi-text); +} + +.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: 3.6em; + resize: vertical; + overflow-y: auto; + padding: 0; + border: 0; + border-radius: 0; + outline: 0; + background: transparent; + color: inherit; + font: inherit; + line-height: 1.6; +} + +.user-edit-actions { + align-self: flex-end; + display: flex; + gap: 6px; + margin-top: 8px; +} + +.user-edit-actions 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-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 { + opacity: 0.45; + pointer-events: none; +} + +.tool-block, +.bash-execution { + position: relative; } -.user-copy-btn, .tool-copy-btn { position: absolute; right: 8px; @@ -2971,19 +3138,12 @@ html.pi-viewport-recovering #pi-web-app { 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 { opacity: 1; } -.tool-block, -.bash-execution { - position: relative; -} - /* Generic extension tools expose their arguments in the header. */ .tool-block .tool-header-arguments { min-width: 0; 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 1f8c177..a41668b 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"; @@ -21,6 +23,14 @@ 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 { + forkSessionTitle, +} from "./user-message-branch.js"; +import { + canEditSession, + resolveEditOutcome, + waitForReplyStart, +} from "./edit-flow.js"; import { emptySessionListPreferences, isSessionArchived, @@ -419,6 +429,186 @@ 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]; + 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) + ? 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. + * + * The SDK has no in-place entry mutation, so the edit is composed from its + * primitives: branch everything before the target message into a new session + * file (predecessor keeps its ids for prefix-cache reuse), open it, send the + * edited text as the fresh user message, then retire and delete the superseded + * session. Editing the first message starts a fresh session. + */ +async function editHistoryMessageFromEntry( + context: vscode.ExtensionContext, + sw: SessionWindow, + entryId: string, + text: string, + originalText?: string, +): 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( + sw.piService.sessionName ?? sw.webviewPanel.summary ?? sw.label, + ); + + 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 }); + 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?: { id?: string; role?: string; content?: unknown }; + }>; + 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 { + 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) { + // 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."); + } + + 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(); + 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.`); + } + + // 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(); + 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."); +} + /** Create a new session window pair. Restore hints prevent transient duplicate rows. */ function createSessionWindow( context: vscode.ExtensionContext, @@ -460,6 +650,24 @@ 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, 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; } @@ -599,7 +807,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); } @@ -805,52 +1020,71 @@ 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, + content?: string, + sourceName?: 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"; 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")); } - // 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."); } + + // 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(); } 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 +1120,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 +1137,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(); @@ -928,7 +1172,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); } @@ -966,6 +1214,67 @@ 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, content: 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, + 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}`); + } + }, + ), + ); + context.subscriptions.push( + vscode.commands.registerCommand( + "pi-on-code.forkHistoryMessage", + 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) { + vscode.window.showErrorMessage("Source session not found."); + return; + } + try { + 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, + forkSessionTitle(title), + typeof content === "string" ? content : undefined, + 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 () => { @@ -1075,6 +1384,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/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/pi-service.ts b/src/pi-service.ts index 6503c6d..d653085 100644 --- a/src/pi-service.ts +++ b/src/pi-service.ts @@ -427,7 +427,29 @@ 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 = ""; + + /** 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 +1323,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, @@ -1407,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", @@ -1625,6 +1677,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 +1732,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/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 new file mode 100644 index 0000000..a7cbc09 --- /dev/null +++ b/src/test/user-message-branch.test.ts @@ -0,0 +1,100 @@ +import * as assert from "node:assert"; +import { readFileSync } from "node:fs"; +import { + forkSessionTitle, + 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("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", () => { + 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 = "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\)/); + + 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, /Selected entry is not a persisted user message\./); + assert.match(extension, /createBranchedSession\(predecessorId\)/); + assert.match(extension, /sendPrompt\(text\)/); + 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\)/); + assert.match(extension, /deleteSessionFileIfPresent\(forkedPath\)/); + assert.match(extension, /Stop the current run before editing a historical message\./); + }); +}); 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 new file mode 100644 index 0000000..a8de43c --- /dev/null +++ b/src/user-message-branch.ts @@ -0,0 +1,53 @@ +/** + * 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, + }; +} + +/** 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 61d6b20..ea6d534 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, content?: string) => void | Promise) | null = null; + private _onForkUserMessage: ((entryId: string, content?: 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, content?: string) => void | Promise) | null) { + this._onEditUserMessage = cb; + } + set onForkUserMessage(cb: ((entryId: string, content?: 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,7 +687,47 @@ 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, + typeof message.content === "string" ? message.content : undefined, + ); + } 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, + typeof message.content === "string" ? message.content : undefined, + ); + } 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 "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); diff --git a/src/webview/handlers/index.ts b/src/webview/handlers/index.ts index 7a05c24..5d8c221 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; @@ -252,9 +256,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) @@ -277,39 +291,93 @@ 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, + // 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-/, ""); + 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); + + host.appendChild(row); } } @@ -333,6 +401,10 @@ export function handleAgentEnd() { toolKeys: Object.keys(state.currentToolBlocks), }); 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(); @@ -569,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(); } } } @@ -589,6 +664,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); @@ -610,13 +688,97 @@ 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 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-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-actions"; + 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(saveButton, cancelButton); + el.appendChild(bar); + area.focus(); + area.select(); + 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, + content: data.content, + }); + } + }); + }); + actions.appendChild(editButton); + } + el.appendChild(actions); } state.chatContainer.appendChild(el); requestLocalMarkdownImages(el); @@ -632,8 +794,17 @@ 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;} + // #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 state.assistantToolCallIds = {}; @@ -980,6 +1151,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(); } @@ -988,6 +1160,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 @@ -1021,6 +1222,10 @@ 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(); } interface HistoryPrependContext { @@ -1160,6 +1365,7 @@ export function handleHistoryPageEnd(data: any) { context.previousScrollTop, ); state.historyLoading = false; + autoFillOlderHistory(); } function submitFollowUpQueue(messages: string[]): void { @@ -2563,9 +2769,56 @@ export function renderInlineCustomMessage(data: any) { scrollToBottom(); } +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"); + 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(); 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; +}