vv_agent/runner/
producer.rs1use 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};
10
11use super::{CheckpointAdmissionSender, NormalizedInput, RunEventStream, Runner};
12
13pub(crate) enum CheckpointStartOutcome {
14 Started {
15 handle: RunHandle,
16 checkpoint: crate::runtime::state_v2::CheckpointV2,
17 },
18 ExistingOwner {
19 checkpoint: crate::runtime::state_v2::CheckpointV2,
20 },
21 TerminalReplay {
22 result: Box<RunResult>,
23 checkpoint: crate::runtime::state_v2::CheckpointV2,
24 },
25}
26
27impl Runner {
28 pub async fn stream(
29 &self,
30 agent: &Agent,
31 input: impl Into<NormalizedInput>,
32 ) -> Result<RunEventStream, String> {
33 self.stream_with_config(agent, input, RunConfig::default())
34 .await
35 }
36
37 pub async fn stream_with_config(
38 &self,
39 agent: &Agent,
40 input: impl Into<NormalizedInput>,
41 config: RunConfig,
42 ) -> Result<RunEventStream, String> {
43 let handle = self.start(agent, input, config).await?;
44 Ok(handle.into_event_stream())
45 }
46
47 pub async fn start(
48 &self,
49 agent: &Agent,
50 input: impl Into<NormalizedInput>,
51 config: RunConfig,
52 ) -> Result<RunHandle, String> {
53 self.start_internal(agent, input.into(), config, None).await
54 }
55
56 pub(crate) async fn start_checkpointed(
57 &self,
58 agent: &Agent,
59 input: impl Into<NormalizedInput>,
60 config: RunConfig,
61 ) -> Result<CheckpointStartOutcome, String> {
62 let checkpoint_config = config
63 .checkpoint_config
64 .clone()
65 .or_else(|| self.default_run_config.checkpoint_config.clone())
66 .ok_or_else(|| {
67 "checkpoint_config_invalid: start_checkpointed requires checkpoint_config"
68 .to_string()
69 })?;
70 checkpoint_config
71 .validate()
72 .map_err(|error| error.to_string())?;
73 let store = checkpoint_config.store.clone().ok_or_else(|| {
74 "checkpoint_store_unavailable: start_checkpointed requires a process-local store"
75 .to_string()
76 })?;
77 let checkpoint_key = checkpoint_config.key.clone().ok_or_else(|| {
78 "checkpoint_key_required: start_checkpointed requires an explicit key".to_string()
79 })?;
80 if let Some(checkpoint) = store
81 .load_checkpoint_v2(&checkpoint_key)
82 .map_err(|error| error.to_string())?
83 {
84 if checkpoint.terminal_result.is_none()
85 && checkpoint
86 .lease_expires_at_ms
87 .is_some_and(|expires_at| expires_at > unix_time_ms())
88 {
89 return Ok(CheckpointStartOutcome::ExistingOwner { checkpoint });
90 }
91 }
92
93 let (admission_sender, admission_receiver) = tokio::sync::oneshot::channel();
94 let handle = self
95 .start_internal(agent, input.into(), config, Some(admission_sender))
96 .await?;
97 match admission_receiver.await {
98 Ok(admission) if admission.terminal_replayed => {
99 let result = handle.result().await?;
100 let checkpoint = store
101 .load_checkpoint_v2(&checkpoint_key)
102 .map_err(|error| error.to_string())?
103 .ok_or_else(|| {
104 "checkpoint_not_found: terminal checkpoint disappeared".to_string()
105 })?;
106 Ok(CheckpointStartOutcome::TerminalReplay {
107 result: Box::new(result),
108 checkpoint,
109 })
110 }
111 Ok(admission) => Ok(CheckpointStartOutcome::Started {
112 handle,
113 checkpoint: admission.checkpoint,
114 }),
115 Err(_) => {
116 let result = handle.result().await;
117 if let Some(checkpoint) = store
118 .load_checkpoint_v2(&checkpoint_key)
119 .map_err(|error| error.to_string())?
120 {
121 if checkpoint.terminal_result.is_none()
122 && checkpoint
123 .lease_expires_at_ms
124 .is_some_and(|expires_at| expires_at > unix_time_ms())
125 {
126 return Ok(CheckpointStartOutcome::ExistingOwner { checkpoint });
127 }
128 }
129 match result {
130 Ok(_) => Err(
131 "checkpoint_admission_missing: checkpointed run completed without admission"
132 .to_string(),
133 ),
134 Err(error) => Err(error),
135 }
136 }
137 }
138 }
139
140 async fn start_internal(
141 &self,
142 agent: &Agent,
143 input: NormalizedInput,
144 mut config: RunConfig,
145 checkpoint_admission_sender: Option<CheckpointAdmissionSender>,
146 ) -> Result<RunHandle, String> {
147 let cancellation_token = config
148 .cancellation_token
149 .clone()
150 .or_else(|| self.default_run_config.cancellation_token.clone())
151 .unwrap_or_default();
152 config.cancellation_token = Some(cancellation_token.clone());
153 let approval_broker = config
154 .approval_broker
155 .clone()
156 .or_else(|| self.default_run_config.approval_broker.clone())
157 .unwrap_or_default();
158 config.approval_broker = Some(approval_broker.clone());
159
160 let (event_sender, _) = broadcast::channel(1024);
161 let event_collector = Arc::new(Mutex::new(Vec::new()));
162 let event_sender_slot: RunEventSenderSlot =
163 Arc::new(Mutex::new(Some(event_sender.clone())));
164 let state = Arc::new(Mutex::new(RunHandleState::running()));
165 let cancel_requested = Arc::new(AtomicBool::new(false));
166 let (completion_sender, completion_receiver) = tokio::sync::watch::channel(false);
167 let runner = self.clone();
168 let agent = agent.clone();
169 let state_for_task = state.clone();
170 let event_collector_for_task = event_collector.clone();
171 let cancellation_token_for_task = cancellation_token.clone();
172 let join = tokio::task::spawn_blocking(move || {
173 struct CompletionGuard {
174 sender: Option<tokio::sync::watch::Sender<bool>>,
175 }
176
177 impl Drop for CompletionGuard {
178 fn drop(&mut self) {
179 if let Some(sender) = self.sender.take() {
180 let _ = sender.send(true);
181 }
182 }
183 }
184
185 let _completion = CompletionGuard {
186 sender: Some(completion_sender),
187 };
188 let result = runner.run_blocking_with_event_sender(
189 &agent,
190 input,
191 config,
192 Some(event_collector_for_task),
193 Some(event_sender),
194 checkpoint_admission_sender,
195 );
196 if let Ok(mut state) = state_for_task.lock() {
197 *state = match &result {
198 Ok(result) if run_result_was_cancelled(result) => {
199 RunHandleState::cancelled_with_reason(
200 result
201 .result()
202 .error
203 .clone()
204 .unwrap_or_else(|| "Operation was cancelled".to_string()),
205 )
206 }
207 Ok(result) => RunHandleState::from_run_result(result),
208 Err(error)
209 if cancellation_token_for_task.is_cancelled()
210 && error.to_ascii_lowercase().contains("cancel") =>
211 {
212 let mut state = RunHandleState::cancelled();
213 state.error = Some(error.clone());
214 state
215 }
216 Err(error) => RunHandleState::failed(error.clone()),
217 };
218 }
219 result
220 });
221 let result = SharedRunResult::new(join);
222 Ok(RunHandle::new(
223 event_sender_slot,
224 event_collector,
225 result,
226 state,
227 cancellation_token,
228 approval_broker,
229 completion_receiver,
230 cancel_requested,
231 ))
232 }
233}
234
235fn unix_time_ms() -> u64 {
236 use std::time::{SystemTime, UNIX_EPOCH};
237
238 SystemTime::now()
239 .duration_since(UNIX_EPOCH)
240 .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
241 .unwrap_or(0)
242}
243
244fn run_result_was_cancelled(result: &RunResult) -> bool {
245 result.status() == crate::types::AgentStatus::Failed
246 && result
247 .result()
248 .error
249 .as_deref()
250 .is_some_and(|error| error.to_ascii_lowercase().contains("cancel"))
251}