stasis/infrastructure/runtime/
grapheme_sdk_workflow_engine.rs

1use async_trait::async_trait;
2use grapheme_sdk::{GraphemeEngine, GraphemeSdkError};
3use serde_json::Value;
4use std::sync::Arc;
5use std::time::Duration;
6use tokio::task;
7use tokio::time;
8
9use crate::domain::errors::{Result, StasisError};
10use crate::ports::outbound::runtime::workflow_engine::{WorkflowEngine, WorkflowExecutionOutput};
11
12pub struct GraphemeSdkWorkflowEngine {
13    engine: Arc<GraphemeEngine>,
14    guardrails: GraphemeWorkflowGuardrails,
15}
16
17const DEFAULT_EXECUTION_TIMEOUT_MS: u64 = 2_000;
18
19#[derive(Clone, Debug)]
20pub struct GraphemeWorkflowGuardrails {
21    pub allowed_imports: Vec<String>,
22    pub max_source_bytes: usize,
23    pub execution_timeout: Duration,
24    pub max_steps: Option<usize>,
25    pub max_call_depth: Option<usize>,
26}
27
28impl Default for GraphemeWorkflowGuardrails {
29    fn default() -> Self {
30        Self {
31            // Allow all built-in Grapheme namespace modules by default.
32            allowed_imports: vec!["grapheme/*".to_string()],
33            max_source_bytes: 128 * 1024,
34            execution_timeout: resolve_execution_timeout_from_env(),
35            max_steps: Some(10_000),
36            max_call_depth: Some(16),
37        }
38    }
39}
40
41fn resolve_execution_timeout_from_env() -> Duration {
42    let raw_value = std::env::var("MEDOUSA_GRAPHEME_EXECUTION_TIMEOUT_MS")
43        .ok()
44        .or_else(|| std::env::var("STASIS_GRAPHEME_EXECUTION_TIMEOUT_MS").ok())
45        .or_else(|| std::env::var("GRAPHEME_EXECUTION_TIMEOUT_MS").ok());
46
47    let timeout_ms = raw_value
48        .as_deref()
49        .map(str::trim)
50        .filter(|value| !value.is_empty())
51        .and_then(|value| value.parse::<u64>().ok())
52        .unwrap_or(DEFAULT_EXECUTION_TIMEOUT_MS);
53
54    Duration::from_millis(timeout_ms)
55}
56
57impl GraphemeSdkWorkflowEngine {
58    pub fn new() -> Self {
59        Self::with_guardrails(GraphemeWorkflowGuardrails::default())
60    }
61
62    pub fn with_guardrails(guardrails: GraphemeWorkflowGuardrails) -> Self {
63        Self {
64            engine: Arc::new(
65                GraphemeEngine::builder()
66                    .with_max_steps(guardrails.max_steps)
67                    .with_max_call_depth(guardrails.max_call_depth)
68                    .build(),
69            ),
70            guardrails,
71        }
72    }
73
74    fn validate_source(&self, source: &str) -> Result<()> {
75        if source.len() > self.guardrails.max_source_bytes {
76            return Err(StasisError::PortFailure(format!(
77                "grapheme policy violation: source size {} exceeds max {} bytes",
78                source.len(),
79                self.guardrails.max_source_bytes
80            )));
81        }
82
83        let imports = Self::extract_imports(source);
84        for import in imports {
85            if !self
86                .guardrails
87                .allowed_imports
88                .iter()
89                .any(|pattern| Self::import_is_allowed(pattern, &import))
90            {
91                return Err(StasisError::PortFailure(format!(
92                    "grapheme policy violation: import '{}' is not allowlisted",
93                    import
94                )));
95            }
96        }
97
98        Ok(())
99    }
100
101    fn import_is_allowed(pattern: &str, import: &str) -> bool {
102        if let Some(prefix) = pattern.strip_suffix('*') {
103            return import.starts_with(prefix);
104        }
105        pattern == import
106    }
107
108    fn extract_imports(source: &str) -> Vec<String> {
109        source
110            .lines()
111            .filter_map(|line| {
112                let trimmed = line.trim();
113                if !trimmed.starts_with("import ") {
114                    return None;
115                }
116
117                let quote = if trimmed.contains('"') { '"' } else { '\'' };
118                let start = trimmed.find(quote)?;
119                let tail = &trimmed[(start + 1)..];
120                let end = tail.find(quote)?;
121                Some(tail[..end].to_string())
122            })
123            .collect()
124    }
125
126    fn map_error(err: GraphemeSdkError) -> StasisError {
127        let msg = err.to_string();
128        if msg.contains("policy:") {
129            return StasisError::PortFailure(format!("grapheme policy violation: {msg}"));
130        }
131
132        StasisError::PortFailure(format!("grapheme sdk execution error: {err}"))
133    }
134}
135
136impl Default for GraphemeSdkWorkflowEngine {
137    fn default() -> Self {
138        Self::new()
139    }
140}
141
142#[async_trait]
143impl WorkflowEngine for GraphemeSdkWorkflowEngine {
144    async fn execute_grapheme_source(
145        &self,
146        source: &str,
147        state_current: Option<&Value>,
148    ) -> Result<WorkflowExecutionOutput> {
149        self.validate_source(source)?;
150
151        if self.guardrails.execution_timeout.is_zero() {
152            return Err(StasisError::PortFailure(
153                "grapheme policy violation: execution timeout must be greater than 0ms".to_string(),
154            ));
155        }
156
157        let source_owned = source.to_string();
158        let state_current_owned = state_current.cloned();
159        let guardrails = self.guardrails.clone();
160        let engine = Arc::clone(&self.engine);
161        let handle = task::spawn_blocking(move || {
162            if let Some(initial_state_current) = state_current_owned {
163                let state_engine = GraphemeEngine::builder()
164                    .with_max_steps(guardrails.max_steps)
165                    .with_max_call_depth(guardrails.max_call_depth)
166                    .with_initial_state_current(initial_state_current)
167                    .build();
168                state_engine.execute_source(&source_owned)
169            } else {
170                engine.execute_source(&source_owned)
171            }
172        });
173        let result = time::timeout(self.guardrails.execution_timeout, handle)
174            .await
175            .map_err(|_| {
176                StasisError::PortFailure(format!(
177                    "grapheme policy violation: execution timed out after {} ms",
178                    self.guardrails.execution_timeout.as_millis()
179                ))
180            })?
181            .map_err(|e| StasisError::PortFailure(format!("grapheme sdk worker join error: {e}")))?
182            .map_err(Self::map_error)?;
183
184        Ok(WorkflowExecutionOutput {
185            run_id: format!("grapheme:{}", result.artifact_id),
186            execution: serde_json::to_value(&result.execution).unwrap_or(Value::Null),
187            final_state: result.final_state,
188        })
189    }
190}
stasis/infrastructure/runtime/grapheme_sdk_workflow_engine.rs

stasis/infrastructure/runtime/
grapheme_sdk_workflow_engine.rs
