Skip to main content

vv_agent/runtime/backends/distributed/
worker.rs

1use std::path::PathBuf;
2use std::sync::{Arc, Condvar, Mutex};
3
4use crate::config::build_vv_llm_from_local_settings;
5use crate::llm::LlmClient;
6use crate::runtime::backends::InlineBackend;
7use crate::runtime::engine::{AgentRuntime, RunEventHandler};
8use crate::workspace::LocalWorkspaceBackend;
9
10use super::capabilities::{DistributedCapabilityRegistry, ResolvedDistributedCapabilities};
11use super::checkpoint_worker::{
12    run_distributed_cycle, DistributedCycleExecutor, DistributedDeliveryMetadata,
13};
14use super::contract::DistributedRunEnvelope;
15use super::dispatch::CycleDispatchResult;
16
17pub(super) struct LeaseHeartbeatStopGuard {
18    stopped: Arc<(Mutex<bool>, Condvar)>,
19}
20
21impl LeaseHeartbeatStopGuard {
22    pub(super) fn new(stopped: Arc<(Mutex<bool>, Condvar)>) -> Self {
23        Self { stopped }
24    }
25}
26
27impl Drop for LeaseHeartbeatStopGuard {
28    fn drop(&mut self) {
29        let (lock, changed) = &*self.stopped;
30        *lock
31            .lock()
32            .unwrap_or_else(std::sync::PoisonError::into_inner) = true;
33        changed.notify_all();
34    }
35}
36
37#[derive(Clone)]
38pub(super) struct LeaseHeartbeatStatus {
39    state: Arc<Mutex<LeaseHeartbeatState>>,
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub(super) enum LeaseCommitPhase {
44    NotStarted,
45    InProgress,
46    Succeeded,
47}
48
49pub(super) struct LeaseHeartbeatFailure {
50    pub(super) renewal: LeaseRenewalFailure,
51    pub(super) renewal_started_during_commit: bool,
52}
53
54struct LeaseHeartbeatState {
55    failure: Option<LeaseHeartbeatFailure>,
56    commit_phase: LeaseCommitPhase,
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub(super) enum LeaseRenewalFailureKind {
61    ActiveClaimLost,
62    ClaimLeaseExpired,
63    Coordination,
64}
65
66pub(super) struct LeaseRenewalFailure {
67    pub(super) kind: LeaseRenewalFailureKind,
68    pub(super) message: String,
69}
70
71impl LeaseRenewalFailure {
72    pub(super) fn active_claim_lost() -> Self {
73        Self {
74            kind: LeaseRenewalFailureKind::ActiveClaimLost,
75            message: "claim is no longer active".to_string(),
76        }
77    }
78
79    pub(super) fn claim_lease_expired() -> Self {
80        Self {
81            kind: LeaseRenewalFailureKind::ClaimLeaseExpired,
82            message: "claim lease expired".to_string(),
83        }
84    }
85
86    pub(super) fn coordination(message: String) -> Self {
87        Self {
88            kind: LeaseRenewalFailureKind::Coordination,
89            message,
90        }
91    }
92}
93
94pub(super) struct LeaseRenewal {
95    pub(super) lease_expires_at_ms: u64,
96    pub(super) effective_lease_ms: u64,
97}
98
99impl LeaseHeartbeatStatus {
100    pub(super) fn new() -> Self {
101        Self {
102            state: Arc::new(Mutex::new(LeaseHeartbeatState {
103                failure: None,
104                commit_phase: LeaseCommitPhase::NotStarted,
105            })),
106        }
107    }
108
109    pub(super) fn commit_phase(&self) -> LeaseCommitPhase {
110        self.state
111            .lock()
112            .unwrap_or_else(std::sync::PoisonError::into_inner)
113            .commit_phase
114    }
115
116    pub(super) fn record(&self, renewal: LeaseRenewalFailure, phase_at_start: LeaseCommitPhase) {
117        let mut state = self
118            .state
119            .lock()
120            .unwrap_or_else(std::sync::PoisonError::into_inner);
121        if state.failure.is_none() {
122            state.failure = Some(LeaseHeartbeatFailure {
123                renewal,
124                renewal_started_during_commit: phase_at_start != LeaseCommitPhase::NotStarted,
125            });
126        }
127    }
128
129    pub(super) fn begin_commit(&self) -> Result<(), String> {
130        let mut state = self
131            .state
132            .lock()
133            .unwrap_or_else(std::sync::PoisonError::into_inner);
134        if let Some(failure) = &state.failure {
135            return Err(format!(
136                "checkpoint lease heartbeat failed: {}",
137                failure.renewal.message
138            ));
139        }
140        state.commit_phase = LeaseCommitPhase::InProgress;
141        Ok(())
142    }
143
144    pub(super) fn mark_commit_succeeded(&self) -> Result<(), String> {
145        let mut state = self
146            .state
147            .lock()
148            .unwrap_or_else(std::sync::PoisonError::into_inner);
149        if state.commit_phase != LeaseCommitPhase::InProgress {
150            return Err("checkpoint commit phase has not started".to_string());
151        }
152        state.commit_phase = LeaseCommitPhase::Succeeded;
153        Ok(())
154    }
155
156    pub(super) fn take(&self) -> (Option<LeaseHeartbeatFailure>, LeaseCommitPhase) {
157        let mut state = self
158            .state
159            .lock()
160            .unwrap_or_else(std::sync::PoisonError::into_inner);
161        (state.failure.take(), state.commit_phase)
162    }
163}
164
165pub(super) struct LeaseOperationResult<T> {
166    pub(super) value: T,
167    pub(super) claim_committed: bool,
168}
169
170impl<T> LeaseOperationResult<T> {
171    pub(super) fn new(value: T, claim_committed: bool) -> Self {
172        Self {
173            value,
174            claim_committed,
175        }
176    }
177}
178
179#[derive(Clone)]
180pub struct DistributedCycleWorker {
181    pub(super) capabilities: DistributedCapabilityRegistry,
182    pub(super) checkpoint_executor: Option<Arc<dyn DistributedCycleExecutor>>,
183}
184
185impl Default for DistributedCycleWorker {
186    fn default() -> Self {
187        Self::new(DistributedCapabilityRegistry::new())
188    }
189}
190
191impl DistributedCycleWorker {
192    pub fn new(capabilities: DistributedCapabilityRegistry) -> Self {
193        Self {
194            capabilities,
195            checkpoint_executor: None,
196        }
197    }
198
199    pub fn with_checkpoint_executor(mut self, executor: Arc<dyn DistributedCycleExecutor>) -> Self {
200        self.checkpoint_executor = Some(executor);
201        self
202    }
203
204    pub fn run_cycle(
205        &self,
206        envelope: DistributedRunEnvelope,
207    ) -> Result<CycleDispatchResult, String> {
208        self.run_cycle_with_delivery(envelope, DistributedDeliveryMetadata::default())
209    }
210
211    pub fn run_cycle_with_delivery(
212        &self,
213        envelope: DistributedRunEnvelope,
214        delivery: DistributedDeliveryMetadata,
215    ) -> Result<CycleDispatchResult, String> {
216        envelope.validate()?;
217        envelope.ensure_not_expired()?;
218        run_distributed_cycle(self, envelope, delivery)
219    }
220}
221
222pub(super) fn lease_expiry_at(
223    now_ms: u64,
224    lease_duration_ms: u64,
225    deadline_unix_ms: Option<u64>,
226) -> Result<u64, String> {
227    let effective_duration = deadline_unix_ms
228        .map(|deadline| lease_duration_ms.min(deadline.saturating_sub(now_ms)))
229        .unwrap_or(lease_duration_ms);
230    now_ms
231        .checked_add(effective_duration)
232        .ok_or_else(|| "checkpoint lease overflow".to_string())
233}
234
235pub(super) fn build_runtime(
236    envelope: &DistributedRunEnvelope,
237    resolved: &ResolvedDistributedCapabilities,
238) -> Result<AgentRuntime<Arc<dyn LlmClient>>, String> {
239    let llm_client = match resolved.llm_client.clone() {
240        Some(client) => client,
241        None => Arc::new(
242            build_vv_llm_from_local_settings(
243                &envelope.recipe.settings_file,
244                &envelope.recipe.backend,
245                &envelope.recipe.model,
246                envelope.recipe.timeout_seconds,
247            )
248            .map_err(|error| error.to_string())?
249            .0,
250        ) as Arc<dyn LlmClient>,
251    };
252    let workspace = PathBuf::from(&envelope.recipe.workspace);
253    let workspace_backend = resolved
254        .workspace_backend
255        .clone()
256        .unwrap_or_else(|| Arc::new(LocalWorkspaceBackend::new(workspace.clone())));
257    let mut runtime = AgentRuntime::new(llm_client)
258        .with_tool_registry(resolved.tool_registry.clone())
259        .with_execution_backend(InlineBackend)
260        .with_default_workspace(workspace)
261        .with_workspace_backend(workspace_backend)
262        .with_settings_file(&envelope.recipe.settings_file)
263        .with_default_backend(&envelope.recipe.backend)
264        .with_hooks(resolved.hooks.clone())
265        .with_after_cycle_hooks(resolved.after_cycle_hooks.clone());
266    if let Some(log_preview_chars) = envelope.recipe.log_preview_chars {
267        runtime = runtime.with_log_preview_chars(log_preview_chars);
268    }
269    runtime.set_tool_policy(resolved.tool_policy.clone());
270    Ok(runtime)
271}
272
273pub(super) fn combined_event_handler(
274    resolved: &ResolvedDistributedCapabilities,
275) -> Option<RunEventHandler> {
276    let mut handlers = resolved.observers.clone();
277    if let Some(event_sink) = &resolved.event_sink {
278        handlers.push(event_sink.clone());
279    }
280    if handlers.is_empty() {
281        return None;
282    }
283    Some(Arc::new(move |event| {
284        for handler in &handlers {
285            handler(event);
286        }
287    }))
288}