Fix branching to restore from checkpoint

2 files changed, 97 insertions, 18 deletions

diff --git a/makima/src/daemon/task/manager.rs b/makima/src/daemon/task/manager.rs
index 74a37bf..86d7e05 100644
--- a/makima/src/daemon/task/manager.rs
+++ b/makima/src/daemon/task/manager.rs
@@ -3734,24 +3734,79 @@ impl TaskManagerInner {
                 // Create worktree - either from scratch or copying from another task
                 let task_name = format!("task-{}", &task_id.to_string()[..8]);
                 let worktree_info = if let Some(from_task_id) = continue_from_task_id {
-                    // Find the source task's worktree path
-                    let source_worktree = self.find_worktree_for_task(from_task_id).await
-                        .map_err(|e| DaemonError::Task(TaskError::SetupFailed(
-                            format!("Cannot continue from task {}: {}", from_task_id, e)
-                        )))?;
+                    // Try to find the source task's worktree path
+                    match self.find_worktree_for_task(from_task_id).await {
+                        Ok(source_worktree) => {
+                            let msg = DaemonMessage::task_output(
+                                task_id,
+                                format!("Continuing from task {} worktree...\n", &from_task_id.to_string()[..8]),
+                                false,
+                            );
+                            let _ = self.ws_tx.send(msg).await;
 
-                    let msg = DaemonMessage::task_output(
-                        task_id,
-                        format!("Continuing from task {} worktree...\n", &from_task_id.to_string()[..8]),
-                        false,
-                    );
-                    let _ = self.ws_tx.send(msg).await;
+                            // Create worktree by copying from source task
+                            self.worktree_manager
+                                .create_worktree_from_task(&source_worktree, task_id, &task_name)
+                                .await
+                                .map_err(|e| DaemonError::Task(TaskError::SetupFailed(e.to_string())))?
+                        }
+                        Err(worktree_err) => {
+                            // Source worktree not found - try to recover using patch data
+                            if let (Some(patch_str), Some(base_sha)) = (&patch_data, &patch_base_sha) {
+                                tracing::info!(
+                                    task_id = %task_id,
+                                    from_task_id = %from_task_id,
+                                    base_sha = %base_sha,
+                                    patch_len = patch_str.len(),
+                                    "Source worktree not found, attempting to restore from patch"
+                                );
 
-                    // Create worktree by copying from source task
-                    self.worktree_manager
-                        .create_worktree_from_task(&source_worktree, task_id, &task_name)
-                        .await
-                        .map_err(|e| DaemonError::Task(TaskError::SetupFailed(e.to_string())))?
+                                let msg = DaemonMessage::task_output(
+                                    task_id,
+                                    format!("Source worktree unavailable, restoring from checkpoint patch...\n"),
+                                    false,
+                                );
+                                let _ = self.ws_tx.send(msg).await;
+
+                                // Decode base64 patch data
+                                match base64::Engine::decode(&base64::engine::general_purpose::STANDARD, patch_str) {
+                                    Ok(patch_bytes) => {
+                                        match self.worktree_manager.restore_from_patch(
+                                            source,
+                                            task_id,
+                                            &task_name,
+                                            base_sha,
+                                            &patch_bytes,
+                                        ).await {
+                                            Ok(worktree_info) => {
+                                                tracing::info!(
+                                                    task_id = %task_id,
+                                                    path = %worktree_info.path.display(),
+                                                    "Successfully restored worktree from patch"
+                                                );
+                                                worktree_info
+                                            }
+                                            Err(e) => {
+                                                return Err(DaemonError::Task(TaskError::SetupFailed(
+                                                    format!("Cannot continue from task {}: {} (patch restore also failed: {})", from_task_id, worktree_err, e)
+                                                )));
+                                            }
+                                        }
+                                    }
+                                    Err(e) => {
+                                        return Err(DaemonError::Task(TaskError::SetupFailed(
+                                            format!("Cannot continue from task {}: {} (patch decode failed: {})", from_task_id, worktree_err, e)
+                                        )));
+                                    }
+                                }
+                            } else {
+                                // No patch data available - fail with original error
+                                return Err(DaemonError::Task(TaskError::SetupFailed(
+                                    format!("Cannot continue from task {}: {}", from_task_id, worktree_err)
+                                )));
+                            }
+                        }
+                    }
                 } else {
                     // Create fresh worktree from repo
                     self.worktree_manager
diff --git a/makima/src/server/handlers/mesh.rs b/makima/src/server/handlers/mesh.rs
index 3d64eb4..545d1ea 100644
--- a/makima/src/server/handlers/mesh.rs
+++ b/makima/src/server/handlers/mesh.rs
@@ -3500,6 +3500,30 @@ pub async fn branch_task(
         }),
     ).await;
 
+    // Fetch latest checkpoint patch from source task for worktree recovery
+    let (patch_data, patch_base_sha) = match repository::get_latest_checkpoint_patch(pool, source_task_id).await {
+        Ok(Some(patch)) => {
+            tracing::info!(
+                source_task_id = %source_task_id,
+                new_task_id = %task.id,
+                patch_size = patch.patch_size_bytes,
+                base_sha = %patch.base_commit_sha,
+                files_count = patch.files_count,
+                "Including checkpoint patch for task branching recovery"
+            );
+            let encoded = base64::engine::general_purpose::STANDARD.encode(&patch.patch_data);
+            (Some(encoded), Some(patch.base_commit_sha))
+        }
+        Ok(None) => {
+            tracing::debug!(source_task_id = %source_task_id, "No checkpoint patch found for branching");
+            (None, None)
+        }
+        Err(e) => {
+            tracing::warn!(source_task_id = %source_task_id, error = %e, "Failed to fetch checkpoint patch for branching");
+            (None, None)
+        }
+    };
+
     // Try to find an available daemon to start the task
     let daemon_id = state.daemon_connections
         .iter()
@@ -3536,8 +3560,8 @@ pub async fn branch_task(
                 autonomous_loop: false,
                 resume_session: message_count > 0, // Resume if we have conversation history
                 conversation_history: updated_task.conversation_state.clone(),
-                patch_data: None,
-                patch_base_sha: None,
+                patch_data,
+                patch_base_sha,
             };
 
             if let Err(e) = state.send_daemon_command(target_daemon_id, command).await {