1pub mod mcp;
13pub mod peer_persister;
14pub mod peer_registry;
15pub mod postgres;
16pub mod router;
17pub mod server_ownership;
18pub mod telemetry;
19pub mod ws_handler;
20pub mod ws_timing;
21
22use std::{
24 collections::HashMap,
25 net::SocketAddr,
26 sync::{
27 Arc, RwLock,
28 atomic::{AtomicBool, Ordering},
29 },
30 time::Duration,
31};
32
33use futures_util::StreamExt;
34pub use myko::server::*;
35use myko::{
36 client::MykoClient, command::CommandContext, request::RequestContext, saga::SagaRegistration,
37 search::SearchIndex, store::StoreRegistry, wire::MEvent,
38};
39pub use peer_persister::PeerPersister;
40pub use server_ownership::ServerOwnershipManager;
41use uuid::Uuid;
42
43use crate::postgres::{
44 CellPostgresConsumer, CellPostgresProducer, PostgresConfig, PostgresHistoryReplayProvider,
45 PostgresHistoryStore, PostgresProducerHandle,
46};
47
48#[derive(Clone)]
50pub struct CellServerConfig {
51 pub bind_addr: SocketAddr,
53 pub tcp_nodelay: bool,
59 pub postgres: Option<PostgresConfig>,
61 pub host_id: Option<Uuid>,
63 pub peer_registry: Option<peer_registry::PeerRegistryConfig>,
65 pub default_persister: Option<Arc<dyn Persister>>,
67 pub persister_overrides: HashMap<String, Arc<dyn Persister>>,
69 pub peer_clients: Option<Arc<dashmap::DashMap<Arc<str>, Arc<MykoClient>>>>,
73}
74
75#[derive(Default)]
77pub struct CellServerBuilder {
78 bind_addr: Option<SocketAddr>,
79 tcp_nodelay: Option<bool>,
80 host_id: Option<Uuid>,
81 postgres: Option<PostgresConfig>,
82 peer_registry: Option<peer_registry::PeerRegistryConfig>,
83 default_persister: Option<Arc<dyn Persister>>,
84 persister_overrides: HashMap<String, Arc<dyn Persister>>,
85 peer_clients: Option<Arc<dashmap::DashMap<Arc<str>, Arc<MykoClient>>>>,
89 after_init: Option<AfterInitCallback>,
90 server_info: Option<mcp::dispatch::ServerInfo>,
94}
95
96type AfterInitCallback = Box<dyn FnOnce(&CellServer) + Send>;
97
98impl CellServerBuilder {
99 pub fn new() -> Self {
101 Self::default()
102 }
103
104 pub fn with_bind_addr(mut self, addr: SocketAddr) -> Self {
106 self.bind_addr = Some(addr);
107 self
108 }
109
110 pub fn with_tcp_nodelay(mut self, enabled: bool) -> Self {
115 self.tcp_nodelay = Some(enabled);
116 self
117 }
118
119 pub fn with_host_id(mut self, id: Uuid) -> Self {
121 self.host_id = Some(id);
122 self
123 }
124
125 pub fn with_postgres(mut self, config: PostgresConfig) -> Self {
127 self.postgres = Some(config);
128 self
129 }
130
131 pub fn with_peer_registry(mut self, config: peer_registry::PeerRegistryConfig) -> Self {
133 self.peer_registry = Some(config);
134 self
135 }
136
137 pub fn with_default_persister(mut self, persister: Arc<dyn Persister>) -> Self {
139 self.default_persister = Some(persister);
140 self
141 }
142
143 pub fn with_persister_override(
145 mut self,
146 entity_type: impl Into<String>,
147 persister: Arc<dyn Persister>,
148 ) -> Self {
149 self.persister_overrides
150 .insert(entity_type.into(), persister);
151 self
152 }
153
154 pub fn with_peer_clients(
160 mut self,
161 peer_clients: Arc<dashmap::DashMap<Arc<str>, Arc<MykoClient>>>,
162 ) -> Self {
163 self.peer_clients = Some(peer_clients);
164 self
165 }
166
167 pub fn after_init(mut self, f: impl FnOnce(&CellServer) + Send + 'static) -> Self {
171 self.after_init = Some(Box::new(f));
172 self
173 }
174
175 pub fn with_server_info(mut self, info: mcp::dispatch::ServerInfo) -> Self {
179 self.server_info = Some(info);
180 self
181 }
182
183 pub fn build(self) -> CellServer {
185 let bind_addr = self
186 .bind_addr
187 .unwrap_or_else(|| "127.0.0.1:5155".parse().unwrap());
188
189 let server_info = Arc::new(self.server_info.unwrap_or_default());
190
191 let mut server = CellServer::new(CellServerConfig {
192 bind_addr,
193 tcp_nodelay: self.tcp_nodelay.unwrap_or(true),
194 postgres: self.postgres,
195 host_id: self.host_id,
196 peer_registry: self.peer_registry,
197 default_persister: self.default_persister,
198 persister_overrides: self.persister_overrides,
199 peer_clients: self.peer_clients,
200 });
201 server.after_init = std::sync::Mutex::new(self.after_init);
202 server.server_info = server_info;
203 server
204 }
205}
206
207pub struct CellServer {
211 pub registry: Arc<StoreRegistry>,
213 pub handler_registry: Arc<HandlerRegistry>,
215 pub relationship_manager: Arc<RelationshipManager>,
217 pub postgres_producer: Option<PostgresProducerHandle>,
219 pub search_index: Arc<SearchIndex>,
221 pub persisters: Arc<PersisterRouter>,
223 pub host_id: Uuid,
225 config: CellServerConfig,
227 _postgres_producer_owner: Option<CellPostgresProducer>,
229 postgres_consumer: Option<CellPostgresConsumer>,
231 ready: Arc<AtomicBool>,
233 peer_registry_instance: RwLock<Option<peer_registry::PeerRegistry>>,
235 peer_clients: Arc<dashmap::DashMap<Arc<str>, Arc<MykoClient>>>,
237 after_init: std::sync::Mutex<Option<AfterInitCallback>>,
239 server_info: Arc<mcp::dispatch::ServerInfo>,
243 saga_event_tx: flume::Sender<MEvent>,
245 saga_event_rx: std::sync::Mutex<Option<flume::Receiver<MEvent>>>,
247 saga_tasks: std::sync::Mutex<Vec<tokio::task::JoinHandle<()>>>,
249 _server_ownership_guard: std::sync::Mutex<Option<hyphae::SubscriptionGuard>>,
251}
252
253impl CellServer {
254 pub fn builder() -> CellServerBuilder {
256 CellServerBuilder::new()
257 }
258
259 pub fn new(config: CellServerConfig) -> Self {
261 let host_id = config.host_id.unwrap_or_else(Uuid::new_v4);
262 let registry = Arc::new(StoreRegistry::new());
263 let handler_registry = Arc::new(HandlerRegistry::new());
264 let relationship_manager = Arc::new(RelationshipManager::new());
265
266 init_client_registry();
268
269 let (saga_event_tx, saga_event_rx) = flume::unbounded::<MEvent>();
270 let (postgres_producer_owner, postgres_producer, postgres_consumer) =
271 if let Some(ref postgres_config) = config.postgres {
272 match CellPostgresProducer::new(postgres_config, host_id) {
273 Ok(producer) => {
274 let handle = producer.handle();
275 let consumer = match CellPostgresConsumer::start(
276 postgres_config,
277 host_id,
278 handler_registry.clone(),
279 registry.clone(),
280 ) {
281 Ok(c) => Some(c),
282 Err(e) => {
283 tracing::error!("Failed to start Postgres consumer: {}", e);
284 None
285 }
286 };
287 (Some(producer), Some(handle), consumer)
288 }
289 Err(e) => {
290 tracing::error!("Failed to create Postgres producer: {}", e);
291 (None, None, None)
292 }
293 }
294 } else {
295 (None, None, None)
296 };
297
298 let ready = Arc::new(AtomicBool::new(postgres_consumer.is_none()));
300
301 let search_index = Arc::new(SearchIndex::new());
303
304 let mut persister_router = PersisterRouter::default();
309 if let Some(default_persister) = config.default_persister.clone() {
310 persister_router.set_default(Some(default_persister));
311 } else if let Some(handle) = postgres_producer.clone() {
312 persister_router.set_default(Some(Arc::new(handle) as Arc<dyn Persister>));
313 }
314 for (entity_type, persister) in &config.persister_overrides {
315 persister_router.set_override(entity_type.clone(), persister.clone());
316 }
317 let persisters = Arc::new(persister_router);
318
319 let peer_clients = config
320 .peer_clients
321 .clone()
322 .unwrap_or_else(|| Arc::new(dashmap::DashMap::new()));
323
324 Self {
325 registry,
326 handler_registry,
327 relationship_manager,
328 postgres_producer,
329 search_index,
330 persisters,
331 host_id,
332 config,
333 _postgres_producer_owner: postgres_producer_owner,
334 postgres_consumer,
335 ready,
336 peer_registry_instance: RwLock::new(None),
337 peer_clients,
338 after_init: std::sync::Mutex::new(None),
339 server_info: Arc::new(mcp::dispatch::ServerInfo::default()),
340 saga_event_tx,
341 saga_event_rx: std::sync::Mutex::new(Some(saga_event_rx)),
342 saga_tasks: std::sync::Mutex::new(Vec::new()),
343 _server_ownership_guard: std::sync::Mutex::new(None),
344 }
345 }
346
347 pub fn start_peer_registry(&self, config: Option<peer_registry::PeerRegistryConfig>) {
349 let peer_config = config.or_else(|| self.config.peer_registry.clone());
350
351 if let Some(peer_config) = peer_config {
352 tracing::info!("Starting peer registry");
353 let pr = peer_registry::PeerRegistry::new(self.ctx(), peer_config);
354 *self.peer_registry_instance.write().unwrap() = Some(pr);
355 }
356 }
357
358 pub fn has_peer_registry(&self) -> bool {
360 self.peer_registry_instance.read().unwrap().is_some()
361 }
362
363 pub fn registry(&self) -> Arc<StoreRegistry> {
365 self.registry.clone()
366 }
367
368 pub fn handler_registry(&self) -> Arc<HandlerRegistry> {
370 self.handler_registry.clone()
371 }
372
373 pub fn server_info(&self) -> Arc<mcp::dispatch::ServerInfo> {
376 self.server_info.clone()
377 }
378
379 pub fn ctx(&self) -> CellServerCtx {
381 let history_replay: Option<Arc<dyn myko::server::HistoryReplayProvider>> =
382 self.config.postgres.as_ref().map(|pg| {
383 Arc::new(PostgresHistoryReplayProvider::new(pg.clone()))
384 as Arc<dyn myko::server::HistoryReplayProvider>
385 });
386 CellServerCtx::new(
387 self.host_id,
388 self.registry.clone(),
389 self.handler_registry.clone(),
390 self.relationship_manager.clone(),
391 self.persisters.clone(),
392 self.search_index.clone(),
393 self.peer_clients.clone(),
394 Some(self.saga_event_tx.clone()),
395 history_replay,
396 )
397 }
398
399 fn start_saga_runtime(&self) {
400 let registrations: Vec<_> = inventory::iter::<SagaRegistration>().collect();
401 if registrations.is_empty() {
402 return;
403 }
404 let Some(rx) = self
405 .saga_event_rx
406 .lock()
407 .expect("saga_event_rx mutex poisoned")
408 .take()
409 else {
410 return;
411 };
412
413 tracing::info!("Starting saga runtime with {} saga(s)", registrations.len());
414
415 struct SagaChannel {
418 tx: flume::Sender<MEvent>,
419 entity_type: &'static str,
420 change_type: myko::event::MEventType,
421 }
422 let mut saga_channels: Vec<SagaChannel> = Vec::new();
423
424 for registration in registrations {
425 let saga = (registration.create)();
426 let saga_name = saga.name().to_string();
427 let (saga_tx, saga_rx) = flume::unbounded::<MEvent>();
428 saga_channels.push(SagaChannel {
429 tx: saga_tx,
430 entity_type: registration.event_entity_type,
431 change_type: registration.event_change_type,
432 });
433 let events: myko::saga::EventStream = Box::pin(futures_util::stream::unfold(
434 saga_rx,
435 move |saga_rx| async move {
436 saga_rx
437 .recv_async()
438 .await
439 .ok()
440 .map(|event| (event, saga_rx))
441 },
442 ));
443
444 let saga_ctx = Arc::new(myko::saga::SagaContext::with_event_sink(
445 self.host_id,
446 self.registry.clone(),
447 self.saga_event_tx.clone(),
448 ));
449 let mut command_stream = saga.build_boxed(events, saga_ctx);
450
451 let host_id = self.host_id;
452 let registry = self.registry.clone();
453 let handler_registry = self.handler_registry.clone();
454 let relationship_manager = self.relationship_manager.clone();
455 let persisters = self.persisters.clone();
456 let search_index = self.search_index.clone();
457 let peer_clients = self.peer_clients.clone();
458 let saga_event_tx = self.saga_event_tx.clone();
459
460 let handle = tokio::spawn(async move {
461 while let Some(command) = command_stream.next().await {
462 let command_name = command.command_name();
463 tracing::debug!("Saga {} executing command {}", saga_name, command_name);
464 let req = Arc::new(RequestContext::internal(
465 Arc::from(Uuid::new_v4().to_string()),
466 host_id,
467 &format!("saga:{saga_name}"),
468 ));
469
470 let cmd_ctx = CommandContext::new(
471 Arc::from(command_name),
472 req,
473 Arc::new(CellServerCtx::new(
474 host_id,
475 registry.clone(),
476 handler_registry.clone(),
477 relationship_manager.clone(),
478 persisters.clone(),
479 search_index.clone(),
480 peer_clients.clone(),
481 Some(saga_event_tx.clone()),
482 None,
483 )),
484 );
485
486 if let Err(err) = command.execute_boxed(cmd_ctx) {
487 tracing::error!(
488 "Saga {} command {} failed: {}",
489 saga_name,
490 command_name,
491 err.message
492 );
493 }
494 }
495 });
496
497 self.saga_tasks
498 .lock()
499 .expect("saga_tasks mutex poisoned")
500 .push(handle);
501 }
502
503 let dispatcher = tokio::spawn(async move {
506 while let Ok(event) = rx.recv_async().await {
507 for ch in &saga_channels {
508 if event.item_type == ch.entity_type && event.change_type == ch.change_type {
509 let _ = ch.tx.send(event.clone());
510 }
511 }
512 }
513 });
514 self.saga_tasks
515 .lock()
516 .expect("saga_tasks mutex poisoned")
517 .push(dispatcher);
518 }
519
520 pub fn postgres_history_store(&self) -> Result<Option<PostgresHistoryStore>, String> {
522 self.config
523 .postgres
524 .clone()
525 .map(PostgresHistoryStore::new)
526 .transpose()
527 }
528
529 pub fn init_postgres_and_wait(&self, timeout: Duration) -> Result<(), String> {
531 if self.config.postgres.is_some() && self.postgres_consumer.is_none() {
532 return Err(
533 "Postgres is configured but the Postgres consumer is not running".to_string(),
534 );
535 }
536
537 if let Some(ref consumer) = self.postgres_consumer {
538 consumer.wait_until_caught_up(timeout)?;
539 self.ready.store(true, Ordering::SeqCst);
540 }
541 Ok(())
542 }
543
544 pub fn establish_relations(&self) {
546 if let Err(e) = self.relationship_manager.establish_relations(&self.ctx()) {
547 tracing::error!("Failed to establish relations: {e}");
548 }
549 }
550
551 pub fn is_ready(&self) -> bool {
553 if let Some(ref consumer) = self.postgres_consumer {
554 if consumer.is_caught_up() {
555 self.ready.store(true, Ordering::SeqCst);
556 return true;
557 }
558 return false;
559 }
560 true
561 }
562
563 pub async fn run(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
565 use tokio::net::TcpListener;
566
567 let entity_types: Vec<&str> = self
569 .handler_registry
570 .entity_types()
571 .map(|t| t.as_ref())
572 .collect();
573 self.persisters
574 .startup_healthcheck(&entity_types)
575 .map_err(|reason| format!("Persister startup healthcheck failed: {reason}"))?;
576
577 if self.config.postgres.is_some() && self.postgres_consumer.is_none() {
578 return Err("Postgres is configured but the Postgres consumer failed to start".into());
579 }
580
581 if self.postgres_consumer.is_some() {
583 tracing::info!("Waiting for Postgres event consumer to catch up...");
584 let timeout = std::time::Duration::from_secs(300);
585 self.init_postgres_and_wait(timeout)
586 .map_err(|reason| format!("Postgres startup catch-up failed: {reason}"))?;
587 tracing::info!("Postgres caught up, ready to accept connections");
588 }
589
590 tracing::info!("Building search index...");
592 self.search_index.build_from_registry(&self.registry);
593
594 tracing::info!("Establishing relations...");
596 self.establish_relations();
597
598 tracing::info!("Checking server-owned item ownership...");
600 if let Err(e) = ServerOwnershipManager::claim_orphaned(&self.ctx()) {
601 tracing::error!("Failed to claim orphaned server-owned items: {}", e);
602 }
603 let ownership_guard = ServerOwnershipManager::watch_peer_deaths(&self.ctx());
604 *self
605 ._server_ownership_guard
606 .lock()
607 .expect("server_ownership_guard mutex poisoned") = Some(ownership_guard);
608
609 if let Some(hook) = self
622 .after_init
623 .lock()
624 .expect("after_init mutex poisoned")
625 .take()
626 {
627 hook(self);
628 }
629
630 self.start_saga_runtime();
631
632 crate::ws_timing::start_periodic_logger();
636
637 myko::server::report_cache_stats::start_periodic_logger();
640
641 myko::server::entity_set_stats::start_periodic_logger();
644
645 myko::search::search_stats::start_periodic_logger();
648
649 crate::telemetry::register_item_count_gauge(self.registry.clone());
653
654 let listener = TcpListener::bind(&self.config.bind_addr).await?;
659 tracing::info!("CellServer listening on {}", self.config.bind_addr);
660 tracing::info!(
661 "Myko gateway: ws://{}/myko | MCP: /myko/mcp (POST + WS + SSE)",
662 self.config.bind_addr
663 );
664
665 if self.config.peer_registry.is_some() {
667 self.start_peer_registry(None);
668 }
669
670 tracing::info!("Server started");
671 self.run_ws_accept_loop(listener).await
672 }
673
674 pub async fn run_ws_loop(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
676 use tokio::net::TcpListener;
677
678 let listener = TcpListener::bind(&self.config.bind_addr).await?;
679 tracing::info!("CellServer listening on {}", self.config.bind_addr);
680 tracing::info!(
681 "Myko gateway: ws://{}/myko | MCP: /myko/mcp (POST + WS + SSE)",
682 self.config.bind_addr
683 );
684 self.run_ws_accept_loop(listener).await
685 }
686
687 async fn run_ws_accept_loop(
688 &self,
689 listener: tokio::net::TcpListener,
690 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
691 let ready = self.ready.clone();
692
693 loop {
694 let (stream, addr) = listener.accept().await?;
695
696 if self.config.tcp_nodelay
705 && let Err(e) = stream.set_nodelay(true)
706 {
707 tracing::warn!("failed to set TCP_NODELAY on connection from {addr}: {e}");
708 }
709
710 if !ready.load(Ordering::SeqCst) {
712 if self.is_ready() {
713 tracing::info!("Server is now ready to accept connections");
714 } else {
715 tracing::warn!(
716 "Rejecting connection from {} - server not ready (durable backend catching up)",
717 addr
718 );
719 drop(stream);
720 continue;
721 }
722 }
723
724 tracing::debug!("New connection from {}", addr);
725
726 let ctx = Arc::new(self.ctx());
727 let server_info = self.server_info.clone();
728
729 tokio::spawn(async move {
730 if let Err(e) = router::route_connection(stream, addr, ctx, server_info).await {
731 tracing::error!("Connection error from {}: {}", addr, e);
732 }
733 });
734 }
735 }
736}
737
738#[cfg(test)]
739mod tests {
740 use super::*;
741
742 #[test]
743 fn test_server_creation() {
744 let config = CellServerConfig {
745 bind_addr: "127.0.0.1:0".parse().unwrap(),
746 tcp_nodelay: true,
747 postgres: None,
748 host_id: None,
749 peer_registry: None,
750 default_persister: None,
751 persister_overrides: HashMap::new(),
752 peer_clients: None,
753 };
754 let server = CellServer::new(config);
755 assert!(Arc::strong_count(&server.registry) >= 1);
756 }
757
758 #[test]
759 fn test_server_with_host_id() {
760 let host_id = Uuid::new_v4();
761 let config = CellServerConfig {
762 bind_addr: "127.0.0.1:0".parse().unwrap(),
763 tcp_nodelay: true,
764 postgres: None,
765 host_id: Some(host_id),
766 peer_registry: None,
767 default_persister: None,
768 persister_overrides: HashMap::new(),
769 peer_clients: None,
770 };
771 let server = CellServer::new(config);
772 assert_eq!(server.host_id, host_id);
773 }
774}