1use crate::agents::AgentRegistry;
7use crate::checkpoint::execution_history::ExecutionHistory;
8use crate::checkpoint::file_state::FileSystemState;
9use crate::checkpoint::state::{
10 AgentConfigSnapshot, CheckpointParams, CliArgsSnapshot, PipelineCheckpoint, PipelinePhase,
11 RebaseState,
12};
13use crate::checkpoint::RunContext;
14use crate::config::{Config, ReviewDepth};
15use crate::executor::ProcessExecutor;
16use crate::logger::Logger;
17use crate::workspace::Workspace;
18use std::sync::Arc;
19
20pub struct CheckpointBuilder {
35 phase: Option<PipelinePhase>,
36 iteration: u32,
37 total_iterations: u32,
38 reviewer_pass: u32,
39 total_reviewer_passes: u32,
40 developer_agent: Option<String>,
41 reviewer_agent: Option<String>,
42 cli_args: Option<CliArgsSnapshot>,
43 developer_agent_config: Option<AgentConfigSnapshot>,
44 reviewer_agent_config: Option<AgentConfigSnapshot>,
45 rebase_state: RebaseState,
46 config_path: Option<std::path::PathBuf>,
47 git_user_name: Option<String>,
48 git_user_email: Option<String>,
49 run_context: Option<RunContext>,
51 execution_history: Option<ExecutionHistory>,
53 prompt_history: Option<std::collections::HashMap<String, String>>,
54 skip_rebase: Option<bool>,
56 executor: Option<Arc<dyn ProcessExecutor>>,
58}
59
60impl Default for CheckpointBuilder {
61 fn default() -> Self {
62 Self::new()
63 }
64}
65
66impl CheckpointBuilder {
67 pub fn new() -> Self {
69 Self {
70 phase: None,
71 iteration: 1,
72 total_iterations: 1,
73 reviewer_pass: 0,
74 total_reviewer_passes: 0,
75 developer_agent: None,
76 reviewer_agent: None,
77 cli_args: None,
78 developer_agent_config: None,
79 reviewer_agent_config: None,
80 rebase_state: RebaseState::default(),
81 config_path: None,
82 git_user_name: None,
83 git_user_email: None,
84 run_context: None,
85 execution_history: None,
86 prompt_history: None,
87 skip_rebase: None,
88 executor: None,
89 }
90 }
91
92 pub fn phase(mut self, phase: PipelinePhase, iteration: u32, total_iterations: u32) -> Self {
94 self.phase = Some(phase);
95 self.iteration = iteration;
96 self.total_iterations = total_iterations;
97 self
98 }
99
100 pub fn reviewer_pass(mut self, pass: u32, total: u32) -> Self {
102 self.reviewer_pass = pass;
103 self.total_reviewer_passes = total;
104 self
105 }
106
107 pub fn agents(mut self, developer: &str, reviewer: &str) -> Self {
109 self.developer_agent = Some(developer.to_string());
110 self.reviewer_agent = Some(reviewer.to_string());
111 self
112 }
113
114 pub fn cli_args(mut self, args: CliArgsSnapshot) -> Self {
116 self.cli_args = Some(args);
117 self
118 }
119
120 pub fn developer_config(mut self, config: AgentConfigSnapshot) -> Self {
122 self.developer_agent_config = Some(config);
123 self
124 }
125
126 pub fn reviewer_config(mut self, config: AgentConfigSnapshot) -> Self {
128 self.reviewer_agent_config = Some(config);
129 self
130 }
131
132 pub fn rebase_state(mut self, state: RebaseState) -> Self {
134 self.rebase_state = state;
135 self
136 }
137
138 pub fn config_path(mut self, path: Option<std::path::PathBuf>) -> Self {
140 self.config_path = path;
141 self
142 }
143
144 pub fn git_identity(mut self, name: Option<&str>, email: Option<&str>) -> Self {
146 self.git_user_name = name.map(String::from);
147 self.git_user_email = email.map(String::from);
148 self
149 }
150
151 pub fn skip_rebase(mut self, value: bool) -> Self {
153 self.skip_rebase = Some(value);
154 self
155 }
156
157 pub fn with_executor(mut self, executor: Arc<dyn ProcessExecutor>) -> Self {
159 self.executor = Some(executor);
160 self
161 }
162
163 pub fn capture_cli_args(mut self, config: &Config) -> Self {
165 let review_depth_str = review_depth_to_string(config.review_depth);
166 let skip_rebase = self.skip_rebase.unwrap_or(false);
167
168 let snapshot = crate::checkpoint::state::CliArgsSnapshotBuilder::new(
169 config.developer_iters,
170 config.reviewer_reviews,
171 review_depth_str,
172 skip_rebase,
173 config.isolation_mode,
174 )
175 .verbosity(config.verbosity as u8)
176 .show_streaming_metrics(config.show_streaming_metrics)
177 .reviewer_json_parser(config.reviewer_json_parser.clone())
178 .build();
179 self.cli_args = Some(snapshot);
180 self
181 }
182
183 pub fn capture_from_context(
188 mut self,
189 config: &Config,
190 registry: &AgentRegistry,
191 developer_name: &str,
192 reviewer_name: &str,
193 logger: &Logger,
194 run_context: &RunContext,
195 ) -> Self {
196 self.run_context = Some(run_context.clone());
198
199 self = self.capture_cli_args(config);
201
202 if let Some(agent_config) = registry.resolve_config(developer_name) {
204 let snapshot = AgentConfigSnapshot::new(
205 developer_name.to_string(),
206 agent_config.cmd.clone(),
207 agent_config.output_flag.clone(),
208 Some(agent_config.yolo_flag.clone()),
209 agent_config.can_commit,
210 )
211 .with_model_override(config.developer_model.clone())
212 .with_provider_override(config.developer_provider.clone())
213 .with_context_level(config.developer_context);
214 self.developer_agent_config = Some(snapshot);
215 self.developer_agent = Some(developer_name.to_string());
216 } else {
217 logger.warn(&format!(
218 "Developer agent '{}' not found in registry",
219 developer_name
220 ));
221 }
222
223 if let Some(agent_config) = registry.resolve_config(reviewer_name) {
225 let snapshot = AgentConfigSnapshot::new(
226 reviewer_name.to_string(),
227 agent_config.cmd.clone(),
228 agent_config.output_flag.clone(),
229 Some(agent_config.yolo_flag.clone()),
230 agent_config.can_commit,
231 )
232 .with_model_override(config.reviewer_model.clone())
233 .with_provider_override(config.reviewer_provider.clone())
234 .with_context_level(config.reviewer_context);
235 self.reviewer_agent_config = Some(snapshot);
236 self.reviewer_agent = Some(reviewer_name.to_string());
237 } else {
238 logger.warn(&format!(
239 "Reviewer agent '{}' not found in registry",
240 reviewer_name
241 ));
242 }
243
244 self.git_user_name = config.git_user_name.clone();
246 self.git_user_email = config.git_user_email.clone();
247
248 self
249 }
250
251 pub fn with_executor_from_context(mut self, executor_arc: Arc<dyn ProcessExecutor>) -> Self {
256 self.executor = Some(executor_arc);
257 self
258 }
259
260 pub fn with_execution_history(mut self, history: ExecutionHistory) -> Self {
265 self.execution_history = Some(history);
266 self
267 }
268
269 pub fn with_prompt_history(
277 mut self,
278 history: std::collections::HashMap<String, String>,
279 ) -> Self {
280 self.prompt_history = if history.is_empty() {
281 None
282 } else {
283 Some(history)
284 };
285 self
286 }
287
288 pub fn build(self) -> Option<PipelineCheckpoint> {
296 self.build_internal(None)
297 }
298
299 pub fn build_with_workspace(self, workspace: &dyn Workspace) -> Option<PipelineCheckpoint> {
309 self.build_internal(Some(workspace))
310 }
311
312 fn build_internal(self, workspace: Option<&dyn Workspace>) -> Option<PipelineCheckpoint> {
314 let phase = self.phase?;
315 let developer_agent = self.developer_agent?;
316 let reviewer_agent = self.reviewer_agent?;
317 let cli_args = self.cli_args?;
318 let developer_config = self.developer_agent_config?;
319 let reviewer_config = self.reviewer_agent_config?;
320
321 let git_user_name = self.git_user_name.as_deref();
322 let git_user_email = self.git_user_email.as_deref();
323
324 let run_context = self.run_context.unwrap_or_default();
326
327 let mut checkpoint = PipelineCheckpoint::from_params(CheckpointParams {
328 phase,
329 iteration: self.iteration,
330 total_iterations: self.total_iterations,
331 reviewer_pass: self.reviewer_pass,
332 total_reviewer_passes: self.total_reviewer_passes,
333 developer_agent: &developer_agent,
334 reviewer_agent: &reviewer_agent,
335 cli_args,
336 developer_agent_config: developer_config,
337 reviewer_agent_config: reviewer_config,
338 rebase_state: self.rebase_state,
339 git_user_name,
340 git_user_email,
341 run_id: &run_context.run_id,
342 parent_run_id: run_context.parent_run_id.as_deref(),
343 resume_count: run_context.resume_count,
344 actual_developer_runs: run_context.actual_developer_runs.max(self.iteration),
345 actual_reviewer_runs: run_context.actual_reviewer_runs.max(self.reviewer_pass),
346 });
347
348 if let Some(path) = self.config_path {
349 checkpoint = checkpoint.with_config(Some(path));
350 }
351
352 checkpoint.execution_history = self.execution_history;
354
355 checkpoint.prompt_history = self.prompt_history;
357
358 let executor_ref = self.executor.as_ref().map(|e| e.as_ref());
362 checkpoint.file_system_state = if let Some(ws) = workspace {
363 let executor = executor_ref.unwrap_or_else(|| {
364 static DEFAULT_EXECUTOR: std::sync::LazyLock<crate::executor::RealProcessExecutor> =
367 std::sync::LazyLock::new(crate::executor::RealProcessExecutor::new);
368 &*DEFAULT_EXECUTOR
369 });
370 Some(FileSystemState::capture_with_workspace(ws, executor))
371 } else {
372 Some(FileSystemState::capture_with_optional_executor_impl(
373 executor_ref,
374 ))
375 };
376
377 checkpoint.env_snapshot =
379 Some(crate::checkpoint::state::EnvironmentSnapshot::capture_current());
380
381 Some(checkpoint)
382 }
383}
384
385fn review_depth_to_string(depth: ReviewDepth) -> Option<String> {
387 match depth {
388 ReviewDepth::Standard => Some("standard".to_string()),
389 ReviewDepth::Comprehensive => Some("comprehensive".to_string()),
390 ReviewDepth::Security => Some("security".to_string()),
391 ReviewDepth::Incremental => Some("incremental".to_string()),
392 }
393}
394
395#[cfg(test)]
396mod tests {
397 use super::review_depth_to_string;
398 use crate::checkpoint::state::{AgentConfigSnapshot, CliArgsSnapshot};
399 use crate::checkpoint::CheckpointBuilder;
400 use crate::checkpoint::PipelinePhase;
401 use crate::config::ReviewDepth;
402
403 #[test]
404 fn test_builder_basic() {
405 let cli_args = CliArgsSnapshot::new(5, 2, None, false, true, 2, false, None);
406 let dev_config =
407 AgentConfigSnapshot::new("dev".into(), "cmd".into(), "-o".into(), None, true);
408 let rev_config =
409 AgentConfigSnapshot::new("rev".into(), "cmd".into(), "-o".into(), None, true);
410
411 let checkpoint = CheckpointBuilder::new()
412 .phase(PipelinePhase::Development, 2, 5)
413 .reviewer_pass(1, 2)
414 .agents("dev", "rev")
415 .cli_args(cli_args)
416 .developer_config(dev_config)
417 .reviewer_config(rev_config)
418 .build()
419 .unwrap();
420
421 assert_eq!(checkpoint.phase, PipelinePhase::Development);
422 assert_eq!(checkpoint.iteration, 2);
423 assert_eq!(checkpoint.total_iterations, 5);
424 assert_eq!(checkpoint.reviewer_pass, 1);
425 assert_eq!(checkpoint.total_reviewer_passes, 2);
426 }
427
428 #[test]
429 fn test_builder_missing_required_field() {
430 let result = CheckpointBuilder::new().build();
432 assert!(result.is_none());
433 }
434
435 #[test]
436 fn test_review_depth_to_string() {
437 assert_eq!(
438 review_depth_to_string(ReviewDepth::Standard),
439 Some("standard".to_string())
440 );
441 assert_eq!(
442 review_depth_to_string(ReviewDepth::Comprehensive),
443 Some("comprehensive".to_string())
444 );
445 assert_eq!(
446 review_depth_to_string(ReviewDepth::Security),
447 Some("security".to_string())
448 );
449 assert_eq!(
450 review_depth_to_string(ReviewDepth::Incremental),
451 Some("incremental".to_string())
452 );
453 }
454
455 #[test]
456 fn test_builder_with_prompt_history() {
457 let cli_args = CliArgsSnapshot::new(5, 2, None, false, true, 2, false, None);
458 let dev_config =
459 AgentConfigSnapshot::new("dev".into(), "cmd".into(), "-o".into(), None, true);
460 let rev_config =
461 AgentConfigSnapshot::new("rev".into(), "cmd".into(), "-o".into(), None, true);
462
463 let mut prompts = std::collections::HashMap::new();
464 prompts.insert(
465 "development_1".to_string(),
466 "Implement feature X".to_string(),
467 );
468
469 let checkpoint = CheckpointBuilder::new()
470 .phase(PipelinePhase::Development, 2, 5)
471 .reviewer_pass(1, 2)
472 .agents("dev", "rev")
473 .cli_args(cli_args)
474 .developer_config(dev_config)
475 .reviewer_config(rev_config)
476 .with_prompt_history(prompts)
477 .build()
478 .unwrap();
479
480 assert_eq!(checkpoint.phase, PipelinePhase::Development);
481 assert!(checkpoint.prompt_history.is_some());
482 let history = checkpoint.prompt_history.as_ref().unwrap();
483 assert_eq!(history.len(), 1);
484 assert_eq!(
485 history.get("development_1"),
486 Some(&"Implement feature X".to_string())
487 );
488 }
489
490 #[test]
491 fn test_builder_with_prompt_history_multiple() {
492 let cli_args = CliArgsSnapshot::new(5, 2, None, false, true, 2, false, None);
493 let dev_config =
494 AgentConfigSnapshot::new("dev".into(), "cmd".into(), "-o".into(), None, true);
495 let rev_config =
496 AgentConfigSnapshot::new("rev".into(), "cmd".into(), "-o".into(), None, true);
497
498 let mut prompts = std::collections::HashMap::new();
499 prompts.insert(
500 "development_1".to_string(),
501 "Implement feature X".to_string(),
502 );
503 prompts.insert("review_1".to_string(), "Review the changes".to_string());
504
505 let checkpoint = CheckpointBuilder::new()
506 .phase(PipelinePhase::Development, 2, 5)
507 .reviewer_pass(1, 2)
508 .agents("dev", "rev")
509 .cli_args(cli_args)
510 .developer_config(dev_config)
511 .reviewer_config(rev_config)
512 .with_prompt_history(prompts)
513 .build()
514 .unwrap();
515
516 assert!(checkpoint.prompt_history.is_some());
517 let history = checkpoint.prompt_history.as_ref().unwrap();
518 assert_eq!(history.len(), 2);
519 assert_eq!(
520 history.get("development_1"),
521 Some(&"Implement feature X".to_string())
522 );
523 assert_eq!(
524 history.get("review_1"),
525 Some(&"Review the changes".to_string())
526 );
527 }
528
529 #[cfg(feature = "test-utils")]
534 mod workspace_tests {
535 use super::*;
536 use crate::workspace::MemoryWorkspace;
537
538 #[test]
539 fn test_builder_with_workspace_captures_file_state() {
540 let workspace =
542 MemoryWorkspace::new_test().with_file("PROMPT.md", "# Test prompt content");
543
544 let cli_args = CliArgsSnapshot::new(5, 2, None, false, true, 2, false, None);
545 let dev_config =
546 AgentConfigSnapshot::new("dev".into(), "cmd".into(), "-o".into(), None, true);
547 let rev_config =
548 AgentConfigSnapshot::new("rev".into(), "cmd".into(), "-o".into(), None, true);
549
550 let checkpoint = CheckpointBuilder::new()
551 .phase(PipelinePhase::Development, 2, 5)
552 .reviewer_pass(1, 2)
553 .agents("dev", "rev")
554 .cli_args(cli_args)
555 .developer_config(dev_config)
556 .reviewer_config(rev_config)
557 .build_with_workspace(&workspace)
558 .unwrap();
559
560 assert!(checkpoint.file_system_state.is_some());
562 let fs_state = checkpoint.file_system_state.as_ref().unwrap();
563
564 assert!(fs_state.files.contains_key("PROMPT.md"));
566 let snapshot = &fs_state.files["PROMPT.md"];
567 assert!(snapshot.exists);
568 assert_eq!(snapshot.size, 21); }
570
571 #[test]
572 fn test_builder_with_workspace_captures_agent_files() {
573 let workspace = MemoryWorkspace::new_test()
575 .with_file("PROMPT.md", "# Test prompt")
576 .with_file(".agent/PLAN.md", "# Plan")
577 .with_file(".agent/ISSUES.md", "# Issues");
578
579 let cli_args = CliArgsSnapshot::new(5, 2, None, false, true, 2, false, None);
580 let dev_config =
581 AgentConfigSnapshot::new("dev".into(), "cmd".into(), "-o".into(), None, true);
582 let rev_config =
583 AgentConfigSnapshot::new("rev".into(), "cmd".into(), "-o".into(), None, true);
584
585 let checkpoint = CheckpointBuilder::new()
586 .phase(PipelinePhase::Review, 2, 5)
587 .reviewer_pass(1, 2)
588 .agents("dev", "rev")
589 .cli_args(cli_args)
590 .developer_config(dev_config)
591 .reviewer_config(rev_config)
592 .build_with_workspace(&workspace)
593 .unwrap();
594
595 let fs_state = checkpoint.file_system_state.as_ref().unwrap();
596
597 assert!(fs_state.files.contains_key(".agent/PLAN.md"));
599 assert!(fs_state.files.contains_key(".agent/ISSUES.md"));
600
601 let plan_snapshot = &fs_state.files[".agent/PLAN.md"];
602 assert!(plan_snapshot.exists);
603
604 let issues_snapshot = &fs_state.files[".agent/ISSUES.md"];
605 assert!(issues_snapshot.exists);
606 }
607 }
608}