From 893db7590206aa7afad622858417a74dabfc3e70 Mon Sep 17 00:00:00 2001
From: Daniel Nouri <daniel.nouri@gmail.com>
Date: Fri, 13 Feb 2026 22:05:01 +0100
Subject: [PATCH] feat: fork from point in chat buffer (f key)

Press f on any user turn in the chat buffer to fork the conversation
from that point.  Uses the get_tree RPC to map visible buffer position
to entry ID, handling compaction (compacted-away turns aren't rendered,
so the last-N visible headings are aligned to the active branch).

Also fixes chat navigation (n/p): the old pattern searched for the
never-present "You:" instead of setext headings, and now recenters
the heading to the top of the window.

Changes:
- Setext heading detection: regex + underline predicate, shared across
  navigation and fork-at-point
- Turn detection: collect headings, map point to 0-based turn index
- Tree helpers: flatten-tree (hash index) and active-branch-user-ids
  (leaf-to-root walk)
- execute-fork: extracted shared fork logic (RPC + state refresh +
  history reload + input prefill), used by both fork-at-point and
  the menu fork selector
- resolve-fork-entry: tree-to-entryId mapping, independently testable
- build-tree test helper: flat specs replace deeply nested tree
  construction in tests

49 new tests covering heading detection, navigation, turn indexing,
tree walk scenarios, and fork-at-point integration.
---
 README.org                          |   7 +-
 pi-coding-agent-menu.el             | 142 +++++++++++++----
 pi-coding-agent-ui.el               |  81 +++++++++-
 test/pi-coding-agent-menu-test.el   | 229 +++++++++++++++++++++++++++-
 test/pi-coding-agent-test-common.el |  68 +++++++++
 test/pi-coding-agent-ui-test.el     | 194 +++++++++++++++++++++++
 6 files changed, 686 insertions(+), 35 deletions(-)

diff --git a/README.org b/README.org
index a55c620..fb11290 100644
--- a/README.org
+++ b/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.
diff --git a/pi-coding-agent-menu.el b/pi-coding-agent-menu.el
index ee91994..98ae864 100644
--- a/pi-coding-agent-menu.el
+++ b/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
 
diff --git a/pi-coding-agent-ui.el b/pi-coding-agent-ui.el
index ddc19de..1f9e71d 100644
--- a/pi-coding-agent-ui.el
+++ b/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,6 +334,7 @@ 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)
@@ -340,16 +342,85 @@ Returns \"text\" for unrecognized extensions to ensure consistent fencing."
     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 ()
