fix: make fork-at-point reliable in deep/large sessions

pi-coding-agent-core.el

@@ -84,14 +84,24 @@ Each process has its own table stored as a process property."
         (process-put process 'pi-coding-agent-pending-requests table)
         table)))
 
+(defun pi-coding-agent--get-pending-command-types (process)
+  "Get or create pending command type table for PROCESS.
+Maps request IDs to command type strings."
+  (or (process-get process 'pi-coding-agent-pending-command-types)
+      (let ((table (make-hash-table :test 'equal)))
+        (process-put process 'pi-coding-agent-pending-command-types table)
+        table)))
+
 (defun pi-coding-agent--rpc-async (process command callback)
   "Send COMMAND to pi PROCESS asynchronously.
 COMMAND is a plist that will be augmented with a unique ID.
 CALLBACK is called with the response plist when received."
   (let* ((id (pi-coding-agent--next-request-id))
          (full-command (plist-put (copy-sequence command) :id id))
-         (pending (pi-coding-agent--get-pending-requests process)))
+         (pending (pi-coding-agent--get-pending-requests process))
+         (pending-types (pi-coding-agent--get-pending-command-types process)))
     (puthash id callback pending)
+    (puthash id (plist-get command :type) pending-types)
     (process-send-string process (pi-coding-agent--encode-command full-command))))
 
 (defun pi-coding-agent--send-extension-ui-response (process response)
@@ -138,15 +148,45 @@ Uses process property for per-process partial output storage."
 
 (defun pi-coding-agent--dispatch-response (proc json)
   "Dispatch JSON response from PROC to callback or event handler.
-If JSON is a response with a known ID, call the stored callback.
-Otherwise, treat it as an event and call the process's handler."
+Response routing order: explicit ID, id-less `:command' match, then
+id-less sole pending request. Non-response JSON is treated as an event."
   (let ((type (plist-get json :type))
         (id (plist-get json :id)))
-    (if (and (equal type "response") id)
-        (let ((pending (pi-coding-agent--get-pending-requests proc)))
-          (when-let ((callback (gethash id pending)))
-            (remhash id pending)
-            (funcall callback json)))
+    (if (equal type "response")
+        (let* ((pending (pi-coding-agent--get-pending-requests proc))
+               (pending-types (pi-coding-agent--get-pending-command-types proc))
+               (dispatch-response
+                (lambda (request-id callback)
+                  (remhash request-id pending)
+                  (remhash request-id pending-types)
+                  (funcall callback json))))
+          (cond
+           ((and id (gethash id pending))
+            (funcall dispatch-response id (gethash id pending)))
+           ((null id)
+            (let ((matched-id nil)
+                  (matched-callback nil)
+                  (matched-count 0)
+                  (command (plist-get json :command)))
+              (when command
+                (maphash (lambda (request-id command-type)
+                           (when (equal command-type command)
+                             (setq matched-count (1+ matched-count))
+                             (when (= matched-count 1)
+                               (setq matched-id request-id
+                                     matched-callback (gethash request-id pending)))))
+                         pending-types))
+              (cond
+               ((and (= matched-count 1) matched-callback)
+                (funcall dispatch-response matched-id matched-callback))
+               ((= (hash-table-count pending) 1)
+                (let (only-id only-callback)
+                  (maphash (lambda (request-id callback)
+                             (setq only-id request-id
+                                   only-callback callback))
+                           pending)
+                  (when only-callback
+                    (funcall dispatch-response only-id only-callback)))))))))
       ;; Call only this process's handler, not all handlers
       (pi-coding-agent--handle-event proc json))))
 
@@ -160,16 +200,19 @@ Calls only the handler registered for this specific process."
 (defun pi-coding-agent--handle-process-exit (proc event)
   "Clean up when pi process PROC exits with EVENT.
 Calls pending request callbacks for this process with an error response
-containing EVENT, then clears this process's pending requests table."
+containing EVENT, then clears this process's pending request tables."
   (let ((pending (process-get proc 'pi-coding-agent-pending-requests))
+        (pending-types (process-get proc 'pi-coding-agent-pending-command-types))
         (error-response (list :type "response"
                               :success :false
                               :error (format "Process exited: %s" (string-trim event)))))
     (when pending
       (maphash (lambda (_id callback)
                  (funcall callback error-response))
                pending)
-      (clrhash pending))))
+      (clrhash pending))
+    (when pending-types
+      (clrhash pending-types))))
 
 (defvar pi-coding-agent-executable)  ; forward decl — core.el cannot require ui.el
 

pi-coding-agent-menu.el

@@ -594,15 +594,25 @@ Optional CUSTOM-INSTRUCTIONS provide guidance for the compaction summary."
 (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))
+Returns a hash table for O(1) lookup by id.
+
+Uses iterative traversal to avoid `max-lisp-eval-depth' errors on deep
+session trees."
+  (let ((index (make-hash-table :test 'equal))
+        (stack nil))
+    ;; Push roots in reverse so popping preserves original order.
+    (let ((i (1- (length nodes))))
+      (while (>= i 0)
+        (push (aref nodes i) stack)
+        (setq i (1- i))))
+    (while stack
+      (let* ((node (pop stack))
+             (children (plist-get node :children)))
+        (puthash (plist-get node :id) node index)
+        (let ((i (1- (length children))))
+          (while (>= i 0)
+            (push (aref children i) stack)
+            (setq i (1- i))))))
     index))
 
 (defun pi-coding-agent--active-branch-user-ids (index leaf-id)
@@ -650,23 +660,20 @@ Shows a selector of user messages and creates a fork from the selected one."
                        (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.
+  "Resolve a fork entry ID from get_fork_messages 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))))
+           (messages (append (plist-get data :messages) nil))
+           ;; Use last N messages to align with visible headings in
+           ;; compacted sessions.
+           (visible-messages (last messages heading-count))
+           (selected (nth ordinal visible-messages))
+           (entry-id (plist-get selected :entryId)))
       (when entry-id
-        (cons entry-id (plist-get node :preview))))))
+        (cons entry-id (pi-coding-agent--format-fork-message selected))))))
 
 (defun pi-coding-agent-fork-at-point ()
   "Fork conversation from the user turn at point.
@@ -689,17 +696,21 @@ a preview, then forks.  Only works when the session is idle."
                 (proc (pi-coding-agent--get-process)))
             (unless proc
               (user-error "Pi: No active process"))
-            (pi-coding-agent--rpc-async proc '(:type "get_tree")
+            (pi-coding-agent--rpc-async proc '(:type "get_fork_messages")
               (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)))))))))))))))
