@feltdb/core 0.5.7 → 0.6.0

package/dist/create/server-source/crates/feltdb-server/src/durable_operations.rs

@@ -0,0 +1,992 @@
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+/// Operation state lifecycle.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum OperationState {
+    /// Operation created, awaiting worker claim.
+    Pending,
+    /// Operation claimed by a worker (lease active).
+    Claimed,
+    /// Operation being executed by worker.
+    Running,
+    /// Operation completed successfully.
+    Completed,
+    /// Operation failed.
+    Failed,
+    /// Operation state unknown (after lease expiration).
+    Unknown,
+}
+
+/// Durable operation for work coordination and recovery.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DurableOperation {
+    /// Unique operation ID.
+    pub operation_id: String,
+
+    /// Type of operation (defines behavior/retry policy).
+    pub operation_type: String,
+
+    /// Tenant ID for isolation.
+    pub tenant_id: String,
+
+    /// Current state of the operation.
+    pub state: OperationState,
+
+    /// Worker claiming the operation (if claimed).
+    pub worker_id: Option<String>,
+
+    /// Lease expiration time (unix timestamp).
+    pub lease_expires_at: u64,
+
+    /// When operation was created.
+    pub created_at: u64,
+
+    /// Last state change time.
+    pub updated_at: u64,
+
+    /// Operation payload (caller-defined work).
+    pub payload: Value,
+
+    /// Result (set on completion).
+    pub result: Option<Value>,
+
+    /// Error reason (set on failure).
+    pub error: Option<String>,
+
+    /// Version for optimistic concurrency control.
+    pub version: u64,
+}
+
+/// Request to create a durable operation.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CreateOperationRequest {
+    pub operation_id: String,
+    pub operation_type: String,
+    pub tenant_id: String,
+    pub payload: Value,
+}
+
+/// Request to claim an operation.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ClaimOperationRequest {
+    pub operation_id: String,
+    pub worker_id: String,
+    pub lease_duration_secs: u64,
+}
+
+/// Response with operation after claiming.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ClaimOperationResponse {
+    pub operation: DurableOperation,
+    pub version: u64,
+}
+
+/// Store for durable operations.
+#[derive(Clone)]
+pub struct DurableOperationStore {
+    operations: Arc<RwLock<HashMap<String, DurableOperation>>>,
+}
+
+impl DurableOperationStore {
+    /// Creates a new durable operation store.
+    pub fn new() -> Self {
+        Self {
+            operations: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    /// Creates a new durable operation.
+    pub fn create_operation(
+        &self,
+        operation_id: String,
+        operation_type: String,
+        tenant_id: String,
+        payload: Value,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        if ops.contains_key(&operation_id) {
+            return Err(format!(
+                "operation already exists: {}",
+                operation_id
+            ));
+        }
+
+        let now = now();
+        let operation = DurableOperation {
+            operation_id: operation_id.clone(),
+            operation_type,
+            tenant_id,
+            state: OperationState::Pending,
+            worker_id: None,
+            lease_expires_at: 0,
+            created_at: now,
+            updated_at: now,
+            payload,
+            result: None,
+            error: None,
+            version: 1,
+        };
+
+        ops.insert(operation_id, operation.clone());
+        Ok(operation)
+    }
+
+    /// Claims an operation for a worker.
+    /// Returns the claimed operation and current version for optimistic concurrency.
+    pub fn claim_operation(
+        &self,
+        operation_id: &str,
+        worker_id: String,
+        lease_duration_secs: u64,
+    ) -> Result<ClaimOperationResponse, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        // Can only claim Pending or Unknown operations
+        if operation.state != OperationState::Pending && operation.state != OperationState::Unknown {
+            return Err(format!(
+                "cannot claim operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        let now = now();
+
+        // Transition to Claimed state
+        operation.state = OperationState::Claimed;
+        operation.worker_id = Some(worker_id);
+        operation.lease_expires_at = now + lease_duration_secs;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(ClaimOperationResponse {
+            operation: operation.clone(),
+            version: operation.version,
+        })
+    }
+
+    /// Renews a worker's lease on an operation (extends expiration).
+    /// Requires the worker to prove ownership.
+    pub fn renew_operation(
+        &self,
+        operation_id: &str,
+        worker_id: &str,
+        lease_duration_secs: u64,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        let now = now();
+
+        // Check if lease has expired (worker fencing)
+        if operation.lease_expires_at < now {
+            return Err(format!(
+                "worker lease expired for operation: {}",
+                operation_id
+            ));
+        }
+
+        // Verify worker ownership
+        if operation.worker_id.as_deref() != Some(worker_id) {
+            return Err(format!(
+                "worker mismatch: expected {:?}, got {}",
+                operation.worker_id, worker_id
+            ));
+        }
+
+        // Extend lease
+        operation.lease_expires_at = now + lease_duration_secs;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Marks an operation as Running.
+    pub fn start_operation(
+        &self,
+        operation_id: &str,
+        worker_id: &str,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        let now = now();
+
+        // Check lease expiration
+        if operation.lease_expires_at < now {
+            return Err("worker lease expired".to_string());
+        }
+
+        // Verify worker ownership
+        if operation.worker_id.as_deref() != Some(worker_id) {
+            return Err("worker mismatch".to_string());
+        }
+
+        // Transition from Claimed to Running
+        if operation.state != OperationState::Claimed {
+            return Err(format!(
+                "cannot start operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        operation.state = OperationState::Running;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Completes an operation with a result.
+    pub fn complete_operation(
+        &self,
+        operation_id: &str,
+        worker_id: &str,
+        result: Value,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        let now = now();
+
+        // Check lease expiration
+        if operation.lease_expires_at < now {
+            return Err("worker lease expired".to_string());
+        }
+
+        // Verify worker ownership
+        if operation.worker_id.as_deref() != Some(worker_id) {
+            return Err("worker mismatch".to_string());
+        }
+
+        // Must be in Running state
+        if operation.state != OperationState::Running {
+            return Err(format!(
+                "cannot complete operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        operation.state = OperationState::Completed;
+        operation.result = Some(result);
+        operation.error = None;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Fails an operation with an error.
+    pub fn fail_operation(
+        &self,
+        operation_id: &str,
+        worker_id: &str,
+        error: String,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        let now = now();
+
+        // Check lease expiration
+        if operation.lease_expires_at < now {
+            return Err("worker lease expired".to_string());
+        }
+
+        // Verify worker ownership
+        if operation.worker_id.as_deref() != Some(worker_id) {
+            return Err("worker mismatch".to_string());
+        }
+
+        // Must be in Claimed or Running state
+        if operation.state != OperationState::Claimed && operation.state != OperationState::Running {
+            return Err(format!(
+                "cannot fail operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        operation.state = OperationState::Failed;
+        operation.error = Some(error);
+        operation.result = None;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Marks an operation as Unknown (when lease expires).
+    pub fn mark_unknown(
+        &self,
+        operation_id: &str,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        let now = now();
+
+        // Can only mark Unknown if in Claimed or Running state
+        if operation.state != OperationState::Claimed && operation.state != OperationState::Running {
+            return Err(format!(
+                "cannot mark Unknown operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        operation.state = OperationState::Unknown;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Reconciles an operation after lease expiration.
+    /// Safely transitions from Unknown back to Pending if reconciliation succeeds.
+    pub fn reconcile_operation(
+        &self,
+        operation_id: &str,
+        expected_worker_id: Option<&str>,
+    ) -> Result<DurableOperation, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let operation = ops
+            .get_mut(operation_id)
+            .ok_or(format!("operation not found: {}", operation_id))?;
+
+        // Can only reconcile if Unknown
+        if operation.state != OperationState::Unknown {
+            return Err(format!(
+                "cannot reconcile operation in state {:?}",
+                operation.state
+            ));
+        }
+
+        // If expected_worker_id provided, verify it matches (to detect multiple workers)
+        if let Some(expected) = expected_worker_id {
+            if operation.worker_id.as_deref() != Some(expected) {
+                return Err(format!(
+                    "worker mismatch during reconciliation: expected {}, got {:?}",
+                    expected, operation.worker_id
+                ));
+            }
+        }
+
+        let now = now();
+
+        // Transition back to Pending for retry
+        operation.state = OperationState::Pending;
+        operation.worker_id = None;
+        operation.lease_expires_at = 0;
+        operation.updated_at = now;
+        operation.version += 1;
+
+        Ok(operation.clone())
+    }
+
+    /// Gets an operation by ID.
+    pub fn get_operation(&self, operation_id: &str) -> Result<DurableOperation, String> {
+        let ops = self
+            .operations
+            .read()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        ops.get(operation_id)
+            .cloned()
+            .ok_or(format!("operation not found: {}", operation_id))
+    }
+
+    /// Lists all operations for a tenant.
+    pub fn list_operations(&self, tenant_id: &str) -> Result<Vec<DurableOperation>, String> {
+        let ops = self
+            .operations
+            .read()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        Ok(ops
+            .values()
+            .filter(|op| op.tenant_id == tenant_id)
+            .cloned()
+            .collect())
+    }
+
+    /// Lists operations by state.
+    pub fn list_by_state(
+        &self,
+        tenant_id: &str,
+        state: OperationState,
+    ) -> Result<Vec<DurableOperation>, String> {
+        let ops = self
+            .operations
+            .read()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        Ok(ops
+            .values()
+            .filter(|op| op.tenant_id == tenant_id && op.state == state)
+            .cloned()
+            .collect())
+    }
+
+    /// Cleanup: Removes completed operations older than cutoff time.
+    pub fn cleanup_completed_before(&self, cutoff_time: u64) -> Result<usize, String> {
+        let mut ops = self
+            .operations
+            .write()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let before_len = ops.len();
+        ops.retain(|_, op| {
+            !(op.state == OperationState::Completed && op.updated_at < cutoff_time)
+        });
+        let after_len = ops.len();
+
+        Ok(before_len - after_len)
+    }
+
+    /// Returns count of operations by state (for diagnostics).
+    pub fn get_state_counts(&self) -> Result<StateCountsResponse, String> {
+        let ops = self
+            .operations
+            .read()
+            .map_err(|_| "store lock poisoned".to_string())?;
+
+        let mut counts = StateCountsResponse::default();
+        for op in ops.values() {
+            match op.state {
+                OperationState::Pending => counts.pending += 1,
+                OperationState::Claimed => counts.claimed += 1,
+                OperationState::Running => counts.running += 1,
+                OperationState::Completed => counts.completed += 1,
+                OperationState::Failed => counts.failed += 1,
+                OperationState::Unknown => counts.unknown += 1,
+            }
+        }
+
+        Ok(counts)
+    }
+}
+
+impl Default for DurableOperationStore {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Response with counts of operations by state.
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct StateCountsResponse {
+    pub pending: usize,
+    pub claimed: usize,
+    pub running: usize,
+    pub completed: usize,
+    pub failed: usize,
+    pub unknown: usize,
+}
+
+fn now() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn create_operation() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        let op = store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        assert_eq!(op.operation_id, "op_123");
+        assert_eq!(op.state, OperationState::Pending);
+        assert_eq!(op.version, 1);
+        assert_eq!(op.payload, payload);
+    }
+
+    #[test]
+    fn reject_duplicate_operation_id() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("first create failed");
+
+        let result = store.create_operation(
+            "op_123".into(),
+            "process".into(),
+            "tenant_prod".into(),
+            payload,
+        );
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn claim_pending_operation() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        let response = store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        assert_eq!(response.operation.state, OperationState::Claimed);
+        assert_eq!(response.operation.worker_id, Some("worker_1".into()));
+        assert_eq!(response.version, 2);
+    }
+
+    #[test]
+    fn reject_claim_non_pending_operation() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("first claim failed");
+
+        let result = store.claim_operation("op_123", "worker_2".into(), 60);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn renew_lease() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let renewed = store
+            .renew_operation("op_123", "worker_1", 120)
+            .expect("renew failed");
+
+        assert_eq!(renewed.version, 3);
+        assert_eq!(renewed.state, OperationState::Claimed);
+    }
+
+    #[test]
+    fn reject_renew_wrong_worker() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let result = store.renew_operation("op_123", "worker_2", 120);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn complete_operation() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        store
+            .start_operation("op_123", "worker_1")
+            .expect("start failed");
+
+        let result_value = serde_json::json!({"result": "success"});
+        let completed = store
+            .complete_operation("op_123", "worker_1", result_value.clone())
+            .expect("complete failed");
+
+        assert_eq!(completed.state, OperationState::Completed);
+        assert_eq!(completed.result, Some(result_value));
+        assert_eq!(completed.error, None);
+    }
+
+    #[test]
+    fn fail_operation() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let failed = store
+            .fail_operation("op_123", "worker_1", "operation timeout".into())
+            .expect("fail failed");
+
+        assert_eq!(failed.state, OperationState::Failed);
+        assert_eq!(failed.error, Some("operation timeout".into()));
+    }
+
+    #[test]
+    fn mark_unknown_and_reconcile() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let unknown = store
+            .mark_unknown("op_123")
+            .expect("mark unknown failed");
+        assert_eq!(unknown.state, OperationState::Unknown);
+
+        let reconciled = store
+            .reconcile_operation("op_123", Some("worker_1"))
+            .expect("reconcile failed");
+
+        assert_eq!(reconciled.state, OperationState::Pending);
+        assert_eq!(reconciled.worker_id, None);
+    }
+
+    #[test]
+    fn worker_fencing_prevents_renew_after_lease_expiration() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_123", "worker_1".into(), 1)
+            .expect("claim failed");
+
+        // Wait for lease to expire
+        std::thread::sleep(std::time::Duration::from_secs(2));
+
+        let result = store.renew_operation("op_123", "worker_1", 60);
+        assert!(result.is_err());
+        assert!(result.unwrap_err().contains("lease expired"));
+    }
+
+    #[test]
+    fn list_operations_by_tenant() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_1".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        store
+            .create_operation(
+                "op_2".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        store
+            .create_operation(
+                "op_3".into(),
+                "process".into(),
+                "tenant_other".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        let ops = store
+            .list_operations("tenant_prod")
+            .expect("list failed");
+        assert_eq!(ops.len(), 2);
+    }
+
+    #[test]
+    fn list_operations_by_state() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_1".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        store
+            .create_operation(
+                "op_2".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_2", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let pending = store
+            .list_by_state("tenant_prod", OperationState::Pending)
+            .expect("list pending failed");
+        assert_eq!(pending.len(), 1);
+
+        let claimed = store
+            .list_by_state("tenant_prod", OperationState::Claimed)
+            .expect("list claimed failed");
+        assert_eq!(claimed.len(), 1);
+    }
+
+    #[test]
+    fn cleanup_completed_operations() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_completed".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_completed", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        store
+            .start_operation("op_completed", "worker_1")
+            .expect("start failed");
+
+        store
+            .complete_operation(
+                "op_completed",
+                "worker_1",
+                serde_json::json!({"result": "done"}),
+            )
+            .expect("complete failed");
+
+        store
+            .create_operation(
+                "op_pending".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        let cutoff = now();
+        let deleted = store
+            .cleanup_completed_before(cutoff + 1)
+            .expect("cleanup failed");
+        assert_eq!(deleted, 1);
+
+        let ops = store
+            .list_operations("tenant_prod")
+            .expect("list failed");
+        assert_eq!(ops.len(), 1);
+        assert_eq!(ops[0].operation_id, "op_pending");
+    }
+
+    #[test]
+    fn get_state_counts() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_1".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload.clone(),
+            )
+            .expect("create failed");
+
+        store
+            .create_operation(
+                "op_2".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        store
+            .claim_operation("op_2", "worker_1".into(), 60)
+            .expect("claim failed");
+
+        let counts = store.get_state_counts().expect("get counts failed");
+        assert_eq!(counts.pending, 1);
+        assert_eq!(counts.claimed, 1);
+    }
+
+    #[test]
+    fn state_transitions_are_safe() {
+        let store = DurableOperationStore::new();
+        let payload = serde_json::json!({"work": "data"});
+
+        store
+            .create_operation(
+                "op_123".into(),
+                "process".into(),
+                "tenant_prod".into(),
+                payload,
+            )
+            .expect("create failed");
+
+        let claim_response = store
+            .claim_operation("op_123", "worker_1".into(), 60)
+            .expect("claim failed");
+        assert_eq!(claim_response.operation.version, 2);
+
+        store
+            .start_operation("op_123", "worker_1")
+            .expect("start failed");
+
+        store
+            .complete_operation(
+                "op_123",
+                "worker_1",
+                serde_json::json!({"result": "done"}),
+            )
+            .expect("complete failed");
+
+        let final_op = store.get_operation("op_123").expect("get failed");
+        assert_eq!(final_op.state, OperationState::Completed);
+        assert_eq!(final_op.version, 4);
+    }
+}