Skip to main content

stasis/application/composition/
runtime_composition.rs

1use std::sync::Arc;
2
3use surrealdb::engine::any::Any;
4use surrealdb::opt::auth::Database;
5use surrealdb::Surreal;
6
7use crate::application::runtime::in_memory_runtime::InMemoryRuntime;
8use crate::application::runtime::surreal_runtime::SurrealRuntime;
9use crate::domain::errors::{Result, StasisError};
10use crate::infrastructure::llm::genai_chat_client::GenaiChatClient;
11use crate::infrastructure::memory::locus_context_reader::LocusContextReader;
12use crate::infrastructure::memory::locus_context_writer::LocusContextWriter;
13use crate::infrastructure::memory::locus_memory_operations::LocusMemoryOperations;
14use crate::infrastructure::memory::locus_node_store_factory::LocusNodeStoreFactory;
15use crate::infrastructure::memory::surreal_identity_memory_store::SurrealIdentityMemoryStore;
16use crate::infrastructure::runtime::endpoint_routing_event_publisher::EndpointRoutingEventPublisher;
17use crate::infrastructure::runtime::grapheme_sdk_workflow_engine::GraphemeSdkWorkflowEngine;
18use crate::infrastructure::runtime::in_memory_cluster_node_store::InMemoryClusterNodeStore;
19use crate::infrastructure::runtime::in_memory_delivery_endpoint_store::InMemoryDeliveryEndpointStore;
20use crate::infrastructure::runtime::in_memory_endpoint_delivery_status_store::InMemoryEndpointDeliveryStatusStore;
21use crate::infrastructure::runtime::in_memory_thread_store::InMemoryThreadStore;
22use crate::infrastructure::runtime::surreal_cluster_node_store::SurrealClusterNodeStore;
23use crate::infrastructure::runtime::surreal_delivery_endpoint_store::SurrealDeliveryEndpointStore;
24use crate::infrastructure::runtime::surreal_endpoint_delivery_status_store::SurrealEndpointDeliveryStatusStore;
25use crate::infrastructure::runtime::surreal_thread_store::SurrealThreadStore;
26use crate::ports::outbound::ai_chat_client::AiChatClient;
27use crate::ports::outbound::memory::memory_context_reader::MemoryContextReader;
28use crate::ports::outbound::memory::memory_context_writer::MemoryContextWriter;
29use crate::ports::outbound::memory::memory_operations::MemoryOperations;
30use crate::ports::outbound::runtime::cluster_node_store::ClusterNodeStore;
31use crate::ports::outbound::runtime::delivery_endpoint_store::DeliveryEndpointStore;
32use crate::ports::outbound::runtime::endpoint_delivery_status_store::EndpointDeliveryStatusStore;
33use crate::ports::outbound::runtime::endpoint_routing_policy::EndpointRoutingPolicy;
34use crate::ports::outbound::runtime::endpoint_transport_publisher::EndpointTransportPublisher;
35use crate::ports::outbound::runtime::thread_store::ThreadStore;
36use crate::ports::outbound::runtime::workflow_engine::WorkflowEngine;
37
38#[derive(Clone, Debug)]
39pub enum RuntimeBackend {
40    InMemory,
41    SurrealMem {
42        namespace: String,
43        database: String,
44        auth: Option<SurrealAuth>,
45    },
46    SurrealWs {
47        endpoint: String,
48        namespace: String,
49        database: String,
50        auth: Option<SurrealAuth>,
51    },
52    SurrealKv {
53        path: String,
54        namespace: String,
55        database: String,
56        auth: Option<SurrealAuth>,
57    },
58}
59
60pub use crate::application::composition::surreal_backend_config::SurrealAuth;
61
62impl RuntimeBackend {
63    pub fn surreal_mem(namespace: impl Into<String>, database: impl Into<String>) -> Self {
64        Self::SurrealMem {
65            namespace: namespace.into(),
66            database: database.into(),
67            auth: None,
68        }
69    }
70
71    pub fn surreal_ws(
72        endpoint: impl Into<String>,
73        namespace: impl Into<String>,
74        database: impl Into<String>,
75    ) -> Self {
76        Self::SurrealWs {
77            endpoint: endpoint.into(),
78            namespace: namespace.into(),
79            database: database.into(),
80            auth: None,
81        }
82    }
83
84    pub fn surreal_kv(
85        path: impl Into<String>,
86        namespace: impl Into<String>,
87        database: impl Into<String>,
88    ) -> Self {
89        Self::SurrealKv {
90            path: path.into(),
91            namespace: namespace.into(),
92            database: database.into(),
93            auth: None,
94        }
95    }
96
97    pub fn with_surreal_auth(mut self, auth: SurrealAuth) -> Self {
98        match &mut self {
99            Self::SurrealMem { auth: slot, .. }
100            | Self::SurrealWs { auth: slot, .. }
101            | Self::SurrealKv { auth: slot, .. } => *slot = Some(auth),
102            Self::InMemory => {}
103        }
104        self
105    }
106}
107
108#[derive(Clone)]
109pub enum RuntimeComposition {
110    InMemory(InMemoryRuntime),
111    Surreal(SurrealRuntime),
112}
113
114pub struct RuntimeFactory;
115
116impl RuntimeFactory {
117    async fn connect_surreal_any(
118        endpoint: &str,
119        namespace: String,
120        database: String,
121        auth: Option<SurrealAuth>,
122    ) -> Result<RuntimeComposition> {
123        let db = Surreal::<Any>::init();
124        db.connect(endpoint)
125            .await
126            .map_err(|e| StasisError::PortFailure(format!("connect surreal db ({endpoint}): {e}")))?;
127
128        if let Some(auth) = auth {
129            db.signin(Database {
130                namespace: namespace.clone(),
131                database: database.clone(),
132                username: auth.username,
133                password: auth.password,
134            })
135            .await
136            .map_err(|e| StasisError::PortFailure(format!("signin surreal db: {e}")))?;
137        }
138
139        db.use_ns(namespace).use_db(database).await.map_err(|e| {
140            StasisError::PortFailure(format!("select surreal namespace/database: {e}"))
141        })?;
142
143        SurrealIdentityMemoryStore::ensure_schema_for_db(&db).await?;
144
145        Ok(RuntimeComposition::Surreal(SurrealRuntime::new(db)))
146    }
147
148    pub async fn build(config: RuntimeBackend) -> Result<RuntimeComposition> {
149        match config {
150            RuntimeBackend::InMemory => Ok(RuntimeComposition::InMemory(InMemoryRuntime::new())),
151            RuntimeBackend::SurrealMem {
152                namespace,
153                database,
154                auth,
155            } => Self::connect_surreal_any("mem://", namespace, database, auth).await,
156            RuntimeBackend::SurrealWs {
157                endpoint,
158                namespace,
159                database,
160                auth,
161            } => Self::connect_surreal_any(&endpoint, namespace, database, auth).await,
162            RuntimeBackend::SurrealKv {
163                path,
164                namespace,
165                database,
166                auth,
167            } => {
168                let endpoint = if path.starts_with("surrealkv://") {
169                    path
170                } else {
171                    format!("surrealkv://{path}")
172                };
173                Self::connect_surreal_any(&endpoint, namespace, database, auth).await
174            }
175        }
176    }
177
178    pub fn from_db(db: Surreal<Any>) -> RuntimeComposition {
179        RuntimeComposition::Surreal(SurrealRuntime::new(db))
180    }
181
182    pub fn default_chat_client() -> Arc<dyn AiChatClient> {
183        Arc::new(GenaiChatClient::from_env())
184    }
185
186    pub fn default_workflow_engine() -> Arc<dyn WorkflowEngine> {
187        Arc::new(GraphemeSdkWorkflowEngine::new())
188    }
189
190    pub async fn ensure_locus_memory_adapters(
191        enable_locus_memory: bool,
192        mut memory_context_reader: Option<Arc<dyn MemoryContextReader>>,
193        mut memory_context_writer: Option<Arc<dyn MemoryContextWriter>>,
194        mut memory_operations: Option<Arc<dyn MemoryOperations>>,
195    ) -> Result<(
196        Option<Arc<dyn MemoryContextReader>>,
197        Option<Arc<dyn MemoryContextWriter>>,
198        Option<Arc<dyn MemoryOperations>>,
199    )> {
200        if enable_locus_memory
201            && (memory_context_reader.is_none()
202                || memory_context_writer.is_none()
203                || memory_operations.is_none())
204        {
205            let store = LocusNodeStoreFactory::in_memory().await?;
206            if memory_context_reader.is_none() {
207                memory_context_reader = Some(Arc::new(LocusContextReader::new(store.clone())));
208            }
209            if memory_context_writer.is_none() {
210                memory_context_writer = Some(Arc::new(LocusContextWriter::new(store.clone())));
211            }
212            if memory_operations.is_none() {
213                memory_operations = Some(Arc::new(LocusMemoryOperations::new(store, None)));
214            }
215        }
216
217        Ok((
218            memory_context_reader,
219            memory_context_writer,
220            memory_operations,
221        ))
222    }
223
224    pub fn resolve_thread_store(
225        runtime: &RuntimeComposition,
226        configured: Option<Arc<dyn ThreadStore>>,
227    ) -> Arc<dyn ThreadStore> {
228        if let Some(store) = configured {
229            return store;
230        }
231
232        match runtime {
233            RuntimeComposition::InMemory(_) => Arc::new(InMemoryThreadStore::default()),
234            RuntimeComposition::Surreal(rt) => Arc::new(SurrealThreadStore::new(rt.job_store.db())),
235        }
236    }
237
238    pub fn resolve_cluster_node_store(
239        runtime: &RuntimeComposition,
240        configured: Option<Arc<dyn ClusterNodeStore>>,
241    ) -> Arc<dyn ClusterNodeStore> {
242        if let Some(store) = configured {
243            return store;
244        }
245
246        match runtime {
247            RuntimeComposition::InMemory(_) => Arc::new(InMemoryClusterNodeStore::default()),
248            RuntimeComposition::Surreal(rt) => {
249                Arc::new(SurrealClusterNodeStore::new(rt.job_store.db()))
250            }
251        }
252    }
253
254    pub fn resolve_delivery_endpoint_store(
255        runtime: &RuntimeComposition,
256        configured: Option<Arc<dyn DeliveryEndpointStore>>,
257    ) -> Arc<dyn DeliveryEndpointStore> {
258        if let Some(store) = configured {
259            return store;
260        }
261
262        match runtime {
263            RuntimeComposition::InMemory(_) => Arc::new(InMemoryDeliveryEndpointStore::default()),
264            RuntimeComposition::Surreal(rt) => {
265                Arc::new(SurrealDeliveryEndpointStore::new(rt.job_store.db()))
266            }
267        }
268    }
269
270    pub fn resolve_endpoint_delivery_status_store(
271        runtime: &RuntimeComposition,
272        configured: Option<Arc<dyn EndpointDeliveryStatusStore>>,
273    ) -> Arc<dyn EndpointDeliveryStatusStore> {
274        if let Some(store) = configured {
275            return store;
276        }
277
278        match runtime {
279            RuntimeComposition::InMemory(_) => {
280                Arc::new(InMemoryEndpointDeliveryStatusStore::default())
281            }
282            RuntimeComposition::Surreal(rt) => {
283                Arc::new(SurrealEndpointDeliveryStatusStore::new(rt.job_store.db()))
284            }
285        }
286    }
287
288    pub fn build_endpoint_routing_publisher(
289        endpoint_store: Arc<dyn DeliveryEndpointStore>,
290        status_store: Arc<dyn EndpointDeliveryStatusStore>,
291        transports: &[Arc<dyn EndpointTransportPublisher>],
292        routing_policy: Option<Arc<dyn EndpointRoutingPolicy>>,
293    ) -> EndpointRoutingEventPublisher {
294        let mut routing_publisher =
295            EndpointRoutingEventPublisher::new(endpoint_store).fail_on_unsupported_protocol(false);
296
297        if transports.is_empty() {
298            routing_publisher = routing_publisher
299                .with_http_webhook_transport()
300                .with_tcp_socket_transport();
301        } else {
302            for transport in transports {
303                routing_publisher = routing_publisher.with_transport_arc(transport.clone());
304            }
305        }
306
307        if let Some(policy) = routing_policy {
308            routing_publisher = routing_publisher.with_routing_policy_arc(policy);
309        }
310
311        routing_publisher.with_status_store_arc(status_store)
312    }
313}