feat: prefill forked session prompt draft

.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.

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");
   });
 });

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");
 }
 

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)[] = [];

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 });

src/server/sessions/piSessionService.ts

@@ -103,7 +103,7 @@ export interface PiSessionRuntime {
   readonly cwd: string;
   readonly session: PiAgentSession;
   setRebindSession(rebindSession?: (session: PiAgentSession) => Promise<void>): 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<void>;
 }
 

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" });
   });

src/server/sessions/sessionCommandService.ts

@@ -33,7 +33,7 @@ export interface CommandSession {
 export interface CommandRuntime<TSession extends CommandSession = CommandSession> {
   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<TSession extends CommandSession = CommandSession> {
@@ -96,7 +96,7 @@ export class SessionCommandService<TSession extends CommandSession = CommandSess
     if (sessionHasActiveWork(active.runtime.session)) return forkActiveUnsupported("fork");
     const result = await active.runtime.fork(value);
     if (result.cancelled) return { type: "done", message: "Fork cancelled" };
-    return { type: "done", message: "Session forked", session: clientSessionFromRuntime(active.runtime) };
+    return { type: "done", message: "Session forked", session: clientSessionFromRuntime(active.runtime), ...promptDraft(result.selectedText) };
   }
 
   private nameSession(active: CommandActiveSession<TSession>, 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<Pick<Extract<ClientCommandResult, { type: "done" }>, "promptDraft">> {
+  return text === undefined ? {} : { promptDraft: text };
+}
+
 function formatSessionStats(session: CommandSession): string {
   const stats = session.getSessionStats();
   return [

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 };