prodigy/cook/execution/mapreduce/
types.rs

1//! Shared types and traits for MapReduce operations
2//!
3//! This module contains type definitions shared across the MapReduce
4//! implementation, promoting consistency and reducing coupling.
5
6use crate::cook::execution::interpolation::InterpolationContext;
7use crate::cook::execution::variable_capture::CaptureConfig;
8use crate::cook::execution::variables::{Variable, VariableContext};
9use crate::cook::orchestrator::ExecutionEnvironment;
10use serde::{Deserialize, Serialize};
11use serde_json::{json, Value};
12use std::collections::HashMap;
13use std::path::PathBuf;
14
15/// Configuration for MapReduce execution
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct MapReduceConfig {
18    /// Input source (file path or command)
19    #[serde(default)]
20    pub input: String,
21    /// JSON path expression to extract work items
22    #[serde(default)]
23    pub json_path: String,
24    /// Maximum number of parallel agents
25    #[serde(default = "default_max_parallel")]
26    pub max_parallel: usize,
27    /// Timeout for individual agent execution (in seconds)
28    pub agent_timeout_secs: Option<u64>,
29    /// Whether to continue on agent failures
30    pub continue_on_failure: bool,
31    /// Batch size for processing work items
32    pub batch_size: Option<usize>,
33    /// Enable checkpoint saving
34    pub enable_checkpoints: bool,
35    /// Maximum number of items to process
36    pub max_items: Option<usize>,
37    /// Number of items to skip
38    pub offset: Option<usize>,
39}
40
41fn default_max_parallel() -> usize {
42    10
43}
44
45impl Default for MapReduceConfig {
46    fn default() -> Self {
47        Self {
48            input: String::new(),
49            json_path: String::new(),
50            max_parallel: default_max_parallel(),
51            agent_timeout_secs: Some(300),
52            continue_on_failure: false,
53            batch_size: None,
54            enable_checkpoints: true,
55            max_items: None,
56            offset: None,
57        }
58    }
59}
60
61/// Setup phase configuration
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SetupPhase {
64    /// Commands to execute during setup
65    pub commands: Vec<crate::cook::workflow::WorkflowStep>,
66    /// Timeout for setup phase (in seconds)
67    /// If None, no timeout is applied
68    pub timeout: Option<u64>,
69    /// Variables to capture from setup commands
70    #[serde(default)]
71    pub capture_outputs: HashMap<String, CaptureConfig>,
72}
73
74/// Map phase configuration
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct MapPhase {
77    /// Input source specification
78    #[serde(flatten)]
79    pub config: MapReduceConfig,
80    /// JSONPath expression for data extraction
81    pub json_path: Option<String>,
82    /// Agent template commands
83    pub agent_template: Vec<crate::cook::workflow::WorkflowStep>,
84    /// Filter expression for work items
85    pub filter: Option<String>,
86    /// Sort expression for work items
87    pub sort_by: Option<String>,
88    /// Maximum items to process
89    pub max_items: Option<usize>,
90    /// Optional distinct field for deduplication
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub distinct: Option<String>,
93    /// Timeout configuration
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub timeout_config: Option<crate::cook::execution::mapreduce::timeout::TimeoutConfig>,
96    /// Workflow environment variables (resolved from env section and command-line args)
97    #[serde(skip)]
98    pub workflow_env: std::collections::HashMap<String, String>,
99}
100
101/// Reduce phase configuration
102#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct ReducePhase {
104    /// Commands to execute during reduction
105    pub commands: Vec<crate::cook::workflow::WorkflowStep>,
106    /// Timeout for reduce phase (in seconds)
107    pub timeout_secs: Option<u64>,
108}
109
110/// Options for resuming MapReduce jobs
111#[derive(Debug, Clone, Serialize, Deserialize, Default)]
112pub struct ResumeOptions {
113    /// Whether to reprocess failed items
114    pub reprocess_failed: bool,
115    /// Maximum parallel agents for resume
116    pub max_parallel: Option<usize>,
117    /// Skip validation of checkpoint
118    pub skip_validation: bool,
119    /// Custom timeout for resumed agents
120    pub agent_timeout_secs: Option<u64>,
121}
122
123/// Result of a resume operation
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct ResumeResult {
126    /// Job ID that was resumed
127    pub job_id: String,
128    /// Checkpoint version resumed from
129    pub resumed_from_version: u32,
130    /// Total number of work items
131    pub total_items: usize,
132    /// Number of already completed items
133    pub already_completed: usize,
134    /// Number of remaining items to process
135    pub remaining_items: usize,
136    /// Final results after resumption
137    pub final_results: Vec<crate::cook::execution::mapreduce::AgentResult>,
138}
139
140/// Context for agent execution
141#[derive(Clone)]
142pub struct AgentContext {
143    /// Unique identifier for this agent
144    pub item_id: String,
145    /// Path to the agent's isolated worktree
146    pub worktree_path: PathBuf,
147    /// Name of the agent's worktree
148    pub worktree_name: String,
149    /// Variables available for interpolation
150    pub variables: HashMap<String, String>,
151    /// Last shell command output
152    pub shell_output: Option<String>,
153    /// Environment for command execution
154    pub environment: ExecutionEnvironment,
155    /// Current retry count for failed commands
156    pub retry_count: u32,
157    /// Captured outputs from previous steps
158    pub captured_outputs: HashMap<String, String>,
159    /// Iteration-specific variables
160    pub iteration_vars: HashMap<String, String>,
161    /// Variable store for structured capture data
162    pub variable_store: crate::cook::workflow::variables::VariableStore,
163}
164
165impl AgentContext {
166    /// Create a new agent context
167    pub fn new(
168        item_id: String,
169        worktree_path: PathBuf,
170        worktree_name: String,
171        environment: ExecutionEnvironment,
172    ) -> Self {
173        Self {
174            item_id,
175            worktree_path,
176            worktree_name,
177            variables: HashMap::new(),
178            shell_output: None,
179            environment,
180            retry_count: 0,
181            captured_outputs: HashMap::new(),
182            iteration_vars: HashMap::new(),
183            variable_store: crate::cook::workflow::variables::VariableStore::new(),
184        }
185    }
186
187    /// Update context with command output
188    pub fn update_with_output(&mut self, output: Option<String>) {
189        if let Some(out) = output {
190            self.variables
191                .insert("shell.output".to_string(), out.clone());
192            self.variables
193                .insert("shell.last_output".to_string(), out.clone());
194            self.shell_output = Some(out);
195        }
196    }
197
198    /// Convert to InterpolationContext
199    pub fn to_interpolation_context(&self) -> InterpolationContext {
200        let mut context = InterpolationContext::new();
201
202        for (key, value) in &self.variables {
203            context.set(key.clone(), Value::String(value.as_str().into()));
204        }
205
206        if let Some(ref output) = self.shell_output {
207            context.set(
208                "shell",
209                json!({
210                    "output": output,
211                    "last_output": output
212                }),
213            );
214        }
215
216        for (key, value) in &self.captured_outputs {
217            context.set(key.clone(), Value::String(value.as_str().into()));
218        }
219
220        for (key, value) in &self.iteration_vars {
221            context.set(key.clone(), Value::String(value.as_str().into()));
222        }
223
224        context
225    }
226
227    /// Convert to enhanced variable context
228    pub async fn to_variable_context(&self) -> VariableContext {
229        let mut context = VariableContext::new();
230
231        for (key, value) in &self.variables {
232            if key.starts_with("map.") {
233                if let Ok(num) = value.parse::<f64>() {
234                    context.set_phase(
235                        key.clone(),
236                        Variable::Static(Value::Number(
237                            serde_json::Number::from_f64(num).unwrap_or(0.into()),
238                        )),
239                    );
240                } else {
241                    context.set_phase(key.clone(), Variable::Static(Value::String(value.clone())));
242                }
243            } else {
244                context.set_phase(key.clone(), Variable::Static(Value::String(value.clone())));
245            }
246        }
247
248        let store_vars = self.variable_store.get_all().await;
249        for (key, captured_value) in store_vars {
250            let value = captured_value.to_json();
251
252            if key.starts_with("map.") {
253                context.set_phase(key.clone(), Variable::Static(value));
254            } else {
255                context.set_local(key.clone(), Variable::Static(value));
256            }
257        }
258
259        if let Some(ref output) = self.shell_output {
260            context.set_phase(
261                "shell",
262                Variable::Static(json!({
263                    "output": output,
264                    "last_output": output
265                })),
266            );
267        }
268
269        for (key, value) in &self.captured_outputs {
270            context.set_local(key.clone(), Variable::Static(Value::String(value.clone())));
271        }
272
273        for (key, value) in &self.iteration_vars {
274            context.set_local(key.clone(), Variable::Static(Value::String(value.clone())));
275        }
276
277        context.set_local(
278            "workflow",
279            Variable::Static(json!({
280                "id": self.item_id.clone(),
281                "worktree": Value::String(self.worktree_name.clone()),
282                "path": self.worktree_path.to_string_lossy()
283            })),
284        );
285
286        context
287    }
288}
289
290/// Trait for components that can be initialized with configuration
291pub trait Configurable {
292    type Config;
293
294    /// Initialize with configuration
295    fn configure(config: Self::Config) -> Self;
296
297    /// Get current configuration
298    fn configuration(&self) -> &Self::Config;
299}
300
301/// Trait for components that can be reset to initial state
302pub trait Resettable {
303    /// Reset to initial state
304    fn reset(&mut self);
305
306    /// Check if component needs reset
307    fn needs_reset(&self) -> bool;
308}
prodigy/cook/execution/mapreduce/types.rs

prodigy/cook/execution/mapreduce/
types.rs
