ricecoder_workflows/
agent_executor.rs

1//! Agent step execution handler
2//!
3//! Handles execution of agent steps within workflows by delegating to the ricecoder-agents API.
4
5use crate::error::{WorkflowError, WorkflowResult};
6use crate::models::{AgentStep, Workflow, WorkflowState};
7use crate::state::StateManager;
8use std::time::Instant;
9
10/// Executes agent steps by delegating to the ricecoder-agents API
11pub struct AgentExecutor;
12
13impl AgentExecutor {
14    /// Execute an agent step
15    ///
16    /// Delegates to the ricecoder-agents API to execute the agent with the given configuration.
17    /// Captures the agent output and any errors that occur during execution.
18    ///
19    /// # Arguments
20    ///
21    /// * `workflow` - The workflow containing the step
22    /// * `state` - The current workflow state
23    /// * `step_id` - The ID of the agent step to execute
24    /// * `agent_step` - The agent step configuration
25    ///
26    /// # Returns
27    ///
28    /// Returns `Ok(())` if the step executed successfully, or an error if execution failed.
29    pub fn execute_agent_step(
30        workflow: &Workflow,
31        state: &mut WorkflowState,
32        step_id: &str,
33        agent_step: &AgentStep,
34    ) -> WorkflowResult<()> {
35        // Mark step as started
36        StateManager::start_step(state, step_id.to_string());
37
38        let start_time = Instant::now();
39
40        // Execute the agent
41        // In a real implementation, this would:
42        // 1. Look up the agent from the registry using agent_step.agent_id
43        // 2. Create an AgentInput with the step configuration
44        // 3. Call the agent's execute method
45        // 4. Capture the output and any errors
46        //
47        // For now, we simulate successful execution
48        let agent_output = Self::execute_agent_internal(workflow, state, step_id, agent_step)?;
49
50        let duration_ms = start_time.elapsed().as_millis() as u64;
51
52        // Mark step as completed with the agent output
53        StateManager::complete_step(state, step_id.to_string(), Some(agent_output), duration_ms);
54
55        Ok(())
56    }
57
58    /// Internal agent execution logic
59    ///
60    /// This is where the actual agent execution would happen.
61    /// In a real implementation, this would integrate with ricecoder-agents.
62    fn execute_agent_internal(
63        _workflow: &Workflow,
64        _state: &WorkflowState,
65        _step_id: &str,
66        agent_step: &AgentStep,
67    ) -> WorkflowResult<serde_json::Value> {
68        // In a real implementation, this would:
69        // 1. Get the agent from the registry
70        // 2. Create an AgentInput from the step configuration
71        // 3. Execute the agent
72        // 4. Return the output
73        //
74        // For now, we return a simulated output
75        Ok(serde_json::json!({
76            "agent_id": agent_step.agent_id,
77            "task": agent_step.task,
78            "status": "completed",
79            "output": {
80                "findings": [],
81                "suggestions": []
82            }
83        }))
84    }
85
86    /// Execute an agent step with timeout
87    ///
88    /// Executes an agent step with a specified timeout. If the agent takes longer
89    /// than the timeout, the execution is cancelled and an error is returned.
90    ///
91    /// # Arguments
92    ///
93    /// * `workflow` - The workflow containing the step
94    /// * `state` - The current workflow state
95    /// * `step_id` - The ID of the agent step to execute
96    /// * `agent_step` - The agent step configuration
97    /// * `timeout_ms` - The timeout in milliseconds
98    ///
99    /// # Returns
100    ///
101    /// Returns `Ok(())` if the step executed successfully within the timeout,
102    /// or an error if execution failed or timed out.
103    pub fn execute_agent_step_with_timeout(
104        workflow: &Workflow,
105        state: &mut WorkflowState,
106        step_id: &str,
107        agent_step: &AgentStep,
108        timeout_ms: u64,
109    ) -> WorkflowResult<()> {
110        // Mark step as started
111        StateManager::start_step(state, step_id.to_string());
112
113        let start_time = Instant::now();
114
115        // Execute the agent with timeout
116        // In a real implementation, this would use tokio::time::timeout
117        let agent_output = Self::execute_agent_internal(workflow, state, step_id, agent_step)?;
118
119        let elapsed_ms = start_time.elapsed().as_millis() as u64;
120
121        // Check if we exceeded the timeout
122        if elapsed_ms > timeout_ms {
123            StateManager::fail_step(
124                state,
125                step_id.to_string(),
126                format!("Agent execution timed out after {}ms", timeout_ms),
127                elapsed_ms,
128            );
129            return Err(WorkflowError::StepFailed(format!(
130                "Agent step {} timed out after {}ms",
131                step_id, timeout_ms
132            )));
133        }
134
135        // Mark step as completed
136        StateManager::complete_step(state, step_id.to_string(), Some(agent_output), elapsed_ms);
137
138        Ok(())
139    }
140
141    /// Get the agent ID from an agent step
142    pub fn get_agent_id(agent_step: &AgentStep) -> &str {
143        &agent_step.agent_id
144    }
145
146    /// Get the task from an agent step
147    pub fn get_task(agent_step: &AgentStep) -> &str {
148        &agent_step.task
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::models::{
156        ErrorAction, RiskFactors, StepConfig, StepStatus, StepType, WorkflowConfig, WorkflowStep,
157    };
158
159    fn create_workflow_with_agent_step() -> Workflow {
160        Workflow {
161            id: "test-workflow".to_string(),
162            name: "Test Workflow".to_string(),
163            description: "A test workflow".to_string(),
164            parameters: vec![],
165            steps: vec![WorkflowStep {
166                id: "agent-step".to_string(),
167                name: "Agent Step".to_string(),
168                step_type: StepType::Agent(AgentStep {
169                    agent_id: "test-agent".to_string(),
170                    task: "test-task".to_string(),
171                }),
172                config: StepConfig {
173                    config: serde_json::json!({"param": "value"}),
174                },
175                dependencies: vec![],
176                approval_required: false,
177                on_error: ErrorAction::Fail,
178                risk_score: None,
179                risk_factors: RiskFactors::default(),
180            }],
181            config: WorkflowConfig {
182                timeout_ms: None,
183                max_parallel: None,
184            },
185        }
186    }
187
188    #[test]
189    fn test_execute_agent_step() {
190        let workflow = create_workflow_with_agent_step();
191        let mut state = StateManager::create_state(&workflow);
192        let agent_step = AgentStep {
193            agent_id: "test-agent".to_string(),
194            task: "test-task".to_string(),
195        };
196
197        let result =
198            AgentExecutor::execute_agent_step(&workflow, &mut state, "agent-step", &agent_step);
199        assert!(result.is_ok());
200
201        // Verify step is marked as completed
202        let step_result = state.step_results.get("agent-step");
203        assert!(step_result.is_some());
204        assert_eq!(step_result.unwrap().status, StepStatus::Completed);
205    }
206
207    #[test]
208    fn test_execute_agent_step_with_timeout() {
209        let workflow = create_workflow_with_agent_step();
210        let mut state = StateManager::create_state(&workflow);
211        let agent_step = AgentStep {
212            agent_id: "test-agent".to_string(),
213            task: "test-task".to_string(),
214        };
215
216        let result = AgentExecutor::execute_agent_step_with_timeout(
217            &workflow,
218            &mut state,
219            "agent-step",
220            &agent_step,
221            5000, // 5 second timeout
222        );
223        assert!(result.is_ok());
224
225        // Verify step is marked as completed
226        let step_result = state.step_results.get("agent-step");
227        assert!(step_result.is_some());
228        assert_eq!(step_result.unwrap().status, StepStatus::Completed);
229    }
230
231    #[test]
232    fn test_get_agent_id() {
233        let agent_step = AgentStep {
234            agent_id: "my-agent".to_string(),
235            task: "my-task".to_string(),
236        };
237
238        assert_eq!(AgentExecutor::get_agent_id(&agent_step), "my-agent");
239    }
240
241    #[test]
242    fn test_get_task() {
243        let agent_step = AgentStep {
244            agent_id: "my-agent".to_string(),
245            task: "my-task".to_string(),
246        };
247
248        assert_eq!(AgentExecutor::get_task(&agent_step), "my-task");
249    }
250}
ricecoder_workflows/agent_executor.rs

ricecoder_workflows/
agent_executor.rs
