Skip to main content

stasis/application/composition/
runtime_composition.rs

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