1use std::sync::atomic::AtomicBool;
2use std::sync::{Arc, Mutex};
3
4use tokio::sync::broadcast;
5
6use crate::agent::Agent;
7use crate::result::RunResult;
8use crate::run_config::RunConfig;
9use crate::run_handle::{RunEventSenderSlot, RunHandle, RunHandleState, SharedRunResult};
10use crate::runtime::backends::{
11 DistributedAdvanceDecision, DistributedRunHandle, RuntimeExecutionBackend,
12};
13
14use super::{CheckpointAdmissionSender, NormalizedInput, RunEventStream, Runner};
15
16pub(crate) enum CheckpointStartOutcome {
17 Started {
18 handle: RunHandle,
19 checkpoint: crate::runtime::state::Checkpoint,
20 },
21 ExistingOwner {
22 checkpoint: crate::runtime::state::Checkpoint,
23 },
24 TerminalReplay {
25 result: Box<RunResult>,
26 checkpoint: crate::runtime::state::Checkpoint,
27 },
28}
29
30impl Runner {
31 pub async fn start_distributed(
32 &self,
33 agent: &Agent,
34 input: impl Into<NormalizedInput>,
35 config: RunConfig,
36 ) -> Result<DistributedRunHandle, String> {
37 self.validate_nonblocking_distributed_config(&config)?;
38 let runner = self.clone();
39 let agent = agent.clone();
40 let input = input.into();
41 tokio::task::spawn_blocking(move || {
42 match runner.run_single_agent_operation(
43 &agent,
44 input,
45 config,
46 Some(Arc::new(Mutex::new(Vec::new()))),
47 None,
48 None,
49 None,
50 Some(super::DistributedRunnerOperation::Start),
51 )? {
52 super::SingleRunExecutionOutcome::DistributedStarted(handle) => Ok(handle),
53 super::SingleRunExecutionOutcome::Completed(_) => {
54 Err("distributed start completed before returning a passive handle".to_string())
55 }
56 }
57 })
58 .await
59 .map_err(|error| format!("distributed start task failed: {error}"))?
60 }
61
62 pub async fn finalize_distributed(
63 &self,
64 agent: &Agent,
65 input: impl Into<NormalizedInput>,
66 decision: DistributedAdvanceDecision,
67 config: RunConfig,
68 ) -> Result<RunResult, String> {
69 self.validate_nonblocking_distributed_config(&config)?;
70 let handle = match &decision {
71 DistributedAdvanceDecision::FinalizeRequired { handle, .. } => handle,
72 _ => {
73 return Err(
74 "distributed finalization requires a FinalizeRequired decision".to_string(),
75 )
76 }
77 };
78 let checkpoint_config = config
79 .checkpoint_config
80 .as_ref()
81 .or(self.default_run_config.checkpoint_config.as_ref())
82 .ok_or_else(|| {
83 "checkpoint_config_invalid: distributed finalization requires checkpoint configuration"
84 .to_string()
85 })?;
86 if checkpoint_config.key.as_deref() != Some(handle.checkpoint_key.as_str()) {
87 return Err(
88 "distributed finalization checkpoint does not match the run handle".to_string(),
89 );
90 }
91 let runner = self.clone();
92 let agent = agent.clone();
93 let input = input.into();
94 tokio::task::spawn_blocking(move || {
95 match runner.run_single_agent_operation(
96 &agent,
97 input,
98 config,
99 Some(Arc::new(Mutex::new(Vec::new()))),
100 None,
101 None,
102 None,
103 Some(super::DistributedRunnerOperation::Finalize(Box::new(
104 decision,
105 ))),
106 )? {
107 super::SingleRunExecutionOutcome::Completed(outcome) => Ok(outcome.result),
108 super::SingleRunExecutionOutcome::DistributedStarted(_) => {
109 Err("distributed finalization returned a passive start handle".to_string())
110 }
111 }
112 })
113 .await
114 .map_err(|error| format!("distributed finalization task failed: {error}"))?
115 }
116
117 fn validate_nonblocking_distributed_config(&self, config: &RunConfig) -> Result<(), String> {
118 let checkpoint_config = config
119 .checkpoint_config
120 .as_ref()
121 .or(self.default_run_config.checkpoint_config.as_ref())
122 .ok_or_else(|| {
123 "checkpoint_config_invalid: nonblocking distributed runs require checkpoint configuration"
124 .to_string()
125 })?;
126 checkpoint_config
127 .validate()
128 .map_err(|error| error.to_string())?;
129 if checkpoint_config.key.is_none() {
130 return Err(
131 "checkpoint_key_required: nonblocking distributed runs require an explicit checkpoint key"
132 .to_string(),
133 );
134 }
135 let backend = config
136 .execution_backend
137 .as_ref()
138 .or(self.default_run_config.execution_backend.as_ref());
139 if !matches!(
140 backend,
141 Some(RuntimeExecutionBackend::Distributed(backend)) if backend.has_nonblocking_driver()
142 ) {
143 return Err(
144 "nonblocking distributed runs require an enqueue-only DistributedBackend"
145 .to_string(),
146 );
147 }
148 if config.approval_provider.is_some()
149 || config.approval_broker.is_some()
150 || self.default_run_config.approval_provider.is_some()
151 || self.default_run_config.approval_broker.is_some()
152 {
153 return Err(
154 "nonblocking distributed runs do not support brokered approval waits".to_string(),
155 );
156 }
157 Ok(())
158 }
159
160 pub async fn stream(
161 &self,
162 agent: &Agent,
163 input: impl Into<NormalizedInput>,
164 ) -> Result<RunEventStream, String> {
165 self.stream_with_config(agent, input, RunConfig::default())
166 .await
167 }
168
169 pub async fn stream_with_config(
170 &self,
171 agent: &Agent,
172 input: impl Into<NormalizedInput>,
173 config: RunConfig,
174 ) -> Result<RunEventStream, String> {
175 let handle = self.start(agent, input, config).await?;
176 Ok(handle.into_event_stream())
177 }
178
179 pub async fn start(
180 &self,
181 agent: &Agent,
182 input: impl Into<NormalizedInput>,
183 config: RunConfig,
184 ) -> Result<RunHandle, String> {
185 self.start_internal(agent, input.into(), config, None).await
186 }
187
188 pub(crate) async fn start_checkpointed(
189 &self,
190 agent: &Agent,
191 input: impl Into<NormalizedInput>,
192 config: RunConfig,
193 ) -> Result<CheckpointStartOutcome, String> {
194 let checkpoint_config = config
195 .checkpoint_config
196 .clone()
197 .or_else(|| self.default_run_config.checkpoint_config.clone())
198 .ok_or_else(|| {
199 "checkpoint_config_invalid: start_checkpointed requires checkpoint_config"
200 .to_string()
201 })?;
202 checkpoint_config
203 .validate()
204 .map_err(|error| error.to_string())?;
205 let store = checkpoint_config.store.clone().ok_or_else(|| {
206 "checkpoint_store_unavailable: start_checkpointed requires a process-local store"
207 .to_string()
208 })?;
209 let checkpoint_key = checkpoint_config.key.clone().ok_or_else(|| {
210 "checkpoint_key_required: start_checkpointed requires an explicit key".to_string()
211 })?;
212 if let Some(checkpoint) = store
213 .load_checkpoint(&checkpoint_key)
214 .map_err(|error| error.to_string())?
215 {
216 if checkpoint.terminal_result.is_none()
217 && checkpoint
218 .lease_expires_at_ms
219 .is_some_and(|expires_at| expires_at > unix_time_ms())
220 {
221 return Ok(CheckpointStartOutcome::ExistingOwner { checkpoint });
222 }
223 }
224
225 let (admission_sender, admission_receiver) = tokio::sync::oneshot::channel();
226 let handle = self
227 .start_internal(agent, input.into(), config, Some(admission_sender))
228 .await?;
229 match admission_receiver.await {
230 Ok(admission) if admission.terminal_replayed => {
231 let result = handle.result().await?;
232 let checkpoint = store
233 .load_checkpoint(&checkpoint_key)
234 .map_err(|error| error.to_string())?
235 .ok_or_else(|| {
236 "checkpoint_not_found: terminal checkpoint disappeared".to_string()
237 })?;
238 Ok(CheckpointStartOutcome::TerminalReplay {
239 result: Box::new(result),
240 checkpoint,
241 })
242 }
243 Ok(admission) => Ok(CheckpointStartOutcome::Started {
244 handle,
245 checkpoint: admission.checkpoint,
246 }),
247 Err(_) => {
248 let result = handle.result().await;
249 if let Some(checkpoint) = store
250 .load_checkpoint(&checkpoint_key)
251 .map_err(|error| error.to_string())?
252 {
253 if checkpoint.terminal_result.is_none()
254 && checkpoint
255 .lease_expires_at_ms
256 .is_some_and(|expires_at| expires_at > unix_time_ms())
257 {
258 return Ok(CheckpointStartOutcome::ExistingOwner { checkpoint });
259 }
260 }
261 match result {
262 Ok(_) => Err(
263 "checkpoint_admission_missing: checkpointed run completed without admission"
264 .to_string(),
265 ),
266 Err(error) => Err(error),
267 }
268 }
269 }
270 }
271
272 async fn start_internal(
273 &self,
274 agent: &Agent,
275 input: NormalizedInput,
276 mut config: RunConfig,
277 checkpoint_admission_sender: Option<CheckpointAdmissionSender>,
278 ) -> Result<RunHandle, String> {
279 let cancellation_token = config
280 .cancellation_token
281 .clone()
282 .or_else(|| self.default_run_config.cancellation_token.clone())
283 .unwrap_or_default();
284 config.cancellation_token = Some(cancellation_token.clone());
285 let approval_broker = config
286 .approval_broker
287 .clone()
288 .or_else(|| self.default_run_config.approval_broker.clone())
289 .unwrap_or_default();
290 config.approval_broker = Some(approval_broker.clone());
291
292 let (event_sender, _) = broadcast::channel(1024);
293 let event_collector = Arc::new(Mutex::new(Vec::new()));
294 let event_sender_slot: RunEventSenderSlot =
295 Arc::new(Mutex::new(Some(event_sender.clone())));
296 let state = Arc::new(Mutex::new(RunHandleState::running()));
297 let cancel_requested = Arc::new(AtomicBool::new(false));
298 let (completion_sender, completion_receiver) = tokio::sync::watch::channel(false);
299 let runner = self.clone();
300 let agent = agent.clone();
301 let state_for_task = state.clone();
302 let event_collector_for_task = event_collector.clone();
303 let cancellation_token_for_task = cancellation_token.clone();
304 let join = tokio::task::spawn_blocking(move || {
305 struct CompletionGuard {
306 sender: Option<tokio::sync::watch::Sender<bool>>,
307 }
308
309 impl Drop for CompletionGuard {
310 fn drop(&mut self) {
311 if let Some(sender) = self.sender.take() {
312 let _ = sender.send(true);
313 }
314 }
315 }
316
317 let _completion = CompletionGuard {
318 sender: Some(completion_sender),
319 };
320 let result = runner.run_blocking_with_event_sender(
321 &agent,
322 input,
323 config,
324 Some(event_collector_for_task),
325 Some(event_sender),
326 checkpoint_admission_sender,
327 );
328 if let Ok(mut state) = state_for_task.lock() {
329 *state = match &result {
330 Ok(result) if run_result_was_cancelled(result) => {
331 RunHandleState::cancelled_with_reason(
332 result
333 .result()
334 .error
335 .clone()
336 .unwrap_or_else(|| "Operation was cancelled".to_string()),
337 )
338 }
339 Ok(result) => RunHandleState::from_run_result(result),
340 Err(error)
341 if cancellation_token_for_task.is_cancelled()
342 && error.to_ascii_lowercase().contains("cancel") =>
343 {
344 let mut state = RunHandleState::cancelled();
345 state.error = Some(error.clone());
346 state
347 }
348 Err(error) => RunHandleState::failed(error.clone()),
349 };
350 }
351 result
352 });
353 let result = SharedRunResult::new(join);
354 Ok(RunHandle::new(
355 event_sender_slot,
356 event_collector,
357 result,
358 state,
359 cancellation_token,
360 approval_broker,
361 completion_receiver,
362 cancel_requested,
363 ))
364 }
365}
366
367fn unix_time_ms() -> u64 {
368 use std::time::{SystemTime, UNIX_EPOCH};
369
370 SystemTime::now()
371 .duration_since(UNIX_EPOCH)
372 .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
373 .unwrap_or(0)
374}
375
376fn run_result_was_cancelled(result: &RunResult) -> bool {
377 result.status() == crate::types::AgentStatus::Failed
378 && result
379 .result()
380 .error
381 .as_deref()
382 .is_some_and(|error| error.to_ascii_lowercase().contains("cancel"))
383}