@@ -357,9 +428,11 @@ Returns \"text\" for unrecognized extensions to ensure consistent fencing."
   (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
diff --git a/test/pi-coding-agent-menu-test.el b/test/pi-coding-agent-menu-test.el
index 9169f7c..2bd5ff1 100644
--- a/test/pi-coding-agent-menu-test.el
+++ b/test/pi-coding-agent-menu-test.el
@@ -320,12 +320,13 @@ Also verifies that the new session-file is stored in state for reload to work."
 ;;; Chat Navigation
 
 (ert-deftest pi-coding-agent-test-chat-has-navigation-keys ()
-  "Chat mode has n/p for navigation and TAB for folding."
+  "Chat mode has n/p for navigation, TAB for folding, f for fork."
   (with-temp-buffer
     (pi-coding-agent-chat-mode)
     (should (eq (key-binding "n") 'pi-coding-agent-next-message))
     (should (eq (key-binding "p") 'pi-coding-agent-previous-message))
-    (should (eq (key-binding (kbd "TAB")) 'pi-coding-agent-toggle-tool-section))))
+    (should (eq (key-binding (kbd "TAB")) 'pi-coding-agent-toggle-tool-section))
+    (should (eq (key-binding "f") 'pi-coding-agent-fork-at-point))))
 
 ;;; Reconnect Tests
 
@@ -633,5 +634,229 @@ Pi v0.51.3+ renamed SlashCommandSource from \"template\" to \"prompt\"."
     (should (pi-coding-agent-test--suffix-key-bound-p "A"))
     (should (pi-coding-agent-test--suffix-key-bound-p "B"))))
 
+;;; Fork at Point
+
+(ert-deftest pi-coding-agent-test-fork-at-point-correct-entry-id ()
+  "Fork-at-point on 2nd You heading forks with the correct entryId."
+  (with-temp-buffer
+    (pi-coding-agent-chat-mode)
+    (let ((pi-coding-agent--status 'idle)
+          (pi-coding-agent--process 'mock-proc)
+          (forked-entry-id nil)
+          (tree-data (pi-coding-agent-test--make-3turn-tree)))
+      (let ((inhibit-read-only t))
+        (pi-coding-agent-test--insert-chat-turns))
+      ;; Navigate to 2nd You heading
+      (goto-char (point-min))
+      (pi-coding-agent-next-message)
+      (pi-coding-agent-next-message)
+      (should (looking-at "You · 10:05"))
+      (cl-letf (((symbol-function 'pi-coding-agent--rpc-async)
+                 (lambda (_proc cmd cb)
+                   (cond
+                    ((equal (plist-get cmd :type) "get_tree")
+                     (funcall cb (list :success t :data tree-data)))
+                    ((equal (plist-get cmd :type) "fork")
+                     (setq forked-entry-id (plist-get cmd :entryId))
+                     (funcall cb '(:success t :data (:text "Second question"))))
+                    ((equal (plist-get cmd :type) "get_state")
+                     (funcall cb '(:success t :data (:sessionFile "/tmp/forked.jsonl"))))
+                    ((equal (plist-get cmd :type) "get_messages")
+                     (funcall cb '(:success t :data (:messages [])))))))
+                ((symbol-function 'y-or-n-p) (lambda (_prompt) t))
+                ((symbol-function 'pi-coding-agent--refresh-header) #'ignore))
+        (pi-coding-agent-fork-at-point))
+      (should (equal forked-entry-id "u2")))))
+
+(ert-deftest pi-coding-agent-test-fork-at-point-confirmation-declined ()
+  "Fork-at-point does nothing when user declines confirmation."
+  (with-temp-buffer
+    (pi-coding-agent-chat-mode)
+    (let ((pi-coding-agent--status 'idle)
+          (pi-coding-agent--process 'mock-proc)
+          (fork-called nil)
+          (tree-data (pi-coding-agent-test--make-3turn-tree)))
+      (let ((inhibit-read-only t))
+        (pi-coding-agent-test--insert-chat-turns))
+      (goto-char (point-min))
+      (pi-coding-agent-next-message)
+      (pi-coding-agent-next-message)
+      (cl-letf (((symbol-function 'pi-coding-agent--rpc-async)
+                 (lambda (_proc cmd cb)
+                   (cond
+                    ((equal (plist-get cmd :type) "get_tree")
+                     (funcall cb (list :success t :data tree-data)))
+                    ((equal (plist-get cmd :type) "fork")
+                     (setq fork-called t)))))
+                ((symbol-function 'y-or-n-p) (lambda (_prompt) nil)))
+        (pi-coding-agent-fork-at-point))
+      (should-not fork-called))))
+
+(ert-deftest pi-coding-agent-test-fork-at-point-no-user-turn ()
+  "Fork-at-point before first You heading shows message, no RPC calls."
+  (with-temp-buffer
+    (pi-coding-agent-chat-mode)
+    (let ((pi-coding-agent--status 'idle)
+          (pi-coding-agent--process 'mock-proc)
+          (rpc-called nil))
+      (let ((inhibit-read-only t))
+        (pi-coding-agent-test--insert-chat-turns))
+      (goto-char (point-min))
+      (cl-letf (((symbol-function 'pi-coding-agent--rpc-async)
+                 (lambda (&rest _) (setq rpc-called t))))
+        (pi-coding-agent-fork-at-point))
+      (should-not rpc-called))))
+
+(ert-deftest pi-coding-agent-test-fork-at-point-streaming-guard ()
+  "Fork-at-point during streaming shows message, no RPC calls."
+  (with-temp-buffer
+    (pi-coding-agent-chat-mode)
+    (let ((pi-coding-agent--status 'streaming)
+          (pi-coding-agent--process 'mock-proc)
+          (rpc-called nil))
+      (let ((inhibit-read-only t))
+        (pi-coding-agent-test--insert-chat-turns))
+      (goto-char (point-min))
+      (pi-coding-agent-next-message)
+      (cl-letf (((symbol-function 'pi-coding-agent--rpc-async)
+                 (lambda (&rest _) (setq rpc-called t))))
+        (pi-coding-agent-fork-at-point))
+      (should-not rpc-called))))
+
+(ert-deftest pi-coding-agent-test-fork-at-point-compaction ()
+  "Fork-at-point with compaction uses last-N to pick correct entry ID."
+  (with-temp-buffer
+    (pi-coding-agent-chat-mode)
+    (let ((pi-coding-agent--status 'idle)
+          (pi-coding-agent--process 'mock-proc)
+          (forked-entry-id nil)
+          (tree-data
+           (pi-coding-agent-test--build-tree
+            '("u1" nil "message" :role "user" :preview "Compacted away")
+            '("a1" nil "message" :role "assistant" :preview "Old answer")
+            '("c1" nil "compaction" :tokensBefore 5000)
+            '("u2" nil "message" :role "user" :preview "After compaction")
+            '("a2" nil "message" :role "assistant" :preview "Response")
+            '("u3" nil "message" :role "user" :preview "Latest")
+            '("a3" nil "message" :role "assistant" :preview "Final"))))
+      ;; Buffer has compaction summary + 2 visible user turns (not 3)
+      (let ((inhibit-read-only t))
+        (insert "Pi 1.0.0\n========\nWelcome\n\n"
+                "Compaction\n==========\nSummary of earlier conversation\n\n"
+                "You · 10:05\n===========\nAfter compaction\n\n"
+                "Assistant\n=========\nResponse\n\n"
+                "You · 10:10\n===========\nLatest\n\n"
+                "Assistant\n=========\nFinal\n"))
+      ;; Navigate to first visible You heading (ordinal 0)
+      (goto-char (point-min))
+      (pi-coding-agent-next-message)
+      (should (looking-at "You · 10:05"))
+      (cl-letf (((symbol-function 'pi-coding-agent--rpc-async)
+                 (lambda (_proc cmd cb)
+                   (cond
+                    ((equal (plist-get cmd :type) "get_tree")
+                     (funcall cb (list :success t :data tree-data)))
+                    ((equal (plist-get cmd :type) "fork")
+                     (setq forked-entry-id (plist-get cmd :entryId))
+                     (funcall cb '(:success t :data (:text "After compaction"))))
+                    ((equal (plist-get cmd :type) "get_state")
+                     (funcall cb '(:success t :data (:sessionFile "/tmp/forked.jsonl"))))
+                    ((equal (plist-get cmd :type) "get_messages")
+                     (funcall cb '(:success t :data (:messages [])))))))
+                ((symbol-function 'y-or-n-p) (lambda (_prompt) t))
+                ((symbol-function 'pi-coding-agent--refresh-header) #'ignore))
+        (pi-coding-agent-fork-at-point))
+      ;; Should fork from u2 (not u1 which was compacted away)
+      (should (equal forked-entry-id "u2")))))
+
+;;; Fork Entry Resolution
+
+(ert-deftest pi-coding-agent-test-resolve-fork-entry-maps-ordinal ()
+  "resolve-fork-entry maps ordinal to correct entry ID and preview."
+  (let* ((tree-data (pi-coding-agent-test--make-3turn-tree))
+         (response (list :success t :data tree-data))
+         (result (pi-coding-agent--resolve-fork-entry response 1 3)))
+    (should (equal (car result) "u2"))
+    (should (equal (cdr result) "Second question"))))
+
+(ert-deftest pi-coding-agent-test-resolve-fork-entry-compaction ()
+  "resolve-fork-entry with compaction uses last-N to skip compacted entries."
+  (let* ((tree-data (pi-coding-agent-test--make-3turn-tree))
+         (response (list :success t :data tree-data))
+         ;; 3 user IDs on path, but only 2 visible headings
+         (result (pi-coding-agent--resolve-fork-entry response 0 2)))
+    ;; Should pick u2 (skipping u1 which would be compacted away)
+    (should (equal (car result) "u2"))))
+
+(ert-deftest pi-coding-agent-test-resolve-fork-entry-failure ()
+  "resolve-fork-entry returns nil on RPC failure."
+  (let ((response '(:success nil :error "Network error")))
+    (should-not (pi-coding-agent--resolve-fork-entry response 0 3))))
+
+;;; Active Branch Tree Walk
+
+(ert-deftest pi-coding-agent-test-active-branch-linear ()
+  "Linear tree: u1 → a1 → u2 → a2 (leaf) returns both user IDs."
+  (let* ((data (pi-coding-agent-test--build-tree
+                '("u1" nil "message" :role "user" :preview "Hello")
+                '("a1" nil "message" :role "assistant" :preview "Hi")
+                '("u2" nil "message" :role "user" :preview "More")
+                '("a2" nil "message" :role "assistant" :preview "Sure")))
+         (index (pi-coding-agent--flatten-tree (plist-get data :tree)))
+         (ids (pi-coding-agent--active-branch-user-ids index "a2")))
+    (should (equal ids '("u1" "u2")))))
+
+(ert-deftest pi-coding-agent-test-active-branch-branched ()
+  "Branched tree: active branch u1 → a1 → u2 → a2, ignores u3 → a3."
+  (let* ((data (pi-coding-agent-test--build-tree
+                '("u1" nil "message" :role "user" :preview "Hello")
+                '("a1" nil "message" :role "assistant" :preview "Hi")
+                '("u2" nil "message" :role "user" :preview "Path A")
+                '("a2" nil "message" :role "assistant" :preview "Sure A")
+                '("u3" "a1" "message" :role "user" :preview "Path B")
+                '("a3" nil "message" :role "assistant" :preview "Sure B")))
+         (index (pi-coding-agent--flatten-tree (plist-get data :tree)))
+         (ids (pi-coding-agent--active-branch-user-ids index "a2")))
+    (should (equal ids '("u1" "u2")))))
+
+(ert-deftest pi-coding-agent-test-active-branch-with-compaction ()
+  "Tree with compaction node: u1 → a1 → compaction → u2 → a2."
+  (let* ((data (pi-coding-agent-test--build-tree
+                '("u1" nil "message" :role "user" :preview "First")
+                '("a1" nil "message" :role "assistant" :preview "Response")
+                '("c1" nil "compaction" :tokensBefore 5000)
+                '("u2" nil "message" :role "user" :preview "After compaction")
+                '("a2" nil "message" :role "assistant" :preview "Still here")))
+         (index (pi-coding-agent--flatten-tree (plist-get data :tree)))
+         (ids (pi-coding-agent--active-branch-user-ids index "a2")))
+    (should (equal ids '("u1" "u2")))))
+
+(ert-deftest pi-coding-agent-test-active-branch-with-metadata ()
+  "Tree with model_change and thinking nodes: only user IDs returned."
+  (let* ((data (pi-coding-agent-test--build-tree
+                '("u1" nil "message" :role "user" :preview "Hello")
+                '("a1" nil "message" :role "assistant" :preview "Hi")
+                '("m1" nil "model_change" :provider "anthropic" :modelId "claude-4")
+                '("t1" nil "thinking_level_change" :thinkingLevel "high")
+                '("u2" nil "message" :role "user" :preview "More")
+                '("a2" nil "message" :role "assistant" :preview "Sure")))
+         (index (pi-coding-agent--flatten-tree (plist-get data :tree)))
+         (ids (pi-coding-agent--active-branch-user-ids index "a2")))
+    (should (equal ids '("u1" "u2")))))
+
+(ert-deftest pi-coding-agent-test-active-branch-empty-tree ()
+  "Empty tree returns empty list."
+  (let* ((index (pi-coding-agent--flatten-tree []))
+         (ids (pi-coding-agent--active-branch-user-ids index nil)))
+    (should (equal ids nil))))
+
+(ert-deftest pi-coding-agent-test-active-branch-nil-leaf ()
+  "Nil leafId returns empty list."
+  (let* ((data (pi-coding-agent-test--build-tree
+                '("u1" nil "message" :role "user" :preview "Hello")))
+         (index (pi-coding-agent--flatten-tree (plist-get data :tree)))
+         (ids (pi-coding-agent--active-branch-user-ids index nil)))
+    (should (equal ids nil))))
+
 (provide 'pi-coding-agent-menu-test)
 ;;; pi-coding-agent-menu-test.el ends here
diff --git a/test/pi-coding-agent-test-common.el b/test/pi-coding-agent-test-common.el
index 53ce2cc..cb2160a 100644
--- a/test/pi-coding-agent-test-common.el
+++ b/test/pi-coding-agent-test-common.el
@@ -138,5 +138,73 @@ Buffers are killed after BODY completes, even on error."
        (ignore-errors (kill-buffer ,buf-a))
        (ignore-errors (kill-buffer ,buf-b)))))
 
+;;;; Tree Fixtures
+
+(defun pi-coding-agent-test--build-tree (&rest specs)
+  "Build a conversation tree from flat node SPECS.
+Each SPEC is (ID PARENT-OVERRIDE TYPE &rest PROPS) where:
+- ID is the node identifier string
+- PARENT-OVERRIDE is nil (auto-chain to previous node) or a parent ID
+- TYPE is \"message\", \"compaction\", \"model_change\", etc.
+- PROPS are keyword plist properties (:role, :preview, etc.)
+First node with nil PARENT-OVERRIDE becomes the root.
+Returns (:tree VECTOR :leafId LAST-ID)."
+  (let ((nodes (make-hash-table :test 'equal))
+        (child-ids (make-hash-table :test 'equal))
+        (roots nil)
+        (prev-id nil)
+        (last-id nil))
+    ;; Pass 1: create nodes, track parent-child relationships
+    (dolist (spec specs)
+      (let* ((id (nth 0 spec))
+             (parent-override (nth 1 spec))
+             (type (nth 2 spec))
+             (props (nthcdr 3 spec))
+             (parent-id (or parent-override prev-id))
+             (node (append (list :id id :type type)
+                           (when parent-id (list :parentId parent-id))
+                           props)))
+        (puthash id node nodes)
+        (if parent-id
+            (puthash parent-id
+                     (append (gethash parent-id child-ids) (list id))
+                     child-ids)
+          (push id roots))
+        (setq prev-id id
+              last-id id)))
+    ;; Pass 2: build nested structure with :children vectors
+    (cl-labels ((build (id)
+                  (let* ((node (gethash id nodes))
+                         (kids (gethash id child-ids))
+                         (child-vec (if kids
+                                        (apply #'vector (mapcar #'build kids))
+                                      [])))
+                    (append node (list :children child-vec)))))
+      (list :tree (apply #'vector (mapcar #'build (nreverse roots)))
+            :leafId last-id))))
+
+(defun pi-coding-agent-test--make-3turn-tree ()
+  "Return tree data for a 3-turn conversation: u1→a1→u2→a2→u3→a3."
+  (pi-coding-agent-test--build-tree
+   '("u1" nil "message" :role "user" :preview "First question")
+   '("a1" nil "message" :role "assistant" :preview "First answer")
+   '("u2" nil "message" :role "user" :preview "Second question")
+   '("a2" nil "message" :role "assistant" :preview "Second answer")
+   '("u3" nil "message" :role "user" :preview "Third question")
+   '("a3" nil "message" :role "assistant" :preview "Third answer")))
+
+;;;; Chat Buffer Fixtures
+
+(defun pi-coding-agent-test--insert-chat-turns ()
+  "Insert a 3-turn chat with setext headings into current buffer.
+Returns the buffer with content ready for navigation tests."
+  (insert "Pi 1.0.0\n========\nWelcome\n\n"
+          "You · 10:00\n===========\nFirst question\n\n"
+          "Assistant\n=========\nFirst answer\n\n"
+          "You · 10:05\n===========\nSecond question\n\n"
+          "Assistant\n=========\nSecond answer\n\n"
+          "You · 10:10\n===========\nThird question\n\n"
+          "Assistant\n=========\nThird answer\n"))
+
 (provide 'pi-coding-agent-test-common)
 ;;; pi-coding-agent-test-common.el ends here
diff --git a/test/pi-coding-agent-ui-test.el b/test/pi-coding-agent-ui-test.el
index 84e6ddc..8cf5fd1 100644
--- a/test/pi-coding-agent-ui-test.el
+++ b/test/pi-coding-agent-ui-test.el
@@ -261,5 +261,199 @@ Buffer is read-only with `inhibit-read-only' used for insertion.
       (kill-ring-save (point-min) (point-max))
       (should (equal (car kill-ring) "Hello **bold** world")))))
 
+;;; Chat Navigation Behavior
+
+(ert-deftest pi-coding-agent-test-next-message-from-top ()
+  "n from point-min reaches first You heading."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:00"))))
+
+(ert-deftest pi-coding-agent-test-next-message-successive ()
+  "Successive n reaches each You heading in order."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:00"))
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:05"))
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:10"))))
+
+(ert-deftest pi-coding-agent-test-next-message-at-last ()
+  "n at last You heading keeps point and shows message."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (pi-coding-agent-next-message)
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:10"))
+    (let ((pos (point)))
+      (pi-coding-agent-next-message)
+      ;; Point stays on the last heading
+      (should (= (point) pos)))))
+
+(ert-deftest pi-coding-agent-test-previous-message-from-last ()
+  "p from last You heading reaches previous."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    ;; Navigate to last heading first
+    (pi-coding-agent-next-message)
+    (pi-coding-agent-next-message)
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:10"))
+    (pi-coding-agent-previous-message)
+    (should (looking-at "You · 10:05"))))
+
+(ert-deftest pi-coding-agent-test-previous-message-at-first ()
+  "p at first You heading keeps point."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (should (looking-at "You · 10:00"))
+    (let ((pos (point)))
+      (pi-coding-agent-previous-message)
+      ;; Point stays on the first heading
+      (should (= (point) pos)))))
+
+;;; Turn Detection
+
+(ert-deftest pi-coding-agent-test-turn-index-on-first-heading ()
+  "Turn index is 0 when point is on first You heading."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+(ert-deftest pi-coding-agent-test-turn-index-in-first-body ()
+  "Turn index is 0 when point is in first user message body."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (forward-line 2) ; skip heading + underline into body
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+(ert-deftest pi-coding-agent-test-turn-index-on-underline ()
+  "Turn index is 0 when point is on === underline of first You."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (forward-line 1) ; on ===
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+(ert-deftest pi-coding-agent-test-turn-index-on-second-heading ()
+  "Turn index is 1 on second You heading."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (pi-coding-agent-next-message)
+    (should (= (pi-coding-agent--user-turn-index-at-point) 1))))
+
+(ert-deftest pi-coding-agent-test-turn-index-on-assistant-heading ()
+  "Turn index is index of preceding You when point is on Assistant heading."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    ;; Navigate to first You, then move into assistant section
+    (pi-coding-agent-next-message)
+    (forward-line 4) ; past heading + underline + body + blank → "Assistant"
+    (should (looking-at "Assistant"))
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+(ert-deftest pi-coding-agent-test-turn-index-in-assistant-body ()
+  "Turn index is index of preceding You when point is in assistant response."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (pi-coding-agent-next-message)
+    (forward-line 6) ; heading + underline + body + blank + Assistant + underline → response
+    (should (looking-at "First answer"))
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+(ert-deftest pi-coding-agent-test-turn-index-before-first-you ()
+  "Turn index is nil before first You heading."
+  (with-temp-buffer
+    (pi-coding-agent-test--insert-chat-turns)
+    (goto-char (point-min))
+    (should-not (pi-coding-agent--user-turn-index-at-point))))
+
+(ert-deftest pi-coding-agent-test-turn-index-empty-buffer ()
+  "Turn index is nil in empty buffer."
+  (with-temp-buffer
+    (should-not (pi-coding-agent--user-turn-index-at-point))))
+
+(ert-deftest pi-coding-agent-test-turn-index-no-false-match ()
+  "Turn index ignores text starting with You without setext underline."
+  (with-temp-buffer
+    (insert "You mentioned something\nRegular text\n\n"
+            "You · 10:00\n===========\nFirst question\n")
+    (goto-char (point-min))
+    ;; Point is on "You mentioned" which has no === underline
+    (should-not (pi-coding-agent--user-turn-index-at-point))
+    ;; Move to the real heading
+    (goto-char (point-max))
+    (should (= (pi-coding-agent--user-turn-index-at-point) 0))))
+
+;;; You Heading Detection
+
+(ert-deftest pi-coding-agent-test-heading-re-matches-plain-you ()
+  "Heading regex matches bare `You' at start of line."
+  (should (string-match-p pi-coding-agent--you-heading-re "You")))
+
+(ert-deftest pi-coding-agent-test-heading-re-matches-you-with-timestamp ()
+  "Heading regex matches `You · 22:10' at start of line."
+  (should (string-match-p pi-coding-agent--you-heading-re "You · 22:10")))
+
+(ert-deftest pi-coding-agent-test-heading-re-rejects-you-colon ()
+  "Heading regex does not match `You:' (old broken pattern)."
+  (should-not (string-match-p pi-coding-agent--you-heading-re "You: hello")))
+
+(ert-deftest pi-coding-agent-test-heading-re-rejects-mid-line ()
+  "Heading regex does not match `You' mid-line."
+  (should-not (string-match-p pi-coding-agent--you-heading-re "  You · 22:10")))
+
+(ert-deftest pi-coding-agent-test-heading-re-rejects-you-prefix ()
+  "Heading regex does not match words starting with You like `Your'."
+  (should-not (string-match-p pi-coding-agent--you-heading-re "Your code is fine")))
+
+(ert-deftest pi-coding-agent-test-at-you-heading-p-true ()
+  "Predicate returns t when on a valid You setext heading."
+  (with-temp-buffer
+    (insert "You · 22:10\n===========\n")
+    (goto-char (point-min))
+    (should (pi-coding-agent--at-you-heading-p))))
+
+(ert-deftest pi-coding-agent-test-at-you-heading-p-no-underline ()
+  "Predicate returns nil when You line lacks setext underline."
+  (with-temp-buffer
+    (insert "You · 22:10\nSome text\n")
+    (goto-char (point-min))
+    (should-not (pi-coding-agent--at-you-heading-p))))
+
+(ert-deftest pi-coding-agent-test-at-you-heading-p-short-underline ()
+  "Predicate returns t with minimum 3-char underline."
+  (with-temp-buffer
+    (insert "You\n===\n")
+    (goto-char (point-min))
+    (should (pi-coding-agent--at-you-heading-p))))
+
+(ert-deftest pi-coding-agent-test-at-you-heading-p-wrong-line ()
+  "Predicate returns nil when not on the heading line."
+  (with-temp-buffer
+    (insert "You · 22:10\n===========\nBody text\n")
+    (goto-char (point-max))
+    (forward-line -1)  ; on "Body text"
+    (should-not (pi-coding-agent--at-you-heading-p))))
+
 (provide 'pi-coding-agent-ui-test)
 ;;; pi-coding-agent-ui-test.el ends here