1 files changed, 115 insertions, 1 deletions

diff --git a/makima/src/server/handlers/contracts.rs b/makima/src/server/handlers/contracts.rs
index 3ce29e1..09f78e6 100644
--- a/makima/src/server/handlers/contracts.rs
+++ b/makima/src/server/handlers/contracts.rs
@@ -425,7 +425,7 @@ pub async fn update_contract(
 
     match repository::update_contract_for_owner(pool, id, auth.owner_id, req).await {
         Ok(Some(contract)) => {
-            // If contract is completed, stop the supervisor task
+            // If contract is completed, stop the supervisor task and clean up worktrees
             if contract.status == "completed" {
                 if let Some(supervisor_task_id) = contract.supervisor_task_id {
                     // Get the supervisor task to find its daemon
@@ -456,6 +456,14 @@ pub async fn update_contract(
                         }
                     }
                 }
+
+                // Clean up all task worktrees for this contract
+                let pool_clone = pool.clone();
+                let state_clone = state.clone();
+                let contract_id = id;
+                tokio::spawn(async move {
+                    cleanup_contract_worktrees(&pool_clone, &state_clone, contract_id).await;
+                });
             }
 
             // Get summary with counts
@@ -548,6 +556,30 @@ pub async fn delete_contract(
             .into_response();
     };
 
+    // First, verify contract exists and belongs to owner
+    match repository::get_contract_for_owner(pool, id, auth.owner_id).await {
+        Ok(Some(_)) => {}
+        Ok(None) => {
+            return (
+                StatusCode::NOT_FOUND,
+                Json(ApiError::new("NOT_FOUND", "Contract not found")),
+            )
+                .into_response();
+        }
+        Err(e) => {
+            tracing::error!("Failed to get contract {}: {}", id, e);
+            return (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(ApiError::new("DB_ERROR", e.to_string())),
+            )
+                .into_response();
+        }
+    }
+
+    // Clean up all task worktrees BEFORE deleting the contract
+    // (because CASCADE delete will remove tasks from DB)
+    cleanup_contract_worktrees(pool, &state, id).await;
+
     match repository::delete_contract_for_owner(pool, id, auth.owner_id).await {
         Ok(true) => StatusCode::NO_CONTENT.into_response(),
         Ok(false) => (
@@ -1318,3 +1350,85 @@ pub async fn get_events(
         }
     }
 }
+
+// =============================================================================
+// Internal Helper Functions
+// =============================================================================
+
+/// Clean up all worktrees for tasks in a contract.
+///
+/// This is called when a contract is completed or deleted to remove
+/// all associated task worktrees from connected daemons.
+async fn cleanup_contract_worktrees(
+    pool: &sqlx::PgPool,
+    state: &SharedState,
+    contract_id: Uuid,
+) {
+    tracing::info!(
+        contract_id = %contract_id,
+        "Cleaning up worktrees for contract tasks"
+    );
+
+    // Get all tasks with worktree info for this contract
+    let tasks = match repository::list_contract_tasks_with_worktree_info(pool, contract_id).await {
+        Ok(tasks) => tasks,
+        Err(e) => {
+            tracing::error!(
+                contract_id = %contract_id,
+                error = %e,
+                "Failed to list tasks for worktree cleanup"
+            );
+            return;
+        }
+    };
+
+    if tasks.is_empty() {
+        tracing::debug!(
+            contract_id = %contract_id,
+            "No tasks with worktrees to clean up"
+        );
+        return;
+    }
+
+    tracing::info!(
+        contract_id = %contract_id,
+        task_count = tasks.len(),
+        "Found tasks with worktrees to clean up"
+    );
+
+    // Send cleanup command to each task's daemon
+    for task in tasks {
+        if let Some(daemon_id) = task.daemon_id {
+            let cmd = crate::server::state::DaemonCommand::CleanupWorktree {
+                task_id: task.id,
+                delete_branch: true, // Delete the branch when contract is done
+            };
+
+            match state.send_daemon_command(daemon_id, cmd).await {
+                Ok(()) => {
+                    tracing::info!(
+                        task_id = %task.id,
+                        daemon_id = %daemon_id,
+                        contract_id = %contract_id,
+                        "Sent worktree cleanup command"
+                    );
+                }
+                Err(e) => {
+                    tracing::warn!(
+                        task_id = %task.id,
+                        daemon_id = %daemon_id,
+                        contract_id = %contract_id,
+                        error = %e,
+                        "Failed to send worktree cleanup command (daemon may be offline)"
+                    );
+                }
+            }
+        } else {
+            tracing::debug!(
+                task_id = %task.id,
+                contract_id = %contract_id,
+                "Task has no daemon assigned, skipping worktree cleanup"
+            );
+        }
+    }
+}