Fork from point in chat buffer (f key)

README.org

@@ -15,6 +15,7 @@ An Emacs frontend for the [[https://shittycodingagent.ai/][pi coding agent]].
 
 - Compose prompts in a full Emacs buffer: multi-line, copy/paste, macros, support for Vi bindings
 - Chat history as a markdown buffer: copy, save, search, navigate
+- Fork the converstaion at any point in the chat buffer (f)
 - Live streaming output as bash commands and tool operations run
 - Syntax-highlighted code blocks and diffs
 - Collapsible tool output with smart preview (expand with TAB)
@@ -120,6 +121,7 @@ For multiple sessions in the same directory, use =C-u M-x pi= to create a named
 | =TAB=         | chat    | Toggle section                    |
 | =S-TAB=       | chat    | Cycle all folds                   |
 | =RET=         | chat    | Visit file at point (other window)|
+| =f=           | chat    | Fork from point                   |
 | =q=           | chat    | Quit session                      |
 
 Press =C-c C-p= to access the full menu with model selection, thinking level,
@@ -185,8 +187,9 @@ When a conversation gets long, the AI's context window fills up. The menu
   preserving key information. Use when context is filling up. If you don't
   compact manually, pi does it automatically when needed.
 
-- *Fork* (=f=): Creates a new session starting from a previous message.
-  Go back to any point and take the conversation in a new direction.
+- *Fork* (=f=): Branches the conversation from any earlier turn.
+  Press =f= on any turn in the chat buffer, or use the menu to pick
+  from a list.
 
 - *Export* (=e=): Saves the conversation as an HTML file for sharing or
   archiving.

pi-coding-agent-menu.el

@@ -567,6 +567,38 @@ Optional CUSTOM-INSTRUCTIONS provide guidance for the compaction summary."
 
 ;;;; Fork
 
+(defun pi-coding-agent--flatten-tree (nodes)
+  "Flatten tree NODES into a hash table mapping id to node plist.
+NODES is a vector of tree node plists, each with `:children' vector.
+Returns a hash table for O(1) lookup by id."
+  (let ((index (make-hash-table :test 'equal)))
+    (cl-labels ((walk (ns)
+                  (seq-doseq (node ns)
+                    (puthash (plist-get node :id) node index)
+                    (let ((children (plist-get node :children)))
+                      (when (and children (> (length children) 0))
+                        (walk children))))))
+      (walk nodes))
+    index))
+
+(defun pi-coding-agent--active-branch-user-ids (index leaf-id)
+  "Return chronological list of user message IDs on the active branch.
+INDEX is a hash table from `pi-coding-agent--flatten-tree'.
+LEAF-ID is the current leaf node ID.  Walk from leaf to root via
+`:parentId', collecting IDs of nodes with type \"message\" and role
+\"user\".  Returns list in root-to-leaf (chronological) order."
+  (when leaf-id
+    (let ((user-ids nil)
+          (current-id leaf-id))
+      (while current-id
+        (let ((node (gethash current-id index)))
+          (when (and node
+                     (equal (plist-get node :type) "message")
+                     (equal (plist-get node :role) "user"))
+            (push (plist-get node :id) user-ids))
+          (setq current-id (and node (plist-get node :parentId)))))
+      user-ids)))
+
 (defun pi-coding-agent--format-fork-message (msg &optional index)
   "Format MSG for display in fork selector.
 MSG is a plist with :entryId and :text.
@@ -593,6 +625,87 @@ Shows a selector of user messages and creates a fork from the selected one."
                              (pi-coding-agent--show-fork-selector proc messages)))
                        (message "Pi: Failed to get fork messages"))))))
 
+(defun pi-coding-agent--resolve-fork-entry (response ordinal heading-count)
+  "Resolve a fork entry ID from get_tree RESPONSE.
+ORDINAL is the 0-based user turn index.  HEADING-COUNT is the number
+of visible You headings in the buffer.  Returns (ENTRY-ID . PREVIEW)
+or nil if the ordinal could not be mapped."
+  (when (plist-get response :success)
+    (let* ((data (plist-get response :data))
+           (tree (plist-get data :tree))
+           (leaf-id (plist-get data :leafId))
+           (index (pi-coding-agent--flatten-tree tree))
+           (all-user-ids (pi-coding-agent--active-branch-user-ids index leaf-id))
+           ;; Take last N to handle compaction (compacted-away
+           ;; user messages at start of path aren't rendered)
+           (visible-ids (last all-user-ids heading-count))
+           (entry-id (nth ordinal visible-ids))
+           (node (and entry-id (gethash entry-id index))))
+      (when entry-id
+        (cons entry-id (plist-get node :preview))))))
+
+(defun pi-coding-agent-fork-at-point ()
+  "Fork conversation from the user turn at point.
+Determines which user message point is in (or after), confirms with
+a preview, then forks.  Only works when the session is idle."
+  (interactive)
+  (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
+    (unless chat-buf
+      (user-error "Pi: No chat buffer"))
+    (with-current-buffer chat-buf
+      (let* ((headings (pi-coding-agent--collect-you-headings))
+             (ordinal (pi-coding-agent--user-turn-index-at-point headings)))
+        (cond
+         ((not (eq pi-coding-agent--status 'idle))
+          (message "Pi: Cannot fork while streaming"))
+         ((not ordinal)
+          (message "Pi: No user message at point"))
+         (t
+          (let ((heading-count (length headings))
+                (proc (pi-coding-agent--get-process)))
+            (unless proc
+              (user-error "Pi: No active process"))
+            (pi-coding-agent--rpc-async proc '(:type "get_tree")
+              (lambda (response)
+                (let ((result (pi-coding-agent--resolve-fork-entry
+                               response ordinal heading-count)))
+                  (cond
+                   ((not result)
+                    (message "Pi: Could not map turn to entry ID"))
+                   ((with-current-buffer chat-buf
+                      (y-or-n-p (format "Fork from: %s? " (or (cdr result) "?"))))
+                    (with-current-buffer chat-buf
+                      (pi-coding-agent--execute-fork proc (car result)))))))))))))))
+
+(defun pi-coding-agent--execute-fork (proc entry-id)
+  "Execute fork to ENTRY-ID via PROC.
+Sends the fork RPC, then on success: refreshes state, reloads history,
+and pre-fills the input buffer with the forked message text.
+Captures chat and input buffers at call time (before the async RPC)."
+  (let ((chat-buf (pi-coding-agent--get-chat-buffer))
+        (input-buf (pi-coding-agent--get-input-buffer)))
+    (pi-coding-agent--rpc-async proc (list :type "fork" :entryId entry-id)
+      (lambda (response)
+        (if (plist-get response :success)
+            (let* ((data (plist-get response :data))
+                   (text (plist-get data :text)))
+              ;; Refresh state to get new session-file
+              (pi-coding-agent--rpc-async proc '(:type "get_state")
+                (lambda (resp)
+                  (pi-coding-agent--apply-state-response chat-buf resp)))
+              ;; Reload and display the forked session
+              (pi-coding-agent--load-session-history
+               proc
+               (lambda (count)
+                 (message "Pi: Branched to new session (%d messages)" count))
+               chat-buf)
+              ;; Pre-fill input with the forked message text
+              (when (buffer-live-p input-buf)
+                (with-current-buffer input-buf
+                  (erase-buffer)
+                  (when text (insert text)))))
+          (message "Pi: Branch failed"))))))
+
 (defun pi-coding-agent--show-fork-selector (proc messages)
   "Show selector for MESSAGES and fork on selection.
 PROC is the pi process.
@@ -613,34 +726,9 @@ MESSAGES is a vector of plists from get_fork_messages."
                                         '(metadata (display-sort-function . identity))
                                       (complete-with-action action choice-strings string pred)))
                                   nil t))
-         (selected (cdr (assoc choice formatted)))
-         ;; Capture buffers before async call (callback runs in arbitrary context)
-         (chat-buf (pi-coding-agent--get-chat-buffer))
-         (input-buf (pi-coding-agent--get-input-buffer)))
+         (selected (cdr (assoc choice formatted))))
     (when selected
-      (let ((entry-id (plist-get selected :entryId)))
-        (pi-coding-agent--rpc-async proc (list :type "fork" :entryId entry-id)
-                       (lambda (response)
-                         (if (plist-get response :success)
-                             (let* ((data (plist-get response :data))
-                                    (text (plist-get data :text)))
-                               ;; Refresh state to get new session-file
-                               (pi-coding-agent--rpc-async proc '(:type "get_state")
-                                 (lambda (resp)
-                                   (pi-coding-agent--apply-state-response chat-buf resp)))
-                               ;; Reload and display the forked session
-                               (pi-coding-agent--load-session-history
-                                proc
-                                (lambda (count)
-                                  (message "Pi: Branched to new session (%d messages)" count))
-                                chat-buf)
-                               ;; Pre-fill input with the selected message text
-                               (when (buffer-live-p input-buf)
-                                 (with-current-buffer input-buf
-                                   (erase-buffer)
-                                   ;; text may be nil if RPC returns null
-                                   (when text (insert text)))))
-                           (message "Pi: Branch failed"))))))))
+      (pi-coding-agent--execute-fork proc (plist-get selected :entryId)))))
 
 ;;;; Custom Commands
 

pi-coding-agent-ui.el

@@ -71,6 +71,7 @@
 (declare-function pi-coding-agent-resume-session "pi-coding-agent-menu")
 (declare-function pi-coding-agent-select-model "pi-coding-agent-menu")
 (declare-function pi-coding-agent-cycle-thinking "pi-coding-agent-menu")
+(declare-function pi-coding-agent-fork-at-point "pi-coding-agent-menu")
 
 ;; Optional: phscroll for horizontal table scrolling
 (require 'phscroll nil t)
@@ -333,33 +334,105 @@ Returns \"text\" for unrecognized extensions to ensure consistent fencing."
     (define-key map (kbd "C-c C-p") #'pi-coding-agent-menu)
     (define-key map (kbd "n") #'pi-coding-agent-next-message)
     (define-key map (kbd "p") #'pi-coding-agent-previous-message)
+    (define-key map (kbd "f") #'pi-coding-agent-fork-at-point)
     (define-key map (kbd "TAB") #'pi-coding-agent-toggle-tool-section)
     (define-key map (kbd "<tab>") #'pi-coding-agent-toggle-tool-section)
     (define-key map (kbd "RET") #'pi-coding-agent-visit-file)
     (define-key map (kbd "<return>") #'pi-coding-agent-visit-file)
     map)
   "Keymap for `pi-coding-agent-chat-mode'.")
 
+;;;; You Heading Detection
+
+(defconst pi-coding-agent--you-heading-re
+  "^You\\( · .*\\)?$"
+  "Regex matching the first line of a user turn setext heading.
+Matches `You' at line start, optionally followed by ` · <timestamp>'.
+Must be verified with `pi-coding-agent--at-you-heading-p' to confirm
+the next line is a setext underline (===), avoiding false matches on
+user message text starting with \"You\".")
+
+(defun pi-coding-agent--at-you-heading-p ()
+  "Return non-nil if current line is a You setext heading.
+Checks that the current line matches `pi-coding-agent--you-heading-re'
+and the next line is a setext underline (three or more `=' characters)."
+  (and (save-excursion
+         (beginning-of-line)
+         (looking-at pi-coding-agent--you-heading-re))
+       (save-excursion
+         (forward-line 1)
+         (looking-at "^=\\{3,\\}$"))))
+
+;;;; Turn Detection
+
+(defun pi-coding-agent--collect-you-headings ()
+  "Return list of buffer positions of all You setext headings.
+Scans from `point-min', returns positions in chronological order."
+  (let (headings)
+    (save-excursion
+      (goto-char (point-min))
+      (while (re-search-forward pi-coding-agent--you-heading-re nil t)
+        (let ((pos (match-beginning 0)))
+          (save-excursion
+            (goto-char pos)
+            (when (pi-coding-agent--at-you-heading-p)
+              (push pos headings))))))
+    (nreverse headings)))
+
+(defun pi-coding-agent--user-turn-index-at-point (&optional headings)
+  "Return 0-based index of the user turn at or before point.
+HEADINGS is an optional pre-computed list from
+`pi-coding-agent--collect-you-headings'; when nil, the buffer is scanned.
+Returns nil if point is before the first You heading."
+  (let ((headings (or headings (pi-coding-agent--collect-you-headings)))
+        (limit (point))
+        (index 0)
+        (result nil))
+    (dolist (h headings)
+      (when (<= h limit)
+        (setq result index))
+      (setq index (1+ index)))
+    result))
+
+;;;; Chat Navigation
+
+(defun pi-coding-agent--find-you-heading (search-fn)
+  "Find the next You setext heading using SEARCH-FN.
+SEARCH-FN is `re-search-forward' or `re-search-backward'.
+Returns the position of the heading line start, or nil if not found."
+  (save-excursion
+    (let ((found nil))
+      (while (and (not found)
+                  (funcall search-fn pi-coding-agent--you-heading-re nil t))
+        (let ((candidate (match-beginning 0)))
+          (save-excursion
+            (goto-char candidate)
+            (when (pi-coding-agent--at-you-heading-p)
+              (setq found candidate)))))
+      found)))
+
 (defun pi-coding-agent-next-message ()
   "Move to the next user message in the chat buffer."
   (interactive)
   (let ((pos (save-excursion
                (forward-line 1)
-               (re-search-forward "^You:" nil t))))
+               (pi-coding-agent--find-you-heading #'re-search-forward))))
     (if pos
         (progn
           (goto-char pos)
-          (beginning-of-line))
+          (when (get-buffer-window) (recenter 0)))
       (message "No more messages"))))
 
 (defun pi-coding-agent-previous-message ()
   "Move to the previous user message in the chat buffer."
   (interactive)
   (let ((pos (save-excursion
                (beginning-of-line)
-               (re-search-backward "^You:" nil t))))
+               (pi-coding-agent--find-you-heading #'re-search-backward))))
     (if pos
-        (goto-char pos)
+        (progn
+          (goto-char pos)
+          (when (get-buffer-window) (recenter 0)))
       (message "No previous message"))))
 
 (defconst pi-coding-agent--blockquote-wrap-prefix