1use 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}