From 9e3d2727311f3ba0a137700c38c32998f364a0ea Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Thu, 21 May 2026 22:56:03 +0200 Subject: [PATCH] feat: prefill forked session prompt draft --- .changeset/fork-draft-editor.md | 5 +++ src/client/src/api/parsers.test.ts | 2 +- src/client/src/api/parsers.ts | 2 +- .../src/controllers/sessionController.test.ts | 32 +++++++++++++++++++ .../src/controllers/sessionController.ts | 3 +- src/server/sessions/piSessionService.ts | 2 +- .../sessions/sessionCommandService.test.ts | 3 +- src/server/sessions/sessionCommandService.ts | 8 +++-- src/shared/apiTypes.ts | 2 +- 9 files changed, 51 insertions(+), 8 deletions(-) create mode 100644 .changeset/fork-draft-editor.md diff --git a/.changeset/fork-draft-editor.md b/.changeset/fork-draft-editor.md new file mode 100644 index 0000000..ea996bf --- /dev/null +++ b/.changeset/fork-draft-editor.md @@ -0,0 +1,5 @@ +--- +"@jmfederico/pi-web": patch +--- + +Prefill the prompt editor with the selected user message after forking a session. diff --git a/src/client/src/api/parsers.test.ts b/src/client/src/api/parsers.test.ts index c9be29c..5d26291 100644 --- a/src/client/src/api/parsers.test.ts +++ b/src/client/src/api/parsers.test.ts @@ -69,7 +69,7 @@ describe("API parsers", () => { it("parses command result variants", () => { expect(parseCommandResult({ type: "unsupported", message: "nope" })).toEqual({ type: "unsupported", message: "nope" }); expect(parseCommandResult({ type: "select", requestId: "r1", title: "Pick", options: [{ value: "v", label: "Label", description: "desc" }] })).toEqual({ type: "select", requestId: "r1", title: "Pick", options: [{ value: "v", label: "Label", description: "desc" }] }); - expect(parseCommandResult({ type: "done", message: "ok" })).toEqual({ type: "done", message: "ok" }); + expect(parseCommandResult({ type: "done", message: "ok", promptDraft: "resend me" })).toEqual({ type: "done", message: "ok", promptDraft: "resend me" }); expect(() => parseCommandResult({ type: "later" })).toThrow("Invalid command result type"); }); }); diff --git a/src/client/src/api/parsers.ts b/src/client/src/api/parsers.ts index 66d950f..c81ff5b 100644 --- a/src/client/src/api/parsers.ts +++ b/src/client/src/api/parsers.ts @@ -403,7 +403,7 @@ export function parseCommandResult(value: unknown): CommandResult { const type = requireString(record, "type"); if (type === "unsupported") return { type, message: requireString(record, "message") }; if (type === "select") return { type, requestId: requireString(record, "requestId"), title: requireString(record, "title"), options: arrayOf(parseCommandOption)(record["options"]) }; - if (type === "done") return { type, ...optionalField("message", optionalString(record, "message")), ...optionalSession(record["session"]) }; + if (type === "done") return { type, ...optionalField("message", optionalString(record, "message")), ...optionalSession(record["session"]), ...optionalField("promptDraft", optionalString(record, "promptDraft")) }; throw new Error("Invalid command result type"); } diff --git a/src/client/src/controllers/sessionController.test.ts b/src/client/src/controllers/sessionController.test.ts index a59c984..c968924 100644 --- a/src/client/src/controllers/sessionController.test.ts +++ b/src/client/src/controllers/sessionController.test.ts @@ -133,6 +133,38 @@ describe("SessionController", () => { expect(urlUpdates).toEqual([{ replace: true }]); }); + it("stores command prompt drafts for replacement sessions before selecting them", async () => { + const storage = new MemoryStorage(); + Object.defineProperty(globalThis, "localStorage", { value: storage, configurable: true }); + + let state: AppState = { + ...initialAppState(), + selectedWorkspace: workspace, + selectedSession: oldSession, + sessions: [oldSession], + commandDialog: { type: "select", requestId: "r1", title: "Fork from message", options: [{ value: "m1", label: "fork me" }] }, + }; + const urlUpdates: unknown[] = []; + const api: typeof defaultApi = { + ...defaultApi, + respondToCommand: () => Promise.resolve({ type: "done", message: "Session forked", session: replacementSession, promptDraft: "fork me" }), + messages: () => Promise.resolve(emptyPage), + status: (sessionId) => Promise.resolve(status(sessionId)), + }; + const controller = new SessionController( + () => state, + (patch) => { state = { ...state, ...patch }; }, + (options) => { urlUpdates.push(options); }, + undefined, + { api, socket: new FakeSocket() }, + ); + + await controller.respondToCommand("r1", "m1"); + + expect(state.commandDialog).toBeUndefined(); + expect(loadDraft(replacementSession.id)).toBe("fork me"); + }); + it("forgets the selected active session when archiving leaves only archived sessions", async () => { let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] }; const urlUpdates: ({ replace?: boolean | undefined } | undefined)[] = []; diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts index 7361572..bbf4b6b 100644 --- a/src/client/src/controllers/sessionController.ts +++ b/src/client/src/controllers/sessionController.ts @@ -1,7 +1,7 @@ import { api as defaultApi, type CommandResult, type SessionActivity, type SessionInfo, type SessionStatus, type ThinkingLevel } from "../api"; import { forgetCachedNewSession, isCachedNewSessionInfo, markCachedNewSessionInfo, rememberCachedNewSession, stripCachedNewSessionMarker } from "../cachedNewSessions"; import { textMessage } from "../chatMessages"; -import { clearDraft, moveDraft } from "../promptDraftStorage"; +import { clearDraft, moveDraft, saveDraft } from "../promptDraftStorage"; import { ChatTranscriptStore } from "../chatTranscriptStore"; import { isShellInput } from "../inputModes"; import { SessionSocket, type GlobalSessionEvent, type SessionUiEvent } from "../sessionSocket"; @@ -410,6 +410,7 @@ export class SessionController { const message = result.type === "unsupported" ? result.message : result.message; if (message !== undefined && message !== "") this.setState({ messages: [...this.getState().messages, textMessage(result.type === "unsupported" ? "system" : "tool", message)] }); if (result.type === "done" && result.session) { + if (result.promptDraft !== undefined) saveDraft(result.session.id, result.promptDraft); const current = this.getState().selectedSession; const sessions = [result.session, ...this.getState().sessions.filter((session) => session.id !== result.session?.id)]; this.setState({ sessions, selectedSession: current?.id === result.session.id ? result.session : current }); diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index afb786e..a4fb1b9 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -103,7 +103,7 @@ export interface PiSessionRuntime { readonly cwd: string; readonly session: PiAgentSession; setRebindSession(rebindSession?: (session: PiAgentSession) => Promise): void; - fork(entryId: string, options?: { position?: "before" | "at" }): Promise<{ cancelled: boolean }>; + fork(entryId: string, options?: { position?: "before" | "at" }): Promise<{ cancelled: boolean; selectedText?: string }>; dispose(): Promise; } diff --git a/src/server/sessions/sessionCommandService.test.ts b/src/server/sessions/sessionCommandService.test.ts index eff4ec0..4ab1b24 100644 --- a/src/server/sessions/sessionCommandService.test.ts +++ b/src/server/sessions/sessionCommandService.test.ts @@ -112,13 +112,14 @@ describe("SessionCommandService", () => { { entryId: "newest", text: "newest message" }, ]), }); + vi.mocked(active.runtime.fork).mockResolvedValueOnce({ cancelled: false, selectedText: "newest message" }); const service = new SessionCommandService(() => getActive(active), vi.fn(), eventPublisher()); const result = await service.run("s1", "/fork"); expect(result).toMatchObject({ type: "select", title: "Fork from message", options: [{ value: "newest" }, { value: "middle" }, { value: "oldest" }] }); if (result.type !== "select") throw new Error("Expected select result"); - await expect(service.respond("s1", result.requestId, "newest")).resolves.toMatchObject({ type: "done", message: "Session forked", session: { id: "s1" } }); + await expect(service.respond("s1", result.requestId, "newest")).resolves.toMatchObject({ type: "done", message: "Session forked", session: { id: "s1" }, promptDraft: "newest message" }); expect(active.runtime.fork).toHaveBeenCalledWith("newest"); await expect(service.respond("s1", result.requestId, "newest")).resolves.toEqual({ type: "unsupported", message: "Command request expired" }); }); diff --git a/src/server/sessions/sessionCommandService.ts b/src/server/sessions/sessionCommandService.ts index 4f97b59..b756ab3 100644 --- a/src/server/sessions/sessionCommandService.ts +++ b/src/server/sessions/sessionCommandService.ts @@ -33,7 +33,7 @@ export interface CommandSession { export interface CommandRuntime { cwd: string; session: TSession; - fork: (entryId: string, options?: { position?: "before" | "at" }) => Promise<{ cancelled: boolean }>; + fork: (entryId: string, options?: { position?: "before" | "at" }) => Promise<{ cancelled: boolean; selectedText?: string }>; } export interface CommandActiveSession { @@ -96,7 +96,7 @@ export class SessionCommandService, name: string): ClientCommandResult { @@ -179,6 +179,10 @@ function forkActiveUnsupported(command: "fork" | "clone"): ClientCommandResult { return { type: "unsupported", message: `Cannot ${command} while the session is active. Stop current activity before ${command === "fork" ? "forking" : "cloning"}.` }; } +function promptDraft(text: string | undefined): Partial, "promptDraft">> { + return text === undefined ? {} : { promptDraft: text }; +} + function formatSessionStats(session: CommandSession): string { const stats = session.getSessionStats(); return [ diff --git a/src/shared/apiTypes.ts b/src/shared/apiTypes.ts index 659c642..2e96af0 100644 --- a/src/shared/apiTypes.ts +++ b/src/shared/apiTypes.ts @@ -272,7 +272,7 @@ export interface MessagePage { } export type CommandResult = - | { type: "done"; message?: string; session?: SessionInfo } + | { type: "done"; message?: string; session?: SessionInfo; promptDraft?: string } | { type: "select"; requestId: string; title: string; options: CommandOption[] } | { type: "unsupported"; message: string };