wip(desktop): progress

2 files changed, 246 insertions, 113 deletions

diff --git a/packages/desktop/src/pages/layout.tsx b/packages/desktop/src/pages/layout.tsx
index 53078e01b..6632abe3a 100644
--- a/packages/desktop/src/pages/layout.tsx
+++ b/packages/desktop/src/pages/layout.tsx
@@ -55,10 +55,32 @@ export default function Layout(props: ParentProps) {
   const dialog = useDialog()
   const command = useCommand()
 
+  function flattenSessions(sessions: Session[]): Session[] {
+    const childrenMap = new Map<string, Session[]>()
+    for (const session of sessions) {
+      if (session.parentID) {
+        const children = childrenMap.get(session.parentID) ?? []
+        children.push(session)
+        childrenMap.set(session.parentID, children)
+      }
+    }
+    const result: Session[] = []
+    function visit(session: Session) {
+      result.push(session)
+      for (const child of childrenMap.get(session.id) ?? []) {
+        visit(child)
+      }
+    }
+    for (const session of sessions) {
+      if (!session.parentID) visit(session)
+    }
+    return result
+  }
+
   const currentSessions = createMemo(() => {
     if (!params.dir) return []
     const directory = base64Decode(params.dir)
-    return globalSync.child(directory)[0].session ?? []
+    return flattenSessions(globalSync.child(directory)[0].session ?? [])
   })
 
   function navigateSessionByOffset(offset: number) {
@@ -98,7 +120,7 @@ export default function Layout(props: ParentProps) {
     const nextProject = projects[nextProjectIndex]
     if (!nextProject) return
 
-    const nextProjectSessions = globalSync.child(nextProject.worktree)[0].session ?? []
+    const nextProjectSessions = flattenSessions(globalSync.child(nextProject.worktree)[0].session ?? [])
     if (nextProjectSessions.length === 0) {
       // Navigate to the project's new session page if no sessions
       navigateToProject(nextProject.worktree)
@@ -375,6 +397,98 @@ export default function Layout(props: ParentProps) {
     )
   }
 
+  const SessionItem = (props: {
+    session: Session
+    slug: string
+    project: Project
+    depth?: number
+    childrenMap: Map<string, Session[]>
+  }): JSX.Element => {
+    const notification = useNotification()
+    const depth = props.depth ?? 0
+    const children = createMemo(() => props.childrenMap.get(props.session.id) ?? [])
+    const updated = createMemo(() => DateTime.fromMillis(props.session.time.updated))
+    const notifications = createMemo(() => notification.session.unseen(props.session.id))
+    const hasError = createMemo(() => notifications().some((n) => n.type === "error"))
+    const isWorking = createMemo(
+      () =>
+        props.session.id !== params.id &&
+        globalSync.child(props.project.worktree)[0].session_status[props.session.id]?.type === "busy",
+    )
+    return (
+      <>
+        <div
+          class="group/session relative w-full pr-2 py-1 rounded-md cursor-default transition-colors
+                 hover:bg-surface-raised-base-hover focus-within:bg-surface-raised-base-hover has-[.active]:bg-surface-raised-base-hover"
+          style={{ "padding-left": `${16 + depth * 12}px` }}
+        >
+          <Tooltip placement="right" value={props.session.title} gutter={10}>
+            <A
+              href={`${props.slug}/session/${props.session.id}`}
+              class="flex flex-col min-w-0 text-left w-full focus:outline-none"
+            >
+              <div class="flex items-center self-stretch gap-6 justify-between transition-[padding] group-hover/session:pr-7 group-focus-within/session:pr-7 group-active/session:pr-7">
+                <span class="text-14-regular text-text-strong overflow-hidden text-ellipsis truncate">
+                  {props.session.title}
+                </span>
+                <div class="shrink-0 group-hover/session:hidden group-active/session:hidden group-focus-within/session:hidden">
+                  <Switch>
+                    <Match when={isWorking()}>
+                      <Spinner class="size-2.5 mr-0.5" />
+                    </Match>
+                    <Match when={hasError()}>
+                      <div class="size-1.5 mr-1.5 rounded-full bg-text-diff-delete-base" />
+                    </Match>
+                    <Match when={notifications().length > 0}>
+                      <div class="size-1.5 mr-1.5 rounded-full bg-text-interactive-base" />
+                    </Match>
+                    <Match when={true}>
+                      <span class="text-12-regular text-text-weak text-right whitespace-nowrap">
+                        {Math.abs(updated().diffNow().as("seconds")) < 60
+                          ? "Now"
+                          : updated()
+                              .toRelative({
+                                style: "short",
+                                unit: ["days", "hours", "minutes"],
+                              })
+                              ?.replace(" ago", "")
+                              ?.replace(/ days?/, "d")
+                              ?.replace(" min.", "m")
+                              ?.replace(" hr.", "h")}
+                      </span>
+                    </Match>
+                  </Switch>
+                </div>
+              </div>
+              <Show when={props.session.summary?.files}>
+                <div class="flex justify-between items-center self-stretch">
+                  <span class="text-12-regular text-text-weak">{`${props.session.summary?.files || "No"} file${props.session.summary?.files !== 1 ? "s" : ""} changed`}</span>
+                  <Show when={props.session.summary}>{(summary) => <DiffChanges changes={summary()} />}</Show>
+                </div>
+              </Show>
+            </A>
+          </Tooltip>
+          <div class="hidden group-hover/session:flex group-active/session:flex group-focus-within/session:flex text-text-base gap-1 items-center absolute top-1 right-1">
+            <Tooltip placement="right" value="Archive session">
+              <IconButton icon="archive" variant="ghost" onClick={() => archiveSession(props.session)} />
+            </Tooltip>
+          </div>
+        </div>
+        <For each={children()}>
+          {(child) => (
+            <SessionItem
+              session={child}
+              slug={props.slug}
+              project={props.project}
+              depth={depth + 1}
+              childrenMap={props.childrenMap}
+            />
+          )}
+        </For>
+      </>
+    )
+  }
+
   const SortableProject = (props: { project: Project & { expanded: boolean } }): JSX.Element => {
     const notification = useNotification()
     const sortable = createSortable(props.project.worktree)
@@ -382,6 +496,18 @@ export default function Layout(props: ParentProps) {
     const name = createMemo(() => getFilename(props.project.worktree))
     const [store, setStore] = globalSync.child(props.project.worktree)
     const sessions = createMemo(() => store.session ?? [])
+    const rootSessions = createMemo(() => sessions().filter((s) => !s.parentID))
+    const childSessionsByParent = createMemo(() => {
+      const map = new Map<string, Session[]>()
+      for (const session of sessions()) {
+        if (session.parentID) {
+          const children = map.get(session.parentID) ?? []
+          children.push(session)
+          map.set(session.parentID, children)
+        }
+      }
+      return map
+    })
     const [expanded, setExpanded] = createSignal(true)
     return (
       // @ts-ignore
@@ -421,78 +547,17 @@ export default function Layout(props: ParentProps) {
               </Button>
               <Collapsible.Content>
                 <nav class="hidden @[4rem]:flex w-full flex-col gap-1.5">
-                  <For each={sessions()}>
-                    {(session) => {
-                      const updated = createMemo(() => DateTime.fromMillis(session.time.updated))
-                      const notifications = createMemo(() => notification.session.unseen(session.id))
-                      const hasError = createMemo(() => notifications().some((n) => n.type === "error"))
-                      const isWorking = createMemo(
-                        () =>
-                          session.id !== params.id &&
-                          globalSync.child(props.project.worktree)[0].session_status[session.id]?.type === "busy",
-                      )
-                      return (
-                        <div
-                          class="group/session relative w-full pl-4 pr-2 py-1 rounded-md cursor-default transition-colors
-                                 hover:bg-surface-raised-base-hover focus-within:bg-surface-raised-base-hover has-[.active]:bg-surface-raised-base-hover"
-                        >
-                          <Tooltip placement="right" value={session.title} gutter={10}>
-                            <A
-                              href={`${slug()}/session/${session.id}`}
-                              class="flex flex-col min-w-0 text-left w-full focus:outline-none"
-                            >
-                              <div class="flex items-center self-stretch gap-6 justify-between transition-[padding] group-hover/session:pr-7 group-focus-within/session:pr-7 group-active/session:pr-7">
-                                <span class="text-14-regular text-text-strong overflow-hidden text-ellipsis truncate">
-                                  {session.title}
-                                </span>
-                                <div class="shrink-0 group-hover/session:hidden group-active/session:hidden group-focus-within/session:hidden">
-                                  <Switch>
-                                    <Match when={isWorking()}>
-                                      <Spinner class="size-2.5 mr-0.5" />
-                                    </Match>
-                                    <Match when={hasError()}>
-                                      <div class="size-1.5 mr-1.5 rounded-full bg-text-diff-delete-base" />
-                                    </Match>
-                                    <Match when={notifications().length > 0}>
-                                      <div class="size-1.5 mr-1.5 rounded-full bg-text-interactive-base" />
-                                    </Match>
-                                    <Match when={true}>
-                                      <span class="text-12-regular text-text-weak text-right whitespace-nowrap">
-                                        {Math.abs(updated().diffNow().as("seconds")) < 60
-                                          ? "Now"
-                                          : updated()
-                                              .toRelative({
-                                                style: "short",
-                                                unit: ["days", "hours", "minutes"],
-                                              })
-                                              ?.replace(" ago", "")
-                                              ?.replace(/ days?/, "d")
-                                              ?.replace(" min.", "m")
-                                              ?.replace(" hr.", "h")}
-                                      </span>
-                                    </Match>
-                                  </Switch>
-                                </div>
-                              </div>
-                              <Show when={session.summary?.files}>
-                                <div class="flex justify-between items-center self-stretch">
-                                  <span class="text-12-regular text-text-weak">{`${session.summary?.files || "No"} file${session.summary?.files !== 1 ? "s" : ""} changed`}</span>
-                                  <Show when={session.summary}>{(summary) => <DiffChanges changes={summary()} />}</Show>
-                                </div>
-                              </Show>
-                            </A>
-                          </Tooltip>
-                          <div class="hidden group-hover/session:flex group-active/session:flex group-focus-within/session:flex text-text-base gap-1 items-center absolute top-1 right-1">
-                            {/* <IconButton icon="dot-grid" variant="ghost" /> */}
-                            <Tooltip placement="right" value="Archive session">
-                              <IconButton icon="archive" variant="ghost" onClick={() => archiveSession(session)} />
-                            </Tooltip>
-                          </div>
-                        </div>
-                      )
-                    }}
+                  <For each={rootSessions()}>
+                    {(session) => (
+                      <SessionItem
+                        session={session}
+                        slug={slug()}
+                        project={props.project}
+                        childrenMap={childSessionsByParent()}
+                      />
+                    )}
                   </For>
-                  <Show when={sessions().length === 0}>
+                  <Show when={rootSessions().length === 0}>
                     <div
                       class="group/session relative w-full pl-4 pr-2 py-1 rounded-md cursor-default transition-colors
                              hover:bg-surface-raised-base-hover focus-within:bg-surface-raised-base-hover has-[.active]:bg-surface-raised-base-hover"
diff --git a/packages/desktop/src/pages/session.tsx b/packages/desktop/src/pages/session.tsx
index 05a9e8a1d..11056a598 100644
--- a/packages/desktop/src/pages/session.tsx
+++ b/packages/desktop/src/pages/session.tsx
@@ -1,4 +1,4 @@
-import { For, onCleanup, onMount, Show, Match, Switch, createResource, createMemo, createEffect } from "solid-js"
+import { For, onCleanup, onMount, Show, Match, Switch, createResource, createMemo, createEffect, on } from "solid-js"
 import { useLocal, type LocalFile } from "@/context/local"
 import { createStore } from "solid-js/store"
 import { PromptInput } from "@/components/prompt-input"
@@ -38,6 +38,9 @@ import { DialogSelectModel } from "@/components/dialog-select-model"
 import { useCommand } from "@/context/command"
 import { useNavigate, useParams } from "@solidjs/router"
 import { AssistantMessage, UserMessage } from "@opencode-ai/sdk/v2"
+import { useSDK } from "@/context/sdk"
+import { usePrompt } from "@/context/prompt"
+import { extractPromptFromParts } from "@/utils/prompt"
 
 export default function Page() {
   const layout = useLayout()
@@ -48,45 +51,56 @@ export default function Page() {
   const command = useCommand()
   const params = useParams()
   const navigate = useNavigate()
+  const sdk = useSDK()
+  const prompt = usePrompt()
 
   const sessionKey = createMemo(() => `${params.dir}${params.id ? "/" + params.id : ""}`)
   const tabs = createMemo(() => layout.tabs(sessionKey()))
 
   const info = createMemo(() => (params.id ? sync.session.get(params.id) : undefined))
+  const revertMessageID = createMemo(() => info()?.revert?.messageID)
   const messages = createMemo(() => (params.id ? (sync.data.message[params.id] ?? []) : []))
   const userMessages = createMemo(() =>
     messages()
       .filter((m) => m.role === "user")
       .sort((a, b) => a.id.localeCompare(b.id)),
   )
-  const lastUserMessage = createMemo(() => userMessages()?.at(-1))
+  // Visible user messages excludes reverted messages (those >= revertMessageID)
+  const visibleUserMessages = createMemo(() => {
+    const revert = revertMessageID()
+    if (!revert) return userMessages()
+    return userMessages().filter((m) => m.id < revert)
+  })
+  const lastUserMessage = createMemo(() => visibleUserMessages()?.at(-1))
 
   const [messageStore, setMessageStore] = createStore<{ messageId?: string }>({})
   const activeMessage = createMemo(() => {
     if (!messageStore.messageId) return lastUserMessage()
-    return userMessages()?.find((m) => m.id === messageStore.messageId)
+    // If the stored message is no longer visible (e.g., was reverted), fall back to last visible
+    const found = visibleUserMessages()?.find((m) => m.id === messageStore.messageId)
+    return found ?? lastUserMessage()
   })
   const setActiveMessage = (message: UserMessage | undefined) => {
     setMessageStore("messageId", message?.id)
   }
 
   function navigateMessageByOffset(offset: number) {
-    const messages = userMessages()
-    if (messages.length === 0) return
+    const msgs = visibleUserMessages()
+    if (msgs.length === 0) return
 
     const current = activeMessage()
-    const currentIndex = current ? messages.findIndex((m) => m.id === current.id) : -1
+    const currentIndex = current ? msgs.findIndex((m) => m.id === current.id) : -1
 
     let targetIndex: number
     if (currentIndex === -1) {
-      targetIndex = offset > 0 ? 0 : messages.length - 1
+      targetIndex = offset > 0 ? 0 : msgs.length - 1
     } else {
       targetIndex = currentIndex + offset
     }
 
-    if (targetIndex < 0 || targetIndex >= messages.length) return
+    if (targetIndex < 0 || targetIndex >= msgs.length) return
 
-    setActiveMessage(messages[targetIndex])
+    setActiveMessage(msgs[targetIndex])
   }
 
   const last = createMemo(
@@ -131,6 +145,24 @@ export default function Page() {
     }
   })
 
+  // Auto-navigate to new messages when they're added
+  // This handles the case after undo + submit where we want to see the new message
+  // We track the last message ID and only navigate when a NEW message is added (ID increases)
+  createEffect(
+    on(
+      () => visibleUserMessages().at(-1)?.id,
+      (lastId, prevLastId) => {
+        // Only navigate if a new message was added (lastId is greater/newer than previous)
+        if (lastId && prevLastId && lastId > prevLastId) {
+          setMessageStore("messageId", undefined)
+        }
+      },
+      { defer: true },
+    ),
+  )
+
+  const status = createMemo(() => sync.data.session_status[params.id ?? ""] ?? { type: "idle" })
+
   command.register(() => [
     {
       id: "session.new",
@@ -226,6 +258,66 @@ export default function Page() {
       slash: "agent",
       onSelect: () => local.agent.move(1),
     },
+    {
+      id: "session.undo",
+      title: "Undo",
+      description: "Undo the last message",
+      category: "Session",
+      keybind: "mod+z",
+      slash: "undo",
+      disabled: !params.id || visibleUserMessages().length === 0,
+      onSelect: async () => {
+        const sessionID = params.id
+        if (!sessionID) return
+        if (status()?.type !== "idle") {
+          await sdk.client.session.abort({ sessionID }).catch(() => {})
+        }
+        const revert = info()?.revert?.messageID
+        // Find the last user message that's not already reverted
+        const message = userMessages().findLast((x) => !revert || x.id < revert)
+        if (!message) return
+        await sdk.client.session.revert({ sessionID, messageID: message.id })
+        // Restore the prompt from the reverted message
+        const parts = sync.data.part[message.id]
+        if (parts) {
+          const restored = extractPromptFromParts(parts)
+          prompt.set(restored)
+        }
+        // Navigate to the message before the reverted one (which will be the new last visible message)
+        const priorMessage = userMessages().findLast((x) => x.id < message.id)
+        setActiveMessage(priorMessage)
+      },
+    },
+    {
+      id: "session.redo",
+      title: "Redo",
+      description: "Redo the last undone message",
+      category: "Session",
+      keybind: "mod+shift+z",
+      slash: "redo",
+      disabled: !params.id || !info()?.revert?.messageID,
+      onSelect: async () => {
+        const sessionID = params.id
+        if (!sessionID) return
+        const revertMessageID = info()?.revert?.messageID
+        if (!revertMessageID) return
+        const nextMessage = userMessages().find((x) => x.id > revertMessageID)
+        if (!nextMessage) {
+          // Full unrevert - restore all messages and navigate to last
+          await sdk.client.session.unrevert({ sessionID })
+          prompt.reset()
+          // Navigate to the last message (the one that was at the revert point)
+          const lastMsg = userMessages().findLast((x) => x.id >= revertMessageID)
+          setActiveMessage(lastMsg)
+          return
+        }
+        // Partial redo - move forward to next message
+        await sdk.client.session.revert({ sessionID, messageID: nextMessage.id })
+        // Navigate to the message before the new revert point
+        const priorMsg = userMessages().findLast((x) => x.id < nextMessage.id)
+        setActiveMessage(priorMsg)
+      },
+    },
   ])
 
   const handleKeyDown = (event: KeyboardEvent) => {
@@ -548,7 +640,7 @@ export default function Page() {
                     <Match when={params.id}>
                       <div class="flex items-start justify-start h-full min-h-0">
                         <SessionMessageRail
-                          messages={userMessages()}
+                          messages={visibleUserMessages()}
                           current={activeMessage()}
                           onMessageSelect={setActiveMessage}
                           wide={wide()}
@@ -556,7 +648,7 @@ export default function Page() {
                         <Show when={activeMessage()}>
                           <SessionTurn
                             sessionID={params.id!}
-                            messageID={activeMessage()?.id!}
+                            messageID={activeMessage()!.id}
                             stepsExpanded={store.stepsExpanded}
                             onStepsExpandedChange={(expanded) => setStore("stepsExpanded", expanded)}
                             classes={{
@@ -564,7 +656,11 @@ export default function Page() {
                               content: "pb-20",
                               container:
                                 "w-full " +
-                                (wide() ? "max-w-146 mx-auto px-6" : userMessages().length > 1 ? "pr-6 pl-18" : "px-6"),
+                                (wide()
+                                  ? "max-w-146 mx-auto px-6"
+                                  : visibleUserMessages().length > 1
+                                    ? "pr-6 pl-18"
+                                    : "px-6"),
                             }}
                           />
                         </Show>
@@ -718,34 +814,6 @@ export default function Page() {
             />
           </div>
         </Show>
-        <div class="hidden shrink-0 w-56 p-2 h-full overflow-y-auto">
-          {/* <FileTree path="" onFileClick={ handleTabClick} /> */}
-        </div>
-        <div class="hidden shrink-0 w-56 p-2">
-          <Show
-            when={local.file.changes().length}
-            fallback={<div class="px-2 text-xs text-text-muted">No changes</div>}
-          >
-            <ul class="">
-              <For each={local.file.changes()}>
-                {(path) => (
-                  <li>
-                    <button
-                      onClick={() => local.file.open(path, { view: "diff-unified", pinned: true })}
-                      class="w-full flex items-center px-2 py-0.5 gap-x-2 text-text-muted grow min-w-0 hover:bg-background-element"
-                    >
-                      <FileIcon node={{ path, type: "file" }} class="shrink-0 size-3" />
-                      <span class="text-xs text-text whitespace-nowrap">{getFilename(path)}</span>
-                      <span class="text-xs text-text-muted/60 whitespace-nowrap truncate min-w-0">
-                        {getDirectory(path)}
-                      </span>
-                    </button>
-                  </li>
-                )}
-              </For>
-            </ul>
-          </Show>
-        </div>
       </div>
       <Show when={layout.terminal.opened()}>
         <div