diff options
Diffstat (limited to 'makima/src/server/handlers')
| -rw-r--r-- | makima/src/server/handlers/mesh.rs | 29 | ||||
| -rw-r--r-- | makima/src/server/handlers/mesh_daemon.rs | 117 | ||||
| -rw-r--r-- | makima/src/server/handlers/mesh_supervisor.rs | 39 |
3 files changed, 178 insertions, 7 deletions
diff --git a/makima/src/server/handlers/mesh.rs b/makima/src/server/handlers/mesh.rs index 240e1f7..3d05f35 100644 --- a/makima/src/server/handlers/mesh.rs +++ b/makima/src/server/handlers/mesh.rs @@ -6,6 +6,7 @@ use axum::{ response::IntoResponse, Json, }; +use base64::Engine; use uuid::Uuid; use crate::db::models::{ @@ -2265,6 +2266,30 @@ pub async fn reassign_task( } }; + // Fetch latest checkpoint patch for worktree recovery during reassignment + let (patch_data, patch_base_sha) = match repository::get_latest_checkpoint_patch(pool, id).await { + Ok(Some(patch)) => { + tracing::info!( + old_task_id = %id, + new_task_id = %new_task.id, + patch_size = patch.patch_size_bytes, + base_sha = %patch.base_commit_sha, + files_count = patch.files_count, + "Including checkpoint patch for task reassignment recovery" + ); + let encoded = base64::engine::general_purpose::STANDARD.encode(&patch.patch_data); + (Some(encoded), Some(patch.base_commit_sha)) + } + Ok(None) => { + tracing::debug!(old_task_id = %id, "No checkpoint patch found for reassignment"); + (None, None) + } + Err(e) => { + tracing::warn!(old_task_id = %id, error = %e, "Failed to fetch checkpoint patch for reassignment"); + (None, None) + } + }; + // Send SpawnTask command to daemon for the new task let command = DaemonCommand::SpawnTask { task_id: new_task.id, @@ -2285,8 +2310,8 @@ pub async fn reassign_task( autonomous_loop: false, resume_session: false, conversation_history: None, - patch_data: None, - patch_base_sha: None, + patch_data, + patch_base_sha, }; tracing::info!( diff --git a/makima/src/server/handlers/mesh_daemon.rs b/makima/src/server/handlers/mesh_daemon.rs index 65db373..53ee806 100644 --- a/makima/src/server/handlers/mesh_daemon.rs +++ b/makima/src/server/handlers/mesh_daemon.rs @@ -291,6 +291,19 @@ pub enum DaemonMessage { success: bool, error: Option<String>, }, + /// Task recovery detected after daemon restart + TaskRecoveryDetected { + #[serde(rename = "taskId")] + task_id: Uuid, + #[serde(rename = "previousState")] + previous_state: String, + #[serde(rename = "worktreeIntact")] + worktree_intact: bool, + #[serde(rename = "worktreePath")] + worktree_path: Option<String>, + #[serde(rename = "needsPatch")] + needs_patch: bool, + }, /// Register a tool key for orchestrator API access RegisterToolKey { #[serde(rename = "taskId")] @@ -990,6 +1003,110 @@ async fn handle_daemon_connection(socket: WebSocket, state: SharedState, auth_re }); } } + Ok(DaemonMessage::TaskRecoveryDetected { + task_id, + previous_state, + worktree_intact, + worktree_path, + needs_patch, + }) => { + tracing::info!( + task_id = %task_id, + previous_state = %previous_state, + worktree_intact = worktree_intact, + worktree_path = ?worktree_path, + needs_patch = needs_patch, + "Task recovery detected after daemon restart" + ); + + // Update task in database based on recovery state + if let Some(ref pool) = state.db_pool { + let pool = pool.clone(); + let state = state.clone(); + tokio::spawn(async move { + if worktree_intact { + // Worktree exists - task can be resumed on this daemon + // Update task status to 'pending' so it can be picked up + match sqlx::query( + r#" + UPDATE tasks + SET status = 'pending', + daemon_id = NULL, + error_message = 'Daemon restarted - task ready for resumption', + interrupted_at = NOW(), + updated_at = NOW() + WHERE id = $1 AND owner_id = $2 + RETURNING id + "#, + ) + .bind(task_id) + .bind(owner_id) + .fetch_optional(&pool) + .await + { + Ok(Some(_)) => { + tracing::info!( + task_id = %task_id, + "Task marked as pending for resumption" + ); + state.broadcast_task_update(TaskUpdateNotification { + task_id, + owner_id: Some(owner_id), + version: 0, + status: "pending".into(), + updated_fields: vec![ + "status".into(), + "daemon_id".into(), + "interrupted_at".into(), + ], + updated_by: "daemon_recovery".into(), + }); + } + Ok(None) => { + tracing::warn!( + task_id = %task_id, + "Task not found during recovery update" + ); + } + Err(e) => { + tracing::error!( + task_id = %task_id, + error = %e, + "Failed to update task during recovery" + ); + } + } + } else { + // Worktree missing - mark for retry with patch restoration + match repository::mark_task_for_retry( + &pool, + task_id, + daemon_uuid, // Mark this daemon as failed + ).await { + Ok(Some(_)) => { + tracing::info!( + task_id = %task_id, + "Task marked for retry (worktree missing)" + ); + } + Ok(None) => { + tracing::warn!( + task_id = %task_id, + "Task not found or exceeded retries" + ); + } + Err(e) => { + tracing::error!( + task_id = %task_id, + error = %e, + "Failed to mark task for retry" + ); + } + } + } + }); + } + } Ok(DaemonMessage::Authenticate { .. }) => { // Already authenticated, ignore } diff --git a/makima/src/server/handlers/mesh_supervisor.rs b/makima/src/server/handlers/mesh_supervisor.rs index 21c9515..1b5e376 100644 --- a/makima/src/server/handlers/mesh_supervisor.rs +++ b/makima/src/server/handlers/mesh_supervisor.rs @@ -279,8 +279,9 @@ async fn verify_supervisor_auth( /// Try to start a pending task on an available daemon. /// Returns Ok(Some(task)) if a task was started, Ok(None) if no tasks could be started. -/// For retried tasks, excludes daemons that previously failed the task. -async fn try_start_pending_task( +/// For retried tasks, excludes daemons that previously failed the task and includes +/// checkpoint patch data for worktree recovery. +pub async fn try_start_pending_task( state: &SharedState, contract_id: Uuid, owner_id: Uuid, @@ -348,6 +349,34 @@ async fn try_start_pending_task( } }; + // For retried tasks, fetch checkpoint patch for worktree recovery + let (patch_data, patch_base_sha) = if task.retry_count > 0 { + // This is a retry - try to restore from checkpoint + match repository::get_latest_checkpoint_patch(pool, task.id).await { + Ok(Some(patch)) => { + tracing::info!( + task_id = %task.id, + retry_count = task.retry_count, + patch_size = patch.patch_size_bytes, + base_sha = %patch.base_commit_sha, + "Including checkpoint patch for task retry recovery" + ); + let encoded = base64::engine::general_purpose::STANDARD.encode(&patch.patch_data); + (Some(encoded), Some(patch.base_commit_sha)) + } + Ok(None) => { + tracing::debug!(task_id = %task.id, "No checkpoint patch found for retry"); + (None, None) + } + Err(e) => { + tracing::warn!(task_id = %task.id, error = %e, "Failed to fetch checkpoint patch for retry"); + (None, None) + } + } + } else { + (None, None) + }; + // Send spawn command let cmd = DaemonCommand::SpawnTask { task_id: updated_task.id, @@ -366,10 +395,10 @@ async fn try_start_pending_task( contract_id: updated_task.contract_id, is_supervisor: false, autonomous_loop: false, - resume_session: false, + resume_session: task.retry_count > 0, // Use --continue for retried tasks conversation_history: None, - patch_data: None, - patch_base_sha: None, + patch_data, + patch_base_sha, }; if let Err(e) = state.send_daemon_command(daemon.id, cmd).await { |