1use std::sync::{Arc, Condvar, Mutex};
2use std::time::Duration;
3
4use crate::checkpoint::{CheckpointConfig, EventCursor, IdempotentRunEventStore};
5use crate::checkpoint::{
6 CheckpointStatus, ClaimMode, OperationKind, OperationState, ReconciliationDecision,
7 ReconciliationDecisionKind, ResumeObservation, ToolIdempotency,
8};
9use crate::event_store::{EventStoreError, RunEventIter, RunEventReplayQuery, RunEventStore};
10use crate::events::RunEvent;
11use crate::runtime::checkpoint_resume::{
12 CheckpointControllerRequest, CheckpointEventSink, CheckpointResumeController,
13};
14use crate::runtime::run_definition::validate_distributed_run_definition;
15use crate::runtime::state::{
16 validate_extension_state_size, Checkpoint, CheckpointStore, ExtensionStateEntry, OperationError,
17};
18use crate::runtime::tool_planner::project_tool_policy;
19use crate::runtime::{CheckpointRuntimeControl, ExecutionContext, RuntimeRunControls};
20use crate::types::AgentResult;
21use crate::types::AgentStatus;
22use crate::{ModelRef, RunContext};
23
24use super::capabilities::ResolvedDistributedCapabilities;
25use super::contract::{now_unix_ms, DistributedCheckpointConfig, DistributedRunEnvelope};
26use super::dispatch::CycleDispatchResult;
27use super::worker::{
28 build_runtime, combined_event_handler, lease_expiry_at, DistributedCycleWorker,
29 LeaseCommitPhase, LeaseHeartbeatStatus, LeaseHeartbeatStopGuard, LeaseOperationResult,
30 LeaseRenewal, LeaseRenewalFailure, LeaseRenewalFailureKind,
31};
32
33mod lease;
34mod recovery;
35mod runtime;
36
37use lease::run_with_checkpoint_lease;
38use recovery::{
39 align_active_claim, commit_cycle, effective_claim_mode, initialize_extensions, load_checkpoint,
40 prepare_terminal_candidate, reconcile_recovery, reconciliation_candidate, snapshot_extensions,
41 suspend_reconciliation, terminal_replay, validate_claimed_resume_attempt,
42 validate_envelope_checkpoint_identity, validate_extension_capabilities,
43 validate_resume_attempt_observation,
44};
45use runtime::run_agent_runtime_cycle;
46
47#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
48pub struct DistributedDeliveryMetadata {
49 pub redelivered: bool,
50 pub attempt: u64,
51}
52
53impl DistributedDeliveryMetadata {
54 pub fn redelivery(attempt: u64) -> Self {
55 Self {
56 redelivered: true,
57 attempt,
58 }
59 }
60
61 pub fn is_redelivery(self) -> bool {
62 self.redelivered || self.attempt > 1
63 }
64}
65
66pub trait DistributedCycleExecutor: Send + Sync {
67 fn execute(
68 &self,
69 envelope: &DistributedRunEnvelope,
70 capabilities: &ResolvedDistributedCapabilities,
71 checkpoint: &mut DistributedCheckpointProgress,
72 ) -> Result<DistributedCycleOutcome, String>;
73}
74
75#[derive(Debug, Clone, PartialEq)]
76pub enum DistributedCycleOutcome {
77 Continue(Checkpoint),
78 ReconciliationRequired(Checkpoint),
79 Terminal(Checkpoint),
80}
81
82pub struct DistributedCheckpointProgress {
83 store: Arc<dyn CheckpointStore>,
84 claim_token: String,
85 checkpoint: Checkpoint,
86}
87
88impl DistributedCheckpointProgress {
89 fn new(store: Arc<dyn CheckpointStore>, claim_token: String, checkpoint: Checkpoint) -> Self {
90 Self {
91 store,
92 claim_token,
93 checkpoint,
94 }
95 }
96
97 pub fn checkpoint(&self) -> &Checkpoint {
98 &self.checkpoint
99 }
100
101 pub fn claim_token(&self) -> &str {
102 &self.claim_token
103 }
104
105 pub fn persist(&mut self, mut snapshot: Checkpoint) -> Result<Checkpoint, String> {
106 align_active_claim(&mut snapshot, &self.checkpoint);
107 let expected_revision = self.checkpoint.revision;
108 if !self
109 .store
110 .progress_checkpoint(snapshot, &self.claim_token, expected_revision)
111 .map_err(|error| error.to_string())?
112 {
113 return Err(format!(
114 "checkpoint progress conflict at revision {expected_revision} for {}",
115 self.checkpoint.checkpoint_key
116 ));
117 }
118 self.reload()?;
119 Ok(self.checkpoint.clone())
120 }
121
122 fn reload(&mut self) -> Result<(), String> {
123 let checkpoint = self
124 .store
125 .load_checkpoint(&self.checkpoint.checkpoint_key)
126 .map_err(|error| error.to_string())?
127 .ok_or_else(|| {
128 format!(
129 "No checkpoint found for key {}",
130 self.checkpoint.checkpoint_key
131 )
132 })?;
133 if checkpoint.claim_token.as_deref() != Some(self.claim_token.as_str()) {
134 return Err(format!(
135 "checkpoint claim changed while progressing {}",
136 checkpoint.checkpoint_key
137 ));
138 }
139 self.checkpoint = checkpoint;
140 Ok(())
141 }
142}
143
144enum RecoveryDisposition {
145 Continue,
146 Suspend,
147 Abort(Box<Checkpoint>),
148}
149
150enum PostCommitAction {
151 Unfinished {
152 revision: u64,
153 cycle_index: u64,
154 },
155 TerminalCandidate {
156 result: Box<AgentResult>,
157 revision: u64,
158 },
159}
160
161pub(super) fn run_distributed_cycle(
162 worker: &DistributedCycleWorker,
163 envelope: DistributedRunEnvelope,
164 delivery: DistributedDeliveryMetadata,
165) -> Result<CycleDispatchResult, String> {
166 let checkpoint_store_ref = envelope
167 .recipe
168 .capabilities
169 .checkpoint_store_ref
170 .as_ref()
171 .ok_or_else(|| "distributed execution requires checkpoint_store_ref".to_string())?;
172 let store = worker
173 .capabilities
174 .resolve_checkpoint_store_required(checkpoint_store_ref)
175 .map_err(|error| error.to_string())?;
176 let config = &envelope.checkpoint_config;
177 let checkpoint_key = config.key.as_str();
178 let now_ms = now_unix_ms()?;
179 let checkpoint = load_checkpoint(store.as_ref(), checkpoint_key)?;
180 validate_envelope_checkpoint_identity(&envelope, &checkpoint)?;
181 validate_distributed_run_definition(&envelope, &checkpoint, None)
182 .map_err(|error| error.to_string())?;
183
184 if checkpoint.terminal_result.is_some() {
185 return terminal_replay(&checkpoint);
186 }
187 if checkpoint.cycle_index >= u64::from(envelope.cycle_index) && checkpoint.claim_token.is_none()
188 {
189 return CycleDispatchResult::committed(checkpoint.cycle_index, checkpoint.revision);
190 }
191 validate_resume_attempt_observation(&envelope, &checkpoint, delivery)?;
192 if checkpoint
193 .lease_expires_at_ms
194 .is_some_and(|lease| lease > now_ms)
195 {
196 return Ok(CycleDispatchResult::pending());
197 }
198
199 let resolved = worker
200 .capabilities
201 .resolve(&envelope.recipe.capabilities)
202 .map_err(|error| error.to_string())?;
203 validate_distributed_run_definition(&envelope, &checkpoint, Some(&resolved))
204 .map_err(|error| error.to_string())?;
205
206 validate_extension_capabilities(config, &resolved)?;
207 if worker.checkpoint_executor.is_none() {
208 return run_agent_runtime_cycle(envelope, delivery, resolved, store, checkpoint);
209 }
210 let executor = worker
211 .checkpoint_executor
212 .clone()
213 .expect("checkpoint executor checked above");
214
215 let claim_mode = effective_claim_mode(&envelope, &checkpoint, delivery, now_ms);
216 let lease_expires_at_ms = lease_expiry_at(
217 now_ms,
218 envelope.lease_duration_ms,
219 envelope.deadline_unix_ms,
220 )?;
221 let claim_token = uuid::Uuid::new_v4().simple().to_string();
222 let resume_attempt_before_claim = checkpoint.resume_attempt;
223 let Some(claimed) = store
224 .claim_checkpoint(
225 checkpoint_key,
226 u64::from(envelope.cycle_index),
227 &claim_token,
228 lease_expires_at_ms,
229 now_ms,
230 claim_mode,
231 )
232 .map_err(|error| format!("retryable distributed delivery conflict: {error}"))?
233 else {
234 let latest = load_checkpoint(store.as_ref(), checkpoint_key)?;
235 validate_envelope_checkpoint_identity(&envelope, &latest)?;
236 if latest.terminal_result.is_some() {
237 return terminal_replay(&latest);
238 }
239 return Ok(CycleDispatchResult::pending());
240 };
241 validate_claimed_resume_attempt(resume_attempt_before_claim, &claimed, claim_mode)?;
242
243 let action = run_with_checkpoint_lease(
244 store.clone(),
245 checkpoint_key,
246 u64::from(envelope.cycle_index),
247 envelope.lease_duration_ms,
248 envelope.deadline_unix_ms,
249 &claim_token,
250 |heartbeat_status| {
251 let result = (|| -> Result<PostCommitAction, String> {
252 let mut progress =
253 DistributedCheckpointProgress::new(store.clone(), claim_token.clone(), claimed);
254 initialize_extensions(config, &resolved, &mut progress)?;
255
256 if claim_mode == ClaimMode::Recovery {
257 match reconcile_recovery(config, &resolved, &mut progress)? {
258 RecoveryDisposition::Continue => {}
259 RecoveryDisposition::Suspend => {
260 suspend_reconciliation(&mut progress, heartbeat_status)?;
261 let checkpoint = load_checkpoint(store.as_ref(), checkpoint_key)?;
262 let result = reconciliation_candidate(&checkpoint)?;
263 return Ok(PostCommitAction::TerminalCandidate {
264 result: Box::new(result),
265 revision: checkpoint.revision,
266 });
267 }
268 RecoveryDisposition::Abort(checkpoint) => {
269 let (result, revision) = prepare_terminal_candidate(
270 *checkpoint,
271 &mut progress,
272 u64::from(envelope.cycle_index),
273 )?;
274 return Ok(PostCommitAction::TerminalCandidate {
275 result: Box::new(result),
276 revision,
277 });
278 }
279 }
280 }
281
282 let outcome = executor.execute(&envelope, &resolved, &mut progress)?;
283 match outcome {
284 DistributedCycleOutcome::Continue(mut checkpoint) => {
285 snapshot_extensions(config, &resolved, &mut checkpoint)?;
286 commit_cycle(
287 checkpoint,
288 &mut progress,
289 heartbeat_status,
290 u64::from(envelope.cycle_index),
291 )?;
292 let committed = load_checkpoint(store.as_ref(), checkpoint_key)?;
293 Ok(PostCommitAction::Unfinished {
294 revision: committed.revision,
295 cycle_index: committed.cycle_index,
296 })
297 }
298 DistributedCycleOutcome::ReconciliationRequired(mut checkpoint) => {
299 snapshot_extensions(config, &resolved, &mut checkpoint)?;
300 align_active_claim(&mut checkpoint, &progress.checkpoint);
301 progress.checkpoint = checkpoint;
302 suspend_reconciliation(&mut progress, heartbeat_status)?;
303 let checkpoint = load_checkpoint(store.as_ref(), checkpoint_key)?;
304 let result = reconciliation_candidate(&checkpoint)?;
305 Ok(PostCommitAction::TerminalCandidate {
306 result: Box::new(result),
307 revision: checkpoint.revision,
308 })
309 }
310 DistributedCycleOutcome::Terminal(mut checkpoint) => {
311 snapshot_extensions(config, &resolved, &mut checkpoint)?;
312 let (result, revision) = prepare_terminal_candidate(
313 checkpoint,
314 &mut progress,
315 u64::from(envelope.cycle_index),
316 )?;
317 Ok(PostCommitAction::TerminalCandidate {
318 result: Box::new(result),
319 revision,
320 })
321 }
322 }
323 })();
324 let claim_committed = match &result {
325 Ok(PostCommitAction::Unfinished { .. }) => true,
326 Ok(PostCommitAction::TerminalCandidate { result, .. }) => {
327 result.status == crate::types::AgentStatus::ReconciliationRequired
328 }
329 Err(_) => false,
330 };
331 LeaseOperationResult::new(result, claim_committed)
332 },
333 )??;
334
335 match action {
336 PostCommitAction::Unfinished {
337 revision,
338 cycle_index,
339 } => CycleDispatchResult::committed(cycle_index, revision),
340 PostCommitAction::TerminalCandidate { result, revision } => {
341 CycleDispatchResult::terminal_candidate(*result, revision)
342 }
343 }
344}