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