Skip to main content

mcpmem_runtime/
lib.rs

1//! Feature-gated runtime role selection and supervision.
2
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6
7use mcpmem_core::LifecycleState;
8use thiserror::Error;
9use tokio::task::{JoinError, JoinSet};
10
11pub type RoleFuture = Pin<Box<dyn Future<Output = Result<(), RuntimeError>> + Send + 'static>>;
12
13#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
14pub enum RuntimeRole {
15    Mcp,
16    Indexer,
17    Webhooks,
18}
19
20impl RuntimeRole {
21    const fn name(self) -> &'static str {
22        match self {
23            Self::Mcp => "mcp",
24            Self::Indexer => "indexer",
25            Self::Webhooks => "webhooks",
26        }
27    }
28
29    const fn is_compiled(self) -> bool {
30        match self {
31            Self::Mcp => true,
32            Self::Indexer => cfg!(feature = "indexer"),
33            Self::Webhooks => cfg!(feature = "webhooks"),
34        }
35    }
36}
37
38impl std::fmt::Display for RuntimeRole {
39    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        formatter.write_str(self.name())
41    }
42}
43
44#[derive(Clone, Debug, Eq, PartialEq, Error)]
45pub enum ConfigError {
46    #[error("at least one runtime role is required")]
47    EmptyRoleSet,
48    #[error("runtime role names cannot be empty")]
49    EmptyRoleName,
50    #[error("unknown runtime role '{0}'")]
51    UnknownRole(String),
52    #[error("runtime role '{0}' was selected more than once")]
53    DuplicateRole(RuntimeRole),
54    #[error("runtime role '{0}' was selected but its Cargo feature is not compiled")]
55    RoleNotCompiled(RuntimeRole),
56}
57
58#[derive(Clone, Debug, Eq, PartialEq)]
59pub struct RoleSet(Vec<RuntimeRole>);
60
61impl RoleSet {
62    pub fn parse_csv(input: &str) -> Result<Self, ConfigError> {
63        if input.trim().is_empty() {
64            return Err(ConfigError::EmptyRoleSet);
65        }
66
67        let roles = input
68            .split(',')
69            .map(str::trim)
70            .map(Self::parse_role)
71            .collect::<Result<Vec<_>, _>>()?;
72
73        for (index, role) in roles.iter().enumerate() {
74            if roles[..index].contains(role) {
75                return Err(ConfigError::DuplicateRole(*role));
76            }
77            if !role.is_compiled() {
78                return Err(ConfigError::RoleNotCompiled(*role));
79            }
80        }
81
82        Ok(Self(roles))
83    }
84
85    pub fn mcp_only() -> Self {
86        Self(vec![RuntimeRole::Mcp])
87    }
88
89    pub fn roles(&self) -> &[RuntimeRole] {
90        &self.0
91    }
92
93    fn parse_role(value: &str) -> Result<RuntimeRole, ConfigError> {
94        match value {
95            "mcp" => Ok(RuntimeRole::Mcp),
96            "indexer" => Ok(RuntimeRole::Indexer),
97            "webhooks" => Ok(RuntimeRole::Webhooks),
98            "" => Err(ConfigError::EmptyRoleName),
99            unknown => Err(ConfigError::UnknownRole(unknown.to_owned())),
100        }
101    }
102}
103
104#[derive(Clone, Debug, Eq, PartialEq)]
105pub struct RoleLifecycle {
106    pub role: RuntimeRole,
107    pub state: LifecycleState,
108}
109
110pub trait RoleService: Send + Sync {
111    fn run(&self) -> RoleFuture;
112}
113
114#[derive(Clone)]
115pub struct AppServices {
116    mcp: Arc<dyn RoleService>,
117    #[cfg(feature = "indexer")]
118    indexer: Arc<dyn RoleService>,
119    #[cfg(feature = "webhooks")]
120    webhooks: Arc<dyn RoleService>,
121}
122
123impl AppServices {
124    pub fn new(mcp: Arc<dyn RoleService>) -> Self {
125        Self {
126            mcp,
127            #[cfg(feature = "indexer")]
128            indexer: Arc::new(NoopService),
129            #[cfg(feature = "webhooks")]
130            webhooks: Arc::new(NoopService),
131        }
132    }
133    #[cfg(feature = "webhooks")]
134    pub fn with_webhooks(mut self, webhooks: Arc<dyn RoleService>) -> Self {
135        self.webhooks = webhooks;
136        self
137    }
138
139    #[cfg(feature = "indexer")]
140    pub fn with_indexer(mut self, indexer: Arc<dyn RoleService>) -> Self {
141        self.indexer = indexer;
142        self
143    }
144}
145
146#[derive(Debug, Error)]
147pub enum RuntimeError {
148    #[error("runtime role '{role}' failed: {message}")]
149    RoleFailed { role: RuntimeRole, message: String },
150    #[error("runtime role task failed: {source}")]
151    TaskFailed {
152        #[source]
153        source: JoinError,
154    },
155}
156
157pub struct RuntimeComposition;
158
159impl RuntimeComposition {
160    pub fn start(roles: RoleSet, services: Arc<AppServices>) -> Result<RunningRoles, RuntimeError> {
161        let RoleSet(roles) = roles;
162        let AppServices {
163            mcp,
164            #[cfg(feature = "indexer")]
165            indexer,
166            #[cfg(feature = "webhooks")]
167            webhooks,
168        } = Arc::unwrap_or_clone(services);
169        let mut tasks = JoinSet::new();
170        let lifecycle = roles
171            .into_iter()
172            .map(|role| {
173                let service: Arc<dyn RoleService> = match role {
174                    RuntimeRole::Mcp => mcp.clone(),
175                    #[cfg(feature = "indexer")]
176                    RuntimeRole::Indexer => indexer.clone(),
177                    #[cfg(not(feature = "indexer"))]
178                    RuntimeRole::Indexer => Arc::new(NoopService),
179                    #[cfg(feature = "webhooks")]
180                    RuntimeRole::Webhooks => webhooks.clone(),
181                    #[cfg(not(feature = "webhooks"))]
182                    RuntimeRole::Webhooks => Arc::new(NoopService),
183                };
184                tasks.spawn(async move { (role, service.run().await) });
185                RoleLifecycle {
186                    role,
187                    state: LifecycleState::Running,
188                }
189            })
190            .collect();
191
192        Ok(RunningRoles { lifecycle, tasks })
193    }
194}
195
196struct NoopService;
197
198impl RoleService for NoopService {
199    fn run(&self) -> RoleFuture {
200        Box::pin(std::future::pending())
201    }
202}
203
204pub struct RunningRoles {
205    lifecycle: Vec<RoleLifecycle>,
206    tasks: JoinSet<(RuntimeRole, Result<(), RuntimeError>)>,
207}
208
209impl RunningRoles {
210    pub fn lifecycle(&self) -> &[RoleLifecycle] {
211        &self.lifecycle
212    }
213
214    pub async fn wait_for_shutdown(mut self) -> Result<(), RuntimeError> {
215        let result = self.tasks.join_next().await;
216        self.tasks.abort_all();
217
218        match result {
219            Some(Ok((_, Ok(())))) | None => Ok(()),
220            Some(Ok((_, Err(error)))) => Err(error),
221            Some(Err(source)) => Err(RuntimeError::TaskFailed { source }),
222        }
223    }
224}