Initial implementation of the agent framework

src/context/manager.rs

@@ -0,0 +1,251 @@
+//! Context manager for handling multiple job contexts.
+
+use std::collections::HashMap;
+
+use tokio::sync::RwLock;
+use uuid::Uuid;
+
+use crate::context::{JobContext, Memory};
+use crate::error::JobError;
+
+/// Manages contexts for multiple concurrent jobs.
+pub struct ContextManager {
+    /// Active job contexts.
+    contexts: RwLock<HashMap<Uuid, JobContext>>,
+    /// Memory for each job.
+    memories: RwLock<HashMap<Uuid, Memory>>,
+    /// Maximum concurrent jobs.
+    max_jobs: usize,
+}
+
+impl ContextManager {
+    /// Create a new context manager.
+    pub fn new(max_jobs: usize) -> Self {
+        Self {
+            contexts: RwLock::new(HashMap::new()),
+            memories: RwLock::new(HashMap::new()),
+            max_jobs,
+        }
+    }
+
+    /// Create a new job context.
+    pub async fn create_job(
+        &self,
+        title: impl Into<String>,
+        description: impl Into<String>,
+    ) -> Result<Uuid, JobError> {
+        let contexts = self.contexts.read().await;
+        let active_count = contexts.values().filter(|c| c.state.is_active()).count();
+
+        if active_count >= self.max_jobs {
+            return Err(JobError::MaxJobsExceeded { max: self.max_jobs });
+        }
+        drop(contexts);
+
+        let context = JobContext::new(title, description);
+        let job_id = context.job_id;
+
+        let memory = Memory::new(job_id);
+
+        self.contexts.write().await.insert(job_id, context);
+        self.memories.write().await.insert(job_id, memory);
+
+        Ok(job_id)
+    }
+
+    /// Get a job context by ID.
+    pub async fn get_context(&self, job_id: Uuid) -> Result<JobContext, JobError> {
+        self.contexts
+            .read()
+            .await
+            .get(&job_id)
+            .cloned()
+            .ok_or(JobError::NotFound { id: job_id })
+    }
+
+    /// Get a mutable reference to update a job context.
+    pub async fn update_context<F, R>(&self, job_id: Uuid, f: F) -> Result<R, JobError>
+    where
+        F: FnOnce(&mut JobContext) -> R,
+    {
+        let mut contexts = self.contexts.write().await;
+        let context = contexts
+            .get_mut(&job_id)
+            .ok_or(JobError::NotFound { id: job_id })?;
+        Ok(f(context))
+    }
+
+    /// Get job memory.
+    pub async fn get_memory(&self, job_id: Uuid) -> Result<Memory, JobError> {
+        self.memories
+            .read()
+            .await
+            .get(&job_id)
+            .cloned()
+            .ok_or(JobError::NotFound { id: job_id })
+    }
+
+    /// Update job memory.
+    pub async fn update_memory<F, R>(&self, job_id: Uuid, f: F) -> Result<R, JobError>
+    where
+        F: FnOnce(&mut Memory) -> R,
+    {
+        let mut memories = self.memories.write().await;
+        let memory = memories
+            .get_mut(&job_id)
+            .ok_or(JobError::NotFound { id: job_id })?;
+        Ok(f(memory))
+    }
+
+    /// List all active job IDs.
+    pub async fn active_jobs(&self) -> Vec<Uuid> {
+        self.contexts
+            .read()
+            .await
+            .iter()
+            .filter(|(_, c)| c.state.is_active())
+            .map(|(id, _)| *id)
+            .collect()
+    }
+
+    /// List all job IDs.
+    pub async fn all_jobs(&self) -> Vec<Uuid> {
+        self.contexts.read().await.keys().cloned().collect()
+    }
+
+    /// Get count of active jobs.
+    pub async fn active_count(&self) -> usize {
+        self.contexts
+            .read()
+            .await
+            .values()
+            .filter(|c| c.state.is_active())
+            .count()
+    }
+
+    /// Remove a completed job (cleanup).
+    pub async fn remove_job(&self, job_id: Uuid) -> Result<(JobContext, Memory), JobError> {
+        let context = self
+            .contexts
+            .write()
+            .await
+            .remove(&job_id)
+            .ok_or(JobError::NotFound { id: job_id })?;
+
+        let memory = self
+            .memories
+            .write()
+            .await
+            .remove(&job_id)
+            .ok_or(JobError::NotFound { id: job_id })?;
+
+        Ok((context, memory))
+    }
+
+    /// Find stuck jobs.
+    pub async fn find_stuck_jobs(&self) -> Vec<Uuid> {
+        self.contexts
+            .read()
+            .await
+            .iter()
+            .filter(|(_, c)| c.state == crate::context::JobState::Stuck)
+            .map(|(id, _)| *id)
+            .collect()
+    }
+
+    /// Get summary of all jobs.
+    pub async fn summary(&self) -> ContextSummary {
+        let contexts = self.contexts.read().await;
+
+        let mut summary = ContextSummary::default();
+        for ctx in contexts.values() {
+            match ctx.state {
+                crate::context::JobState::Pending => summary.pending += 1,
+                crate::context::JobState::InProgress => summary.in_progress += 1,
+                crate::context::JobState::Completed => summary.completed += 1,
+                crate::context::JobState::Submitted => summary.submitted += 1,
+                crate::context::JobState::Accepted => summary.accepted += 1,
+                crate::context::JobState::Failed => summary.failed += 1,
+                crate::context::JobState::Stuck => summary.stuck += 1,
+                crate::context::JobState::Cancelled => summary.cancelled += 1,
+            }
+        }
+
+        summary.total = contexts.len();
+        summary
+    }
+}
+
+impl Default for ContextManager {
+    fn default() -> Self {
+        Self::new(10)
+    }
+}
+
+/// Summary of all job contexts.
+#[derive(Debug, Default)]
+pub struct ContextSummary {
+    pub total: usize,
+    pub pending: usize,
+    pub in_progress: usize,
+    pub completed: usize,
+    pub submitted: usize,
+    pub accepted: usize,
+    pub failed: usize,
+    pub stuck: usize,
+    pub cancelled: usize,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_create_job() {
+        let manager = ContextManager::new(5);
+        let job_id = manager.create_job("Test", "Description").await.unwrap();
+
+        let context = manager.get_context(job_id).await.unwrap();
+        assert_eq!(context.title, "Test");
+    }
+
+    #[tokio::test]
+    async fn test_max_jobs_limit() {
+        let manager = ContextManager::new(2);
+
+        manager.create_job("Job 1", "Desc").await.unwrap();
+        manager.create_job("Job 2", "Desc").await.unwrap();
+
+        // Start the jobs to make them active
+        for job_id in manager.all_jobs().await {
+            manager
+                .update_context(job_id, |ctx| {
+                    ctx.transition_to(crate::context::JobState::InProgress, None)
+                })
+                .await
+                .unwrap()
+                .unwrap();
+        }
+
+        // Third job should fail
+        let result = manager.create_job("Job 3", "Desc").await;
+        assert!(matches!(result, Err(JobError::MaxJobsExceeded { max: 2 })));
+    }
+
+    #[tokio::test]
+    async fn test_update_context() {
+        let manager = ContextManager::new(5);
+        let job_id = manager.create_job("Test", "Desc").await.unwrap();
+
+        manager
+            .update_context(job_id, |ctx| {
+                ctx.transition_to(crate::context::JobState::InProgress, None)
+            })
+            .await
+            .unwrap()
+            .unwrap();
+
+        let context = manager.get_context(job_id).await.unwrap();
+        assert_eq!(context.state, crate::context::JobState::InProgress);
+    }
+}

