feat(app): restore to message and fork session (#17092)

5 files changed, 256 insertions, 4 deletions

diff --git a/packages/app/src/i18n/en.ts b/packages/app/src/i18n/en.ts
index 97a572f1c..c87e7cb9d 100644
--- a/packages/app/src/i18n/en.ts
+++ b/packages/app/src/i18n/en.ts
@@ -530,6 +530,11 @@ export const dict = {
   "session.todo.title": "Todos",
   "session.todo.collapse": "Collapse",
   "session.todo.expand": "Expand",
+  "session.revertDock.summary.one": "{{count}} rolled back message",
+  "session.revertDock.summary.other": "{{count}} rolled back messages",
+  "session.revertDock.collapse": "Collapse rolled back messages",
+  "session.revertDock.expand": "Expand rolled back messages",
+  "session.revertDock.restore": "Restore message",
 
   "session.new.title": "Build anything",
   "session.new.worktree.main": "Main branch",
diff --git a/packages/app/src/pages/session.tsx b/packages/app/src/pages/session.tsx
index a5c7bf90b..1b62b9429 100644
--- a/packages/app/src/pages/session.tsx
+++ b/packages/app/src/pages/session.tsx
@@ -43,6 +43,7 @@ import { SessionSidePanel } from "@/pages/session/session-side-panel"
 import { TerminalPanel } from "@/pages/session/terminal-panel"
 import { useSessionCommands } from "@/pages/session/use-session-commands"
 import { useSessionHashScroll } from "@/pages/session/use-session-hash-scroll"
+import { extractPromptFromParts } from "@/utils/prompt"
 import { same } from "@/utils/same"
 import { formatServerError } from "@/utils/server-errors"
 
@@ -286,6 +287,7 @@ export default function Page() {
   const [ui, setUi] = createStore({
     git: false,
     pendingMessage: undefined as string | undefined,
+    restoring: undefined as string | undefined,
     reviewSnap: false,
     scrollGesture: 0,
     scroll: {
@@ -1179,6 +1181,110 @@ export default function Page() {
     scroller: () => scroller,
   })
 
+  const draft = (id: string) =>
+    extractPromptFromParts(sync.data.part[id] ?? [], {
+      directory: sdk.directory,
+      attachmentName: language.t("common.attachment"),
+    })
+
+  const line = (id: string) => {
+    const text = draft(id)
+      .map((part) => (part.type === "image" ? `[image:${part.filename}]` : part.content))
+      .join("")
+      .replace(/\s+/g, " ")
+      .trim()
+    if (text) return text
+    return `[${language.t("common.attachment")}]`
+  }
+
+  const fail = (err: unknown) => {
+    showToast({
+      variant: "error",
+      title: language.t("common.requestFailed"),
+      description: formatServerError(err, language.t),
+    })
+  }
+
+  const busy = (sessionID: string) => {
+    if (sync.data.session_status[sessionID]?.type !== "idle") return true
+    return (sync.data.message[sessionID] ?? []).some(
+      (item) => item.role === "assistant" && typeof item.time.completed !== "number",
+    )
+  }
+
+  const halt = (sessionID: string) =>
+    busy(sessionID) ? sdk.client.session.abort({ sessionID }).catch(() => {}) : Promise.resolve()
+
+  const fork = (input: { sessionID: string; messageID: string }) => {
+    const value = draft(input.messageID)
+    return sdk.client.session
+      .fork(input)
+      .then((result) => {
+        const next = result.data
+        if (!next) {
+          showToast({
+            variant: "error",
+            title: language.t("common.requestFailed"),
+          })
+          return
+        }
+        navigate(`/${base64Encode(sdk.directory)}/session/${next.id}`)
+        requestAnimationFrame(() => {
+          prompt.set(value)
+        })
+      })
+      .catch(fail)
+  }
+
+  const revert = (input: { sessionID: string; messageID: string }) => {
+    const value = draft(input.messageID)
+    return halt(input.sessionID)
+      .then(() => sdk.client.session.revert(input))
+      .then(() => {
+        prompt.set(value)
+      })
+      .catch(fail)
+  }
+
+  const restore = (id: string) => {
+    const sessionID = params.id
+    if (!sessionID || ui.restoring) return
+
+    const next = userMessages().find((item) => item.id > id)
+    setUi("restoring", id)
+
+    const task = !next
+      ? halt(sessionID)
+          .then(() => sdk.client.session.unrevert({ sessionID }))
+          .then(() => {
+            prompt.reset()
+          })
+      : halt(sessionID)
+          .then(() =>
+            sdk.client.session.revert({
+              sessionID,
+              messageID: next.id,
+            }),
+          )
+          .then(() => {
+            prompt.set(draft(next.id))
+          })
+
+    return task.catch(fail).finally(() => {
+      setUi("restoring", (value) => (value === id ? undefined : value))
+    })
+  }
+
+  const rolled = createMemo(() => {
+    const id = revertMessageID()
+    if (!id) return []
+    return userMessages()
+      .filter((item) => item.id >= id)
+      .map((item) => ({ id: item.id, text: line(item.id) }))
+  })
+
+  const actions = { fork, revert }
+
   createResizeObserver(
     () => promptDock,
     ({ height }) => {
@@ -1268,6 +1374,7 @@ export default function Page() {
                       loadingClass: "px-4 py-4 text-text-weak",
                       emptyClass: "h-full pb-64 -mt-4 flex flex-col items-center justify-center text-center gap-6",
                     })}
+                    actions={actions}
                     scroll={ui.scroll}
                     onResumeScroll={resumeScroll}
                     setScrollRef={setScrollRef}
@@ -1333,6 +1440,15 @@ export default function Page() {
               resumeScroll()
             }}
             onResponseSubmit={resumeScroll}
+            revert={
+              rolled().length > 0
+                ? {
+                    items: rolled(),
+                    restoring: ui.restoring,
+                    onRestore: restore,
+                  }
+                : undefined
+            }
             setPromptDockRef={(el) => {
               promptDock = el
             }}
diff --git a/packages/app/src/pages/session/composer/session-composer-region.tsx b/packages/app/src/pages/session/composer/session-composer-region.tsx
index 93ea3d465..08746b51a 100644
--- a/packages/app/src/pages/session/composer/session-composer-region.tsx
+++ b/packages/app/src/pages/session/composer/session-composer-region.tsx
@@ -8,6 +8,7 @@ import { usePrompt } from "@/context/prompt"
 import { getSessionHandoff, setSessionHandoff } from "@/pages/session/handoff"
 import { SessionPermissionDock } from "@/pages/session/composer/session-permission-dock"
 import { SessionQuestionDock } from "@/pages/session/composer/session-question-dock"
+import { SessionRevertDock } from "@/pages/session/composer/session-revert-dock"
 import type { SessionComposerState } from "@/pages/session/composer/session-composer-state"
 import { SessionTodoDock } from "@/pages/session/composer/session-todo-dock"
 
@@ -20,6 +21,11 @@ export function SessionComposerRegion(props: {
   onNewSessionWorktreeReset: () => void
   onSubmit: () => void
   onResponseSubmit: () => void
+  revert?: {
+    items: { id: string; text: string }[]
+    restoring?: string
+    onRestore: (id: string) => void
+  }
   setPromptDockRef: (el: HTMLDivElement) => void
   visualDuration?: number
   bounce?: number
@@ -116,6 +122,8 @@ export function SessionComposerRegion(props: {
   const value = createMemo(() => Math.max(0, Math.min(1, progress())))
   const [height, setHeight] = createSignal(320)
   const dock = createMemo(() => (gate.ready && props.state.dock()) || value() > 0.001)
+  const rolled = createMemo(() => (props.revert?.items.length ? props.revert : undefined))
+  const lift = createMemo(() => (rolled() ? 18 : 36 * value()))
   const full = createMemo(() => Math.max(78, height()))
   const [contentRef, setContentRef] = createSignal<HTMLDivElement>()
 
@@ -170,9 +178,22 @@ export function SessionComposerRegion(props: {
           <Show
             when={prompt.ready()}
             fallback={
-              <div class="w-full min-h-32 md:min-h-40 rounded-md border border-border-weak-base bg-background-base/50 px-4 py-3 text-text-weak whitespace-pre-wrap pointer-events-none">
-                {handoffPrompt() || language.t("prompt.loading")}
-              </div>
+              <>
+                <Show when={rolled()} keyed>
+                  {(revert) => (
+                    <div class="pb-2">
+                      <SessionRevertDock
+                        items={revert.items}
+                        restoring={revert.restoring}
+                        onRestore={revert.onRestore}
+                      />
+                    </div>
+                  )}
+                </Show>
+                <div class="w-full min-h-32 md:min-h-40 rounded-md border border-border-weak-base bg-background-base/50 px-4 py-3 text-text-weak whitespace-pre-wrap pointer-events-none">
+                  {handoffPrompt() || language.t("prompt.loading")}
+                </div>
+              </>
             }
           >
             <Show when={dock()}>
@@ -209,12 +230,23 @@ export function SessionComposerRegion(props: {
                 </div>
               </div>
             </Show>
+            <Show when={rolled()} keyed>
+              {(revert) => (
+                <div
+                  style={{
+                    "margin-top": `${-36 * value()}px`,
+                  }}
+                >
+                  <SessionRevertDock items={revert.items} restoring={revert.restoring} onRestore={revert.onRestore} />
+                </div>
+              )}
+            </Show>
             <div
               classList={{
                 "relative z-10": true,
               }}
               style={{
-                "margin-top": `${-36 * value()}px`,
+                "margin-top": `${-lift()}px`,
               }}
             >
               <PromptInput
diff --git a/packages/app/src/pages/session/composer/session-revert-dock.tsx b/packages/app/src/pages/session/composer/session-revert-dock.tsx
new file mode 100644
index 000000000..ce7644b67
--- /dev/null
+++ b/packages/app/src/pages/session/composer/session-revert-dock.tsx
@@ -0,0 +1,92 @@
+import { For, Show, createMemo } from "solid-js"
+import { createStore } from "solid-js/store"
+import { Button } from "@opencode-ai/ui/button"
+import { DockTray } from "@opencode-ai/ui/dock-surface"
+import { IconButton } from "@opencode-ai/ui/icon-button"
+import { useLanguage } from "@/context/language"
+
+export function SessionRevertDock(props: {
+  items: { id: string; text: string }[]
+  restoring?: string
+  onRestore: (id: string) => void
+}) {
+  const language = useLanguage()
+  const [store, setStore] = createStore({
+    collapsed: false,
+  })
+
+  const toggle = () => setStore("collapsed", (value) => !value)
+  const total = createMemo(() => props.items.length)
+  const label = createMemo(() =>
+    language.t(total() === 1 ? "session.revertDock.summary.one" : "session.revertDock.summary.other", {
+      count: total(),
+    }),
+  )
+  const preview = createMemo(() => props.items[0]?.text ?? "")
+
+  return (
+    <DockTray data-component="session-revert-dock">
+      <div
+        class="pl-3 pr-2 py-2 flex items-center gap-2"
+        role="button"
+        tabIndex={0}
+        onClick={toggle}
+        onKeyDown={(event) => {
+          if (event.key !== "Enter" && event.key !== " ") return
+          event.preventDefault()
+          toggle()
+        }}
+      >
+        <span class="shrink-0 text-14-regular text-text-strong cursor-default">{label()}</span>
+        <Show when={store.collapsed && preview()}>
+          <span class="min-w-0 flex-1 truncate text-14-regular text-text-base cursor-default">{preview()}</span>
+        </Show>
+        <div class="ml-auto shrink-0">
+          <IconButton
+            data-collapsed={store.collapsed ? "true" : "false"}
+            icon="chevron-down"
+            size="normal"
+            variant="ghost"
+            style={{ transform: `rotate(${store.collapsed ? 180 : 0}deg)` }}
+            onMouseDown={(event) => {
+              event.preventDefault()
+              event.stopPropagation()
+            }}
+            onClick={(event) => {
+              event.stopPropagation()
+              toggle()
+            }}
+            aria-label={
+              store.collapsed ? language.t("session.revertDock.expand") : language.t("session.revertDock.collapse")
+            }
+          />
+        </div>
+      </div>
+
+      <Show when={store.collapsed}>
+        <div class="h-5" aria-hidden="true" />
+      </Show>
+
+      <Show when={!store.collapsed}>
+        <div class="px-3 pb-11 flex flex-col gap-1.5 max-h-42 overflow-y-auto no-scrollbar">
+          <For each={props.items}>
+            {(item) => (
+              <div class="flex items-center gap-2 min-w-0 rounded-[10px] border border-border-weak-base bg-background-stronger px-2.5 py-2">
+                <span class="min-w-0 flex-1 truncate text-13-regular text-text-strong">{item.text}</span>
+                <Button
+                  size="small"
+                  variant="secondary"
+                  class="shrink-0"
+                  disabled={!!props.restoring}
+                  onClick={() => props.onRestore(item.id)}
+                >
+                  {language.t("session.revertDock.restore")}
+                </Button>
+              </div>
+            )}
+          </For>
+        </div>
+      </Show>
+    </DockTray>
+  )
+}
diff --git a/packages/app/src/pages/session/message-timeline.tsx b/packages/app/src/pages/session/message-timeline.tsx
index 223539092..cb9c61918 100644
--- a/packages/app/src/pages/session/message-timeline.tsx
+++ b/packages/app/src/pages/session/message-timeline.tsx
@@ -36,6 +36,11 @@ type MessageComment = {
 const emptyMessages: MessageType[] = []
 const idle = { type: "idle" as const }
 
+type UserActions = {
+  fork?: (input: { sessionID: string; messageID: string }) => Promise<void> | void
+  revert?: (input: { sessionID: string; messageID: string }) => Promise<void> | void
+}
+
 const messageComments = (parts: Part[]): MessageComment[] =>
   parts.flatMap((part) => {
     if (part.type !== "text" || !(part as TextPart).synthetic) return []
@@ -186,6 +191,7 @@ function createTimelineStaging(input: TimelineStageInput) {
 export function MessageTimeline(props: {
   mobileChanges: boolean
   mobileFallback: JSX.Element
+  actions?: UserActions
   scroll: { overflow: boolean; bottom: boolean }
   onResumeScroll: () => void
   setScrollRef: (el: HTMLDivElement | undefined) => void
@@ -805,6 +811,7 @@ export function MessageTimeline(props: {
                       <SessionTurn
                         sessionID={sessionID() ?? ""}
                         messageID={messageID}
+                        actions={props.actions}
                         active={active()}
                         queued={queued()}
                         status={active() ? sessionStatus() : undefined}