feat: prefill forked session prompt draft

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