Fix branching to restore from checkpoint

1 files changed, 71 insertions, 16 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