1use std::collections::{HashMap, VecDeque};
2use std::fmt;
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::{Arc, Mutex, RwLock};
7use std::thread::JoinHandle;
8use std::time::{Duration, Instant};
9
10use serde::{Deserialize, Serialize};
11use serde_json::{json, Value};
12use tokio::sync::{mpsc, Notify};
13
14use crate::{
15 EffectRequest, EffectResult, LocalProviderBinding, RuntimeManifest, RuntimeRelease,
16 RuntimeSnapshot, RuntimeTraceRecord,
17};
18
19const SNAPSHOT_VERSION: u32 = 1;
20const DEFAULT_TIMEOUT_MS: u64 = 30_000;
21const MAX_TIMEOUT_MS: u64 = 120_000;
22const DEFAULT_EFFECT_LIMIT: usize = 64;
23const MAX_IN_FLIGHT_INVOCATIONS: usize = 64;
24const MAX_RETAINED_INVOCATIONS: usize = 256;
25const MAX_RETAINED_INVOCATION_EVENTS: usize = 256;
26const MAX_COALESCED_EVENT_BYTES: usize = 64 * 1024;
27const WORKER_QUEUE_CAPACITY: usize = 64;
28
29pub type ProviderFuture<'a> =
31 Pin<Box<dyn Future<Output = Result<ProviderResponse, RuntimeError>> + Send + 'a>>;
32
33#[derive(Clone, PartialEq, Serialize, Deserialize)]
35#[serde(tag = "format", content = "value", rename_all = "camelCase")]
36pub enum InvocationData {
37 Json(Value),
38 Binary(Vec<u8>),
39}
40
41impl fmt::Debug for InvocationData {
42 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
43 match self {
44 Self::Json(_) => formatter.write_str("InvocationData::Json([REDACTED])"),
45 Self::Binary(bytes) => formatter
46 .debug_tuple("InvocationData::Binary")
47 .field(&format_args!("{} bytes", bytes.len()))
48 .finish(),
49 }
50 }
51}
52
53impl Default for InvocationData {
54 fn default() -> Self {
55 Self::Json(Value::Null)
56 }
57}
58
59#[derive(Clone, PartialEq, Serialize, Deserialize)]
61#[serde(rename_all = "camelCase")]
62pub struct InvocationInput {
63 pub endpoint: String,
64 #[serde(default = "default_session_id")]
65 pub session_id: String,
66 #[serde(default)]
67 pub data: InvocationData,
68 #[serde(default)]
69 pub metadata: Value,
70}
71
72impl InvocationInput {
73 pub fn json(endpoint: impl Into<String>, data: Value) -> Self {
74 Self {
75 endpoint: endpoint.into(),
76 session_id: default_session_id(),
77 data: InvocationData::Json(data),
78 metadata: Value::Object(Default::default()),
79 }
80 }
81
82 pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
83 self.session_id = session_id.into();
84 self
85 }
86}
87
88impl fmt::Debug for InvocationInput {
89 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
90 formatter
91 .debug_struct("InvocationInput")
92 .field("endpoint", &self.endpoint)
93 .field("session_id", &self.session_id)
94 .field("data", &"[REDACTED]")
95 .field("metadata", &"[REDACTED]")
96 .finish()
97 }
98}
99
100#[derive(Clone, PartialEq, Serialize, Deserialize)]
102#[serde(rename_all = "camelCase")]
103pub struct AgentDefinition {
104 pub id: String,
105 pub name: String,
106 pub provider: String,
107 pub capabilities: Vec<String>,
108 #[serde(default)]
109 pub metadata: Value,
110}
111
112impl fmt::Debug for AgentDefinition {
113 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
114 formatter
115 .debug_struct("AgentDefinition")
116 .field("id", &self.id)
117 .field("name", &self.name)
118 .field("provider", &self.provider)
119 .field("capabilities", &self.capabilities)
120 .field("metadata", &"[REDACTED]")
121 .finish()
122 }
123}
124
125#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
127#[serde(rename_all = "camelCase")]
128pub struct EndpointDefinition {
129 pub name: String,
130 pub agent: String,
131 pub capability: String,
132 #[serde(default = "default_timeout_ms")]
133 pub timeout_ms: u64,
134}
135
136#[derive(Clone)]
138pub struct ProviderRequest {
139 pub project_id: String,
140 pub endpoint: String,
141 pub session_id: String,
142 pub agent: AgentDefinition,
143 pub capability: String,
144 pub data: InvocationData,
145 pub metadata: Value,
146 pub snapshot: RuntimeSnapshot,
147}
148
149impl fmt::Debug for ProviderRequest {
150 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
151 formatter
152 .debug_struct("ProviderRequest")
153 .field("project_id", &self.project_id)
154 .field("endpoint", &self.endpoint)
155 .field("session_id", &self.session_id)
156 .field("agent", &self.agent.id)
157 .field("capability", &self.capability)
158 .field("data", &"[REDACTED]")
159 .field("metadata", &"[REDACTED]")
160 .field("snapshot_revision", &self.snapshot.revision)
161 .finish()
162 }
163}
164
165#[derive(Clone, PartialEq, Serialize, Deserialize)]
167#[serde(rename_all = "camelCase")]
168pub struct ProviderResponse {
169 #[serde(default)]
170 pub data: InvocationData,
171 #[serde(default)]
172 pub metadata: Value,
173 #[serde(default, skip_serializing_if = "Option::is_none")]
174 pub state: Option<Value>,
175}
176
177impl ProviderResponse {
178 pub fn json(data: Value) -> Self {
179 Self {
180 data: InvocationData::Json(data),
181 metadata: Value::Object(Default::default()),
182 state: None,
183 }
184 }
185}
186
187impl fmt::Debug for ProviderResponse {
188 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
189 formatter
190 .debug_struct("ProviderResponse")
191 .field("data", &"[REDACTED]")
192 .field("metadata", &"[REDACTED]")
193 .field("state", &self.state.as_ref().map(|_| "[REDACTED]"))
194 .finish()
195 }
196}
197
198#[derive(Clone, Default)]
200pub struct CancellationToken {
201 inner: Arc<CancellationState>,
202}
203
204#[derive(Default)]
205struct CancellationState {
206 cancelled: std::sync::atomic::AtomicBool,
207 notify: Notify,
208}
209
210impl CancellationToken {
211 pub fn cancel(&self) {
212 if !self.inner.cancelled.swap(true, Ordering::AcqRel) {
213 self.inner.notify.notify_waiters();
214 }
215 }
216
217 pub fn is_cancelled(&self) -> bool {
218 self.inner.cancelled.load(Ordering::Acquire)
219 }
220
221 pub async fn cancelled(&self) {
222 if self.is_cancelled() {
223 return;
224 }
225 self.inner.notify.notified().await;
226 }
227}
228
229impl fmt::Debug for CancellationToken {
230 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
231 formatter
232 .debug_struct("CancellationToken")
233 .field("cancelled", &self.is_cancelled())
234 .finish()
235 }
236}
237
238#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
240#[serde(rename_all = "camelCase")]
241pub enum InvocationEventKind {
242 Started,
243 OutputDelta,
244 Completed,
245 Failed,
246 Cancelled,
247}
248
249#[derive(Clone, PartialEq, Serialize, Deserialize)]
251#[serde(rename_all = "camelCase")]
252pub struct InvocationEvent {
253 pub sequence: u64,
254 pub kind: InvocationEventKind,
255 #[serde(default, skip_serializing_if = "Option::is_none")]
256 pub data: Option<InvocationData>,
257 #[serde(default, skip_serializing_if = "Option::is_none")]
258 pub error: Option<String>,
259}
260
261impl fmt::Debug for InvocationEvent {
262 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
263 formatter
264 .debug_struct("InvocationEvent")
265 .field("sequence", &self.sequence)
266 .field("kind", &self.kind)
267 .field("data", &self.data.as_ref().map(|_| "[REDACTED]"))
268 .field("error", &self.error.as_ref().map(|_| "[REDACTED]"))
269 .finish()
270 }
271}
272
273#[derive(Clone)]
279pub struct ProviderEventSink {
280 output_delta: Arc<dyn Fn(InvocationData) + Send + Sync>,
281}
282
283impl ProviderEventSink {
284 fn new(output_delta: impl Fn(InvocationData) + Send + Sync + 'static) -> Self {
285 Self {
286 output_delta: Arc::new(output_delta),
287 }
288 }
289
290 pub fn discard() -> Self {
291 Self::new(|_data| {})
292 }
293
294 pub fn output_delta(&self, data: InvocationData) {
295 (self.output_delta)(data);
296 }
297}
298
299impl fmt::Debug for ProviderEventSink {
300 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
301 formatter.write_str("ProviderEventSink")
302 }
303}
304
305pub trait AgentProvider: Send + Sync + 'static {
311 fn supports(&self, capability: &str) -> bool;
312
313 fn invoke<'a>(
314 &'a self,
315 request: ProviderRequest,
316 cancellation: CancellationToken,
317 ) -> ProviderFuture<'a>;
318
319 fn invoke_with_events<'a>(
320 &'a self,
321 request: ProviderRequest,
322 cancellation: CancellationToken,
323 _events: ProviderEventSink,
324 ) -> ProviderFuture<'a> {
325 self.invoke(request, cancellation)
326 }
327}
328
329pub trait RuntimeStore: Send + Sync + 'static {
334 fn load(
335 &self,
336 project_id: &str,
337 session_id: &str,
338 ) -> Result<Option<RuntimeSnapshot>, RuntimeError>;
339
340 fn save(
341 &self,
342 project_id: &str,
343 session_id: &str,
344 snapshot: &RuntimeSnapshot,
345 ) -> Result<(), RuntimeError>;
346
347 fn save_release(&self, _release: &RuntimeRelease) -> Result<(), RuntimeError> {
348 Err(RuntimeError::store(
349 "this runtime store does not support releases".to_string(),
350 ))
351 }
352
353 fn load_release(
354 &self,
355 _project_id: &str,
356 _version: u64,
357 ) -> Result<Option<RuntimeRelease>, RuntimeError> {
358 Ok(None)
359 }
360
361 fn list_releases(&self, _project_id: &str) -> Result<Vec<RuntimeRelease>, RuntimeError> {
362 Ok(Vec::new())
363 }
364
365 fn active_release(&self, _project_id: &str) -> Result<Option<u64>, RuntimeError> {
366 Ok(None)
367 }
368
369 fn set_active_release(&self, _project_id: &str, _version: u64) -> Result<(), RuntimeError> {
370 Err(RuntimeError::store(
371 "this runtime store does not support releases".to_string(),
372 ))
373 }
374
375 fn save_local_provider_binding(
376 &self,
377 _project_id: &str,
378 _binding: &LocalProviderBinding,
379 ) -> Result<(), RuntimeError> {
380 Err(RuntimeError::store(
381 "this runtime store does not support provider bindings".to_string(),
382 ))
383 }
384
385 fn local_provider_bindings(
386 &self,
387 _project_id: &str,
388 ) -> Result<Vec<LocalProviderBinding>, RuntimeError> {
389 Ok(Vec::new())
390 }
391
392 fn enqueue_trace(&self, _trace: &RuntimeTraceRecord) -> Result<(), RuntimeError> {
393 Ok(())
394 }
395
396 fn pending_traces(&self, _limit: usize) -> Result<Vec<RuntimeTraceRecord>, RuntimeError> {
397 Ok(Vec::new())
398 }
399
400 fn acknowledge_traces(&self, _trace_ids: &[String]) -> Result<(), RuntimeError> {
401 Ok(())
402 }
403}
404
405#[derive(Default)]
407pub struct MemoryRuntimeStore {
408 snapshots: RwLock<HashMap<(String, String), RuntimeSnapshot>>,
409 releases: RwLock<HashMap<(String, u64), RuntimeRelease>>,
410 active_releases: RwLock<HashMap<String, u64>>,
411 provider_bindings: RwLock<HashMap<(String, String), LocalProviderBinding>>,
412 trace_outbox: RwLock<VecDeque<RuntimeTraceRecord>>,
413}
414
415impl RuntimeStore for MemoryRuntimeStore {
416 fn load(
417 &self,
418 project_id: &str,
419 session_id: &str,
420 ) -> Result<Option<RuntimeSnapshot>, RuntimeError> {
421 let snapshots = self.snapshots.read().map_err(|_| RuntimeError::Internal)?;
422 Ok(snapshots
423 .get(&(project_id.to_string(), session_id.to_string()))
424 .cloned())
425 }
426
427 fn save(
428 &self,
429 project_id: &str,
430 session_id: &str,
431 snapshot: &RuntimeSnapshot,
432 ) -> Result<(), RuntimeError> {
433 let mut snapshots = self.snapshots.write().map_err(|_| RuntimeError::Internal)?;
434 snapshots.insert(
435 (project_id.to_string(), session_id.to_string()),
436 snapshot.clone(),
437 );
438 Ok(())
439 }
440
441 fn save_release(&self, release: &RuntimeRelease) -> Result<(), RuntimeError> {
442 release.validate()?;
443 let key = (release.manifest.project_id.clone(), release.version);
444 let mut releases = self.releases.write().map_err(|_| RuntimeError::Internal)?;
445 if let Some(existing) = releases.get(&key) {
446 if existing != release {
447 return Err(RuntimeError::store(
448 "runtime release versions are immutable".to_string(),
449 ));
450 }
451 return Ok(());
452 }
453 releases.insert(key, release.clone());
454 Ok(())
455 }
456
457 fn load_release(
458 &self,
459 project_id: &str,
460 version: u64,
461 ) -> Result<Option<RuntimeRelease>, RuntimeError> {
462 Ok(self
463 .releases
464 .read()
465 .map_err(|_| RuntimeError::Internal)?
466 .get(&(project_id.to_string(), version))
467 .cloned())
468 }
469
470 fn list_releases(&self, project_id: &str) -> Result<Vec<RuntimeRelease>, RuntimeError> {
471 let mut releases = self
472 .releases
473 .read()
474 .map_err(|_| RuntimeError::Internal)?
475 .iter()
476 .filter(|((stored_project_id, _), _)| stored_project_id == project_id)
477 .map(|(_, release)| release.clone())
478 .collect::<Vec<_>>();
479 releases.sort_by_key(|release| std::cmp::Reverse(release.version));
480 Ok(releases)
481 }
482
483 fn active_release(&self, project_id: &str) -> Result<Option<u64>, RuntimeError> {
484 Ok(self
485 .active_releases
486 .read()
487 .map_err(|_| RuntimeError::Internal)?
488 .get(project_id)
489 .copied())
490 }
491
492 fn set_active_release(&self, project_id: &str, version: u64) -> Result<(), RuntimeError> {
493 if !self
494 .releases
495 .read()
496 .map_err(|_| RuntimeError::Internal)?
497 .contains_key(&(project_id.to_string(), version))
498 {
499 return Err(RuntimeError::store("runtime release was not found"));
500 }
501 self.active_releases
502 .write()
503 .map_err(|_| RuntimeError::Internal)?
504 .insert(project_id.to_string(), version);
505 Ok(())
506 }
507
508 fn save_local_provider_binding(
509 &self,
510 project_id: &str,
511 binding: &LocalProviderBinding,
512 ) -> Result<(), RuntimeError> {
513 self.provider_bindings
514 .write()
515 .map_err(|_| RuntimeError::Internal)?
516 .insert(
517 (project_id.to_string(), binding.provider_id.clone()),
518 binding.clone(),
519 );
520 Ok(())
521 }
522
523 fn local_provider_bindings(
524 &self,
525 project_id: &str,
526 ) -> Result<Vec<LocalProviderBinding>, RuntimeError> {
527 let mut bindings = self
528 .provider_bindings
529 .read()
530 .map_err(|_| RuntimeError::Internal)?
531 .iter()
532 .filter(|((stored_project_id, _), _)| stored_project_id == project_id)
533 .map(|(_, binding)| binding.clone())
534 .collect::<Vec<_>>();
535 bindings.sort_by(|left, right| left.provider_id.cmp(&right.provider_id));
536 Ok(bindings)
537 }
538
539 fn enqueue_trace(&self, trace: &RuntimeTraceRecord) -> Result<(), RuntimeError> {
540 const MAX_MEMORY_TRACES: usize = 1_000;
541 let mut traces = self
542 .trace_outbox
543 .write()
544 .map_err(|_| RuntimeError::Internal)?;
545 if traces.iter().any(|stored| stored.id == trace.id) {
546 return Ok(());
547 }
548 traces.push_back(trace.clone());
549 while traces.len() > MAX_MEMORY_TRACES {
550 traces.pop_front();
551 }
552 Ok(())
553 }
554
555 fn pending_traces(&self, limit: usize) -> Result<Vec<RuntimeTraceRecord>, RuntimeError> {
556 Ok(self
557 .trace_outbox
558 .read()
559 .map_err(|_| RuntimeError::Internal)?
560 .iter()
561 .take(limit)
562 .cloned()
563 .collect())
564 }
565
566 fn acknowledge_traces(&self, trace_ids: &[String]) -> Result<(), RuntimeError> {
567 self.trace_outbox
568 .write()
569 .map_err(|_| RuntimeError::Internal)?
570 .retain(|trace| !trace_ids.contains(&trace.id));
571 Ok(())
572 }
573}
574
575#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
577#[serde(rename_all = "camelCase")]
578pub struct InvocationTraceEvent {
579 pub name: String,
580 pub status: String,
581 pub duration_ms: u64,
582 #[serde(default)]
583 pub attributes: Value,
584}
585
586#[derive(Clone, PartialEq, Serialize, Deserialize)]
588#[serde(rename_all = "camelCase")]
589pub struct InvocationOutput {
590 pub invocation_id: String,
591 pub project_id: String,
592 pub endpoint: String,
593 pub session_id: String,
594 pub agent: String,
595 pub provider: String,
596 pub capability: String,
597 pub data: InvocationData,
598 #[serde(default)]
599 pub metadata: Value,
600 pub snapshot: RuntimeSnapshot,
601 pub trace: Vec<InvocationTraceEvent>,
602}
603
604impl fmt::Debug for InvocationOutput {
605 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
606 formatter
607 .debug_struct("InvocationOutput")
608 .field("invocation_id", &self.invocation_id)
609 .field("project_id", &self.project_id)
610 .field("endpoint", &self.endpoint)
611 .field("session_id", &self.session_id)
612 .field("agent", &self.agent)
613 .field("provider", &self.provider)
614 .field("capability", &self.capability)
615 .field("data", &"[REDACTED]")
616 .field("metadata", &"[REDACTED]")
617 .field("snapshot_revision", &self.snapshot.revision)
618 .field("trace_count", &self.trace.len())
619 .finish()
620 }
621}
622
623#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
625pub struct InvocationHandle(pub String);
626
627#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
629#[serde(rename_all = "camelCase")]
630pub enum InvocationStatus {
631 Pending,
632 Running,
633 Completed,
634 Failed,
635 Cancelled,
636}
637
638#[derive(Clone, PartialEq, Serialize, Deserialize)]
640#[serde(rename_all = "camelCase")]
641pub struct InvocationPoll {
642 pub handle: InvocationHandle,
643 pub status: InvocationStatus,
644 #[serde(default, skip_serializing_if = "Option::is_none")]
645 pub output: Option<InvocationOutput>,
646 #[serde(default, skip_serializing_if = "Option::is_none")]
647 pub error: Option<String>,
648}
649
650impl fmt::Debug for InvocationPoll {
651 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
652 formatter
653 .debug_struct("InvocationPoll")
654 .field("handle", &self.handle)
655 .field("status", &self.status)
656 .field("output", &self.output.as_ref().map(|_| "[REDACTED]"))
657 .field("error", &self.error.as_ref().map(|_| "[REDACTED]"))
658 .finish()
659 }
660}
661
662#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
664#[serde(rename_all = "camelCase")]
665pub struct EffectExecution {
666 pub results: Vec<EffectResult>,
667 pub unhandled: Vec<EffectRequest>,
668}
669
670pub enum RuntimeError {
672 InvalidDefinition(String),
673 EndpointNotFound(String),
674 AgentNotFound(String),
675 ProviderNotFound(String),
676 CapabilityUnavailable {
677 provider: String,
678 capability: String,
679 },
680 Timeout(u64),
681 Cancelled,
682 Unavailable(String),
683 Backpressure(String),
684 Provider {
685 provider: String,
686 message: String,
687 },
688 Store(String),
689 Snapshot(String),
690 EffectLimitExceeded(usize),
691 InvocationNotFound(String),
692 Internal,
693}
694
695impl RuntimeError {
696 pub fn provider(provider: impl Into<String>, message: impl Into<String>) -> Self {
697 Self::Provider {
698 provider: provider.into(),
699 message: message.into(),
700 }
701 }
702
703 pub fn store(message: impl Into<String>) -> Self {
704 Self::Store(message.into())
705 }
706
707 pub fn public_message(&self) -> String {
708 match self {
709 Self::Provider { provider, .. } => {
710 format!("provider {provider} request failed")
711 }
712 Self::Store(_) => "runtime state could not be persisted".to_string(),
713 Self::Snapshot(_) => "runtime snapshot is invalid".to_string(),
714 Self::Unavailable(_) => "provider is not available".to_string(),
715 Self::Backpressure(_) => "runtime is busy".to_string(),
716 _ => self.to_string(),
717 }
718 }
719}
720
721impl fmt::Debug for RuntimeError {
722 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
723 match self {
724 Self::InvalidDefinition(_) => formatter.write_str("InvalidDefinition([REDACTED])"),
725 Self::EndpointNotFound(endpoint) => formatter
726 .debug_tuple("EndpointNotFound")
727 .field(endpoint)
728 .finish(),
729 Self::AgentNotFound(agent) => {
730 formatter.debug_tuple("AgentNotFound").field(agent).finish()
731 }
732 Self::ProviderNotFound(provider) => formatter
733 .debug_tuple("ProviderNotFound")
734 .field(provider)
735 .finish(),
736 Self::CapabilityUnavailable {
737 provider,
738 capability,
739 } => formatter
740 .debug_struct("CapabilityUnavailable")
741 .field("provider", provider)
742 .field("capability", capability)
743 .finish(),
744 Self::Timeout(timeout) => formatter.debug_tuple("Timeout").field(timeout).finish(),
745 Self::Cancelled => formatter.write_str("Cancelled"),
746 Self::Unavailable(_) => formatter.write_str("Unavailable([REDACTED])"),
747 Self::Backpressure(_) => formatter.write_str("Backpressure([REDACTED])"),
748 Self::Provider { provider, .. } => formatter
749 .debug_struct("Provider")
750 .field("provider", provider)
751 .field("message", &"[REDACTED]")
752 .finish(),
753 Self::Store(_) => formatter.write_str("Store([REDACTED])"),
754 Self::Snapshot(_) => formatter.write_str("Snapshot([REDACTED])"),
755 Self::EffectLimitExceeded(limit) => formatter
756 .debug_tuple("EffectLimitExceeded")
757 .field(limit)
758 .finish(),
759 Self::InvocationNotFound(handle) => formatter
760 .debug_tuple("InvocationNotFound")
761 .field(handle)
762 .finish(),
763 Self::Internal => formatter.write_str("Internal"),
764 }
765 }
766}
767
768impl fmt::Display for RuntimeError {
769 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
770 match self {
771 Self::InvalidDefinition(message) => {
772 write!(formatter, "invalid runtime definition: {message}")
773 }
774 Self::EndpointNotFound(endpoint) => {
775 write!(formatter, "endpoint {endpoint} is not registered")
776 }
777 Self::AgentNotFound(agent) => write!(formatter, "agent {agent} is not registered"),
778 Self::ProviderNotFound(provider) => {
779 write!(formatter, "provider {provider} is not registered")
780 }
781 Self::CapabilityUnavailable {
782 provider,
783 capability,
784 } => write!(
785 formatter,
786 "provider {provider} does not support capability {capability}"
787 ),
788 Self::Timeout(timeout_ms) => {
789 write!(
790 formatter,
791 "agent invocation timed out after {timeout_ms} ms"
792 )
793 }
794 Self::Cancelled => formatter.write_str("agent invocation was cancelled"),
795 Self::Unavailable(message) => {
796 write!(formatter, "provider is not available: {message}")
797 }
798 Self::Backpressure(message) => write!(formatter, "runtime is busy: {message}"),
799 Self::Provider { provider, message } => {
800 write!(formatter, "provider {provider} failed: {message}")
801 }
802 Self::Store(message) => write!(formatter, "runtime store failed: {message}"),
803 Self::Snapshot(message) => write!(formatter, "runtime snapshot failed: {message}"),
804 Self::EffectLimitExceeded(limit) => {
805 write!(formatter, "runtime effect limit {limit} was exceeded")
806 }
807 Self::InvocationNotFound(handle) => {
808 write!(formatter, "invocation {handle} was not found")
809 }
810 Self::Internal => formatter.write_str("runtime internal error"),
811 }
812 }
813}
814
815impl std::error::Error for RuntimeError {}
816
817#[derive(Default)]
818struct RuntimeRegistry {
819 providers: HashMap<String, Arc<dyn AgentProvider>>,
820 agents: HashMap<String, AgentDefinition>,
821 endpoints: HashMap<String, EndpointDefinition>,
822}
823
824struct RuntimeCore {
825 project_id: String,
826 registry: RwLock<RuntimeRegistry>,
827 manifest: RwLock<Option<RuntimeManifest>>,
828 store: Arc<dyn RuntimeStore>,
829 sessions: RwLock<HashMap<String, RuntimeSnapshot>>,
830 session_locks: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
831 invocations: Mutex<InvocationRegistry>,
832 next_invocation: AtomicU64,
833}
834
835struct InvocationEntry {
836 poll: InvocationPoll,
837 cancellation: CancellationToken,
838 events: VecDeque<InvocationEvent>,
839 next_event_sequence: u64,
840}
841
842#[derive(Default)]
843struct InvocationRegistry {
844 entries: HashMap<String, InvocationEntry>,
845 terminal_order: VecDeque<String>,
846 active_count: usize,
847}
848
849impl InvocationRegistry {
850 fn insert(
851 &mut self,
852 handle: InvocationHandle,
853 cancellation: CancellationToken,
854 ) -> Result<(), RuntimeError> {
855 if self.active_count >= MAX_IN_FLIGHT_INVOCATIONS {
856 return Err(RuntimeError::Backpressure(
857 "too many invocations are already running".to_string(),
858 ));
859 }
860 self.entries.insert(
861 handle.0.clone(),
862 InvocationEntry {
863 poll: InvocationPoll {
864 handle,
865 status: InvocationStatus::Pending,
866 output: None,
867 error: None,
868 },
869 cancellation,
870 events: VecDeque::new(),
871 next_event_sequence: 1,
872 },
873 );
874 self.active_count += 1;
875 Ok(())
876 }
877
878 fn update(
879 &mut self,
880 handle: &InvocationHandle,
881 status: InvocationStatus,
882 output: Option<InvocationOutput>,
883 error: Option<String>,
884 ) {
885 let Some(entry) = self.entries.get_mut(&handle.0) else {
886 return;
887 };
888 if is_terminal_status(entry.poll.status) {
889 return;
890 }
891 entry.poll.status = status;
892 entry.poll.output = output;
893 entry.poll.error = error;
894 let event = match status {
895 InvocationStatus::Pending => None,
896 InvocationStatus::Running => Some((InvocationEventKind::Started, None, None)),
897 InvocationStatus::Completed => Some((
898 InvocationEventKind::Completed,
899 entry.poll.output.as_ref().map(|value| value.data.clone()),
900 None,
901 )),
902 InvocationStatus::Failed => {
903 Some((InvocationEventKind::Failed, None, entry.poll.error.clone()))
904 }
905 InvocationStatus::Cancelled => Some((InvocationEventKind::Cancelled, None, None)),
906 };
907 if let Some((kind, data, error)) = event {
908 entry.push_event(kind, data, error);
909 }
910 if is_terminal_status(status) {
911 self.active_count = self.active_count.saturating_sub(1);
912 self.terminal_order.push_back(handle.0.clone());
913 self.evict_old_terminal_entries();
914 }
915 }
916
917 fn remove(&mut self, handle: &InvocationHandle) {
918 if let Some(entry) = self.entries.remove(&handle.0) {
919 if !is_terminal_status(entry.poll.status) {
920 self.active_count = self.active_count.saturating_sub(1);
921 }
922 }
923 self.terminal_order.retain(|stored| stored != &handle.0);
924 }
925
926 fn take(&mut self, handle: &InvocationHandle) -> Result<InvocationPoll, RuntimeError> {
927 let poll = self
928 .entries
929 .get(&handle.0)
930 .map(|entry| entry.poll.clone())
931 .ok_or_else(|| RuntimeError::InvocationNotFound(handle.0.clone()))?;
932 if is_terminal_status(poll.status) {
933 self.remove(handle);
934 }
935 Ok(poll)
936 }
937
938 fn push_output_delta(&mut self, handle: &InvocationHandle, data: InvocationData) {
939 if let Some(entry) = self.entries.get_mut(&handle.0) {
940 if entry.poll.status != InvocationStatus::Running {
941 return;
942 }
943 entry.push_event(InvocationEventKind::OutputDelta, Some(data), None);
944 }
945 }
946
947 fn drain_events(
948 &mut self,
949 handle: &InvocationHandle,
950 ) -> Result<Vec<InvocationEvent>, RuntimeError> {
951 let entry = self
952 .entries
953 .get_mut(&handle.0)
954 .ok_or_else(|| RuntimeError::InvocationNotFound(handle.0.clone()))?;
955 Ok(entry.events.drain(..).collect())
956 }
957
958 fn evict_old_terminal_entries(&mut self) {
959 while self.terminal_order.len() > MAX_RETAINED_INVOCATIONS {
960 if let Some(handle) = self.terminal_order.pop_front() {
961 self.entries.remove(&handle);
962 }
963 }
964 }
965}
966
967impl InvocationEntry {
968 fn push_event(
969 &mut self,
970 kind: InvocationEventKind,
971 data: Option<InvocationData>,
972 error: Option<String>,
973 ) {
974 if kind == InvocationEventKind::OutputDelta {
975 if let (
976 Some(InvocationEvent {
977 kind: InvocationEventKind::OutputDelta,
978 data: Some(previous),
979 ..
980 }),
981 Some(next),
982 ) = (self.events.back_mut(), data.as_ref())
983 {
984 if merge_invocation_data(previous, next) {
985 return;
986 }
987 }
988 }
989 self.events.push_back(InvocationEvent {
990 sequence: self.next_event_sequence,
991 kind,
992 data,
993 error,
994 });
995 self.next_event_sequence = self.next_event_sequence.saturating_add(1);
996 while self.events.len() > MAX_RETAINED_INVOCATION_EVENTS {
997 self.events.pop_front();
998 }
999 }
1000}
1001
1002impl RuntimeCore {
1003 async fn invoke(
1004 self: &Arc<Self>,
1005 invocation_id: String,
1006 input: InvocationInput,
1007 cancellation: CancellationToken,
1008 ) -> Result<InvocationOutput, RuntimeError> {
1009 let endpoint = input.endpoint.clone();
1010 let started = Instant::now();
1011 let result = self
1012 .invoke_provider(invocation_id.clone(), input, cancellation)
1013 .await;
1014 let created_at_ms = crate::unix_time_ms();
1015 let trace = match &result {
1016 Ok(output) => RuntimeTraceRecord {
1017 id: format!("trace-{created_at_ms}-{invocation_id}"),
1018 project_id: self.project_id.clone(),
1019 invocation_id,
1020 endpoint,
1021 agent: Some(output.agent.clone()),
1022 provider: Some(output.provider.clone()),
1023 capability: Some(output.capability.clone()),
1024 status: "completed".to_string(),
1025 duration_ms: duration_ms(started.elapsed()),
1026 created_at_ms,
1027 },
1028 Err(error) => RuntimeTraceRecord {
1029 id: format!("trace-{created_at_ms}-{invocation_id}"),
1030 project_id: self.project_id.clone(),
1031 invocation_id,
1032 endpoint,
1033 agent: None,
1034 provider: None,
1035 capability: None,
1036 status: match error {
1037 RuntimeError::Cancelled => "cancelled",
1038 _ => "error",
1039 }
1040 .to_string(),
1041 duration_ms: duration_ms(started.elapsed()),
1042 created_at_ms,
1043 },
1044 };
1045 let _ = self.store.enqueue_trace(&trace);
1046 result
1047 }
1048
1049 async fn invoke_provider(
1050 self: &Arc<Self>,
1051 invocation_id: String,
1052 input: InvocationInput,
1053 cancellation: CancellationToken,
1054 ) -> Result<InvocationOutput, RuntimeError> {
1055 validate_identifier("endpoint", &input.endpoint)?;
1056 validate_identifier("session", &input.session_id)?;
1057 let session_lock = {
1058 let mut locks = self
1059 .session_locks
1060 .lock()
1061 .map_err(|_| RuntimeError::Internal)?;
1062 Arc::clone(
1063 locks
1064 .entry(input.session_id.clone())
1065 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
1066 )
1067 };
1068 let _session_guard = session_lock.lock().await;
1069 let (endpoint, agent, provider) = {
1070 let registry = self.registry.read().map_err(|_| RuntimeError::Internal)?;
1071 let endpoint = registry
1072 .endpoints
1073 .get(&input.endpoint)
1074 .cloned()
1075 .ok_or_else(|| RuntimeError::EndpointNotFound(input.endpoint.clone()))?;
1076 let agent = registry
1077 .agents
1078 .get(&endpoint.agent)
1079 .cloned()
1080 .ok_or_else(|| RuntimeError::AgentNotFound(endpoint.agent.clone()))?;
1081 let provider = registry
1082 .providers
1083 .get(&agent.provider)
1084 .cloned()
1085 .ok_or_else(|| RuntimeError::ProviderNotFound(agent.provider.clone()))?;
1086 (endpoint, agent, provider)
1087 };
1088 if !agent
1089 .capabilities
1090 .iter()
1091 .any(|capability| capability == &endpoint.capability)
1092 || !provider.supports(&endpoint.capability)
1093 {
1094 return Err(RuntimeError::CapabilityUnavailable {
1095 provider: agent.provider.clone(),
1096 capability: endpoint.capability,
1097 });
1098 }
1099 if cancellation.is_cancelled() {
1100 return Err(RuntimeError::Cancelled);
1101 }
1102
1103 let snapshot = self.load_snapshot(&input.session_id)?;
1104 let request = ProviderRequest {
1105 project_id: self.project_id.clone(),
1106 endpoint: endpoint.name.clone(),
1107 session_id: input.session_id.clone(),
1108 agent: agent.clone(),
1109 capability: endpoint.capability.clone(),
1110 data: input.data,
1111 metadata: input.metadata,
1112 snapshot: snapshot.clone(),
1113 };
1114 let started = Instant::now();
1115 let events = self.provider_event_sink(&InvocationHandle(invocation_id.clone()));
1116 let provider_call = provider.invoke_with_events(request, cancellation.clone(), events);
1117 let response = tokio::select! {
1118 _ = cancellation.cancelled() => return Err(RuntimeError::Cancelled),
1119 response = tokio::time::timeout(Duration::from_millis(endpoint.timeout_ms), provider_call) => {
1120 response.map_err(|_| RuntimeError::Timeout(endpoint.timeout_ms))??
1121 }
1122 };
1123 if cancellation.is_cancelled() {
1124 return Err(RuntimeError::Cancelled);
1125 }
1126
1127 let next_snapshot = RuntimeSnapshot {
1128 revision: snapshot.revision.saturating_add(1),
1129 state: response.state.unwrap_or(snapshot.state),
1130 };
1131 self.store
1132 .save(&self.project_id, &input.session_id, &next_snapshot)?;
1133 self.sessions
1134 .write()
1135 .map_err(|_| RuntimeError::Internal)?
1136 .insert(input.session_id.clone(), next_snapshot.clone());
1137 Ok(InvocationOutput {
1138 invocation_id,
1139 project_id: self.project_id.clone(),
1140 endpoint: endpoint.name,
1141 session_id: input.session_id,
1142 agent: agent.id,
1143 provider: agent.provider,
1144 capability: endpoint.capability,
1145 data: response.data,
1146 metadata: response.metadata,
1147 snapshot: next_snapshot,
1148 trace: vec![InvocationTraceEvent {
1149 name: "provider.invoke".to_string(),
1150 status: "completed".to_string(),
1151 duration_ms: duration_ms(started.elapsed()),
1152 attributes: json!({
1153 "endpoint": input.endpoint,
1154 }),
1155 }],
1156 })
1157 }
1158
1159 fn load_snapshot(&self, session_id: &str) -> Result<RuntimeSnapshot, RuntimeError> {
1160 if let Some(snapshot) = self
1161 .sessions
1162 .read()
1163 .map_err(|_| RuntimeError::Internal)?
1164 .get(session_id)
1165 .cloned()
1166 {
1167 return Ok(snapshot);
1168 }
1169 let snapshot = self
1170 .store
1171 .load(&self.project_id, session_id)?
1172 .unwrap_or_default();
1173 self.sessions
1174 .write()
1175 .map_err(|_| RuntimeError::Internal)?
1176 .insert(session_id.to_string(), snapshot.clone());
1177 Ok(snapshot)
1178 }
1179
1180 fn next_invocation_id(&self) -> String {
1181 let id = self.next_invocation.fetch_add(1, Ordering::Relaxed);
1182 format!("invocation-{id}")
1183 }
1184
1185 fn update_poll(
1186 &self,
1187 handle: &InvocationHandle,
1188 status: InvocationStatus,
1189 output: Option<InvocationOutput>,
1190 error: Option<String>,
1191 ) {
1192 if let Ok(mut invocations) = self.invocations.lock() {
1193 invocations.update(handle, status, output, error);
1194 }
1195 }
1196
1197 fn provider_event_sink(self: &Arc<Self>, handle: &InvocationHandle) -> ProviderEventSink {
1198 let core = Arc::clone(self);
1199 let handle = handle.clone();
1200 ProviderEventSink::new(move |data| {
1201 if let Ok(mut invocations) = core.invocations.lock() {
1202 invocations.push_output_delta(&handle, data);
1203 }
1204 })
1205 }
1206}
1207
1208fn merge_invocation_data(previous: &mut InvocationData, next: &InvocationData) -> bool {
1209 match (previous, next) {
1210 (
1211 InvocationData::Json(Value::String(previous)),
1212 InvocationData::Json(Value::String(next)),
1213 ) if previous.len().saturating_add(next.len()) <= MAX_COALESCED_EVENT_BYTES => {
1214 previous.push_str(next);
1215 true
1216 }
1217 (InvocationData::Binary(previous), InvocationData::Binary(next))
1218 if previous.len().saturating_add(next.len()) <= MAX_COALESCED_EVENT_BYTES =>
1219 {
1220 previous.extend_from_slice(next);
1221 true
1222 }
1223 _ => false,
1224 }
1225}
1226
1227enum WorkerCommand {
1228 Start {
1229 handle: InvocationHandle,
1230 input: InvocationInput,
1231 cancellation: CancellationToken,
1232 },
1233}
1234
1235struct RuntimeWorker {
1236 sender: Mutex<Option<mpsc::Sender<WorkerCommand>>>,
1237 thread: Mutex<Option<JoinHandle<()>>>,
1238}
1239
1240impl RuntimeWorker {
1241 fn spawn(core: Arc<RuntimeCore>) -> Result<Self, RuntimeError> {
1242 let (sender, mut receiver) = mpsc::channel(WORKER_QUEUE_CAPACITY);
1243 let thread = std::thread::Builder::new()
1244 .name(format!("vifu-runtime-{}", core.project_id))
1245 .spawn(move || {
1246 let Ok(runtime) = tokio::runtime::Builder::new_current_thread()
1247 .enable_time()
1248 .build()
1249 else {
1250 return;
1251 };
1252 runtime.block_on(async move {
1253 while let Some(command) = receiver.recv().await {
1254 match command {
1255 WorkerCommand::Start {
1256 handle,
1257 input,
1258 cancellation,
1259 } => {
1260 let invocation_core = Arc::clone(&core);
1261 tokio::spawn(async move {
1262 invocation_core.update_poll(
1263 &handle,
1264 InvocationStatus::Running,
1265 None,
1266 None,
1267 );
1268 let result = invocation_core
1269 .invoke(handle.0.clone(), input, cancellation)
1270 .await;
1271 match result {
1272 Ok(output) => invocation_core.update_poll(
1273 &handle,
1274 InvocationStatus::Completed,
1275 Some(output),
1276 None,
1277 ),
1278 Err(RuntimeError::Cancelled) => invocation_core
1279 .update_poll(
1280 &handle,
1281 InvocationStatus::Cancelled,
1282 None,
1283 None,
1284 ),
1285 Err(error) => invocation_core.update_poll(
1286 &handle,
1287 InvocationStatus::Failed,
1288 None,
1289 Some(error.public_message()),
1290 ),
1291 }
1292 });
1293 }
1294 }
1295 }
1296 });
1297 })
1298 .map_err(|_error| RuntimeError::Internal)?;
1299 Ok(Self {
1300 sender: Mutex::new(Some(sender)),
1301 thread: Mutex::new(Some(thread)),
1302 })
1303 }
1304
1305 fn send(&self, command: WorkerCommand) -> Result<(), RuntimeError> {
1306 self.sender
1307 .lock()
1308 .map_err(|_| RuntimeError::Internal)?
1309 .as_ref()
1310 .ok_or(RuntimeError::Internal)?
1311 .try_send(command)
1312 .map_err(|error| match error {
1313 mpsc::error::TrySendError::Full(_) => {
1314 RuntimeError::Backpressure("invocation queue is full".to_string())
1315 }
1316 mpsc::error::TrySendError::Closed(_) => RuntimeError::Internal,
1317 })
1318 }
1319}
1320
1321impl Drop for RuntimeWorker {
1322 fn drop(&mut self) {
1323 if let Ok(sender) = self.sender.get_mut() {
1324 sender.take();
1325 }
1326 if let Ok(thread) = self.thread.get_mut() {
1327 if let Some(thread) = thread.take() {
1328 let _ = thread.join();
1329 }
1330 }
1331 }
1332}
1333
1334#[derive(Clone)]
1340pub struct VifuRuntime {
1341 core: Arc<RuntimeCore>,
1342 worker: Arc<Mutex<Option<RuntimeWorker>>>,
1343}
1344
1345impl fmt::Debug for VifuRuntime {
1346 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1347 let counts = self.core.registry.read().ok().map(|registry| {
1348 (
1349 registry.providers.len(),
1350 registry.agents.len(),
1351 registry.endpoints.len(),
1352 )
1353 });
1354 formatter
1355 .debug_struct("VifuRuntime")
1356 .field("project_id", &self.core.project_id)
1357 .field("resource_counts", &counts)
1358 .finish()
1359 }
1360}
1361
1362impl VifuRuntime {
1363 pub fn new(project_id: impl Into<String>) -> Result<Self, RuntimeError> {
1364 Self::with_store(project_id, Arc::new(MemoryRuntimeStore::default()))
1365 }
1366
1367 pub fn with_store(
1368 project_id: impl Into<String>,
1369 store: Arc<dyn RuntimeStore>,
1370 ) -> Result<Self, RuntimeError> {
1371 let project_id = project_id.into();
1372 validate_identifier("project", &project_id)?;
1373 let core = Arc::new(RuntimeCore {
1374 project_id,
1375 registry: RwLock::new(RuntimeRegistry::default()),
1376 manifest: RwLock::new(None),
1377 store,
1378 sessions: RwLock::new(HashMap::new()),
1379 session_locks: Mutex::new(HashMap::new()),
1380 invocations: Mutex::new(InvocationRegistry::default()),
1381 next_invocation: AtomicU64::new(1),
1382 });
1383 let worker = Arc::new(Mutex::new(None));
1384 Ok(Self { core, worker })
1385 }
1386
1387 pub fn project_id(&self) -> &str {
1388 &self.core.project_id
1389 }
1390
1391 pub fn register_provider(
1392 &self,
1393 name: impl Into<String>,
1394 provider: Arc<dyn AgentProvider>,
1395 ) -> Result<(), RuntimeError> {
1396 let name = name.into();
1397 validate_identifier("provider", &name)?;
1398 self.core
1399 .registry
1400 .write()
1401 .map_err(|_| RuntimeError::Internal)?
1402 .providers
1403 .insert(name, provider);
1404 Ok(())
1405 }
1406
1407 pub fn register_agent(&self, mut agent: AgentDefinition) -> Result<(), RuntimeError> {
1408 validate_identifier("agent", &agent.id)?;
1409 validate_identifier("provider", &agent.provider)?;
1410 if agent.name.trim().is_empty() || agent.capabilities.is_empty() {
1411 return Err(RuntimeError::InvalidDefinition(
1412 "agent name and at least one capability are required".to_string(),
1413 ));
1414 }
1415 for capability in &mut agent.capabilities {
1416 *capability = capability.trim().to_ascii_lowercase();
1417 validate_identifier("capability", capability)?;
1418 }
1419 agent.capabilities.sort();
1420 agent.capabilities.dedup();
1421 let mut registry = self
1422 .core
1423 .registry
1424 .write()
1425 .map_err(|_| RuntimeError::Internal)?;
1426 if !registry.providers.contains_key(&agent.provider) {
1427 return Err(RuntimeError::ProviderNotFound(agent.provider));
1428 }
1429 registry.agents.insert(agent.id.clone(), agent);
1430 Ok(())
1431 }
1432
1433 pub fn register_endpoint(&self, mut endpoint: EndpointDefinition) -> Result<(), RuntimeError> {
1434 validate_identifier("endpoint", &endpoint.name)?;
1435 validate_identifier("agent", &endpoint.agent)?;
1436 endpoint.capability = endpoint.capability.trim().to_ascii_lowercase();
1437 validate_identifier("capability", &endpoint.capability)?;
1438 if !(1..=MAX_TIMEOUT_MS).contains(&endpoint.timeout_ms) {
1439 return Err(RuntimeError::InvalidDefinition(format!(
1440 "endpoint timeout must be between 1 and {MAX_TIMEOUT_MS} ms"
1441 )));
1442 }
1443 let mut registry = self
1444 .core
1445 .registry
1446 .write()
1447 .map_err(|_| RuntimeError::Internal)?;
1448 let agent = registry
1449 .agents
1450 .get(&endpoint.agent)
1451 .ok_or_else(|| RuntimeError::AgentNotFound(endpoint.agent.clone()))?;
1452 if !agent
1453 .capabilities
1454 .iter()
1455 .any(|capability| capability == &endpoint.capability)
1456 {
1457 return Err(RuntimeError::CapabilityUnavailable {
1458 provider: agent.provider.clone(),
1459 capability: endpoint.capability,
1460 });
1461 }
1462 registry.endpoints.insert(endpoint.name.clone(), endpoint);
1463 Ok(())
1464 }
1465
1466 pub fn agent_definitions(&self) -> Result<Vec<AgentDefinition>, RuntimeError> {
1467 let mut agents = self
1468 .core
1469 .registry
1470 .read()
1471 .map_err(|_| RuntimeError::Internal)?
1472 .agents
1473 .values()
1474 .cloned()
1475 .collect::<Vec<_>>();
1476 agents.sort_by(|left, right| left.id.cmp(&right.id));
1477 Ok(agents)
1478 }
1479
1480 pub fn endpoint_definitions(&self) -> Result<Vec<EndpointDefinition>, RuntimeError> {
1481 let mut endpoints = self
1482 .core
1483 .registry
1484 .read()
1485 .map_err(|_| RuntimeError::Internal)?
1486 .endpoints
1487 .values()
1488 .cloned()
1489 .collect::<Vec<_>>();
1490 endpoints.sort_by(|left, right| left.name.cmp(&right.name));
1491 Ok(endpoints)
1492 }
1493
1494 pub fn apply_manifest(&self, manifest: RuntimeManifest) -> Result<(), RuntimeError> {
1497 manifest.validate()?;
1498 if manifest.project_id != self.core.project_id {
1499 return Err(RuntimeError::InvalidDefinition(
1500 "runtime manifest belongs to another project".to_string(),
1501 ));
1502 }
1503 let mut registry = self
1504 .core
1505 .registry
1506 .write()
1507 .map_err(|_| RuntimeError::Internal)?;
1508 for requirement in &manifest.providers {
1509 let provider = registry
1510 .providers
1511 .get(&requirement.id)
1512 .ok_or_else(|| RuntimeError::ProviderNotFound(requirement.id.clone()))?;
1513 for capability in &requirement.capabilities {
1514 if !provider.supports(capability) {
1515 return Err(RuntimeError::CapabilityUnavailable {
1516 provider: requirement.id.clone(),
1517 capability: capability.clone(),
1518 });
1519 }
1520 }
1521 }
1522 registry.agents = manifest
1523 .agents
1524 .iter()
1525 .cloned()
1526 .map(|agent| (agent.id.clone(), agent))
1527 .collect();
1528 registry.endpoints = manifest
1529 .endpoints
1530 .iter()
1531 .cloned()
1532 .map(|endpoint| (endpoint.name.clone(), endpoint))
1533 .collect();
1534 *self
1535 .core
1536 .manifest
1537 .write()
1538 .map_err(|_| RuntimeError::Internal)? = Some(manifest);
1539 Ok(())
1540 }
1541
1542 pub fn current_manifest(&self) -> Result<Option<RuntimeManifest>, RuntimeError> {
1543 Ok(self
1544 .core
1545 .manifest
1546 .read()
1547 .map_err(|_| RuntimeError::Internal)?
1548 .clone())
1549 }
1550
1551 pub fn install_release(&self, release: &RuntimeRelease) -> Result<(), RuntimeError> {
1552 release.validate()?;
1553 if release.manifest.project_id != self.core.project_id {
1554 return Err(RuntimeError::InvalidDefinition(
1555 "runtime release belongs to another project".to_string(),
1556 ));
1557 }
1558 self.core.store.save_release(release)
1559 }
1560
1561 pub fn releases(&self) -> Result<Vec<RuntimeRelease>, RuntimeError> {
1562 self.core.store.list_releases(&self.core.project_id)
1563 }
1564
1565 pub fn active_release_version(&self) -> Result<Option<u64>, RuntimeError> {
1566 self.core.store.active_release(&self.core.project_id)
1567 }
1568
1569 pub fn activate_release(&self, version: u64) -> Result<RuntimeRelease, RuntimeError> {
1570 let release = self
1571 .core
1572 .store
1573 .load_release(&self.core.project_id, version)?
1574 .ok_or_else(|| RuntimeError::store("runtime release was not found"))?;
1575 self.apply_manifest(release.manifest.clone())?;
1576 self.core
1577 .store
1578 .set_active_release(&self.core.project_id, version)?;
1579 Ok(release)
1580 }
1581
1582 pub fn restore_active_release(&self) -> Result<Option<RuntimeRelease>, RuntimeError> {
1583 self.active_release_version()?
1584 .map(|version| self.activate_release(version))
1585 .transpose()
1586 }
1587
1588 pub fn bootstrap_release(
1589 &self,
1590 manifest: RuntimeManifest,
1591 ) -> Result<RuntimeRelease, RuntimeError> {
1592 if let Some(active) = self.restore_active_release()? {
1593 return Ok(active);
1594 }
1595 let release = RuntimeRelease::new(1, manifest)?;
1596 self.install_release(&release)?;
1597 self.activate_release(release.version)
1598 }
1599
1600 pub fn save_local_provider_binding(
1601 &self,
1602 binding: &LocalProviderBinding,
1603 ) -> Result<(), RuntimeError> {
1604 validate_identifier("provider", &binding.provider_id)?;
1605 self.core
1606 .store
1607 .save_local_provider_binding(&self.core.project_id, binding)
1608 }
1609
1610 pub fn local_provider_bindings(&self) -> Result<Vec<LocalProviderBinding>, RuntimeError> {
1611 self.core
1612 .store
1613 .local_provider_bindings(&self.core.project_id)
1614 }
1615
1616 pub fn pending_traces(&self, limit: usize) -> Result<Vec<RuntimeTraceRecord>, RuntimeError> {
1617 self.core.store.pending_traces(limit.min(1_000))
1618 }
1619
1620 pub fn acknowledge_traces(&self, trace_ids: &[String]) -> Result<(), RuntimeError> {
1621 self.core.store.acknowledge_traces(trace_ids)
1622 }
1623
1624 pub fn session(&self, session_id: impl Into<String>) -> Result<RuntimeSession, RuntimeError> {
1625 let session_id = session_id.into();
1626 validate_identifier("session", &session_id)?;
1627 Ok(RuntimeSession {
1628 runtime: self.clone(),
1629 session_id,
1630 })
1631 }
1632
1633 pub async fn invoke(&self, input: InvocationInput) -> Result<InvocationOutput, RuntimeError> {
1634 let invocation_id = self.core.next_invocation_id();
1635 self.core
1636 .invoke(invocation_id, input, CancellationToken::default())
1637 .await
1638 }
1639
1640 pub fn start_invoke(&self, input: InvocationInput) -> Result<InvocationHandle, RuntimeError> {
1641 validate_identifier("endpoint", &input.endpoint)?;
1642 validate_identifier("session", &input.session_id)?;
1643 let handle = InvocationHandle(self.core.next_invocation_id());
1644 let cancellation = CancellationToken::default();
1645 let mut worker = self.worker.lock().map_err(|_| RuntimeError::Internal)?;
1646 if worker.is_none() {
1647 *worker = Some(RuntimeWorker::spawn(Arc::clone(&self.core))?);
1648 }
1649 self.core
1650 .invocations
1651 .lock()
1652 .map_err(|_| RuntimeError::Internal)?
1653 .insert(handle.clone(), cancellation.clone())?;
1654 let send_result =
1655 worker
1656 .as_ref()
1657 .ok_or(RuntimeError::Internal)?
1658 .send(WorkerCommand::Start {
1659 handle: handle.clone(),
1660 input,
1661 cancellation,
1662 });
1663 if let Err(error) = send_result {
1664 self.core
1665 .invocations
1666 .lock()
1667 .map_err(|_| RuntimeError::Internal)?
1668 .remove(&handle);
1669 return Err(error);
1670 }
1671 Ok(handle)
1672 }
1673
1674 pub fn poll_invocation(
1675 &self,
1676 handle: &InvocationHandle,
1677 ) -> Result<InvocationPoll, RuntimeError> {
1678 self.core
1679 .invocations
1680 .lock()
1681 .map_err(|_| RuntimeError::Internal)?
1682 .entries
1683 .get(&handle.0)
1684 .map(|entry| entry.poll.clone())
1685 .ok_or_else(|| RuntimeError::InvocationNotFound(handle.0.clone()))
1686 }
1687
1688 pub fn drain_invocation_events(
1690 &self,
1691 handle: &InvocationHandle,
1692 ) -> Result<Vec<InvocationEvent>, RuntimeError> {
1693 self.core
1694 .invocations
1695 .lock()
1696 .map_err(|_| RuntimeError::Internal)?
1697 .drain_events(handle)
1698 }
1699
1700 pub fn take_invocation(
1705 &self,
1706 handle: &InvocationHandle,
1707 ) -> Result<InvocationPoll, RuntimeError> {
1708 self.core
1709 .invocations
1710 .lock()
1711 .map_err(|_| RuntimeError::Internal)?
1712 .take(handle)
1713 }
1714
1715 pub fn cancel_invocation(&self, handle: &InvocationHandle) -> Result<(), RuntimeError> {
1716 let cancellation = self
1717 .core
1718 .invocations
1719 .lock()
1720 .map_err(|_| RuntimeError::Internal)?
1721 .entries
1722 .get(&handle.0)
1723 .map(|entry| entry.cancellation.clone())
1724 .ok_or_else(|| RuntimeError::InvocationNotFound(handle.0.clone()))?;
1725 cancellation.cancel();
1726 self.core
1727 .update_poll(handle, InvocationStatus::Cancelled, None, None);
1728 Ok(())
1729 }
1730
1731 pub async fn execute_effects(
1732 &self,
1733 effects: Vec<EffectRequest>,
1734 ) -> Result<EffectExecution, RuntimeError> {
1735 self.execute_effects_with_limit(effects, DEFAULT_EFFECT_LIMIT)
1736 .await
1737 }
1738
1739 pub async fn execute_effects_with_limit(
1740 &self,
1741 effects: Vec<EffectRequest>,
1742 limit: usize,
1743 ) -> Result<EffectExecution, RuntimeError> {
1744 if effects.len() > limit {
1745 return Err(RuntimeError::EffectLimitExceeded(limit));
1746 }
1747 let mut results = Vec::new();
1748 let mut unhandled = Vec::new();
1749 for effect in effects {
1750 if effect.kind != "agent.invoke" {
1751 unhandled.push(effect);
1752 continue;
1753 }
1754 let input = serde_json::from_value::<InvocationInput>(effect.payload.clone())
1755 .map_err(|error| RuntimeError::InvalidDefinition(error.to_string()))?;
1756 let result = self.invoke(input).await;
1757 match result {
1758 Ok(output) => results.push(EffectResult {
1759 effect_id: effect.id,
1760 succeeded: true,
1761 output: serde_json::to_value(output)
1762 .map_err(|_error| RuntimeError::Internal)?,
1763 }),
1764 Err(error) => results.push(EffectResult {
1765 effect_id: effect.id,
1766 succeeded: false,
1767 output: json!({ "error": error.public_message() }),
1768 }),
1769 }
1770 }
1771 Ok(EffectExecution { results, unhandled })
1772 }
1773
1774 pub fn export_snapshot(&self) -> Result<Vec<u8>, RuntimeError> {
1775 let snapshot = PortableProjectSnapshot {
1776 version: SNAPSHOT_VERSION,
1777 project_id: self.core.project_id.clone(),
1778 sessions: self
1779 .core
1780 .sessions
1781 .read()
1782 .map_err(|_| RuntimeError::Internal)?
1783 .clone(),
1784 };
1785 serde_json::to_vec(&snapshot).map_err(|error| RuntimeError::Snapshot(error.to_string()))
1786 }
1787
1788 pub fn restore_snapshot(&self, bytes: &[u8]) -> Result<(), RuntimeError> {
1789 let snapshot = serde_json::from_slice::<PortableProjectSnapshot>(bytes)
1790 .map_err(|error| RuntimeError::Snapshot(error.to_string()))?;
1791 if snapshot.version != SNAPSHOT_VERSION || snapshot.project_id != self.core.project_id {
1792 return Err(RuntimeError::Snapshot(
1793 "snapshot version or project does not match".to_string(),
1794 ));
1795 }
1796 for (session_id, state) in &snapshot.sessions {
1797 validate_identifier("session", session_id)?;
1798 self.core
1799 .store
1800 .save(&self.core.project_id, session_id, state)?;
1801 }
1802 *self
1803 .core
1804 .sessions
1805 .write()
1806 .map_err(|_| RuntimeError::Internal)? = snapshot.sessions;
1807 Ok(())
1808 }
1809}
1810
1811#[derive(Clone, Debug)]
1813pub struct RuntimeSession {
1814 runtime: VifuRuntime,
1815 session_id: String,
1816}
1817
1818impl RuntimeSession {
1819 pub fn id(&self) -> &str {
1820 &self.session_id
1821 }
1822
1823 pub async fn invoke(
1824 &self,
1825 mut input: InvocationInput,
1826 ) -> Result<InvocationOutput, RuntimeError> {
1827 input.session_id.clone_from(&self.session_id);
1828 self.runtime.invoke(input).await
1829 }
1830
1831 pub fn start_invoke(
1832 &self,
1833 mut input: InvocationInput,
1834 ) -> Result<InvocationHandle, RuntimeError> {
1835 input.session_id.clone_from(&self.session_id);
1836 self.runtime.start_invoke(input)
1837 }
1838}
1839
1840#[derive(Serialize, Deserialize)]
1841#[serde(rename_all = "camelCase")]
1842struct PortableProjectSnapshot {
1843 version: u32,
1844 project_id: String,
1845 sessions: HashMap<String, RuntimeSnapshot>,
1846}
1847
1848fn default_session_id() -> String {
1849 "default".to_string()
1850}
1851
1852const fn default_timeout_ms() -> u64 {
1853 DEFAULT_TIMEOUT_MS
1854}
1855
1856fn validate_identifier(kind: &str, value: &str) -> Result<(), RuntimeError> {
1857 if value.is_empty()
1858 || value.len() > 128
1859 || !value
1860 .bytes()
1861 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b':'))
1862 {
1863 return Err(RuntimeError::InvalidDefinition(format!(
1864 "{kind} must be a portable identifier"
1865 )));
1866 }
1867 Ok(())
1868}
1869
1870fn duration_ms(duration: Duration) -> u64 {
1871 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1872}
1873
1874const fn is_terminal_status(status: InvocationStatus) -> bool {
1875 matches!(
1876 status,
1877 InvocationStatus::Completed | InvocationStatus::Failed | InvocationStatus::Cancelled
1878 )
1879}
1880
1881#[cfg(test)]
1882mod tests {
1883 use super::*;
1884
1885 struct TestProvider {
1886 fail: bool,
1887 delay: Duration,
1888 }
1889
1890 impl TestProvider {
1891 fn immediate() -> Self {
1892 Self {
1893 fail: false,
1894 delay: Duration::ZERO,
1895 }
1896 }
1897 }
1898
1899 impl AgentProvider for TestProvider {
1900 fn supports(&self, capability: &str) -> bool {
1901 matches!(capability, "chat" | "speech" | "transcription")
1902 }
1903
1904 fn invoke<'a>(
1905 &'a self,
1906 request: ProviderRequest,
1907 cancellation: CancellationToken,
1908 ) -> ProviderFuture<'a> {
1909 Box::pin(async move {
1910 if !self.delay.is_zero() {
1911 tokio::select! {
1912 _ = tokio::time::sleep(self.delay) => {}
1913 _ = cancellation.cancelled() => {
1914 return Err(RuntimeError::Cancelled);
1915 }
1916 }
1917 }
1918 if self.fail {
1919 return Err(RuntimeError::provider(
1920 request.agent.provider,
1921 "synthetic provider failure",
1922 ));
1923 }
1924 Ok(ProviderResponse {
1925 data: match request.data {
1926 InvocationData::Json(data) => InvocationData::Json(json!({
1927 "capability": request.capability,
1928 "input": data,
1929 })),
1930 InvocationData::Binary(bytes) => InvocationData::Binary(bytes),
1931 },
1932 metadata: json!({}),
1933 state: Some(json!({
1934 "lastEndpoint": request.endpoint,
1935 "previousRevision": request.snapshot.revision,
1936 })),
1937 })
1938 })
1939 }
1940 }
1941
1942 struct StreamingTestProvider;
1943
1944 impl AgentProvider for StreamingTestProvider {
1945 fn supports(&self, capability: &str) -> bool {
1946 capability == "chat"
1947 }
1948
1949 fn invoke<'a>(
1950 &'a self,
1951 request: ProviderRequest,
1952 cancellation: CancellationToken,
1953 ) -> ProviderFuture<'a> {
1954 self.invoke_with_events(request, cancellation, ProviderEventSink::new(|_data| {}))
1955 }
1956
1957 fn invoke_with_events<'a>(
1958 &'a self,
1959 _request: ProviderRequest,
1960 cancellation: CancellationToken,
1961 events: ProviderEventSink,
1962 ) -> ProviderFuture<'a> {
1963 Box::pin(async move {
1964 if cancellation.is_cancelled() {
1965 return Err(RuntimeError::Cancelled);
1966 }
1967 events.output_delta(InvocationData::Json(Value::String("Hello".to_string())));
1968 events.output_delta(InvocationData::Json(Value::String(", world".to_string())));
1969 Ok(ProviderResponse::json(json!({ "text": "Hello, world" })))
1970 })
1971 }
1972 }
1973
1974 fn configured_runtime(provider: Arc<dyn AgentProvider>) -> VifuRuntime {
1975 let runtime = VifuRuntime::new("test-project").expect("runtime should start");
1976 runtime
1977 .register_provider("test-provider", provider)
1978 .expect("provider should register");
1979 runtime
1980 .register_agent(AgentDefinition {
1981 id: "guide".to_string(),
1982 name: "Guide".to_string(),
1983 provider: "test-provider".to_string(),
1984 capabilities: vec![
1985 "chat".to_string(),
1986 "speech".to_string(),
1987 "transcription".to_string(),
1988 ],
1989 metadata: json!({ "public": true }),
1990 })
1991 .expect("agent should register");
1992 for capability in ["chat", "speech", "transcription"] {
1993 runtime
1994 .register_endpoint(EndpointDefinition {
1995 name: capability.to_string(),
1996 agent: "guide".to_string(),
1997 capability: capability.to_string(),
1998 timeout_ms: 500,
1999 })
2000 .expect("endpoint should register");
2001 }
2002 runtime
2003 }
2004
2005 #[tokio::test(flavor = "current_thread")]
2006 async fn embedded_runtime_invokes_chat_speech_and_transcription_without_a_server() {
2007 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2008
2009 for capability in ["chat", "speech", "transcription"] {
2010 let output = runtime
2011 .invoke(InvocationInput::json(
2012 capability,
2013 json!({ "message": capability }),
2014 ))
2015 .await
2016 .expect("endpoint should invoke");
2017 assert_eq!(output.capability, capability);
2018 }
2019 }
2020
2021 #[tokio::test(flavor = "current_thread")]
2022 async fn runtime_sessions_keep_independent_durable_state() {
2023 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2024 let first = runtime
2025 .session("player-one")
2026 .expect("first session should open");
2027 let second = runtime
2028 .session("player-two")
2029 .expect("second session should open");
2030
2031 let first_output = first
2032 .invoke(InvocationInput::json("chat", json!({ "text": "one" })))
2033 .await
2034 .expect("first session should invoke");
2035 let second_output = second
2036 .invoke(InvocationInput::json("chat", json!({ "text": "two" })))
2037 .await
2038 .expect("second session should invoke");
2039
2040 assert_eq!(first_output.snapshot.revision, 1);
2041 assert_eq!(second_output.snapshot.revision, 1);
2042 }
2043
2044 #[tokio::test(flavor = "current_thread")]
2045 async fn concurrent_calls_serialize_state_updates_for_one_session() {
2046 let runtime = configured_runtime(Arc::new(TestProvider {
2047 fail: false,
2048 delay: Duration::from_millis(5),
2049 }));
2050 let first = runtime.invoke(
2051 InvocationInput::json("chat", json!({ "text": "one" })).with_session("shared-session"),
2052 );
2053 let second = runtime.invoke(
2054 InvocationInput::json("chat", json!({ "text": "two" })).with_session("shared-session"),
2055 );
2056
2057 let (first, second) = tokio::join!(first, second);
2058 let mut revisions = [
2059 first
2060 .expect("first invocation should complete")
2061 .snapshot
2062 .revision,
2063 second
2064 .expect("second invocation should complete")
2065 .snapshot
2066 .revision,
2067 ];
2068 revisions.sort_unstable();
2069 assert_eq!(revisions, [1, 2]);
2070 }
2071
2072 #[tokio::test(flavor = "current_thread")]
2073 async fn runtime_round_trips_binary_provider_results() {
2074 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2075 let output = runtime
2076 .invoke(InvocationInput {
2077 endpoint: "speech".to_string(),
2078 session_id: "audio-session".to_string(),
2079 data: InvocationData::Binary(vec![1, 2, 3, 4]),
2080 metadata: json!({}),
2081 })
2082 .await
2083 .expect("binary invocation should complete");
2084
2085 assert_eq!(output.data, InvocationData::Binary(vec![1, 2, 3, 4]));
2086 }
2087
2088 #[tokio::test(flavor = "current_thread")]
2089 async fn runtime_times_out_slow_providers() {
2090 let runtime = VifuRuntime::new("timeout-project").expect("runtime should start");
2091 runtime
2092 .register_provider(
2093 "slow",
2094 Arc::new(TestProvider {
2095 fail: false,
2096 delay: Duration::from_millis(100),
2097 }),
2098 )
2099 .expect("provider should register");
2100 runtime
2101 .register_agent(AgentDefinition {
2102 id: "slow-agent".to_string(),
2103 name: "Slow agent".to_string(),
2104 provider: "slow".to_string(),
2105 capabilities: vec!["chat".to_string()],
2106 metadata: json!({}),
2107 })
2108 .expect("agent should register");
2109 runtime
2110 .register_endpoint(EndpointDefinition {
2111 name: "slow-chat".to_string(),
2112 agent: "slow-agent".to_string(),
2113 capability: "chat".to_string(),
2114 timeout_ms: 10,
2115 })
2116 .expect("endpoint should register");
2117
2118 let error = runtime
2119 .invoke(InvocationInput::json("slow-chat", json!({})))
2120 .await
2121 .expect_err("slow invocation should time out");
2122 assert!(matches!(error, RuntimeError::Timeout(10)));
2123 }
2124
2125 #[test]
2126 fn game_loop_api_starts_polls_and_cancels_invocations() {
2127 let runtime = configured_runtime(Arc::new(TestProvider {
2128 fail: false,
2129 delay: Duration::from_secs(5),
2130 }));
2131 let handle = runtime
2132 .start_invoke(InvocationInput::json("chat", json!({})))
2133 .expect("invocation should start");
2134 let running_deadline = Instant::now() + Duration::from_secs(1);
2135 loop {
2136 let poll = runtime
2137 .poll_invocation(&handle)
2138 .expect("invocation should remain pollable");
2139 if poll.status == InvocationStatus::Running {
2140 break;
2141 }
2142 assert!(
2143 Instant::now() < running_deadline,
2144 "invocation did not start"
2145 );
2146 std::thread::sleep(Duration::from_millis(5));
2147 }
2148 runtime
2149 .cancel_invocation(&handle)
2150 .expect("invocation should cancel");
2151
2152 let deadline = Instant::now() + Duration::from_secs(1);
2153 loop {
2154 let poll = runtime
2155 .poll_invocation(&handle)
2156 .expect("invocation should remain pollable");
2157 if poll.status == InvocationStatus::Cancelled {
2158 break;
2159 }
2160 assert!(
2161 Instant::now() < deadline,
2162 "cancelled provider did not observe cancellation"
2163 );
2164 std::thread::sleep(Duration::from_millis(5));
2165 }
2166 }
2167
2168 #[test]
2169 fn game_loop_poll_returns_the_same_provider_result_shape_as_async_invoke() {
2170 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2171 let handle = runtime
2172 .start_invoke(
2173 InvocationInput::json("chat", json!({ "text": "hello" }))
2174 .with_session("poll-session"),
2175 )
2176 .expect("invocation should start");
2177 let deadline = Instant::now() + Duration::from_secs(1);
2178 let output = loop {
2179 let poll = runtime
2180 .poll_invocation(&handle)
2181 .expect("invocation should remain pollable");
2182 if let Some(output) = poll.output {
2183 break output;
2184 }
2185 assert!(
2186 !matches!(
2187 poll.status,
2188 InvocationStatus::Failed | InvocationStatus::Cancelled
2189 ),
2190 "invocation unexpectedly failed: {poll:?}"
2191 );
2192 assert!(Instant::now() < deadline, "invocation did not complete");
2193 std::thread::sleep(Duration::from_millis(5));
2194 };
2195
2196 assert_eq!(
2197 output.data,
2198 InvocationData::Json(json!({
2199 "capability": "chat",
2200 "input": { "text": "hello" },
2201 }))
2202 );
2203 }
2204
2205 #[test]
2206 fn game_loop_api_drains_ordered_streaming_events() {
2207 let runtime = configured_runtime(Arc::new(StreamingTestProvider));
2208 let handle = runtime
2209 .start_invoke(InvocationInput::json("chat", json!({})))
2210 .expect("invocation should start");
2211 let deadline = Instant::now() + Duration::from_secs(1);
2212 loop {
2213 let poll = runtime
2214 .poll_invocation(&handle)
2215 .expect("invocation should remain pollable");
2216 if poll.status == InvocationStatus::Completed {
2217 break;
2218 }
2219 assert!(Instant::now() < deadline, "invocation did not complete");
2220 std::thread::sleep(Duration::from_millis(5));
2221 }
2222
2223 let events = runtime
2224 .drain_invocation_events(&handle)
2225 .expect("events should be available");
2226 assert_eq!(
2227 events.iter().map(|event| event.kind).collect::<Vec<_>>(),
2228 vec![
2229 InvocationEventKind::Started,
2230 InvocationEventKind::OutputDelta,
2231 InvocationEventKind::Completed,
2232 ]
2233 );
2234 assert_eq!(
2235 events[1].data,
2236 Some(InvocationData::Json(Value::String(
2237 "Hello, world".to_string()
2238 )))
2239 );
2240 }
2241
2242 #[test]
2243 fn invocation_registry_ignores_output_after_terminal_event() {
2244 let handle = InvocationHandle("late-output".to_string());
2245 let mut registry = InvocationRegistry::default();
2246 registry
2247 .insert(handle.clone(), CancellationToken::default())
2248 .expect("invocation should be registered");
2249 registry.update(&handle, InvocationStatus::Running, None, None);
2250 registry.update(
2251 &handle,
2252 InvocationStatus::Failed,
2253 None,
2254 Some("provider failed".to_string()),
2255 );
2256
2257 registry.push_output_delta(
2258 &handle,
2259 InvocationData::Json(Value::String("too late".to_string())),
2260 );
2261
2262 let events = registry
2263 .drain_events(&handle)
2264 .expect("events should remain available");
2265 assert_eq!(
2266 events.iter().map(|event| event.kind).collect::<Vec<_>>(),
2267 vec![InvocationEventKind::Started, InvocationEventKind::Failed]
2268 );
2269 }
2270
2271 #[test]
2272 fn taking_a_terminal_invocation_releases_its_result() {
2273 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2274 let handle = runtime
2275 .start_invoke(InvocationInput::json("chat", json!({})))
2276 .expect("invocation should start");
2277 let deadline = Instant::now() + Duration::from_secs(1);
2278 loop {
2279 let poll = runtime
2280 .take_invocation(&handle)
2281 .expect("invocation should remain available until terminal");
2282 if is_terminal_status(poll.status) {
2283 break;
2284 }
2285 assert!(Instant::now() < deadline, "invocation did not complete");
2286 std::thread::sleep(Duration::from_millis(5));
2287 }
2288
2289 assert!(matches!(
2290 runtime.poll_invocation(&handle),
2291 Err(RuntimeError::InvocationNotFound(_))
2292 ));
2293 }
2294
2295 #[test]
2296 fn game_loop_api_applies_backpressure_to_excess_invocations() {
2297 let runtime = configured_runtime(Arc::new(TestProvider {
2298 fail: false,
2299 delay: Duration::from_secs(5),
2300 }));
2301 let handles = (0..MAX_IN_FLIGHT_INVOCATIONS)
2302 .map(|index| {
2303 runtime
2304 .start_invoke(
2305 InvocationInput::json("chat", json!({}))
2306 .with_session(format!("session-{index}")),
2307 )
2308 .expect("invocation within the bound should start")
2309 })
2310 .collect::<Vec<_>>();
2311
2312 let error = runtime
2313 .start_invoke(
2314 InvocationInput::json("chat", json!({})).with_session("one-session-too-many"),
2315 )
2316 .expect_err("invocations above the bound should be rejected");
2317 assert!(matches!(error, RuntimeError::Backpressure(_)));
2318
2319 for handle in handles {
2320 runtime
2321 .cancel_invocation(&handle)
2322 .expect("test invocation should cancel");
2323 }
2324 }
2325
2326 #[test]
2327 fn terminal_invocation_history_is_bounded() {
2328 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2329 let first = runtime
2330 .start_invoke(
2331 InvocationInput::json("chat", json!({})).with_session("retained-session-0"),
2332 )
2333 .expect("first invocation should start");
2334 let mut last = first.clone();
2335 for index in 0..=MAX_RETAINED_INVOCATIONS {
2336 let handle = if index == 0 {
2337 first.clone()
2338 } else {
2339 runtime
2340 .start_invoke(
2341 InvocationInput::json("chat", json!({}))
2342 .with_session(format!("retained-session-{index}")),
2343 )
2344 .expect("invocation should start")
2345 };
2346 let deadline = Instant::now() + Duration::from_secs(1);
2347 loop {
2348 let poll = runtime
2349 .poll_invocation(&handle)
2350 .expect("latest invocation should remain available");
2351 if is_terminal_status(poll.status) {
2352 break;
2353 }
2354 assert!(Instant::now() < deadline, "invocation did not complete");
2355 std::thread::sleep(Duration::from_millis(2));
2356 }
2357 last = handle;
2358 }
2359
2360 assert!(matches!(
2361 runtime.poll_invocation(&first),
2362 Err(RuntimeError::InvocationNotFound(_))
2363 ));
2364 assert!(runtime.poll_invocation(&last).is_ok());
2365 }
2366
2367 #[tokio::test(flavor = "current_thread")]
2368 async fn runtime_executes_agent_effects_and_returns_custom_effects_to_the_host() {
2369 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2370 let execution = runtime
2371 .execute_effects(vec![
2372 EffectRequest {
2373 id: "agent-effect".to_string(),
2374 kind: "agent.invoke".to_string(),
2375 payload: serde_json::to_value(InvocationInput::json(
2376 "chat",
2377 json!({ "text": "hello" }),
2378 ))
2379 .unwrap(),
2380 },
2381 EffectRequest {
2382 id: "host-effect".to_string(),
2383 kind: "game.play_animation".to_string(),
2384 payload: json!({ "name": "wave" }),
2385 },
2386 ])
2387 .await
2388 .expect("effects should execute");
2389
2390 assert_eq!(execution.results.len(), 1);
2391 assert_eq!(execution.unhandled[0].kind, "game.play_animation");
2392 }
2393
2394 #[tokio::test(flavor = "current_thread")]
2395 async fn runtime_rejects_effect_batches_above_the_bound() {
2396 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2397 let effects = (0..3)
2398 .map(|index| EffectRequest {
2399 id: format!("effect-{index}"),
2400 kind: "host.effect".to_string(),
2401 payload: json!({}),
2402 })
2403 .collect();
2404
2405 let error = runtime
2406 .execute_effects_with_limit(effects, 2)
2407 .await
2408 .expect_err("oversized effect batch should fail");
2409 assert!(matches!(error, RuntimeError::EffectLimitExceeded(2)));
2410 }
2411
2412 #[tokio::test(flavor = "current_thread")]
2413 async fn snapshots_restore_session_state_without_runtime_definitions_or_secrets() {
2414 let secret = "synthetic-secret-must-not-leak";
2415 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2416 runtime
2417 .invoke(
2418 InvocationInput::json("chat", json!({ "text": "hello" }))
2419 .with_session("saved-session"),
2420 )
2421 .await
2422 .expect("invocation should create state");
2423 let bytes = runtime.export_snapshot().expect("snapshot should export");
2424 assert!(!String::from_utf8_lossy(&bytes).contains(secret));
2425
2426 let restored = configured_runtime(Arc::new(TestProvider::immediate()));
2427 restored
2428 .restore_snapshot(&bytes)
2429 .expect("snapshot should restore");
2430 let output = restored
2431 .invoke(
2432 InvocationInput::json("chat", json!({ "text": "again" }))
2433 .with_session("saved-session"),
2434 )
2435 .await
2436 .expect("restored session should invoke");
2437
2438 assert_eq!(output.snapshot.revision, 2);
2439 }
2440
2441 #[test]
2442 fn debug_output_redacts_payloads_provider_errors_and_snapshots() {
2443 let secret = "synthetic-secret-must-not-leak";
2444 let input = InvocationInput::json("chat", json!({ "secret": secret }));
2445 let error = RuntimeError::provider("test-provider", secret);
2446 let runtime = configured_runtime(Arc::new(TestProvider::immediate()));
2447
2448 assert!(!format!("{input:?}").contains(secret));
2449 assert!(!format!("{error:?}").contains(secret));
2450 assert!(!format!("{runtime:?}").contains(secret));
2451 }
2452}