+                (if (not (plist-get response :success))
+                    (if-let ((error-text (plist-get response :error)))
+                        (message "Pi: Failed to get fork messages: %s" error-text)
+                      (message "Pi: Failed to get fork messages"))
+                  (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.

test/pi-coding-agent-core-test.el

@@ -126,15 +126,19 @@
 ;;;; Process Cleanup Tests
 
 (ert-deftest pi-coding-agent-test-process-exit-clears-pending ()
-  "Process exit clears that process's pending requests."
+  "Process exit clears pending request state."
   (let ((pi-coding-agent--request-id-counter 0)
         (fake-proc (start-process "cat" nil "cat")))
     (unwind-protect
-        (let ((pending (pi-coding-agent--get-pending-requests fake-proc)))
+        (let ((pending (pi-coding-agent--get-pending-requests fake-proc))
+              (pending-types (pi-coding-agent--get-pending-command-types fake-proc)))
           (puthash "req_1" #'ignore pending)
           (puthash "req_2" #'ignore pending)
+          (puthash "req_1" "get_tree" pending-types)
+          (puthash "req_2" "get_state" pending-types)
           (pi-coding-agent--handle-process-exit fake-proc "finished\n")
-          (should (= (hash-table-count pending) 0)))
+          (should (= (hash-table-count pending) 0))
+          (should (= (hash-table-count pending-types) 0)))
       (ignore-errors (delete-process fake-proc)))))
 
 (ert-deftest pi-coding-agent-test-process-exit-calls-callbacks-with-error ()
@@ -175,13 +179,52 @@
           (should (null (gethash "req_1" pending))))
       (delete-process fake-proc))))
 
+(ert-deftest pi-coding-agent-test-dispatch-idless-response-to-sole-pending ()
+  "Id-less response routes to sole pending callback."
+  (let ((received nil)
+        (fake-proc (start-process "cat" nil "cat")))
+    (unwind-protect
+        (let ((pending (pi-coding-agent--get-pending-requests fake-proc)))
+          (puthash "req_1" (lambda (r) (setq received r)) pending)
+          (pi-coding-agent--dispatch-response
+           fake-proc
+           '(:type "response" :command "get_tree" :success nil :error "Unknown command: get_tree"))
+          (should (equal (plist-get received :error) "Unknown command: get_tree"))
+          (should (= (hash-table-count pending) 0)))
+      (delete-process fake-proc))))
+
+(ert-deftest pi-coding-agent-test-dispatch-idless-response-matches-command ()
+  "Id-less response with :command routes to matching request."
+  (let ((pi-coding-agent--request-id-counter 0)
+        (received-tree nil)
+        (received-state nil)
+        (fake-proc (start-process "cat" nil "cat")))
+    (unwind-protect
+        (cl-letf (((symbol-function 'process-send-string) #'ignore))
+          (pi-coding-agent--rpc-async
+           fake-proc
+           '(:type "get_tree")
+           (lambda (response)
+             (setq received-tree response)))
+          (pi-coding-agent--rpc-async
+           fake-proc
+           '(:type "get_state")
+           (lambda (response)
+             (setq received-state response)))
+          (pi-coding-agent--dispatch-response
+           fake-proc
+           '(:type "response" :command "get_tree" :success nil :error "Unknown command: get_tree"))
+          (should (equal (plist-get received-tree :error) "Unknown command: get_tree"))
+          (should-not received-state)
+          (should (= (hash-table-count (pi-coding-agent--get-pending-requests fake-proc)) 1)))
+      (delete-process fake-proc))))
+
 (ert-deftest pi-coding-agent-test-dispatch-event-calls-handler ()
-  "Non-response messages call process's handler."
+  "Events call the process handler."
   (let ((event-received nil)
         (fake-proc (start-process "cat" nil "cat")))
     (unwind-protect
         (progn
-          ;; Register a handler on the process
           (process-put fake-proc 'pi-coding-agent-display-handler
                        (lambda (e) (setq event-received e)))
           (pi-coding-agent--dispatch-response fake-proc '(:type "agent_start"))
@@ -190,7 +233,7 @@
       (delete-process fake-proc))))
 
 (ert-deftest pi-coding-agent-test-dispatch-unknown-id-no-crash ()
-  "Response with unknown ID is handled gracefully."
+  "Unknown response IDs do not crash."
   (let ((fake-proc (start-process "cat" nil "cat")))
     (unwind-protect
         (should (null (pi-coding-agent--dispatch-response fake-proc '(:type "response" :id "unknown" :success t))))