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::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(Root {
130                username: auth.username,
131                password: auth.password,
132            })
133            .await
134            .map_err(|e| StasisError::PortFailure(format!("signin surreal db: {e}")))?;
135        }
136
137        db.use_ns(namespace).use_db(database).await.map_err(|e| {
138            StasisError::PortFailure(format!("select surreal namespace/database: {e}"))
139        })?;
140
141        SurrealIdentityMemoryStore::ensure_schema_for_db(&db).await?;
142
143        Ok(RuntimeComposition::Surreal(SurrealRuntime::new(db)))
144    }
145
146    pub async fn build(config: RuntimeBackend) -> Result<RuntimeComposition> {
147        match config {
148            RuntimeBackend::InMemory => Ok(RuntimeComposition::InMemory(InMemoryRuntime::new())),
149            RuntimeBackend::SurrealMem {
150                namespace,
151                database,
152                auth,
153            } => Self::connect_surreal_any("mem://", namespace, database, auth).await,
154            RuntimeBackend::SurrealWs {
155                endpoint,
156                namespace,
157                database,
158                auth,
159            } => Self::connect_surreal_any(&endpoint, namespace, database, auth).await,
160            RuntimeBackend::SurrealKv {
161                path,
162                namespace,
163                database,
164                auth,
165            } => {
166                let endpoint = if path.starts_with("surrealkv://") {
167                    path
168                } else {
169                    format!("surrealkv://{path}")
170                };
171                Self::connect_surreal_any(&endpoint, namespace, database, auth).await
172            }
173        }
174    }
175
176    pub fn from_db(db: Surreal<Any>) -> RuntimeComposition {
177        RuntimeComposition::Surreal(SurrealRuntime::new(db))
178    }
179
180    pub fn default_chat_client() -> Arc<dyn AiChatClient> {
181        Arc::new(GenaiChatClient::from_env())
182    }
183
184    pub fn default_workflow_engine() -> Arc<dyn WorkflowEngine> {
185        Arc::new(GraphemeSdkWorkflowEngine::new())
186    }
187
188    pub async fn ensure_locus_memory_adapters(
189        enable_locus_memory: bool,
190        mut memory_context_reader: Option<Arc<dyn MemoryContextReader>>,
191        mut memory_context_writer: Option<Arc<dyn MemoryContextWriter>>,
192        mut memory_operations: Option<Arc<dyn MemoryOperations>>,
193    ) -> Result<(
194        Option<Arc<dyn MemoryContextReader>>,
195        Option<Arc<dyn MemoryContextWriter>>,
196        Option<Arc<dyn MemoryOperations>>,
197    )> {
198        if enable_locus_memory
199            && (memory_context_reader.is_none()
200                || memory_context_writer.is_none()
201                || memory_operations.is_none())
202        {
203            let store = LocusNodeStoreFactory::in_memory().await?;
204            if memory_context_reader.is_none() {
205                memory_context_reader = Some(Arc::new(LocusContextReader::new(store.clone())));
206            }
207            if memory_context_writer.is_none() {
208                memory_context_writer = Some(Arc::new(LocusContextWriter::new(store.clone())));
209            }
210            if memory_operations.is_none() {
211                memory_operations = Some(Arc::new(LocusMemoryOperations::new(store, None)));
212            }
213        }
214
215        Ok((
216            memory_context_reader,
217            memory_context_writer,
218            memory_operations,
219        ))
220    }
221
222    pub fn resolve_thread_store(
223        runtime: &RuntimeComposition,
224        configured: Option<Arc<dyn ThreadStore>>,
225    ) -> Arc<dyn ThreadStore> {
226        if let Some(store) = configured {
227            return store;
228        }
229
230        match runtime {
231            RuntimeComposition::InMemory(_) => Arc::new(InMemoryThreadStore::default()),
232            RuntimeComposition::Surreal(rt) => Arc::new(SurrealThreadStore::new(rt.job_store.db())),
233        }
234    }
235
236    pub fn resolve_cluster_node_store(
237        runtime: &RuntimeComposition,
238        configured: Option<Arc<dyn ClusterNodeStore>>,
239    ) -> Arc<dyn ClusterNodeStore> {
240        if let Some(store) = configured {
241            return store;
242        }
243
244        match runtime {
245            RuntimeComposition::InMemory(_) => Arc::new(InMemoryClusterNodeStore::default()),
246            RuntimeComposition::Surreal(rt) => {
247                Arc::new(SurrealClusterNodeStore::new(rt.job_store.db()))
248            }
249        }
250    }
251
252    pub fn resolve_delivery_endpoint_store(
253        runtime: &RuntimeComposition,
254        configured: Option<Arc<dyn DeliveryEndpointStore>>,
255    ) -> Arc<dyn DeliveryEndpointStore> {
256        if let Some(store) = configured {
257            return store;
258        }
259
260        match runtime {
261            RuntimeComposition::InMemory(_) => Arc::new(InMemoryDeliveryEndpointStore::default()),
262            RuntimeComposition::Surreal(rt) => {
263                Arc::new(SurrealDeliveryEndpointStore::new(rt.job_store.db()))
264            }
265        }
266    }
267
268    pub fn resolve_endpoint_delivery_status_store(
269        runtime: &RuntimeComposition,
270        configured: Option<Arc<dyn EndpointDeliveryStatusStore>>,
271    ) -> Arc<dyn EndpointDeliveryStatusStore> {
272        if let Some(store) = configured {
273            return store;
274        }
275
276        match runtime {
277            RuntimeComposition::InMemory(_) => {
278                Arc::new(InMemoryEndpointDeliveryStatusStore::default())
279            }
280            RuntimeComposition::Surreal(rt) => {
281                Arc::new(SurrealEndpointDeliveryStatusStore::new(rt.job_store.db()))
282            }
283        }
284    }
285
286    pub fn build_endpoint_routing_publisher(
287        endpoint_store: Arc<dyn DeliveryEndpointStore>,
288        status_store: Arc<dyn EndpointDeliveryStatusStore>,
289        transports: &[Arc<dyn EndpointTransportPublisher>],
290        routing_policy: Option<Arc<dyn EndpointRoutingPolicy>>,
291    ) -> EndpointRoutingEventPublisher {
292        let mut routing_publisher =
293            EndpointRoutingEventPublisher::new(endpoint_store).fail_on_unsupported_protocol(false);
294
295        if transports.is_empty() {
296            routing_publisher = routing_publisher
297                .with_http_webhook_transport()
298                .with_tcp_socket_transport();
299        } else {
300            for transport in transports {
301                routing_publisher = routing_publisher.with_transport_arc(transport.clone());
302            }
303        }
304
305        if let Some(policy) = routing_policy {
306            routing_publisher = routing_publisher.with_routing_policy_arc(policy);
307        }
308
309        routing_publisher.with_status_store_arc(status_store)
310    }
311}