1use crate::contract::{
2 RuntimeActionRequest, RuntimeApplyRequest, RuntimeCapabilities, RuntimeExecRequest,
3 RuntimeExecResult, RuntimeInspection, RuntimeLogChunk, RuntimeLogQuery, RuntimeObservation,
4 RuntimeRemoval, RuntimeUnitState,
5};
6use crate::{
7 RuntimeActionKind, RuntimeClient, RuntimeClock, RuntimeDriver, RuntimeError,
8 RuntimeRequestKind, RuntimeRequestReceipt, RuntimeRequestState, RuntimeResult,
9 RuntimeStateStore, SystemRuntimeClock,
10};
11use async_trait::async_trait;
12use std::future::Future;
13use std::sync::Arc;
14use std::time::Duration;
15
16pub struct ManagedRuntimeClient {
18 state: Arc<dyn RuntimeStateStore>,
19 driver: Arc<dyn RuntimeDriver>,
20 clock: Arc<dyn RuntimeClock>,
21}
22
23impl ManagedRuntimeClient {
24 pub fn new(state: Arc<dyn RuntimeStateStore>, driver: Arc<dyn RuntimeDriver>) -> Self {
25 Self::with_clock(state, driver, Arc::new(SystemRuntimeClock))
26 }
27
28 pub fn with_clock(
29 state: Arc<dyn RuntimeStateStore>,
30 driver: Arc<dyn RuntimeDriver>,
31 clock: Arc<dyn RuntimeClock>,
32 ) -> Self {
33 Self {
34 state,
35 driver,
36 clock,
37 }
38 }
39
40 async fn checked_capabilities(&self) -> RuntimeResult<RuntimeCapabilities> {
41 let capabilities = self.driver.capabilities().await?;
42 capabilities.validate().map_err(RuntimeError::Protocol)?;
43 if &capabilities.provider_id != self.driver.provider_id() {
44 return Err(RuntimeError::Protocol(format!(
45 "Runtime driver {:?} reported capabilities for {:?}",
46 self.driver.provider_id().as_str(),
47 capabilities.provider_id.as_str()
48 )));
49 }
50 Ok(capabilities)
51 }
52
53 fn check_deadline(&self, deadline_at_ms: Option<u64>) -> RuntimeResult<()> {
54 if deadline_at_ms.is_some_and(|deadline| deadline <= self.clock.now_ms()) {
55 return Err(RuntimeError::DeadlineExceeded(
56 "request expired before provider dispatch".into(),
57 ));
58 }
59 Ok(())
60 }
61
62 async fn bounded<T, F>(
63 &self,
64 deadline_at_ms: Option<u64>,
65 stage: &'static str,
66 future: F,
67 ) -> RuntimeResult<T>
68 where
69 T: Send,
70 F: Future<Output = RuntimeResult<T>> + Send,
71 {
72 let Some(deadline_at_ms) = deadline_at_ms else {
73 return future.await;
74 };
75 let now_ms = self.clock.now_ms();
76 let Some(remaining_ms) = deadline_at_ms.checked_sub(now_ms) else {
77 return Err(RuntimeError::DeadlineExceeded(format!(
78 "request expired before {stage}"
79 )));
80 };
81 if remaining_ms == 0 {
82 return Err(RuntimeError::DeadlineExceeded(format!(
83 "request expired before {stage}"
84 )));
85 }
86 tokio::time::timeout(Duration::from_millis(remaining_ms), future)
87 .await
88 .map_err(|_| {
89 RuntimeError::DeadlineExceeded(format!("request deadline elapsed during {stage}"))
90 })?
91 }
92
93 fn exec_deadline(&self, request: &RuntimeExecRequest, started_at_ms: u64) -> u64 {
94 let relative = started_at_ms.saturating_add(request.timeout_ms);
95 request
96 .deadline_at_ms
97 .map_or(relative, |absolute| absolute.min(relative))
98 }
99
100 async fn matching_receipt(
101 &self,
102 unit_id: &str,
103 generation: u64,
104 request_id: &str,
105 kind: RuntimeRequestKind,
106 request_digest: &str,
107 ) -> RuntimeResult<Option<RuntimeRequestReceipt>> {
108 let receipt = match self.state.load_request(unit_id, request_id).await {
109 Ok(receipt) => receipt,
110 Err(RuntimeError::RequestNotFound { .. }) => return Ok(None),
111 Err(error) => return Err(error),
112 };
113 receipt.validate().map_err(RuntimeError::Protocol)?;
114 if receipt.unit_id != unit_id || receipt.request_id != request_id {
115 return Err(RuntimeError::Protocol(
116 "Runtime request receipt storage key mismatch".into(),
117 ));
118 }
119 if receipt.generation != generation
120 || receipt.kind != kind
121 || receipt.request_digest != request_digest
122 {
123 return Err(RuntimeError::RequestConflict {
124 request_id: request_id.into(),
125 });
126 }
127 Ok(Some(receipt))
128 }
129
130 async fn completed_replay(
131 &self,
132 unit_id: &str,
133 generation: u64,
134 request_id: &str,
135 kind: RuntimeRequestKind,
136 request_digest: &str,
137 ) -> RuntimeResult<Option<RuntimeRequestReceipt>> {
138 Ok(self
139 .matching_receipt(unit_id, generation, request_id, kind, request_digest)
140 .await?
141 .filter(|receipt| receipt.state == RuntimeRequestState::Completed))
142 }
143}
144
145#[async_trait]
146impl RuntimeClient for ManagedRuntimeClient {
147 async fn capabilities(&self) -> RuntimeResult<RuntimeCapabilities> {
148 self.checked_capabilities().await
149 }
150
151 async fn apply(&self, request: &RuntimeApplyRequest) -> RuntimeResult<RuntimeObservation> {
152 request.validate().map_err(RuntimeError::InvalidRequest)?;
153 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
154 if self
155 .completed_replay(
156 &request.spec.unit_id,
157 request.spec.generation,
158 &request.request_id,
159 RuntimeRequestKind::Apply,
160 &request_digest,
161 )
162 .await?
163 .is_some()
164 {
165 let _lease = self
166 .state
167 .acquire_operation_lease(&request.spec.unit_id)
168 .await?;
169 let reservation = self
170 .state
171 .reserve_apply(request, self.clock.now_ms())
172 .await?;
173 if reservation.dispatch {
174 return Err(RuntimeError::Protocol(
175 "completed apply receipt regressed to pending".into(),
176 ));
177 }
178 return reservation.receipt.observation.ok_or_else(|| {
179 RuntimeError::Protocol("completed apply receipt has no observation".into())
180 });
181 }
182 self.check_deadline(request.deadline_at_ms)?;
183 let capabilities = self
184 .bounded(
185 request.deadline_at_ms,
186 "capability query",
187 self.checked_capabilities(),
188 )
189 .await?;
190 let missing = capabilities
191 .missing_for(&request.spec)
192 .map_err(RuntimeError::InvalidRequest)?;
193 if !missing.is_empty() {
194 return Err(RuntimeError::UnsupportedCapabilities(missing));
195 }
196
197 let _lease = self
198 .bounded(
199 request.deadline_at_ms,
200 "operation lease wait",
201 self.state.acquire_operation_lease(&request.spec.unit_id),
202 )
203 .await?;
204 self.check_deadline(request.deadline_at_ms)?;
205
206 let reservation = self
207 .state
208 .reserve_apply(request, self.clock.now_ms())
209 .await?;
210 if !reservation.dispatch {
211 return reservation.receipt.observation.ok_or_else(|| {
212 RuntimeError::Protocol("completed apply receipt has no observation".into())
213 });
214 }
215
216 let observation = self
217 .bounded(
218 request.deadline_at_ms,
219 "provider apply",
220 self.driver
221 .apply(&request.spec, &reservation.record.observation),
222 )
223 .await?;
224 observation
225 .validate_against(&request.spec)
226 .map_err(RuntimeError::Protocol)?;
227 ensure_apply_result(&request.spec, &observation)?;
228 Ok(self
229 .state
230 .update_observation(Some(&request.request_id), &observation)
231 .await?
232 .observation)
233 }
234
235 async fn inspect(&self, unit_id: &str) -> RuntimeResult<RuntimeInspection> {
236 let _lease = self.state.acquire_operation_lease(unit_id).await?;
237 let record = match self.state.load(unit_id).await {
238 Ok(record) => record,
239 Err(RuntimeError::NotFound { .. }) => {
240 return Ok(RuntimeInspection::NotFound {
241 schema: RuntimeInspection::SCHEMA.into(),
242 unit_id: unit_id.into(),
243 last_generation: None,
244 });
245 }
246 Err(error) => return Err(error),
247 };
248 if record.removed_at_ms.is_some() {
249 return Ok(RuntimeInspection::NotFound {
250 schema: RuntimeInspection::SCHEMA.into(),
251 unit_id: unit_id.into(),
252 last_generation: Some(record.spec.generation),
253 });
254 }
255 if record.observation.state.is_terminal() {
256 return Ok(RuntimeInspection::Found {
257 schema: RuntimeInspection::SCHEMA.into(),
258 observation: Box::new(record.observation),
259 });
260 }
261
262 let inspection = self.driver.inspect(&record).await?;
263 inspection.validate().map_err(RuntimeError::Protocol)?;
264 match inspection {
265 RuntimeInspection::Found { observation, .. } => {
266 observation
267 .validate_against(&record.spec)
268 .map_err(RuntimeError::Protocol)?;
269 let record = self
270 .state
271 .update_observation(None, observation.as_ref())
272 .await?;
273 Ok(RuntimeInspection::Found {
274 schema: RuntimeInspection::SCHEMA.into(),
275 observation: Box::new(record.observation),
276 })
277 }
278 RuntimeInspection::NotFound { .. } => {
279 let mut unknown = record.observation;
280 unknown.state = RuntimeUnitState::Unknown;
281 unknown.observed_at_ms = unknown.observed_at_ms.max(self.clock.now_ms());
282 unknown.finished_at_ms = None;
283 unknown.health = None;
284 unknown.outputs.clear();
285 unknown.failure = None;
286 let record = self.state.update_observation(None, &unknown).await?;
287 Ok(RuntimeInspection::Found {
288 schema: RuntimeInspection::SCHEMA.into(),
289 observation: Box::new(record.observation),
290 })
291 }
292 }
293 }
294
295 async fn stop(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeInspection> {
296 request.validate().map_err(RuntimeError::InvalidRequest)?;
297 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
298 if self
299 .completed_replay(
300 &request.unit_id,
301 request.generation,
302 &request.request_id,
303 RuntimeRequestKind::Stop,
304 &request_digest,
305 )
306 .await?
307 .is_some()
308 {
309 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
310 let reservation = self
311 .state
312 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
313 .await?;
314 if reservation.dispatch {
315 return Err(RuntimeError::Protocol(
316 "completed stop receipt regressed to pending".into(),
317 ));
318 }
319 return Ok(RuntimeInspection::Found {
320 schema: RuntimeInspection::SCHEMA.into(),
321 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
322 RuntimeError::Protocol("completed stop receipt has no observation".into())
323 })?),
324 });
325 }
326 self.check_deadline(request.deadline_at_ms)?;
327 let capabilities = self
328 .bounded(
329 request.deadline_at_ms,
330 "capability query",
331 self.checked_capabilities(),
332 )
333 .await?;
334 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Stop) {
335 return Err(RuntimeError::UnsupportedCapabilities(vec![
336 "feature:Stop".into()
337 ]));
338 }
339 let _lease = self
340 .bounded(
341 request.deadline_at_ms,
342 "operation lease wait",
343 self.state.acquire_operation_lease(&request.unit_id),
344 )
345 .await?;
346 self.check_deadline(request.deadline_at_ms)?;
347 let reservation = self
348 .state
349 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
350 .await?;
351 if !reservation.dispatch {
352 return Ok(RuntimeInspection::Found {
353 schema: RuntimeInspection::SCHEMA.into(),
354 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
355 RuntimeError::Protocol("completed stop receipt has no observation".into())
356 })?),
357 });
358 }
359 let observation = self
360 .bounded(
361 request.deadline_at_ms,
362 "provider stop",
363 self.driver.stop(&reservation.record, request),
364 )
365 .await?;
366 observation
367 .validate_against(&reservation.record.spec)
368 .map_err(RuntimeError::Protocol)?;
369 ensure_stop_result(&reservation.record.observation, &observation)?;
370 let record = self
371 .state
372 .update_observation(Some(&request.request_id), &observation)
373 .await?;
374 Ok(RuntimeInspection::Found {
375 schema: RuntimeInspection::SCHEMA.into(),
376 observation: Box::new(record.observation),
377 })
378 }
379
380 async fn remove(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeRemoval> {
381 request.validate().map_err(RuntimeError::InvalidRequest)?;
382 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
383 if self
384 .completed_replay(
385 &request.unit_id,
386 request.generation,
387 &request.request_id,
388 RuntimeRequestKind::Remove,
389 &request_digest,
390 )
391 .await?
392 .is_some()
393 {
394 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
395 let reservation = self
396 .state
397 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
398 .await?;
399 if reservation.dispatch {
400 return Err(RuntimeError::Protocol(
401 "completed remove receipt regressed to pending".into(),
402 ));
403 }
404 return reservation.receipt.removal.ok_or_else(|| {
405 RuntimeError::Protocol("completed remove receipt has no removal".into())
406 });
407 }
408 self.check_deadline(request.deadline_at_ms)?;
409 let capabilities = self
410 .bounded(
411 request.deadline_at_ms,
412 "capability query",
413 self.checked_capabilities(),
414 )
415 .await?;
416 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Remove) {
417 return Err(RuntimeError::UnsupportedCapabilities(vec![
418 "feature:Remove".into(),
419 ]));
420 }
421 let _lease = self
422 .bounded(
423 request.deadline_at_ms,
424 "operation lease wait",
425 self.state.acquire_operation_lease(&request.unit_id),
426 )
427 .await?;
428 self.check_deadline(request.deadline_at_ms)?;
429 let reservation = self
430 .state
431 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
432 .await?;
433 if !reservation.dispatch {
434 return reservation.receipt.removal.ok_or_else(|| {
435 RuntimeError::Protocol("completed remove receipt has no removal".into())
436 });
437 }
438 let removal = self
439 .bounded(
440 request.deadline_at_ms,
441 "provider remove",
442 self.driver.remove(&reservation.record, request),
443 )
444 .await?;
445 removal.validate().map_err(RuntimeError::Protocol)?;
446 if removal.request_id != request.request_id
447 || removal.unit_id != request.unit_id
448 || removal.generation != request.generation
449 {
450 return Err(RuntimeError::Protocol(
451 "provider removal changed immutable request identity".into(),
452 ));
453 }
454 self.state.complete_removal(&removal).await?;
455 Ok(removal)
456 }
457
458 async fn logs(&self, query: &RuntimeLogQuery) -> RuntimeResult<Vec<RuntimeLogChunk>> {
459 query.validate().map_err(RuntimeError::InvalidRequest)?;
460 let capabilities = self.checked_capabilities().await?;
461 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Logs) {
462 return Err(RuntimeError::UnsupportedCapabilities(vec![
463 "feature:Logs".into()
464 ]));
465 }
466 let _lease = self.state.acquire_operation_lease(&query.unit_id).await?;
467 let record = self.state.load(&query.unit_id).await?;
468 ensure_current_generation(&record, query.generation)?;
469 let chunks = self.driver.logs(&record, query).await?;
470 for chunk in &chunks {
471 chunk.validate().map_err(RuntimeError::Protocol)?;
472 }
473 if chunks
474 .windows(2)
475 .any(|pair| pair[0].sequence >= pair[1].sequence)
476 {
477 return Err(RuntimeError::Protocol(
478 "provider returned unordered log chunks".into(),
479 ));
480 }
481 Ok(chunks)
482 }
483
484 async fn exec(&self, request: &RuntimeExecRequest) -> RuntimeResult<RuntimeExecResult> {
485 request.validate().map_err(RuntimeError::InvalidRequest)?;
486 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
487 let operation_started_at_ms = self.clock.now_ms();
488 let existing = self
489 .matching_receipt(
490 &request.unit_id,
491 request.generation,
492 &request.request_id,
493 RuntimeRequestKind::Exec,
494 &request_digest,
495 )
496 .await?;
497 if existing
498 .as_ref()
499 .is_some_and(|receipt| receipt.state == RuntimeRequestState::Completed)
500 {
501 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
502 let reservation = self
503 .state
504 .reserve_exec(request, operation_started_at_ms)
505 .await?;
506 if reservation.dispatch {
507 return Err(RuntimeError::Protocol(
508 "completed exec receipt regressed to pending".into(),
509 ));
510 }
511 return reservation.receipt.exec_result.ok_or_else(|| {
512 RuntimeError::Protocol("completed exec receipt has no result".into())
513 });
514 }
515 let deadline_at_ms = existing
516 .as_ref()
517 .and_then(|receipt| receipt.deadline_at_ms)
518 .unwrap_or_else(|| self.exec_deadline(request, operation_started_at_ms));
519 let deadline_at_ms = Some(deadline_at_ms);
520 self.check_deadline(deadline_at_ms)?;
521 let capabilities = self
522 .bounded(
523 deadline_at_ms,
524 "capability query",
525 self.checked_capabilities(),
526 )
527 .await?;
528 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Exec) {
529 return Err(RuntimeError::UnsupportedCapabilities(vec![
530 "feature:Exec".into()
531 ]));
532 }
533 let _lease = self
534 .bounded(
535 deadline_at_ms,
536 "operation lease wait",
537 self.state.acquire_operation_lease(&request.unit_id),
538 )
539 .await?;
540 self.check_deadline(deadline_at_ms)?;
541 let reservation = self
542 .state
543 .reserve_exec(request, operation_started_at_ms)
544 .await?;
545 if !reservation.dispatch {
546 return reservation.receipt.exec_result.ok_or_else(|| {
547 RuntimeError::Protocol("completed exec receipt has no result".into())
548 });
549 }
550 let deadline_at_ms = reservation.receipt.deadline_at_ms.ok_or_else(|| {
551 RuntimeError::Protocol("pending exec receipt has no effective deadline".into())
552 })?;
553 let deadline_at_ms = Some(deadline_at_ms);
554 self.check_deadline(deadline_at_ms)?;
555 let mut provider_request = request.clone();
556 provider_request.deadline_at_ms = deadline_at_ms;
557 let result = self
558 .bounded(
559 deadline_at_ms,
560 "provider exec",
561 self.driver.exec(&reservation.record, &provider_request),
562 )
563 .await?;
564 result.validate().map_err(RuntimeError::Protocol)?;
565 result
566 .observation
567 .validate_against(&reservation.record.spec)
568 .map_err(RuntimeError::Protocol)?;
569 if result.request_id != request.request_id
570 || result.observation.unit_id != request.unit_id
571 || result.observation.generation != request.generation
572 {
573 return Err(RuntimeError::Protocol(
574 "provider exec changed immutable request identity".into(),
575 ));
576 }
577 self.state.complete_exec(&result).await?;
578 Ok(result)
579 }
580}
581
582fn ensure_current_generation(
583 record: &crate::RuntimeUnitRecord,
584 requested: u64,
585) -> RuntimeResult<()> {
586 if record.removed_at_ms.is_some() {
587 return Err(RuntimeError::NotFound {
588 unit_id: record.spec.unit_id.clone(),
589 });
590 }
591 if requested < record.spec.generation {
592 return Err(RuntimeError::StaleGeneration {
593 unit_id: record.spec.unit_id.clone(),
594 requested,
595 current: record.spec.generation,
596 });
597 }
598 if requested != record.spec.generation {
599 return Err(RuntimeError::GenerationConflict {
600 unit_id: record.spec.unit_id.clone(),
601 generation: requested,
602 });
603 }
604 Ok(())
605}
606
607fn ensure_apply_result(
608 spec: &crate::contract::RuntimeUnitSpec,
609 observation: &RuntimeObservation,
610) -> RuntimeResult<()> {
611 let allowed = match spec.class {
612 crate::contract::RuntimeUnitClass::Task => matches!(
613 observation.state,
614 RuntimeUnitState::Succeeded | RuntimeUnitState::Failed
615 ),
616 crate::contract::RuntimeUnitClass::Service => matches!(
617 observation.state,
618 RuntimeUnitState::Running
619 | RuntimeUnitState::Stopped
620 | RuntimeUnitState::Failed
621 | RuntimeUnitState::Unknown
622 ),
623 };
624 if !allowed {
625 return Err(RuntimeError::Protocol(format!(
626 "provider apply returned invalid {:?} result {:?}",
627 spec.class, observation.state
628 )));
629 }
630 if observation.provider_resource_id.is_none() || observation.provider_build.is_none() {
631 return Err(RuntimeError::Protocol(
632 "provider apply returned an observation without provider identity".into(),
633 ));
634 }
635 Ok(())
636}
637
638fn ensure_stop_result(
639 current: &RuntimeObservation,
640 observation: &RuntimeObservation,
641) -> RuntimeResult<()> {
642 if observation.state == RuntimeUnitState::Stopped
643 || observation.state == RuntimeUnitState::Unknown
644 || current.state.is_terminal() && observation == current
645 {
646 return Ok(());
647 }
648 Err(RuntimeError::Protocol(format!(
649 "provider stop returned nonterminal state {:?}",
650 observation.state
651 )))
652}