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 #[cfg(feature = "inspector")]
253 _inspector: hyphae::server::InspectorServer,
254}
255
256impl CellServer {
257 pub fn builder() -> CellServerBuilder {
259 CellServerBuilder::new()
260 }
261
262 pub fn new(config: CellServerConfig) -> Self {
264 let host_id = config.host_id.unwrap_or_else(Uuid::new_v4);
265 let registry = Arc::new(StoreRegistry::new());
266 let handler_registry = Arc::new(HandlerRegistry::new());
267 let relationship_manager = Arc::new(RelationshipManager::new());
268
269 init_client_registry();
271
272 let (saga_event_tx, saga_event_rx) = flume::unbounded::<MEvent>();
273 let (postgres_producer_owner, postgres_producer, postgres_consumer) =
274 if let Some(ref postgres_config) = config.postgres {
275 match CellPostgresProducer::new(postgres_config, host_id) {
276 Ok(producer) => {
277 let handle = producer.handle();
278 let consumer = match CellPostgresConsumer::start(
279 postgres_config,
280 host_id,
281 handler_registry.clone(),
282 registry.clone(),
283 ) {
284 Ok(c) => Some(c),
285 Err(e) => {
286 tracing::error!("Failed to start Postgres consumer: {}", e);
287 None
288 }
289 };
290 (Some(producer), Some(handle), consumer)
291 }
292 Err(e) => {
293 tracing::error!("Failed to create Postgres producer: {}", e);
294 (None, None, None)
295 }
296 }
297 } else {
298 (None, None, None)
299 };
300
301 let ready = Arc::new(AtomicBool::new(postgres_consumer.is_none()));
303
304 let search_index = Arc::new(SearchIndex::new());
306
307 let mut persister_router = PersisterRouter::default();
312 if let Some(default_persister) = config.default_persister.clone() {
313 persister_router.set_default(Some(default_persister));
314 } else if let Some(handle) = postgres_producer.clone() {
315 persister_router.set_default(Some(Arc::new(handle) as Arc<dyn Persister>));
316 }
317 for (entity_type, persister) in &config.persister_overrides {
318 persister_router.set_override(entity_type.clone(), persister.clone());
319 }
320 let persisters = Arc::new(persister_router);
321
322 #[cfg(feature = "inspector")]
324 let inspector = hyphae::server::start_server("myko");
325 #[cfg(feature = "inspector")]
326 tracing::info!("Hyphae inspector on port {}", inspector.port());
327
328 let peer_clients = config
329 .peer_clients
330 .clone()
331 .unwrap_or_else(|| Arc::new(dashmap::DashMap::new()));
332
333 Self {
334 registry,
335 handler_registry,
336 relationship_manager,
337 postgres_producer,
338 search_index,
339 persisters,
340 host_id,
341 config,
342 _postgres_producer_owner: postgres_producer_owner,
343 postgres_consumer,
344 ready,
345 peer_registry_instance: RwLock::new(None),
346 peer_clients,
347 after_init: std::sync::Mutex::new(None),
348 server_info: Arc::new(mcp::dispatch::ServerInfo::default()),
349 saga_event_tx,
350 saga_event_rx: std::sync::Mutex::new(Some(saga_event_rx)),
351 saga_tasks: std::sync::Mutex::new(Vec::new()),
352 _server_ownership_guard: std::sync::Mutex::new(None),
353 #[cfg(feature = "inspector")]
354 _inspector: inspector,
355 }
356 }
357
358 pub fn start_peer_registry(&self, config: Option<peer_registry::PeerRegistryConfig>) {
360 let peer_config = config.or_else(|| self.config.peer_registry.clone());
361
362 if let Some(peer_config) = peer_config {
363 tracing::info!("Starting peer registry");
364 let pr = peer_registry::PeerRegistry::new(self.ctx(), peer_config);
365 *self.peer_registry_instance.write().unwrap() = Some(pr);
366 }
367 }
368
369 pub fn has_peer_registry(&self) -> bool {
371 self.peer_registry_instance.read().unwrap().is_some()
372 }
373
374 pub fn registry(&self) -> Arc<StoreRegistry> {
376 self.registry.clone()
377 }
378
379 pub fn handler_registry(&self) -> Arc<HandlerRegistry> {
381 self.handler_registry.clone()
382 }
383
384 pub fn server_info(&self) -> Arc<mcp::dispatch::ServerInfo> {
387 self.server_info.clone()
388 }
389
390 pub fn ctx(&self) -> CellServerCtx {
392 let history_replay: Option<Arc<dyn myko::server::HistoryReplayProvider>> =
393 self.config.postgres.as_ref().map(|pg| {
394 Arc::new(PostgresHistoryReplayProvider::new(pg.clone()))
395 as Arc<dyn myko::server::HistoryReplayProvider>
396 });
397 CellServerCtx::new(
398 self.host_id,
399 self.registry.clone(),
400 self.handler_registry.clone(),
401 self.relationship_manager.clone(),
402 self.persisters.clone(),
403 self.search_index.clone(),
404 self.peer_clients.clone(),
405 Some(self.saga_event_tx.clone()),
406 history_replay,
407 )
408 }
409
410 fn start_saga_runtime(&self) {
411 let registrations: Vec<_> = inventory::iter::<SagaRegistration>().collect();
412 if registrations.is_empty() {
413 return;
414 }
415 let Some(rx) = self
416 .saga_event_rx
417 .lock()
418 .expect("saga_event_rx mutex poisoned")
419 .take()
420 else {
421 return;
422 };
423
424 tracing::info!("Starting saga runtime with {} saga(s)", registrations.len());
425
426 struct SagaChannel {
429 tx: flume::Sender<MEvent>,
430 entity_type: &'static str,
431 change_type: myko::event::MEventType,
432 }
433 let mut saga_channels: Vec<SagaChannel> = Vec::new();
434
435 for registration in registrations {
436 let saga = (registration.create)();
437 let saga_name = saga.name().to_string();
438 let (saga_tx, saga_rx) = flume::unbounded::<MEvent>();
439 saga_channels.push(SagaChannel {
440 tx: saga_tx,
441 entity_type: registration.event_entity_type,
442 change_type: registration.event_change_type,
443 });
444 let events: myko::saga::EventStream = Box::pin(futures_util::stream::unfold(
445 saga_rx,
446 move |saga_rx| async move {
447 saga_rx
448 .recv_async()
449 .await
450 .ok()
451 .map(|event| (event, saga_rx))
452 },
453 ));
454
455 let saga_ctx = Arc::new(myko::saga::SagaContext::with_event_sink(
456 self.host_id,
457 self.registry.clone(),
458 self.saga_event_tx.clone(),
459 ));
460 let mut command_stream = saga.build_boxed(events, saga_ctx);
461
462 let host_id = self.host_id;
463 let registry = self.registry.clone();
464 let handler_registry = self.handler_registry.clone();
465 let relationship_manager = self.relationship_manager.clone();
466 let persisters = self.persisters.clone();
467 let search_index = self.search_index.clone();
468 let peer_clients = self.peer_clients.clone();
469 let saga_event_tx = self.saga_event_tx.clone();
470
471 let handle = tokio::spawn(async move {
472 while let Some(command) = command_stream.next().await {
473 let command_name = command.command_name();
474 tracing::debug!("Saga {} executing command {}", saga_name, command_name);
475 let req = Arc::new(RequestContext::internal(
476 Arc::from(Uuid::new_v4().to_string()),
477 host_id,
478 &format!("saga:{saga_name}"),
479 ));
480
481 let cmd_ctx = CommandContext::new(
482 Arc::from(command_name),
483 req,
484 Arc::new(CellServerCtx::new(
485 host_id,
486 registry.clone(),
487 handler_registry.clone(),
488 relationship_manager.clone(),
489 persisters.clone(),
490 search_index.clone(),
491 peer_clients.clone(),
492 Some(saga_event_tx.clone()),
493 None,
494 )),
495 );
496
497 if let Err(err) = command.execute_boxed(cmd_ctx) {
498 tracing::error!(
499 "Saga {} command {} failed: {}",
500 saga_name,
501 command_name,
502 err.message
503 );
504 }
505 }
506 });
507
508 self.saga_tasks
509 .lock()
510 .expect("saga_tasks mutex poisoned")
511 .push(handle);
512 }
513
514 let dispatcher = tokio::spawn(async move {
517 while let Ok(event) = rx.recv_async().await {
518 for ch in &saga_channels {
519 if event.item_type == ch.entity_type && event.change_type == ch.change_type {
520 let _ = ch.tx.send(event.clone());
521 }
522 }
523 }
524 });
525 self.saga_tasks
526 .lock()
527 .expect("saga_tasks mutex poisoned")
528 .push(dispatcher);
529 }
530
531 pub fn postgres_history_store(&self) -> Result<Option<PostgresHistoryStore>, String> {
533 self.config
534 .postgres
535 .clone()
536 .map(PostgresHistoryStore::new)
537 .transpose()
538 }
539
540 pub fn init_postgres_and_wait(&self, timeout: Duration) -> Result<(), String> {
542 if self.config.postgres.is_some() && self.postgres_consumer.is_none() {
543 return Err(
544 "Postgres is configured but the Postgres consumer is not running".to_string(),
545 );
546 }
547
548 if let Some(ref consumer) = self.postgres_consumer {
549 consumer.wait_until_caught_up(timeout)?;
550 self.ready.store(true, Ordering::SeqCst);
551 }
552 Ok(())
553 }
554
555 pub fn establish_relations(&self) {
557 if let Err(e) = self.relationship_manager.establish_relations(&self.ctx()) {
558 tracing::error!("Failed to establish relations: {e}");
559 }
560 }
561
562 pub fn is_ready(&self) -> bool {
564 if let Some(ref consumer) = self.postgres_consumer {
565 if consumer.is_caught_up() {
566 self.ready.store(true, Ordering::SeqCst);
567 return true;
568 }
569 return false;
570 }
571 true
572 }
573
574 pub async fn run(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
576 use tokio::net::TcpListener;
577
578 let entity_types: Vec<&str> = self
580 .handler_registry
581 .entity_types()
582 .map(|t| t.as_ref())
583 .collect();
584 self.persisters
585 .startup_healthcheck(&entity_types)
586 .map_err(|reason| format!("Persister startup healthcheck failed: {reason}"))?;
587
588 if self.config.postgres.is_some() && self.postgres_consumer.is_none() {
589 return Err("Postgres is configured but the Postgres consumer failed to start".into());
590 }
591
592 if self.postgres_consumer.is_some() {
594 tracing::info!("Waiting for Postgres event consumer to catch up...");
595 let timeout = std::time::Duration::from_secs(300);
596 self.init_postgres_and_wait(timeout)
597 .map_err(|reason| format!("Postgres startup catch-up failed: {reason}"))?;
598 tracing::info!("Postgres caught up, ready to accept connections");
599 }
600
601 tracing::info!("Building search index...");
603 self.search_index.build_from_registry(&self.registry);
604
605 tracing::info!("Establishing relations...");
607 self.establish_relations();
608
609 tracing::info!("Checking server-owned item ownership...");
611 if let Err(e) = ServerOwnershipManager::claim_orphaned(&self.ctx()) {
612 tracing::error!("Failed to claim orphaned server-owned items: {}", e);
613 }
614 let ownership_guard = ServerOwnershipManager::watch_peer_deaths(&self.ctx());
615 *self
616 ._server_ownership_guard
617 .lock()
618 .expect("server_ownership_guard mutex poisoned") = Some(ownership_guard);
619
620 if let Some(hook) = self
633 .after_init
634 .lock()
635 .expect("after_init mutex poisoned")
636 .take()
637 {
638 hook(self);
639 }
640
641 self.start_saga_runtime();
642
643 crate::ws_timing::start_periodic_logger();
647
648 myko::server::report_cache_stats::start_periodic_logger();
651
652 myko::server::entity_set_stats::start_periodic_logger();
655
656 myko::search::search_stats::start_periodic_logger();
659
660 crate::telemetry::register_item_count_gauge(self.registry.clone());
664
665 let listener = TcpListener::bind(&self.config.bind_addr).await?;
670 tracing::info!("CellServer listening on {}", self.config.bind_addr);
671 tracing::info!(
672 "Myko gateway: ws://{}/myko | MCP: /myko/mcp (POST + WS + SSE)",
673 self.config.bind_addr
674 );
675
676 if self.config.peer_registry.is_some() {
678 self.start_peer_registry(None);
679 }
680
681 tracing::info!("Server started");
682 self.run_ws_accept_loop(listener).await
683 }
684
685 pub async fn run_ws_loop(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
687 use tokio::net::TcpListener;
688
689 let listener = TcpListener::bind(&self.config.bind_addr).await?;
690 tracing::info!("CellServer listening on {}", self.config.bind_addr);
691 tracing::info!(
692 "Myko gateway: ws://{}/myko | MCP: /myko/mcp (POST + WS + SSE)",
693 self.config.bind_addr
694 );
695 self.run_ws_accept_loop(listener).await
696 }
697
698 async fn run_ws_accept_loop(
699 &self,
700 listener: tokio::net::TcpListener,
701 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
702 let ready = self.ready.clone();
703
704 loop {
705 let (stream, addr) = listener.accept().await?;
706
707 if self.config.tcp_nodelay
716 && let Err(e) = stream.set_nodelay(true)
717 {
718 tracing::warn!("failed to set TCP_NODELAY on connection from {addr}: {e}");
719 }
720
721 if !ready.load(Ordering::SeqCst) {
723 if self.is_ready() {
724 tracing::info!("Server is now ready to accept connections");
725 } else {
726 tracing::warn!(
727 "Rejecting connection from {} - server not ready (durable backend catching up)",
728 addr
729 );
730 drop(stream);
731 continue;
732 }
733 }
734
735 tracing::debug!("New connection from {}", addr);
736
737 let ctx = Arc::new(self.ctx());
738 let server_info = self.server_info.clone();
739
740 tokio::spawn(async move {
741 if let Err(e) = router::route_connection(stream, addr, ctx, server_info).await {
742 tracing::error!("Connection error from {}: {}", addr, e);
743 }
744 });
745 }
746 }
747}
748
749#[cfg(test)]
750mod tests {
751 use super::*;
752
753 #[test]
754 fn test_server_creation() {
755 let config = CellServerConfig {
756 bind_addr: "127.0.0.1:0".parse().unwrap(),
757 tcp_nodelay: true,
758 postgres: None,
759 host_id: None,
760 peer_registry: None,
761 default_persister: None,
762 persister_overrides: HashMap::new(),
763 peer_clients: None,
764 };
765 let server = CellServer::new(config);
766 assert!(Arc::strong_count(&server.registry) >= 1);
767 }
768
769 #[test]
770 fn test_server_with_host_id() {
771 let host_id = Uuid::new_v4();
772 let config = CellServerConfig {
773 bind_addr: "127.0.0.1:0".parse().unwrap(),
774 tcp_nodelay: true,
775 postgres: None,
776 host_id: Some(host_id),
777 peer_registry: None,
778 default_persister: None,
779 persister_overrides: HashMap::new(),
780 peer_clients: None,
781 };
782 let server = CellServer::new(config);
783 assert_eq!(server.host_id, host_id);
784 }
785}