From 2b4deeb4199b2a3f6414811792dfa1d832ff1260 Mon Sep 17 00:00:00 2001
From: 2witstudios <2witstudios@gmail.com>
Date: Sat, 7 Mar 2026 21:57:57 -0600
Subject: [PATCH] fix: Store session handle before manifest write in
 handle_resume

handle_spawn correctly stores the session handle BEFORE writing the
manifest because ManifestWatcher in Swift fires on manifest write and
immediately tries to attach. handle_resume had this ordering reversed,
creating a race where the Swift app would try to attach before the
session existed in the map. Also fixes the rollback path to remove
the session from the map on manifest update failure, matching the
pattern in handle_spawn.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/pu-engine/src/engine.rs | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/crates/pu-engine/src/engine.rs b/crates/pu-engine/src/engine.rs
index 33fdcf3..e89e54c 100644
--- a/crates/pu-engine/src/engine.rs
+++ b/crates/pu-engine/src/engine.rs
@@ -1404,6 +1404,14 @@ impl Engine {
 
         let pid = handle.pid;
 
+        // Store handle in session map BEFORE writing manifest.
+        // ManifestWatcher in Swift fires on manifest write and immediately
+        // tries to attach — the session must already be in the map.
+        self.sessions
+            .lock()
+            .await
+            .insert(agent_id.to_string(), handle);
+
         // 6. Update manifest: Suspended → Running, new PID
         let aid = agent_id.to_string();
         let sid = session_id.clone();
@@ -1427,23 +1435,19 @@ impl Engine {
         .unwrap_or_else(|e| Err(PuError::Io(std::io::Error::other(e))));
 
         if let Err(e) = manifest_result {
-            // Rollback: kill the resumed process
-            self.pty_host
-                .kill(&handle, Duration::from_secs(2))
-                .await
-                .ok();
+            // Rollback: remove session and kill process
+            if let Some(handle) = self.sessions.lock().await.remove(agent_id) {
+                self.pty_host
+                    .kill(&handle, Duration::from_secs(2))
+                    .await
+                    .ok();
+            }
             return Response::Error {
                 code: "RESUME_FAILED".into(),
                 message: format!("failed to update manifest: {e}"),
             };
         }
 
-        // 7. Store handle in session map
-        self.sessions
-            .lock()
-            .await
-            .insert(agent_id.to_string(), handle);
-
         self.notify_status_change(project_root).await;
 
         Response::ResumeResult {