src/context/memory.rs

@@ -0,0 +1,293 @@
+//! Memory management for job contexts.
+
+use std::time::Duration;
+
+use chrono::{DateTime, Utc};
+use rust_decimal::Decimal;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+use crate::llm::ChatMessage;
+
+/// A record of an action taken during job execution.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ActionRecord {
+    /// Unique action ID.
+    pub id: Uuid,
+    /// Sequence number within the job.
+    pub sequence: u32,
+    /// Tool that was used.
+    pub tool_name: String,
+    /// Input parameters.
+    pub input: serde_json::Value,
+    /// Raw output (before sanitization).
+    pub output_raw: Option<String>,
+    /// Sanitized output.
+    pub output_sanitized: Option<serde_json::Value>,
+    /// Any sanitization warnings.
+    pub sanitization_warnings: Vec<String>,
+    /// Cost of the action.
+    pub cost: Option<Decimal>,
+    /// Duration of the action.
+    pub duration: Duration,
+    /// Whether the action succeeded.
+    pub success: bool,
+    /// Error message if failed.
+    pub error: Option<String>,
+    /// When the action was executed.
+    pub executed_at: DateTime<Utc>,
+}
+
+impl ActionRecord {
+    /// Create a new action record.
+    pub fn new(sequence: u32, tool_name: impl Into<String>, input: serde_json::Value) -> Self {
+        Self {
+            id: Uuid::new_v4(),
+            sequence,
+            tool_name: tool_name.into(),
+            input,
+            output_raw: None,
+            output_sanitized: None,
+            sanitization_warnings: Vec::new(),
+            cost: None,
+            duration: Duration::ZERO,
+            success: false,
+            error: None,
+            executed_at: Utc::now(),
+        }
+    }
+
+    /// Mark the action as successful.
+    pub fn succeed(
+        mut self,
+        output_raw: Option<String>,
+        output_sanitized: serde_json::Value,
+        duration: Duration,
+    ) -> Self {
+        self.success = true;
+        self.output_raw = output_raw;
+        self.output_sanitized = Some(output_sanitized);
+        self.duration = duration;
+        self
+    }
+
+    /// Mark the action as failed.
+    pub fn fail(mut self, error: impl Into<String>, duration: Duration) -> Self {
+        self.success = false;
+        self.error = Some(error.into());
+        self.duration = duration;
+        self
+    }
+
+    /// Add sanitization warnings.
+    pub fn with_warnings(mut self, warnings: Vec<String>) -> Self {
+        self.sanitization_warnings = warnings;
+        self
+    }
+
+    /// Set the cost.
+    pub fn with_cost(mut self, cost: Decimal) -> Self {
+        self.cost = Some(cost);
+        self
+    }
+}
+
+/// Conversation history.
+#[derive(Debug, Clone, Default)]
+pub struct ConversationMemory {
+    /// Messages in the conversation.
+    messages: Vec<ChatMessage>,
+    /// Maximum messages to keep.
+    max_messages: usize,
+}
+
+impl ConversationMemory {
+    /// Create a new conversation memory.
+    pub fn new(max_messages: usize) -> Self {
+        Self {
+            messages: Vec::new(),
+            max_messages,
+        }
+    }
+
+    /// Add a message.
+    pub fn add(&mut self, message: ChatMessage) {
+        self.messages.push(message);
+
+        // Trim old messages if needed (keeping system message if present)
+        while self.messages.len() > self.max_messages {
+            // Don't remove system messages
+            if self.messages.first().map(|m| m.role) == Some(crate::llm::Role::System) {
+                if self.messages.len() > 1 {
+                    self.messages.remove(1);
+                } else {
+                    break;
+                }
+            } else {
+                self.messages.remove(0);
+            }
+        }
+    }
+
+    /// Get all messages.
+    pub fn messages(&self) -> &[ChatMessage] {
+        &self.messages
+    }
+
+    /// Get the last N messages.
+    pub fn last_n(&self, n: usize) -> &[ChatMessage] {
+        let start = self.messages.len().saturating_sub(n);
+        &self.messages[start..]
+    }
+
+    /// Clear the conversation.
+    pub fn clear(&mut self) {
+        self.messages.clear();
+    }
+
+    /// Get message count.
+    pub fn len(&self) -> usize {
+        self.messages.len()
+    }
+
+    /// Check if empty.
+    pub fn is_empty(&self) -> bool {
+        self.messages.is_empty()
+    }
+}
+
+/// Combined memory for a job.
+#[derive(Debug, Clone)]
+pub struct Memory {
+    /// Job ID.
+    pub job_id: Uuid,
+    /// Conversation history.
+    pub conversation: ConversationMemory,
+    /// Action history.
+    pub actions: Vec<ActionRecord>,
+    /// Next action sequence number.
+    next_sequence: u32,
+}
+
+impl Memory {
+    /// Create a new memory instance.
+    pub fn new(job_id: Uuid) -> Self {
+        Self {
+            job_id,
+            conversation: ConversationMemory::new(100),
+            actions: Vec::new(),
+            next_sequence: 0,
+        }
+    }
+
+    /// Add a conversation message.
+    pub fn add_message(&mut self, message: ChatMessage) {
+        self.conversation.add(message);
+    }
+
+    /// Create a new action record.
+    pub fn create_action(
+        &mut self,
+        tool_name: impl Into<String>,
+        input: serde_json::Value,
+    ) -> ActionRecord {
+        let seq = self.next_sequence;
+        self.next_sequence += 1;
+        ActionRecord::new(seq, tool_name, input)
+    }
+
+    /// Record a completed action.
+    pub fn record_action(&mut self, action: ActionRecord) {
+        self.actions.push(action);
+    }
+
+    /// Get total cost of all actions.
+    pub fn total_cost(&self) -> Decimal {
+        self.actions
+            .iter()
+            .filter_map(|a| a.cost)
+            .fold(Decimal::ZERO, |acc, c| acc + c)
+    }
+
+    /// Get total duration of all actions.
+    pub fn total_duration(&self) -> Duration {
+        self.actions
+            .iter()
+            .map(|a| a.duration)
+            .fold(Duration::ZERO, |acc, d| acc + d)
+    }
+
+    /// Get successful action count.
+    pub fn successful_actions(&self) -> usize {
+        self.actions.iter().filter(|a| a.success).count()
+    }
+
+    /// Get failed action count.
+    pub fn failed_actions(&self) -> usize {
+        self.actions.iter().filter(|a| !a.success).count()
+    }
+
+    /// Get the last action.
+    pub fn last_action(&self) -> Option<&ActionRecord> {
+        self.actions.last()
+    }
+
+    /// Get actions by tool name.
+    pub fn actions_by_tool(&self, tool_name: &str) -> Vec<&ActionRecord> {
+        self.actions
+            .iter()
+            .filter(|a| a.tool_name == tool_name)
+            .collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_action_record() {
+        let action = ActionRecord::new(0, "test", serde_json::json!({"key": "value"}));
+        assert_eq!(action.sequence, 0);
+        assert!(!action.success);
+
+        let action = action.succeed(
+            Some("raw".to_string()),
+            serde_json::json!({"result": "ok"}),
+            Duration::from_millis(100),
+        );
+        assert!(action.success);
+    }
+
+    #[test]
+    fn test_conversation_memory() {
+        let mut memory = ConversationMemory::new(3);
+        memory.add(ChatMessage::user("Hello"));
+        memory.add(ChatMessage::assistant("Hi"));
+        memory.add(ChatMessage::user("How are you?"));
+        memory.add(ChatMessage::assistant("Good!"));
+
+        assert_eq!(memory.len(), 3); // Oldest removed
+    }
+
+    #[test]
+    fn test_memory_totals() {
+        let mut memory = Memory::new(Uuid::new_v4());
+
+        let action1 = memory
+            .create_action("tool1", serde_json::json!({}))
+            .succeed(None, serde_json::json!({}), Duration::from_secs(1))
+            .with_cost(Decimal::new(10, 1));
+        memory.record_action(action1);
+
+        let action2 = memory
+            .create_action("tool2", serde_json::json!({}))
+            .succeed(None, serde_json::json!({}), Duration::from_secs(2))
+            .with_cost(Decimal::new(20, 1));
+        memory.record_action(action2);
+
+        assert_eq!(memory.total_cost(), Decimal::new(30, 1));
+        assert_eq!(memory.total_duration(), Duration::from_secs(3));
+        assert_eq!(memory.successful_actions(), 2);
+    }
+}

src/context/mod.rs

@@ -0,0 +1,15 @@
+//! Per-job context isolation and state management.
+//!
+//! Each job runs with its own isolated context that includes:
+//! - Conversation history
+//! - Action history
+//! - State machine
+//! - Resource tracking
+
+mod manager;
+mod memory;
+mod state;
+
+pub use manager::ContextManager;
+pub use memory::{ActionRecord, ConversationMemory, Memory};
+pub use state::{JobContext, JobState, StateTransition};

src/context/state.rs

@@ -0,0 +1,276 @@
+//! Job state machine.
+
+use std::time::Duration;
+
+use chrono::{DateTime, Utc};
+use rust_decimal::Decimal;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+/// State of a job.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum JobState {
+    /// Job is waiting to be started.
+    Pending,
+    /// Job is currently being worked on.
+    InProgress,
+    /// Job work is complete, awaiting submission.
+    Completed,
+    /// Job has been submitted for review.
+    Submitted,
+    /// Job was accepted/paid.
+    Accepted,
+    /// Job failed and cannot be completed.
+    Failed,
+    /// Job is stuck and needs repair.
+    Stuck,
+    /// Job was cancelled.
+    Cancelled,
+}
+
+impl JobState {
+    /// Check if this state allows transitioning to another state.
+    pub fn can_transition_to(&self, target: JobState) -> bool {
+        use JobState::*;
+
+        matches!(
+            (self, target),
+            // From Pending
+            (Pending, InProgress) | (Pending, Cancelled) |
+            // From InProgress
+            (InProgress, Completed) | (InProgress, Failed) |
+            (InProgress, Stuck) | (InProgress, Cancelled) |
+            // From Completed
+            (Completed, Submitted) | (Completed, Failed) |
+            // From Submitted
+            (Submitted, Accepted) | (Submitted, Failed) |
+            // From Stuck (can recover or fail)
+            (Stuck, InProgress) | (Stuck, Failed) | (Stuck, Cancelled)
+        )
+    }
+
+    /// Check if this is a terminal state.
+    pub fn is_terminal(&self) -> bool {
+        matches!(self, Self::Accepted | Self::Failed | Self::Cancelled)
+    }
+
+    /// Check if the job is active (not terminal).
+    pub fn is_active(&self) -> bool {
+        !self.is_terminal()
+    }
+}
+
+impl std::fmt::Display for JobState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let s = match self {
+            Self::Pending => "pending",
+            Self::InProgress => "in_progress",
+            Self::Completed => "completed",
+            Self::Submitted => "submitted",
+            Self::Accepted => "accepted",
+            Self::Failed => "failed",
+            Self::Stuck => "stuck",
+            Self::Cancelled => "cancelled",
+        };
+        write!(f, "{}", s)
+    }
+}
+
+/// A state transition event.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StateTransition {
+    /// Previous state.
+    pub from: JobState,
+    /// New state.
+    pub to: JobState,
+    /// When the transition occurred.
+    pub timestamp: DateTime<Utc>,
+    /// Reason for the transition.
+    pub reason: Option<String>,
+}
+
+/// Context for a running job.
+#[derive(Debug, Clone)]
+pub struct JobContext {
+    /// Unique job ID.
+    pub job_id: Uuid,
+    /// Current state.
+    pub state: JobState,
+    /// Conversation ID if linked to a conversation.
+    pub conversation_id: Option<Uuid>,
+    /// Job title.
+    pub title: String,
+    /// Job description.
+    pub description: String,
+    /// Job category.
+    pub category: Option<String>,
+    /// Budget amount (if from marketplace).
+    pub budget: Option<Decimal>,
+    /// Budget token (e.g., "NEAR", "USD").
+    pub budget_token: Option<String>,
+    /// Our bid amount.
+    pub bid_amount: Option<Decimal>,
+    /// Estimated cost to complete.
+    pub estimated_cost: Option<Decimal>,
+    /// Estimated time to complete.
+    pub estimated_duration: Option<Duration>,
+    /// Actual cost so far.
+    pub actual_cost: Decimal,
+    /// When the job was created.
+    pub created_at: DateTime<Utc>,
+    /// When the job was started.
+    pub started_at: Option<DateTime<Utc>>,
+    /// When the job was completed.
+    pub completed_at: Option<DateTime<Utc>>,
+    /// Number of repair attempts.
+    pub repair_attempts: u32,
+    /// State transition history.
+    pub transitions: Vec<StateTransition>,
+    /// Metadata.
+    pub metadata: serde_json::Value,
+}
+
+impl JobContext {
+    /// Create a new job context.
+    pub fn new(title: impl Into<String>, description: impl Into<String>) -> Self {
+        Self {
+            job_id: Uuid::new_v4(),
+            state: JobState::Pending,
+            conversation_id: None,
+            title: title.into(),
+            description: description.into(),
+            category: None,
+            budget: None,
+            budget_token: None,
+            bid_amount: None,
+            estimated_cost: None,
+            estimated_duration: None,
+            actual_cost: Decimal::ZERO,
+            created_at: Utc::now(),
+            started_at: None,
+            completed_at: None,
+            repair_attempts: 0,
+            transitions: Vec::new(),
+            metadata: serde_json::Value::Null,
+        }
+    }
+
+    /// Transition to a new state.
+    pub fn transition_to(
+        &mut self,
+        new_state: JobState,
+        reason: Option<String>,
+    ) -> Result<(), String> {
+        if !self.state.can_transition_to(new_state) {
+            return Err(format!(
+                "Cannot transition from {} to {}",
+                self.state, new_state
+            ));
+        }
+
+        let transition = StateTransition {
+            from: self.state,
+            to: new_state,
+            timestamp: Utc::now(),
+            reason,
+        };
+
+        self.transitions.push(transition);
+        self.state = new_state;
+
+        // Update timestamps
+        match new_state {
+            JobState::InProgress if self.started_at.is_none() => {
+                self.started_at = Some(Utc::now());
+            }
+            JobState::Completed | JobState::Accepted | JobState::Failed | JobState::Cancelled => {
+                self.completed_at = Some(Utc::now());
+            }
+            _ => {}
+        }
+
+        Ok(())
+    }
+
+    /// Add to the actual cost.
+    pub fn add_cost(&mut self, cost: Decimal) {
+        self.actual_cost += cost;
+    }
+
+    /// Get the duration since the job started.
+    pub fn elapsed(&self) -> Option<Duration> {
+        self.started_at.map(|start| {
+            let end = self.completed_at.unwrap_or_else(Utc::now);
+            let duration = end.signed_duration_since(start);
+            Duration::from_secs(duration.num_seconds().max(0) as u64)
+        })
+    }
+
+    /// Mark the job as stuck.
+    pub fn mark_stuck(&mut self, reason: impl Into<String>) -> Result<(), String> {
+        self.transition_to(JobState::Stuck, Some(reason.into()))
+    }
+
+    /// Attempt to recover from stuck state.
+    pub fn attempt_recovery(&mut self) -> Result<(), String> {
+        if self.state != JobState::Stuck {
+            return Err("Job is not stuck".to_string());
+        }
+        self.repair_attempts += 1;
+        self.transition_to(JobState::InProgress, Some("Recovery attempt".to_string()))
+    }
+}
+
+impl Default for JobContext {
+    fn default() -> Self {
+        Self::new("Untitled", "No description")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_state_transitions() {
+        assert!(JobState::Pending.can_transition_to(JobState::InProgress));
+        assert!(JobState::InProgress.can_transition_to(JobState::Completed));
+        assert!(!JobState::Completed.can_transition_to(JobState::Pending));
+        assert!(!JobState::Accepted.can_transition_to(JobState::InProgress));
+    }
+
+    #[test]
+    fn test_terminal_states() {
+        assert!(JobState::Accepted.is_terminal());
+        assert!(JobState::Failed.is_terminal());
+        assert!(JobState::Cancelled.is_terminal());
+        assert!(!JobState::InProgress.is_terminal());
+    }
+
+    #[test]
+    fn test_job_context_transitions() {
+        let mut ctx = JobContext::new("Test", "Test job");
+        assert_eq!(ctx.state, JobState::Pending);
+
+        ctx.transition_to(JobState::InProgress, None).unwrap();
+        assert_eq!(ctx.state, JobState::InProgress);
+        assert!(ctx.started_at.is_some());
+
+        ctx.transition_to(JobState::Completed, Some("Done".to_string()))
+            .unwrap();
+        assert_eq!(ctx.state, JobState::Completed);
+    }
+
+    #[test]
+    fn test_stuck_recovery() {
+        let mut ctx = JobContext::new("Test", "Test job");
+        ctx.transition_to(JobState::InProgress, None).unwrap();
+        ctx.mark_stuck("Timed out").unwrap();
+        assert_eq!(ctx.state, JobState::Stuck);
+
+        ctx.attempt_recovery().unwrap();
+        assert_eq!(ctx.state, JobState::InProgress);
+        assert_eq!(ctx.repair_attempts, 1);
+    }
+}