1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::{Arc, Mutex};
3
4use tokio::sync::broadcast;
5use tokio::task::JoinHandle;
6
7use crate::approval::{ApprovalBroker, ApprovalError};
8use crate::events::{RunEvent, RunEventPayload};
9use crate::result::{RunResult, RunState};
10use crate::runner::{NormalizedInput, RunEventStream};
11use crate::runtime::CancellationToken;
12use crate::tools::ApprovalDecision;
13use crate::types::AgentStatus;
14
15pub(crate) type RunEventSenderSlot = Arc<Mutex<Option<broadcast::Sender<RunEvent>>>>;
16pub(crate) type RunCompletionSignal = tokio::sync::watch::Receiver<bool>;
17type RunJoinHandle = JoinHandle<Result<RunResult, String>>;
18type SharedJoinHandle = Arc<tokio::sync::Mutex<Option<RunJoinHandle>>>;
19type CachedRunResult = Arc<tokio::sync::OnceCell<Result<RunResult, String>>>;
20
21pub(crate) trait RunHandleController: Send + Sync {
22 fn steer(&self, message: String) -> Result<(), String>;
23 fn follow_up(&self, message: String) -> Result<(), String>;
24}
25
26#[derive(Default)]
27struct ControllerBinding {
28 generation: u64,
29 active: Option<(u64, Arc<dyn RunHandleController>)>,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum RunHandleStatus {
34 Running,
35 Completed,
36 Failed,
37 Cancelled,
38 WaitUser,
39 MaxCycles,
40 ReconciliationRequired,
41}
42
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct RunHandleState {
45 pub status: RunHandleStatus,
46 pub done: bool,
47 pub cancelled: bool,
48 pub error: Option<String>,
49}
50
51impl RunHandleState {
52 pub fn running() -> Self {
53 Self {
54 status: RunHandleStatus::Running,
55 done: false,
56 cancelled: false,
57 error: None,
58 }
59 }
60
61 pub fn completed() -> Self {
62 Self {
63 status: RunHandleStatus::Completed,
64 done: true,
65 cancelled: false,
66 error: None,
67 }
68 }
69
70 pub fn from_agent_status(status: AgentStatus) -> Self {
71 match status {
72 AgentStatus::Completed => Self::completed(),
73 AgentStatus::WaitUser => Self {
74 status: RunHandleStatus::WaitUser,
75 done: true,
76 cancelled: false,
77 error: None,
78 },
79 AgentStatus::MaxCycles => Self {
80 status: RunHandleStatus::MaxCycles,
81 done: true,
82 cancelled: false,
83 error: None,
84 },
85 AgentStatus::Failed => Self {
86 status: RunHandleStatus::Failed,
87 done: true,
88 cancelled: false,
89 error: None,
90 },
91 AgentStatus::ReconciliationRequired => Self {
92 status: RunHandleStatus::ReconciliationRequired,
93 done: true,
94 cancelled: false,
95 error: None,
96 },
97 AgentStatus::Pending | AgentStatus::Running => Self::running(),
98 }
99 }
100
101 pub(crate) fn from_run_result(result: &RunResult) -> Self {
102 match result.status() {
103 AgentStatus::Failed => Self {
104 status: RunHandleStatus::Failed,
105 done: true,
106 cancelled: false,
107 error: result.result().error.clone(),
108 },
109 status => Self::from_agent_status(status),
110 }
111 }
112
113 pub fn failed(error: impl Into<String>) -> Self {
114 Self {
115 status: RunHandleStatus::Failed,
116 done: true,
117 cancelled: false,
118 error: Some(error.into()),
119 }
120 }
121
122 pub fn cancelled() -> Self {
123 Self::cancelled_with_reason("run cancelled")
124 }
125
126 pub fn cancelled_with_reason(reason: impl Into<String>) -> Self {
127 Self {
128 status: RunHandleStatus::Cancelled,
129 done: true,
130 cancelled: true,
131 error: Some(reason.into()),
132 }
133 }
134}
135
136#[derive(Clone)]
137pub(crate) struct SharedRunResult {
138 join: SharedJoinHandle,
139 cached: CachedRunResult,
140}
141
142impl SharedRunResult {
143 pub(crate) fn new(join: RunJoinHandle) -> Self {
144 Self {
145 join: Arc::new(tokio::sync::Mutex::new(Some(join))),
146 cached: Arc::new(tokio::sync::OnceCell::new()),
147 }
148 }
149
150 pub(crate) async fn wait(&self) -> Result<RunResult, String> {
151 self.cached
152 .get_or_init(|| async {
153 let join = self
154 .join
155 .lock()
156 .await
157 .take()
158 .ok_or_else(|| "run result task is unavailable".to_string())?;
159 join.await
160 .map_err(|error| format!("run task failed: {error}"))?
161 })
162 .await
163 .clone()
164 }
165}
166
167#[derive(Clone)]
168pub struct RunHandle {
169 sender: RunEventSenderSlot,
170 events: Arc<Mutex<Vec<RunEvent>>>,
171 result: SharedRunResult,
172 state: Arc<Mutex<RunHandleState>>,
173 cancellation_token: CancellationToken,
174 approval_broker: ApprovalBroker,
175 completion: RunCompletionSignal,
176 cancel_requested: Arc<AtomicBool>,
177 controller: Arc<Mutex<ControllerBinding>>,
178}
179
180impl RunHandle {
181 #[allow(clippy::too_many_arguments)]
182 pub(crate) fn new(
183 sender: RunEventSenderSlot,
184 events: Arc<Mutex<Vec<RunEvent>>>,
185 result: SharedRunResult,
186 state: Arc<Mutex<RunHandleState>>,
187 cancellation_token: CancellationToken,
188 approval_broker: ApprovalBroker,
189 completion: RunCompletionSignal,
190 cancel_requested: Arc<AtomicBool>,
191 ) -> Self {
192 Self {
193 sender,
194 events,
195 result,
196 state,
197 cancellation_token,
198 approval_broker,
199 completion,
200 cancel_requested,
201 controller: Arc::new(Mutex::new(ControllerBinding::default())),
202 }
203 }
204
205 pub fn events(&self) -> RunEventStream {
206 let receiver = self
207 .sender
208 .lock()
209 .ok()
210 .and_then(|sender| sender.as_ref().map(broadcast::Sender::subscribe));
211 RunEventStream::from_live(
212 receiver,
213 Some(self.result.clone()),
214 self.events.clone(),
215 self.completion.clone(),
216 )
217 }
218
219 pub(crate) fn into_event_stream(self) -> RunEventStream {
220 let receiver = self
221 .sender
222 .lock()
223 .ok()
224 .and_then(|sender| sender.as_ref().map(broadcast::Sender::subscribe));
225 RunEventStream::from_live(
226 receiver,
227 Some(self.result.clone()),
228 self.events.clone(),
229 self.completion.clone(),
230 )
231 }
232
233 pub async fn result(&self) -> Result<RunResult, String> {
234 self.result.wait().await
235 }
236
237 pub fn state(&self) -> RunHandleState {
238 let events = self
239 .events
240 .lock()
241 .unwrap_or_else(std::sync::PoisonError::into_inner);
242 let has_active_sub_runs = !active_sub_run_ids(&events).is_empty();
243 drop(events);
244 let mut state = self
245 .state
246 .lock()
247 .map(|state| state.clone())
248 .unwrap_or_else(|_| RunHandleState::failed("run handle state lock poisoned"));
249 if has_active_sub_runs {
250 return RunHandleState {
251 status: RunHandleStatus::Running,
252 done: false,
253 cancelled: self.cancel_requested.load(Ordering::SeqCst)
254 || self.cancellation_token.is_cancelled(),
255 error: None,
256 };
257 }
258 if !state.done
259 && (self.cancel_requested.load(Ordering::SeqCst)
260 || self.cancellation_token.is_cancelled())
261 {
262 state.cancelled = true;
263 }
264 if state.done && self.cancel_requested.load(Ordering::SeqCst) {
265 return RunHandleState::cancelled_with_reason(
266 self.cancellation_token
267 .reason()
268 .unwrap_or_else(|| "Operation was cancelled".to_string()),
269 );
270 }
271 state
272 }
273
274 pub fn cancel(&self) -> bool {
275 self.cancel_with_reason("Run was cancelled.")
276 }
277
278 pub fn cancel_with_reason(&self, reason: impl Into<String>) -> bool {
279 let reason = reason.into();
280 if self.cancellation_token.is_cancelled() {
281 return false;
282 }
283 let events = self
284 .events
285 .lock()
286 .unwrap_or_else(std::sync::PoisonError::into_inner);
287 let active_sub_runs = active_sub_run_ids(&events);
288 let state = self
289 .state
290 .lock()
291 .unwrap_or_else(std::sync::PoisonError::into_inner);
292 if (state.done || terminal_event_seen(&events)) && active_sub_runs.is_empty() {
293 return false;
294 }
295 if self
296 .cancel_requested
297 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
298 .is_err()
299 {
300 return false;
301 }
302 drop(state);
303 drop(events);
304 self.cancellation_token.cancel_with_reason(reason);
305 true
306 }
307
308 pub fn cancellation_token(&self) -> &CancellationToken {
309 &self.cancellation_token
310 }
311
312 pub fn steer(&self, message: impl Into<String>) -> Result<(), String> {
313 let message = message.into();
314 self.with_controller("steer", |controller| controller.steer(message))
315 }
316
317 pub fn follow_up(&self, message: impl Into<String>) -> Result<(), String> {
318 let message = message.into();
319 self.with_controller("follow_up", |controller| controller.follow_up(message))
320 }
321
322 pub async fn resume(&self, state: RunState) -> Result<RunResult, String> {
323 let origin_runner = state
324 .result()
325 .resume_context()
326 .map(|context| context.runner.clone())
327 .ok_or_else(|| "run state does not include resume context".to_string())?;
328 origin_runner.resume(state).await
329 }
330
331 pub async fn resume_with_input(
332 &self,
333 state: RunState,
334 input: impl Into<NormalizedInput>,
335 ) -> Result<RunResult, String> {
336 let origin_runner = state
337 .result()
338 .resume_context()
339 .map(|context| context.runner.clone())
340 .ok_or_else(|| "run state does not include resume context".to_string())?;
341 origin_runner.resume_with_input(state, input).await
342 }
343
344 pub async fn approve(
345 &self,
346 request_id: impl AsRef<str>,
347 decision: ApprovalDecision,
348 ) -> Result<(), ApprovalError> {
349 self.approval_broker.resolve(request_id, decision)
350 }
351
352 pub fn approval_broker(&self) -> &ApprovalBroker {
353 &self.approval_broker
354 }
355
356 pub(crate) fn attach_controller(&self, controller: Arc<dyn RunHandleController>) -> u64 {
357 let mut binding = self
358 .controller
359 .lock()
360 .unwrap_or_else(std::sync::PoisonError::into_inner);
361 binding.generation = binding.generation.wrapping_add(1).max(1);
362 let generation = binding.generation;
363 binding.active = Some((generation, controller));
364 generation
365 }
366
367 pub(crate) fn detach_controller(&self, generation: u64) -> bool {
368 let mut binding = self
369 .controller
370 .lock()
371 .unwrap_or_else(std::sync::PoisonError::into_inner);
372 if binding
373 .active
374 .as_ref()
375 .is_some_and(|(active_generation, _)| *active_generation == generation)
376 {
377 binding.active = None;
378 return true;
379 }
380 false
381 }
382
383 fn with_controller<T>(
384 &self,
385 method: &str,
386 callback: impl FnOnce(&dyn RunHandleController) -> Result<T, String>,
387 ) -> Result<T, String> {
388 let binding = self
389 .controller
390 .lock()
391 .unwrap_or_else(std::sync::PoisonError::into_inner);
392 let controller = binding
393 .active
394 .as_ref()
395 .map(|(_, controller)| controller.as_ref())
396 .ok_or_else(|| {
397 format!(
398 "RunHandle.{method}() is only available when the handle is attached to an interactive session."
399 )
400 })?;
401 callback(controller)
402 }
403}
404
405fn terminal_event_seen(events: &[RunEvent]) -> bool {
406 let mut active_handoffs = 0usize;
407 let mut terminal_seen = false;
408 let mut terminal_seen_during_handoff = false;
409 for event in events {
410 match event.payload() {
411 RunEventPayload::HandoffStarted { .. } => {
412 active_handoffs += 1;
413 terminal_seen = false;
414 terminal_seen_during_handoff = false;
415 }
416 RunEventPayload::HandoffCompleted { .. } => {
417 active_handoffs = active_handoffs.saturating_sub(1);
418 if event
419 .metadata()
420 .get("chain_continues")
421 .and_then(serde_json::Value::as_bool)
422 .unwrap_or(false)
423 {
424 terminal_seen = false;
425 terminal_seen_during_handoff = false;
426 } else if active_handoffs == 0 && terminal_seen_during_handoff {
427 terminal_seen = true;
428 terminal_seen_during_handoff = false;
429 }
430 }
431 RunEventPayload::RunCompleted { .. }
432 | RunEventPayload::RunFailed { .. }
433 | RunEventPayload::RunCancelled { .. } => {
434 if active_handoffs > 0 {
435 terminal_seen_during_handoff = true;
436 } else {
437 terminal_seen = true;
438 }
439 }
440 _ => {}
441 }
442 }
443 terminal_seen
444}
445
446pub(crate) fn active_sub_run_ids(events: &[RunEvent]) -> std::collections::HashSet<String> {
447 let mut active = std::collections::HashSet::new();
448 for event in events {
449 match event.payload() {
450 RunEventPayload::SubRunStarted { .. } => {
451 active.insert(event.run_id().to_string());
452 }
453 RunEventPayload::SubRunCompleted { .. } => {
454 active.remove(event.run_id());
455 }
456 _ => {}
457 }
458 }
459 active
460}