1 files changed, 64 insertions, 0 deletions
diff --git a/makima/src/server/state.rs b/makima/src/server/state.rs
index 5b75281..32c0af3 100644
--- a/makima/src/server/state.rs
+++ b/makima/src/server/state.rs
@@ -797,6 +797,70 @@ impl AppState {
         self.question_responses.remove(&question_id);
     }
 
+    /// Remove all pending questions for a specific task.
+    ///
+    /// This should be called when a task is deleted to clean up orphaned questions.
+    /// Returns the number of questions removed.
+    pub fn remove_pending_questions_for_task(&self, task_id: Uuid) -> usize {
+        // Collect question IDs to remove
+        let question_ids: Vec<Uuid> = self
+            .pending_questions
+            .iter()
+            .filter(|entry| entry.value().task_id == task_id)
+            .map(|entry| entry.value().question_id)
+            .collect();
+
+        let count = question_ids.len();
+
+        // Remove pending questions and their responses
+        for question_id in question_ids {
+            self.pending_questions.remove(&question_id);
+            self.question_responses.remove(&question_id);
+        }
+
+        if count > 0 {
+            tracing::info!(
+                task_id = %task_id,
+                count = count,
+                "Cleaned up pending questions for deleted task"
+            );
+        }
+
+        count
+    }
+
+    /// Remove all pending questions for a specific contract.
+    ///
+    /// This should be called when a contract is deleted to clean up orphaned questions.
+    /// Returns the number of questions removed.
+    pub fn remove_pending_questions_for_contract(&self, contract_id: Uuid) -> usize {
+        // Collect question IDs to remove
+        let question_ids: Vec<Uuid> = self
+            .pending_questions
+            .iter()
+            .filter(|entry| entry.value().contract_id == contract_id)
+            .map(|entry| entry.value().question_id)
+            .collect();
+
+        let count = question_ids.len();
+
+        // Remove pending questions and their responses
+        for question_id in question_ids {
+            self.pending_questions.remove(&question_id);
+            self.question_responses.remove(&question_id);
+        }
+
+        if count > 0 {
+            tracing::info!(
+                contract_id = %contract_id,
+                count = count,
+                "Cleaned up pending questions for deleted contract"
+            );
+        }
+
+        count
+    }
+
     /// Register a new daemon connection.
     ///
     /// Returns the connection_id for later reference.