vv_agent/
run_handle.rs

1use std::sync::{Arc, Mutex};
2
3use tokio::sync::broadcast;
4use tokio::task::JoinHandle;
5
6use crate::approval::{ApprovalBroker, ApprovalError};
7use crate::events::RunEvent;
8use crate::result::RunResult;
9use crate::runner::RunEventStream;
10use crate::runtime::CancellationToken;
11use crate::tools::ApprovalDecision;
12
13pub(crate) type RunEventSenderSlot = Arc<Mutex<Option<broadcast::Sender<RunEvent>>>>;
14type RunJoinHandle = JoinHandle<Result<RunResult, String>>;
15type SharedJoinHandle = Arc<tokio::sync::Mutex<Option<RunJoinHandle>>>;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub enum RunHandleStatus {
19    Running,
20    Completed,
21    Failed,
22    Cancelled,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct RunHandleState {
27    pub status: RunHandleStatus,
28    pub done: bool,
29    pub cancelled: bool,
30    pub error: Option<String>,
31}
32
33impl RunHandleState {
34    pub fn running() -> Self {
35        Self {
36            status: RunHandleStatus::Running,
37            done: false,
38            cancelled: false,
39            error: None,
40        }
41    }
42
43    pub fn completed() -> Self {
44        Self {
45            status: RunHandleStatus::Completed,
46            done: true,
47            cancelled: false,
48            error: None,
49        }
50    }
51
52    pub fn failed(error: impl Into<String>) -> Self {
53        Self {
54            status: RunHandleStatus::Failed,
55            done: true,
56            cancelled: false,
57            error: Some(error.into()),
58        }
59    }
60
61    pub fn cancelled() -> Self {
62        Self {
63            status: RunHandleStatus::Cancelled,
64            done: true,
65            cancelled: true,
66            error: Some("run cancelled".to_string()),
67        }
68    }
69}
70
71#[derive(Clone)]
72pub(crate) struct SharedRunResult {
73    join: SharedJoinHandle,
74}
75
76impl SharedRunResult {
77    pub(crate) fn new(join: RunJoinHandle) -> Self {
78        Self {
79            join: Arc::new(tokio::sync::Mutex::new(Some(join))),
80        }
81    }
82
83    pub(crate) async fn wait(&self) -> Result<RunResult, String> {
84        let join = self
85            .join
86            .lock()
87            .await
88            .take()
89            .ok_or_else(|| "run result already taken".to_string())?;
90        join.await
91            .map_err(|error| format!("run task failed: {error}"))?
92    }
93}
94
95#[derive(Clone)]
96pub struct RunHandle {
97    sender: RunEventSenderSlot,
98    events: Arc<Mutex<Vec<RunEvent>>>,
99    result: SharedRunResult,
100    state: Arc<Mutex<RunHandleState>>,
101    cancellation_token: CancellationToken,
102    approval_broker: ApprovalBroker,
103}
104
105impl RunHandle {
106    pub(crate) fn new(
107        sender: RunEventSenderSlot,
108        events: Arc<Mutex<Vec<RunEvent>>>,
109        result: SharedRunResult,
110        state: Arc<Mutex<RunHandleState>>,
111        cancellation_token: CancellationToken,
112        approval_broker: ApprovalBroker,
113    ) -> Self {
114        Self {
115            sender,
116            events,
117            result,
118            state,
119            cancellation_token,
120            approval_broker,
121        }
122    }
123
124    pub fn events(&self) -> RunEventStream {
125        let receiver = self
126            .sender
127            .lock()
128            .ok()
129            .and_then(|sender| sender.as_ref().map(broadcast::Sender::subscribe));
130        RunEventStream::from_live(receiver, Some(self.result.clone()), self.event_snapshot())
131    }
132
133    pub(crate) fn into_event_stream(self) -> RunEventStream {
134        let receiver = self
135            .sender
136            .lock()
137            .ok()
138            .and_then(|sender| sender.as_ref().map(broadcast::Sender::subscribe));
139        RunEventStream::from_live(receiver, Some(self.result.clone()), self.event_snapshot())
140    }
141
142    pub async fn result(&self) -> Result<RunResult, String> {
143        self.result.wait().await
144    }
145
146    pub fn state(&self) -> RunHandleState {
147        self.state
148            .lock()
149            .map(|state| state.clone())
150            .unwrap_or_else(|_| RunHandleState::failed("run handle state lock poisoned"))
151    }
152
153    pub fn cancel(&self) {
154        self.cancellation_token.cancel();
155        if let Ok(mut state) = self.state.lock() {
156            if !state.done {
157                *state = RunHandleState::cancelled();
158            }
159        }
160    }
161
162    pub fn cancellation_token(&self) -> &CancellationToken {
163        &self.cancellation_token
164    }
165
166    pub async fn approve(
167        &self,
168        request_id: impl AsRef<str>,
169        decision: ApprovalDecision,
170    ) -> Result<(), ApprovalError> {
171        self.approval_broker.resolve(request_id, decision)
172    }
173
174    pub fn approval_broker(&self) -> &ApprovalBroker {
175        &self.approval_broker
176    }
177
178    fn event_snapshot(&self) -> Vec<RunEvent> {
179        self.events
180            .lock()
181            .map(|events| events.clone())
182            .unwrap_or_default()
183    }
184}
vv_agent/run_handle.rs

vv_agent/
run_handle.rs
