1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use sha2::{Digest, Sha256};
9use thiserror::Error;
10use tokio::sync::Mutex;
11
12use crate::ConnectorDescription;
13
14pub const DEFAULT_MAX_EXECUTIONS: usize = 50;
16pub const MAX_DURABLE_VALUE_BYTES: usize = 1_000_000;
18pub const DEFAULT_PAUSED_TTL_MS: u64 = 24 * 60 * 60 * 1_000;
20pub const DEFAULT_MAX_EVENTS: usize = 1_000;
22
23static EXECUTION_SEQUENCE: AtomicU64 = AtomicU64::new(0);
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct CapabilitySnapshot {
28 pub connectors: Vec<ConnectorDescription>,
30 pub fingerprint: String,
32}
33
34impl CapabilitySnapshot {
35 pub fn new(connectors: Vec<ConnectorDescription>) -> Result<Self, String> {
37 let bytes = serde_json::to_vec(&connectors).map_err(|error| error.to_string())?;
38 Ok(Self {
39 connectors,
40 fingerprint: format!("{:x}", Sha256::digest(bytes)),
41 })
42 }
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ArtifactRef {
48 pub id: String,
50 pub execution_id: String,
52 pub bytes: usize,
54 pub preview: String,
56}
57
58#[async_trait]
60pub trait ArtifactStore: Send + Sync {
61 async fn put(&self, execution_id: &str, value: &Value) -> Result<ArtifactRef, String>;
63 async fn get(&self, execution_id: &str, artifact_id: &str) -> Result<Option<Value>, String>;
65 async fn delete_execution(&self, execution_id: &str) -> Result<(), String>;
67}
68
69#[derive(Clone)]
70struct MemoryArtifact {
71 execution_id: String,
72 value: Value,
73}
74
75#[derive(Default)]
77pub struct MemoryArtifactStore {
78 values: Mutex<BTreeMap<String, MemoryArtifact>>,
79}
80
81#[async_trait]
82impl ArtifactStore for MemoryArtifactStore {
83 async fn put(&self, execution_id: &str, value: &Value) -> Result<ArtifactRef, String> {
84 let bytes = serde_json::to_vec(value).map_err(|error| error.to_string())?;
85 let id = format!(
86 "{:x}",
87 Sha256::digest([execution_id.as_bytes(), &bytes].concat())
88 );
89 self.values.lock().await.insert(
90 id.clone(),
91 MemoryArtifact {
92 execution_id: execution_id.to_string(),
93 value: value.clone(),
94 },
95 );
96 Ok(ArtifactRef {
97 id,
98 execution_id: execution_id.to_string(),
99 bytes: bytes.len(),
100 preview: truncate_bytes(String::from_utf8_lossy(&bytes).into_owned(), 512),
101 })
102 }
103
104 async fn get(&self, execution_id: &str, artifact_id: &str) -> Result<Option<Value>, String> {
105 Ok(self
106 .values
107 .lock()
108 .await
109 .get(artifact_id)
110 .filter(|artifact| artifact.execution_id == execution_id)
111 .map(|artifact| artifact.value.clone()))
112 }
113
114 async fn delete_execution(&self, execution_id: &str) -> Result<(), String> {
115 self.values
116 .lock()
117 .await
118 .retain(|_, artifact| artifact.execution_id != execution_id);
119 Ok(())
120 }
121}
122
123#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
125#[serde(rename_all = "snake_case")]
126pub enum ExecutionStatus {
127 Running,
129 Paused,
131 Completed,
133 Error,
135 Rejected,
137 RolledBack,
139 Cancelled,
141}
142
143impl ExecutionStatus {
144 fn terminal(self) -> bool {
145 matches!(
146 self,
147 Self::Completed | Self::Error | Self::Rejected | Self::RolledBack | Self::Cancelled
148 )
149 }
150}
151
152#[derive(Debug, Clone, Serialize, Deserialize)]
154#[serde(tag = "type", rename_all = "snake_case")]
155pub enum ExecutionEvent {
156 Status {
158 status: ExecutionStatus,
160 at: u64,
162 },
163 Chunk {
165 data: Value,
167 at: u64,
169 },
170 Log {
172 level: String,
174 message: String,
176 at: u64,
178 },
179 Progress {
181 message: String,
183 fraction: Option<f64>,
185 at: u64,
187 },
188}
189
190#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
192#[serde(rename_all = "snake_case")]
193pub enum LogEntryState {
194 Pending,
196 Executing,
198 Applied,
200 Reverted,
202}
203
204#[derive(Debug, Clone, Serialize, Deserialize)]
206pub struct LogEntry {
207 pub seq: u64,
209 pub connector: String,
211 pub method: String,
213 pub arguments: Value,
215 pub result: Option<Value>,
217 pub requires_approval: bool,
219 pub ephemeral: bool,
221 pub state: LogEntryState,
223}
224
225#[derive(Debug, Clone, Serialize, Deserialize)]
227pub struct ExecutionState {
228 pub id: String,
230 pub code: String,
232 pub status: ExecutionStatus,
234 pub log: Vec<LogEntry>,
236 pub result: Option<Value>,
238 pub error: Option<String>,
240 pub logs: Vec<String>,
242 pub connectors: Vec<String>,
244 #[serde(default)]
246 pub capabilities: Option<CapabilitySnapshot>,
247 #[serde(default)]
249 pub events: Vec<ExecutionEvent>,
250 pub created_at: u64,
252 pub updated_at: u64,
254}
255
256#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
258pub struct PendingAction {
259 pub execution_id: String,
261 pub seq: u64,
263 pub connector: String,
265 pub method: String,
267 pub arguments: Value,
269}
270
271#[derive(Debug, Clone, PartialEq)]
273pub enum ToolDecision {
274 Replay(Value),
276 Execute(u64),
278 Pause(u64),
280}
281
282#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
284pub struct Snippet {
285 pub name: String,
287 pub description: String,
289 pub code: String,
291 pub saved_at: u64,
293 pub input_schema: Option<Value>,
295 pub connectors: Vec<String>,
297}
298
299#[async_trait]
301pub trait RuntimeStore: Send + Sync {
302 async fn get_execution(&self, id: &str) -> Result<Option<ExecutionState>, String>;
304 async fn put_execution(&self, execution: &ExecutionState) -> Result<(), String>;
306 async fn list_executions(&self) -> Result<Vec<ExecutionState>, String>;
308 async fn delete_execution(&self, id: &str) -> Result<(), String>;
310 async fn get_snippet(&self, name: &str) -> Result<Option<Snippet>, String>;
312 async fn put_snippet(&self, snippet: &Snippet) -> Result<(), String>;
314 async fn list_snippets(&self) -> Result<Vec<Snippet>, String>;
316 async fn delete_snippet(&self, name: &str) -> Result<bool, String>;
318}
319
320#[derive(Default)]
321struct MemoryState {
322 executions: BTreeMap<String, ExecutionState>,
323 snippets: BTreeMap<String, Snippet>,
324}
325
326#[derive(Default)]
328pub struct MemoryStore {
329 state: Mutex<MemoryState>,
330}
331
332#[async_trait]
333impl RuntimeStore for MemoryStore {
334 async fn get_execution(&self, id: &str) -> Result<Option<ExecutionState>, String> {
335 Ok(self.state.lock().await.executions.get(id).cloned())
336 }
337
338 async fn put_execution(&self, execution: &ExecutionState) -> Result<(), String> {
339 self.state
340 .lock()
341 .await
342 .executions
343 .insert(execution.id.clone(), execution.clone());
344 Ok(())
345 }
346
347 async fn list_executions(&self) -> Result<Vec<ExecutionState>, String> {
348 Ok(self
349 .state
350 .lock()
351 .await
352 .executions
353 .values()
354 .cloned()
355 .collect())
356 }
357
358 async fn delete_execution(&self, id: &str) -> Result<(), String> {
359 self.state.lock().await.executions.remove(id);
360 Ok(())
361 }
362
363 async fn get_snippet(&self, name: &str) -> Result<Option<Snippet>, String> {
364 Ok(self.state.lock().await.snippets.get(name).cloned())
365 }
366
367 async fn put_snippet(&self, snippet: &Snippet) -> Result<(), String> {
368 self.state
369 .lock()
370 .await
371 .snippets
372 .insert(snippet.name.clone(), snippet.clone());
373 Ok(())
374 }
375
376 async fn list_snippets(&self) -> Result<Vec<Snippet>, String> {
377 Ok(self.state.lock().await.snippets.values().cloned().collect())
378 }
379
380 async fn delete_snippet(&self, name: &str) -> Result<bool, String> {
381 Ok(self.state.lock().await.snippets.remove(name).is_some())
382 }
383}
384
385#[derive(Debug, Error)]
387pub enum RuntimeError {
388 #[error("Execution \"{0}\" not found")]
390 MissingExecution(String),
391 #[error("{0}")]
393 InvalidState(String),
394 #[error("{0}")]
396 DurableValue(String),
397 #[error("Runtime store failed: {0}")]
399 Store(String),
400}
401
402pub struct CodeModeRuntime {
404 store: Arc<dyn RuntimeStore>,
405 artifacts: Arc<dyn ArtifactStore>,
406 gate: Mutex<()>,
407}
408
409impl CodeModeRuntime {
410 pub fn new(store: Arc<dyn RuntimeStore>) -> Self {
412 Self::with_artifacts(store, Arc::new(MemoryArtifactStore::default()))
413 }
414
415 pub fn with_artifacts(store: Arc<dyn RuntimeStore>, artifacts: Arc<dyn ArtifactStore>) -> Self {
417 Self {
418 store,
419 artifacts,
420 gate: Mutex::new(()),
421 }
422 }
423
424 pub async fn begin(
426 &self,
427 code: impl Into<String>,
428 connectors: Vec<String>,
429 now: u64,
430 ) -> Result<String, RuntimeError> {
431 self.begin_inner(code.into(), connectors, None, now).await
432 }
433
434 pub async fn begin_with_capabilities(
436 &self,
437 code: &str,
438 capabilities: CapabilitySnapshot,
439 now: u64,
440 ) -> Result<String, RuntimeError> {
441 let connectors = capabilities
442 .connectors
443 .iter()
444 .map(|connector| connector.name.clone())
445 .collect();
446 self.begin_inner(code.to_string(), connectors, Some(capabilities), now)
447 .await
448 }
449
450 async fn begin_inner(
451 &self,
452 code: String,
453 connectors: Vec<String>,
454 capabilities: Option<CapabilitySnapshot>,
455 now: u64,
456 ) -> Result<String, RuntimeError> {
457 let _guard = self.gate.lock().await;
458 ensure_size("The execution code", &Value::String(code.clone()))?;
459 let sequence = EXECUTION_SEQUENCE.fetch_add(1, Ordering::Relaxed);
460 let id = format!("exec_{now:016}_{sequence:016x}");
461 self.store
462 .put_execution(&ExecutionState {
463 id: id.clone(),
464 code,
465 status: ExecutionStatus::Running,
466 log: Vec::new(),
467 result: None,
468 error: None,
469 logs: Vec::new(),
470 connectors,
471 capabilities,
472 events: vec![ExecutionEvent::Status {
473 status: ExecutionStatus::Running,
474 at: now,
475 }],
476 created_at: now,
477 updated_at: now,
478 })
479 .await
480 .map_err(RuntimeError::Store)?;
481 self.prune_locked(DEFAULT_MAX_EXECUTIONS, Some(&id)).await?;
482 Ok(id)
483 }
484
485 pub async fn resume(&self, id: &str, now: u64) -> Result<ExecutionState, RuntimeError> {
487 let _guard = self.gate.lock().await;
488 let mut execution = self.require(id).await?;
489 if execution.status != ExecutionStatus::Paused {
490 return Err(RuntimeError::InvalidState(format!(
491 "Execution \"{id}\" is {:?}, not paused",
492 execution.status
493 )));
494 }
495 execution.status = ExecutionStatus::Running;
496 push_event(
497 &mut execution,
498 ExecutionEvent::Status {
499 status: ExecutionStatus::Running,
500 at: now,
501 },
502 );
503 execution.updated_at = now;
504 self.save(&execution).await?;
505 Ok(execution)
506 }
507
508 pub async fn cancel(&self, id: &str, now: u64) -> Result<bool, RuntimeError> {
510 let _guard = self.gate.lock().await;
511 let mut execution = self.require(id).await?;
512 if execution.status.terminal() {
513 return Ok(false);
514 }
515 execution.status = ExecutionStatus::Cancelled;
516 execution.error = Some("Execution cancelled".to_string());
517 execution.updated_at = now;
518 push_event(
519 &mut execution,
520 ExecutionEvent::Status {
521 status: ExecutionStatus::Cancelled,
522 at: now,
523 },
524 );
525 self.save(&execution).await?;
526 Ok(true)
527 }
528
529 pub async fn event(&self, id: &str, event: ExecutionEvent) -> Result<(), RuntimeError> {
531 let _guard = self.gate.lock().await;
532 let mut execution = self.require(id).await?;
533 push_event(&mut execution, event);
534 self.save(&execution).await
535 }
536
537 pub async fn artifact(
539 &self,
540 execution_id: &str,
541 artifact_id: &str,
542 ) -> Result<Option<Value>, RuntimeError> {
543 self.artifacts
544 .get(execution_id, artifact_id)
545 .await
546 .map_err(RuntimeError::Store)
547 }
548
549 async fn spill(&self, execution_id: &str, value: Value) -> Result<Value, RuntimeError> {
550 if serialized_size(&value) <= MAX_DURABLE_VALUE_BYTES {
551 return Ok(value);
552 }
553 let artifact = self
554 .artifacts
555 .put(execution_id, &value)
556 .await
557 .map_err(RuntimeError::Store)?;
558 Ok(serde_json::json!({ "$artifact": artifact }))
559 }
560
561 async fn rehydrate(&self, execution_id: &str, value: Value) -> Result<Value, RuntimeError> {
562 let Some(artifact) = artifact_ref(&value)? else {
563 return Ok(value);
564 };
565 self.artifacts
566 .get(execution_id, &artifact.id)
567 .await
568 .map_err(RuntimeError::Store)?
569 .ok_or_else(|| {
570 RuntimeError::Store(format!("Artifact \"{}\" is unavailable", artifact.id))
571 })
572 }
573
574 async fn rehydrate_execution(
575 &self,
576 mut execution: ExecutionState,
577 ) -> Result<ExecutionState, RuntimeError> {
578 for entry in &mut execution.log {
579 entry.arguments = self
580 .rehydrate(&execution.id, entry.arguments.clone())
581 .await?;
582 if let Some(result) = entry.result.take() {
583 entry.result = Some(self.rehydrate(&execution.id, result).await?);
584 }
585 }
586 if let Some(result) = execution.result.take() {
587 execution.result = Some(self.rehydrate(&execution.id, result).await?);
588 }
589 Ok(execution)
590 }
591
592 async fn rehydrate_entry(
593 &self,
594 execution_id: &str,
595 mut entry: LogEntry,
596 ) -> Result<LogEntry, RuntimeError> {
597 entry.arguments = self.rehydrate(execution_id, entry.arguments).await?;
598 if let Some(result) = entry.result.take() {
599 entry.result = Some(self.rehydrate(execution_id, result).await?);
600 }
601 Ok(entry)
602 }
603
604 async fn rehydrate_pending(
605 &self,
606 execution_id: &str,
607 mut action: PendingAction,
608 ) -> Result<PendingAction, RuntimeError> {
609 action.arguments = self.rehydrate(execution_id, action.arguments).await?;
610 Ok(action)
611 }
612
613 #[allow(clippy::too_many_arguments)]
615 pub async fn decide(
616 &self,
617 id: &str,
618 seq: u64,
619 connector: &str,
620 method: &str,
621 arguments: Value,
622 requires_approval: bool,
623 ephemeral: bool,
624 now: u64,
625 ) -> Result<ToolDecision, RuntimeError> {
626 let _guard = self.gate.lock().await;
627 let mut execution = self.require(id).await?;
628 if execution.status != ExecutionStatus::Running {
629 return Ok(ToolDecision::Pause(seq));
630 }
631 if requires_approval && ephemeral {
632 return Err(RuntimeError::InvalidState(format!(
633 "{connector}.{method} cannot require approval and re-execute on replay"
634 )));
635 }
636 let arguments = self.spill(id, arguments).await?;
637 if let Some(entry) = execution.log.iter_mut().find(|entry| entry.seq == seq) {
638 if entry.connector != connector
639 || entry.method != method
640 || stable_json(&entry.arguments) != stable_json(&arguments)
641 {
642 execution.status = ExecutionStatus::Error;
643 execution.error = Some(format!(
644 "Replay diverged at step {seq}: expected {}.{}, got {connector}.{method}. \
645 Code must be deterministic up to tool calls and steps.",
646 entry.connector, entry.method
647 ));
648 execution.updated_at = now;
649 self.save(&execution).await?;
650 return Ok(ToolDecision::Pause(seq));
651 }
652 return match entry.state {
653 LogEntryState::Applied if !entry.ephemeral => Ok(ToolDecision::Replay(
654 self.rehydrate(id, entry.result.clone().unwrap_or(Value::Null))
655 .await?,
656 )),
657 LogEntryState::Pending | LogEntryState::Executing | LogEntryState::Applied => {
658 entry.state = LogEntryState::Executing;
659 execution.updated_at = now;
660 self.save(&execution).await?;
661 Ok(ToolDecision::Execute(seq))
662 }
663 LogEntryState::Reverted => {
664 entry.state = if requires_approval {
665 LogEntryState::Pending
666 } else {
667 LogEntryState::Executing
668 };
669 entry.arguments = arguments;
670 entry.result = None;
671 entry.requires_approval = requires_approval;
672 entry.ephemeral = ephemeral;
673 if requires_approval {
674 execution.status = ExecutionStatus::Paused;
675 push_event(
676 &mut execution,
677 ExecutionEvent::Status {
678 status: ExecutionStatus::Paused,
679 at: now,
680 },
681 );
682 }
683 execution.updated_at = now;
684 self.save(&execution).await?;
685 Ok(if requires_approval {
686 ToolDecision::Pause(seq)
687 } else {
688 ToolDecision::Execute(seq)
689 })
690 }
691 };
692 }
693 execution.log.push(LogEntry {
694 seq,
695 connector: connector.to_string(),
696 method: method.to_string(),
697 arguments,
698 result: None,
699 requires_approval,
700 ephemeral,
701 state: if requires_approval {
702 LogEntryState::Pending
703 } else {
704 LogEntryState::Executing
705 },
706 });
707 if requires_approval {
708 execution.status = ExecutionStatus::Paused;
709 push_event(
710 &mut execution,
711 ExecutionEvent::Status {
712 status: ExecutionStatus::Paused,
713 at: now,
714 },
715 );
716 }
717 execution.updated_at = now;
718 self.save(&execution).await?;
719 Ok(if requires_approval {
720 ToolDecision::Pause(seq)
721 } else {
722 ToolDecision::Execute(seq)
723 })
724 }
725
726 pub async fn record_result(
728 &self,
729 id: &str,
730 seq: u64,
731 result: Value,
732 now: u64,
733 ) -> Result<(), RuntimeError> {
734 let _guard = self.gate.lock().await;
735 let mut execution = self.require(id).await?;
736 let result = self.spill(id, result).await?;
737 if execution.status != ExecutionStatus::Running {
738 return Err(RuntimeError::InvalidState(format!(
739 "Execution \"{id}\" is {:?}, not running",
740 execution.status
741 )));
742 }
743 let Some(entry) = execution.log.iter_mut().find(|entry| entry.seq == seq) else {
744 return Err(RuntimeError::InvalidState(format!(
745 "No log entry at step {seq}"
746 )));
747 };
748 if entry.state != LogEntryState::Executing {
749 return Err(RuntimeError::InvalidState(format!(
750 "Log entry at step {seq} is {:?}, not executing",
751 entry.state
752 )));
753 }
754 if !entry.ephemeral {
755 entry.result = Some(result);
756 }
757 entry.state = LogEntryState::Applied;
758 execution.updated_at = now;
759 self.save(&execution).await
760 }
761
762 pub async fn complete(
764 &self,
765 id: &str,
766 result: Value,
767 logs: Vec<String>,
768 now: u64,
769 ) -> Result<(), RuntimeError> {
770 let _guard = self.gate.lock().await;
771 let mut execution = self.require(id).await?;
772 let result = self.spill(id, result).await?;
773 execution.status = ExecutionStatus::Completed;
774 push_event(
775 &mut execution,
776 ExecutionEvent::Status {
777 status: ExecutionStatus::Completed,
778 at: now,
779 },
780 );
781 execution.result = Some(result);
782 execution.logs = bounded_logs(logs);
783 execution.updated_at = now;
784 self.save(&execution).await
785 }
786
787 pub async fn fail(
789 &self,
790 id: &str,
791 error: impl Into<String>,
792 logs: Vec<String>,
793 now: u64,
794 ) -> Result<(), RuntimeError> {
795 let _guard = self.gate.lock().await;
796 let mut execution = self.require(id).await?;
797 execution.status = ExecutionStatus::Error;
798 push_event(
799 &mut execution,
800 ExecutionEvent::Status {
801 status: ExecutionStatus::Error,
802 at: now,
803 },
804 );
805 execution.error = Some(truncate_bytes(error.into(), MAX_DURABLE_VALUE_BYTES));
806 execution.logs = bounded_logs(logs);
807 execution.updated_at = now;
808 self.save(&execution).await
809 }
810
811 pub async fn approve(&self, id: &str, seq: u64, now: u64) -> Result<bool, RuntimeError> {
813 let _guard = self.gate.lock().await;
814 let mut execution = self.require(id).await?;
815 if execution.status != ExecutionStatus::Paused {
816 return Ok(false);
817 }
818 let Some(entry) = execution
819 .log
820 .iter()
821 .find(|entry| entry.seq == seq && entry.state == LogEntryState::Pending)
822 else {
823 return Ok(false);
824 };
825 let _ = entry;
826 execution.status = ExecutionStatus::Running;
827 execution.updated_at = now;
828 self.save(&execution).await?;
829 Ok(true)
830 }
831
832 pub async fn reject(&self, id: &str, seq: u64, now: u64) -> Result<bool, RuntimeError> {
834 let _guard = self.gate.lock().await;
835 let mut execution = self.require(id).await?;
836 let Some(entry) = execution
837 .log
838 .iter_mut()
839 .find(|entry| entry.seq == seq && entry.state == LogEntryState::Pending)
840 else {
841 return Ok(false);
842 };
843 entry.state = LogEntryState::Reverted;
844 let error = format!(
845 "Action {}.{} rejected by user",
846 entry.connector, entry.method
847 );
848 execution.status = ExecutionStatus::Rejected;
849 push_event(
850 &mut execution,
851 ExecutionEvent::Status {
852 status: ExecutionStatus::Rejected,
853 at: now,
854 },
855 );
856 execution.error = Some(error);
857 execution.updated_at = now;
858 self.save(&execution).await?;
859 Ok(true)
860 }
861
862 pub async fn pending(&self, id: Option<&str>) -> Result<Vec<PendingAction>, RuntimeError> {
864 let executions = if let Some(id) = id {
865 vec![self.require(id).await?]
866 } else {
867 self.store
868 .list_executions()
869 .await
870 .map_err(RuntimeError::Store)?
871 };
872 let mut result = Vec::new();
873 for execution in executions
874 .into_iter()
875 .filter(|execution| execution.status == ExecutionStatus::Paused)
876 {
877 for entry in execution
878 .log
879 .into_iter()
880 .filter(|entry| entry.state == LogEntryState::Pending)
881 {
882 let action = PendingAction {
883 execution_id: execution.id.clone(),
884 seq: entry.seq,
885 connector: entry.connector,
886 method: entry.method,
887 arguments: entry.arguments,
888 };
889 result.push(self.rehydrate_pending(&execution.id, action).await?);
890 }
891 }
892 result.sort_by(|left, right| {
893 right
894 .execution_id
895 .cmp(&left.execution_id)
896 .then_with(|| left.seq.cmp(&right.seq))
897 });
898 Ok(result)
899 }
900
901 pub async fn actions_to_revert(&self, id: &str) -> Result<Vec<LogEntry>, RuntimeError> {
903 let execution = self.require(id).await?;
904 let mut actions = Vec::new();
905 for entry in execution.log.into_iter().filter(|entry| {
906 entry.state == LogEntryState::Applied && !entry.ephemeral && entry.connector != "__step"
907 }) {
908 actions.push(self.rehydrate_entry(&execution.id, entry).await?);
909 }
910 actions.sort_by_key(|entry| std::cmp::Reverse(entry.seq));
911 Ok(actions)
912 }
913
914 pub async fn mark_reverted(&self, id: &str, seq: u64, now: u64) -> Result<(), RuntimeError> {
916 let _guard = self.gate.lock().await;
917 let mut execution = self.require(id).await?;
918 if let Some(entry) = execution.log.iter_mut().find(|entry| entry.seq == seq) {
919 entry.state = LogEntryState::Reverted;
920 execution.updated_at = now;
921 self.save(&execution).await?;
922 }
923 Ok(())
924 }
925
926 pub async fn finish_rollback(&self, id: &str, now: u64) -> Result<(), RuntimeError> {
928 let _guard = self.gate.lock().await;
929 let mut execution = self.require(id).await?;
930 execution.status = ExecutionStatus::RolledBack;
931 push_event(
932 &mut execution,
933 ExecutionEvent::Status {
934 status: ExecutionStatus::RolledBack,
935 at: now,
936 },
937 );
938 execution.updated_at = now;
939 self.save(&execution).await
940 }
941
942 pub async fn expire(&self, now: u64, max_age_ms: u64) -> Result<Vec<String>, RuntimeError> {
944 let _guard = self.gate.lock().await;
945 let mut expired = Vec::new();
946 for mut execution in self
947 .store
948 .list_executions()
949 .await
950 .map_err(RuntimeError::Store)?
951 {
952 if !matches!(
953 execution.status,
954 ExecutionStatus::Paused | ExecutionStatus::Running
955 ) || now.saturating_sub(execution.updated_at) <= max_age_ms
956 {
957 continue;
958 }
959 execution.status = if execution.status == ExecutionStatus::Paused {
960 ExecutionStatus::Rejected
961 } else {
962 ExecutionStatus::Error
963 };
964 execution.error = Some(if execution.status == ExecutionStatus::Rejected {
965 "Expired awaiting approval".to_string()
966 } else {
967 "Expired while running - the host never completed the pass".to_string()
968 });
969 for entry in &mut execution.log {
970 if entry.state == LogEntryState::Pending {
971 entry.state = LogEntryState::Reverted;
972 }
973 }
974 execution.updated_at = now;
975 self.save(&execution).await?;
976 expired.push(execution.id);
977 }
978 Ok(expired)
979 }
980
981 pub async fn execution(&self, id: &str) -> Result<Option<ExecutionState>, RuntimeError> {
983 let Some(execution) = self
984 .store
985 .get_execution(id)
986 .await
987 .map_err(RuntimeError::Store)?
988 else {
989 return Ok(None);
990 };
991 self.rehydrate_execution(execution).await.map(Some)
992 }
993
994 pub async fn execution_snapshot(
996 &self,
997 id: &str,
998 ) -> Result<Option<ExecutionState>, RuntimeError> {
999 self.store
1000 .get_execution(id)
1001 .await
1002 .map_err(RuntimeError::Store)
1003 }
1004
1005 pub async fn executions(&self) -> Result<Vec<ExecutionState>, RuntimeError> {
1007 let mut executions = self
1008 .store
1009 .list_executions()
1010 .await
1011 .map_err(RuntimeError::Store)?;
1012 executions.sort_by_key(|execution| {
1013 (
1014 std::cmp::Reverse(execution.created_at),
1015 std::cmp::Reverse(execution.id.clone()),
1016 )
1017 });
1018 Ok(executions)
1019 }
1020
1021 pub async fn prune(&self, keep: usize) -> Result<usize, RuntimeError> {
1023 let _guard = self.gate.lock().await;
1024 self.prune_locked(keep, None).await
1025 }
1026
1027 pub async fn save_snippet(
1029 &self,
1030 name: impl Into<String>,
1031 description: impl Into<String>,
1032 execution_id: &str,
1033 input_schema: Option<Value>,
1034 now: u64,
1035 ) -> Result<Snippet, RuntimeError> {
1036 let execution = self.require(execution_id).await?;
1037 if execution.status != ExecutionStatus::Completed {
1038 return Err(RuntimeError::InvalidState(
1039 "Only completed executions can be saved as snippets".to_string(),
1040 ));
1041 }
1042 let snippet = Snippet {
1043 name: name.into(),
1044 description: description.into(),
1045 code: execution.code,
1046 saved_at: now,
1047 input_schema,
1048 connectors: execution.connectors,
1049 };
1050 self.store
1051 .put_snippet(&snippet)
1052 .await
1053 .map_err(RuntimeError::Store)?;
1054 Ok(snippet)
1055 }
1056
1057 pub async fn snippets(&self) -> Result<Vec<Snippet>, RuntimeError> {
1059 let mut snippets = self
1060 .store
1061 .list_snippets()
1062 .await
1063 .map_err(RuntimeError::Store)?;
1064 snippets.sort_by(|left, right| left.name.cmp(&right.name));
1065 Ok(snippets)
1066 }
1067
1068 pub async fn delete_snippet(&self, name: &str) -> Result<bool, RuntimeError> {
1070 self.store
1071 .delete_snippet(name)
1072 .await
1073 .map_err(RuntimeError::Store)
1074 }
1075
1076 async fn require(&self, id: &str) -> Result<ExecutionState, RuntimeError> {
1077 self.store
1078 .get_execution(id)
1079 .await
1080 .map_err(RuntimeError::Store)?
1081 .ok_or_else(|| RuntimeError::MissingExecution(id.to_string()))
1082 }
1083
1084 async fn save(&self, execution: &ExecutionState) -> Result<(), RuntimeError> {
1085 self.store
1086 .put_execution(execution)
1087 .await
1088 .map_err(RuntimeError::Store)
1089 }
1090
1091 async fn prune_locked(
1092 &self,
1093 keep: usize,
1094 excluding: Option<&str>,
1095 ) -> Result<usize, RuntimeError> {
1096 let mut terminal = self
1097 .store
1098 .list_executions()
1099 .await
1100 .map_err(RuntimeError::Store)?
1101 .into_iter()
1102 .filter(|execution| {
1103 execution.status.terminal()
1104 && excluding.is_none_or(|id| execution.id.as_str() != id)
1105 })
1106 .collect::<Vec<_>>();
1107 terminal.sort_by_key(|execution| std::cmp::Reverse(execution.created_at));
1108 let remove = terminal.len().saturating_sub(keep);
1109 for execution in terminal.into_iter().rev().take(remove) {
1110 self.store
1111 .delete_execution(&execution.id)
1112 .await
1113 .map_err(RuntimeError::Store)?;
1114 self.artifacts
1115 .delete_execution(&execution.id)
1116 .await
1117 .map_err(RuntimeError::Store)?;
1118 }
1119 Ok(remove)
1120 }
1121}
1122
1123fn ensure_size(what: &str, value: &Value) -> Result<(), RuntimeError> {
1124 let size = serialized_size(value);
1125 if size > MAX_DURABLE_VALUE_BYTES {
1126 Err(RuntimeError::DurableValue(format!(
1127 "{what} is too large to record durably ({size} bytes > \
1128 {MAX_DURABLE_VALUE_BYTES} byte limit). Write large data to a file or workspace \
1129 instead and pass or return a small reference."
1130 )))
1131 } else {
1132 Ok(())
1133 }
1134}
1135
1136fn serialized_size(value: &Value) -> usize {
1137 serde_json::to_vec(value).map_or(usize::MAX, |value| value.len())
1138}
1139
1140fn stable_json(value: &Value) -> String {
1141 match value {
1142 Value::Object(values) => {
1143 let mut values = values.iter().collect::<Vec<_>>();
1144 values.sort_by_key(|(key, _)| *key);
1145 format!(
1146 "{{{}}}",
1147 values
1148 .into_iter()
1149 .map(|(key, value)| format!(
1150 "{}:{}",
1151 serde_json::to_string(key).unwrap(),
1152 stable_json(value)
1153 ))
1154 .collect::<Vec<_>>()
1155 .join(",")
1156 )
1157 }
1158 Value::Array(values) => format!(
1159 "[{}]",
1160 values.iter().map(stable_json).collect::<Vec<_>>().join(",")
1161 ),
1162 value => serde_json::to_string(value).unwrap_or_default(),
1163 }
1164}
1165
1166fn bounded_logs(logs: Vec<String>) -> Vec<String> {
1167 let mut size = 2;
1168 logs.into_iter()
1169 .take_while(|log| {
1170 size += log.len() + 3;
1171 size <= MAX_DURABLE_VALUE_BYTES
1172 })
1173 .collect()
1174}
1175
1176fn push_event(execution: &mut ExecutionState, event: ExecutionEvent) {
1177 execution.events.push(event);
1178 if execution.events.len() > DEFAULT_MAX_EVENTS {
1179 execution
1180 .events
1181 .drain(..execution.events.len() - DEFAULT_MAX_EVENTS);
1182 }
1183}
1184
1185fn artifact_ref(value: &Value) -> Result<Option<ArtifactRef>, RuntimeError> {
1186 let Value::Object(values) = value else {
1187 return Ok(None);
1188 };
1189 if values.len() != 1 {
1190 return Ok(None);
1191 }
1192 let Some(artifact) = values.get("$artifact") else {
1193 return Ok(None);
1194 };
1195 serde_json::from_value::<ArtifactRef>(artifact.clone())
1196 .map(Some)
1197 .map_err(|error| RuntimeError::Store(error.to_string()))
1198}
1199
1200fn truncate_bytes(mut value: String, limit: usize) -> String {
1201 if value.len() <= limit {
1202 return value;
1203 }
1204 let mut end = limit;
1205 while !value.is_char_boundary(end) {
1206 end -= 1;
1207 }
1208 value.truncate(end);
1209 value
1210}
1211
1212#[cfg(test)]
1213mod tests {
1214 use serde_json::json;
1215
1216 use super::*;
1217
1218 fn runtime() -> CodeModeRuntime {
1219 CodeModeRuntime::new(Arc::new(MemoryStore::default()))
1220 }
1221
1222 fn runtime_parts() -> (CodeModeRuntime, Arc<MemoryStore>, Arc<MemoryArtifactStore>) {
1223 let store = Arc::new(MemoryStore::default());
1224 let artifacts = Arc::new(MemoryArtifactStore::default());
1225 (
1226 CodeModeRuntime::with_artifacts(store.clone(), artifacts.clone()),
1227 store,
1228 artifacts,
1229 )
1230 }
1231
1232 fn artifact_id(value: &Value) -> String {
1233 serde_json::from_value::<ArtifactRef>(value["$artifact"].clone())
1234 .unwrap()
1235 .id
1236 }
1237
1238 #[tokio::test]
1239 async fn pauses_approves_and_replays_without_reapplying() {
1240 let runtime = runtime();
1241 let id = runtime
1242 .begin("async () => {}", vec!["db".into()], 1)
1243 .await
1244 .unwrap();
1245 assert_eq!(
1246 runtime
1247 .decide(&id, 0, "db", "write", json!({"id": 1}), true, false, 2)
1248 .await
1249 .unwrap(),
1250 ToolDecision::Pause(0)
1251 );
1252 assert!(runtime.approve(&id, 0, 3).await.unwrap());
1253 assert_eq!(
1254 runtime
1255 .decide(&id, 0, "db", "write", json!({"id": 1}), true, false, 4)
1256 .await
1257 .unwrap(),
1258 ToolDecision::Execute(0)
1259 );
1260 runtime
1261 .record_result(&id, 0, json!({"ok": true}), 5)
1262 .await
1263 .unwrap();
1264 assert_eq!(
1265 runtime
1266 .decide(&id, 0, "db", "write", json!({"id": 1}), true, false, 6)
1267 .await
1268 .unwrap(),
1269 ToolDecision::Replay(json!({"ok": true}))
1270 );
1271 }
1272
1273 #[tokio::test]
1274 async fn rejects_results_for_pending_approval_entries() {
1275 let runtime = runtime();
1276 let id = runtime
1277 .begin("async () => {}", vec!["db".into()], 1)
1278 .await
1279 .unwrap();
1280 runtime
1281 .decide(&id, 0, "db", "write", json!({"id": 1}), true, false, 2)
1282 .await
1283 .unwrap();
1284
1285 let error = runtime
1286 .record_result(&id, 0, json!({"forged": true}), 3)
1287 .await
1288 .unwrap_err();
1289
1290 assert!(matches!(error, RuntimeError::InvalidState(_)));
1291 let execution = runtime.execution(&id).await.unwrap().unwrap();
1292 assert_eq!(execution.status, ExecutionStatus::Paused);
1293 assert_eq!(execution.log[0].state, LogEntryState::Pending);
1294 }
1295
1296 #[tokio::test]
1297 async fn rollback_targets_only_logged_connector_actions() {
1298 let runtime = runtime();
1299 let id = runtime
1300 .begin("async () => {}", vec!["db".into()], 1)
1301 .await
1302 .unwrap();
1303 runtime
1304 .decide(&id, 0, "db", "read", json!({}), false, true, 2)
1305 .await
1306 .unwrap();
1307 runtime.record_result(&id, 0, json!(1), 3).await.unwrap();
1308 runtime
1309 .decide(&id, 1, "db", "write", json!({}), false, false, 4)
1310 .await
1311 .unwrap();
1312 runtime.record_result(&id, 1, json!(2), 5).await.unwrap();
1313
1314 let actions = runtime.actions_to_revert(&id).await.unwrap();
1315
1316 assert_eq!(actions.len(), 1);
1317 assert_eq!(actions[0].seq, 1);
1318 }
1319
1320 #[tokio::test]
1321 async fn divergence_is_terminal() {
1322 let runtime = runtime();
1323 let id = runtime.begin("async () => {}", vec![], 1).await.unwrap();
1324 runtime
1325 .decide(&id, 0, "db", "read", json!({"id": 1}), false, false, 2)
1326 .await
1327 .unwrap();
1328 runtime.record_result(&id, 0, json!(1), 3).await.unwrap();
1329 assert_eq!(
1330 runtime
1331 .decide(&id, 0, "db", "read", json!({"id": 2}), false, false, 4)
1332 .await
1333 .unwrap(),
1334 ToolDecision::Pause(0)
1335 );
1336 assert_eq!(
1337 runtime.execution(&id).await.unwrap().unwrap().status,
1338 ExecutionStatus::Error
1339 );
1340 }
1341
1342 #[tokio::test]
1343 async fn spills_and_rehydrates_oversized_replay_values() {
1344 let (runtime, store, _) = runtime_parts();
1345 let id = runtime
1346 .begin("async () => {}", vec!["db".into()], 1)
1347 .await
1348 .unwrap();
1349 runtime
1350 .decide(&id, 0, "db", "read", json!({}), false, false, 2)
1351 .await
1352 .unwrap();
1353 let value = json!("x".repeat(MAX_DURABLE_VALUE_BYTES + 1));
1354 runtime
1355 .record_result(&id, 0, value.clone(), 3)
1356 .await
1357 .unwrap();
1358 let raw = store.get_execution(&id).await.unwrap().unwrap();
1359 assert!(raw.log[0].result.as_ref().unwrap()["$artifact"].is_object());
1360 let execution = runtime.execution(&id).await.unwrap().unwrap();
1361 assert_eq!(execution.log[0].result, Some(value.clone()));
1362 assert_eq!(
1363 runtime
1364 .decide(&id, 0, "db", "read", json!({}), false, false, 4)
1365 .await
1366 .unwrap(),
1367 ToolDecision::Replay(value)
1368 );
1369 }
1370
1371 #[tokio::test]
1372 async fn spills_and_rehydrates_oversized_arguments() {
1373 let (runtime, store, _) = runtime_parts();
1374 let id = runtime
1375 .begin("async () => {}", vec!["db".into()], 1)
1376 .await
1377 .unwrap();
1378 let arguments = json!({"payload": "x".repeat(MAX_DURABLE_VALUE_BYTES + 1)});
1379
1380 assert_eq!(
1381 runtime
1382 .decide(&id, 0, "db", "write", arguments.clone(), false, false, 2)
1383 .await
1384 .unwrap(),
1385 ToolDecision::Execute(0)
1386 );
1387
1388 let raw = store.get_execution(&id).await.unwrap().unwrap();
1389 assert!(raw.log[0].arguments["$artifact"].is_object());
1390 let execution = runtime.execution(&id).await.unwrap().unwrap();
1391 assert_eq!(execution.log[0].arguments, arguments);
1392 }
1393
1394 #[tokio::test]
1395 async fn pending_actions_rehydrate_oversized_arguments() {
1396 let runtime = runtime();
1397 let id = runtime
1398 .begin("async () => {}", vec!["db".into()], 1)
1399 .await
1400 .unwrap();
1401 let arguments = json!({"payload": "x".repeat(MAX_DURABLE_VALUE_BYTES + 1)});
1402
1403 assert_eq!(
1404 runtime
1405 .decide(&id, 0, "db", "write", arguments.clone(), true, false, 2)
1406 .await
1407 .unwrap(),
1408 ToolDecision::Pause(0)
1409 );
1410
1411 let pending = runtime.pending(Some(&id)).await.unwrap();
1412 assert_eq!(pending[0].arguments, arguments);
1413 }
1414
1415 #[tokio::test]
1416 async fn retrieves_oversized_final_result_by_execution_and_artifact_ids() {
1417 let (runtime, store, _) = runtime_parts();
1418 let id = runtime.begin("async () => {}", vec![], 1).await.unwrap();
1419 let value = json!("x".repeat(MAX_DURABLE_VALUE_BYTES + 1));
1420
1421 runtime
1422 .complete(&id, value.clone(), Vec::new(), 2)
1423 .await
1424 .unwrap();
1425
1426 let raw = store.get_execution(&id).await.unwrap().unwrap();
1427 let artifact_id = artifact_id(raw.result.as_ref().unwrap());
1428 assert_eq!(
1429 runtime.artifact(&id, &artifact_id).await.unwrap(),
1430 Some(value.clone())
1431 );
1432 assert_eq!(
1433 runtime.execution(&id).await.unwrap().unwrap().result,
1434 Some(value)
1435 );
1436 }
1437
1438 #[tokio::test]
1439 async fn denies_artifact_lookup_for_wrong_execution_owner() {
1440 let (runtime, store, _) = runtime_parts();
1441 let owner = runtime.begin("async () => {}", vec![], 1).await.unwrap();
1442 let other = runtime.begin("async () => {}", vec![], 2).await.unwrap();
1443 let value = json!("x".repeat(MAX_DURABLE_VALUE_BYTES + 1));
1444 runtime
1445 .complete(&owner, value, Vec::new(), 3)
1446 .await
1447 .unwrap();
1448 let raw = store.get_execution(&owner).await.unwrap().unwrap();
1449 let artifact_id = artifact_id(raw.result.as_ref().unwrap());
1450
1451 assert_eq!(runtime.artifact(&other, &artifact_id).await.unwrap(), None);
1452 }
1453
1454 #[tokio::test]
1455 async fn missing_artifact_fails_rehydrated_execution_lookup() {
1456 let (runtime, _, artifacts) = runtime_parts();
1457 let id = runtime.begin("async () => {}", vec![], 1).await.unwrap();
1458 let value = json!("x".repeat(MAX_DURABLE_VALUE_BYTES + 1));
1459 runtime.complete(&id, value, Vec::new(), 2).await.unwrap();
1460
1461 artifacts.delete_execution(&id).await.unwrap();
1462 let error = runtime.execution(&id).await.unwrap_err();
1463
1464 assert!(matches!(error, RuntimeError::Store(message) if message.contains("unavailable")));
1465 }
1466
1467 #[tokio::test]
1468 async fn prune_deletes_artifacts_for_removed_executions() {
1469 let (runtime, store, _) = runtime_parts();
1470 let id = runtime.begin("async () => {}", vec![], 1).await.unwrap();
1471 let value = json!("x".repeat(MAX_DURABLE_VALUE_BYTES + 1));
1472 runtime
1473 .complete(&id, value.clone(), Vec::new(), 2)
1474 .await
1475 .unwrap();
1476 let raw = store.get_execution(&id).await.unwrap().unwrap();
1477 let artifact_id = artifact_id(raw.result.as_ref().unwrap());
1478 assert_eq!(
1479 runtime.artifact(&id, &artifact_id).await.unwrap(),
1480 Some(value)
1481 );
1482
1483 assert_eq!(runtime.prune(0).await.unwrap(), 1);
1484
1485 assert_eq!(runtime.artifact(&id, &artifact_id).await.unwrap(), None);
1486 assert!(runtime.execution(&id).await.unwrap().is_none());
1487 }
1488}