1use std::sync::Arc;
7
8use anyhow::Result;
9use locus_core_rs::domain::contracts::EmbeddingProvider;
10use locus_core_rs::{
11 CalibrationService, EmbeddingMigrationService, InMemoryNodeStore, MonthlyRollupService,
12 MoodCatalogService, NodeStore, NodeStoreInitializer, NodeValidator, StoreContextService,
13 SttpNodeParser, SurrealDbNodeStore, SurrealDbRuntimeOptions, TreeSitterValidator,
14};
15use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters};
16use rmcp::{ServerHandler, ServiceExt, tool, tool_handler, tool_router};
17use schemars::JsonSchema;
18use serde::Deserialize;
19
20mod composition;
21mod shared;
22mod tools;
23
24use composition::{
25 RuntimeSurrealDbClient, build_embedding_provider, init_logging, load_surreal_settings,
26 resolve_parser_profile, runtime_args,
27};
28
29pub(crate) use shared::{
30 avec_to_json, expanded_limit, filter_nodes_by_context_keywords, infer_store_error_code,
31 mode_to_string, normalize_context_keywords, normalize_tiers, parse_migration_mode,
32 parse_utc_optional, parse_utc_required, schema_first_guidance_payload,
33 strict_typed_ir_profile_name, sttp_node_to_json, to_json_string, tool_error,
34 validate_batch_size, validate_limit, validate_max_nodes,
35};
36
37#[derive(Clone)]
38pub(crate) struct SttpMcpServer {
39 pub(crate) node_store: Arc<dyn NodeStore>,
40 pub(crate) calibration: Arc<CalibrationService>,
41 pub(crate) store_context: Arc<StoreContextService>,
42 pub(crate) embedding_migration: Arc<EmbeddingMigrationService>,
43 pub(crate) embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
44 pub(crate) moods: Arc<MoodCatalogService>,
45 pub(crate) monthly_rollup: Arc<MonthlyRollupService>,
46 #[allow(dead_code)]
47 tool_router: ToolRouter<Self>,
48}
49
50impl SttpMcpServer {
51 fn new(
52 node_store: Arc<dyn NodeStore>,
53 calibration: Arc<CalibrationService>,
54 store_context: Arc<StoreContextService>,
55 embedding_migration: Arc<EmbeddingMigrationService>,
56 embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
57 moods: Arc<MoodCatalogService>,
58 monthly_rollup: Arc<MonthlyRollupService>,
59 ) -> Self {
60 Self {
61 node_store,
62 calibration,
63 store_context,
64 embedding_migration,
65 embedding_provider,
66 moods,
67 monthly_rollup,
68 tool_router: Self::tool_router(),
69 }
70 }
71
72 pub(crate) async fn embed_context_keywords(&self, keywords: &[String]) -> Option<Vec<f32>> {
73 let provider = self.embedding_provider.as_ref()?;
74 let prompt = keywords.join(" ");
75 let prompt = prompt.trim();
76
77 if prompt.is_empty() {
78 return None;
79 }
80
81 provider
82 .embed_async(prompt)
83 .await
84 .ok()
85 .filter(|vector| !vector.is_empty())
86 }
87}
88
89#[tool_router]
90impl SttpMcpServer {
91 #[tool(
92 name = "get_schema",
93 description = "Get a canonical example of what an STTP node should look like before storage."
94 )]
95 async fn get_schema(&self) -> String {
96 tools::get_schema::execute().await
97 }
98
99 #[tool(
100 name = "calibrate_session",
101 description = "Call this at session start and after heavy reasoning work to measure current AVEC drift. Use it to compare your current cognitive state against prior calibration for the same session before storing or retrieving memory. On first calibration, name the session id something similar to the topic of the conversation if no session id was provided by user."
102 )]
103 async fn calibrate_session(
104 &self,
105 Parameters(request): Parameters<CalibrateSessionRequest>,
106 ) -> String {
107 tools::calibrate_session::execute(self, request).await
108 }
109
110 #[tool(
111 name = "store_context",
112 description = "Call this when context should be preserved to memory. Store a complete valid STTP node so future retrieval can rehydrate prior reasoning state, decisions, and confidence signals. If no session id provided by user, use something that the user can semantically relate to the conversation for better retrieval."
113 )]
114 async fn store_context(&self, Parameters(request): Parameters<StoreContextRequest>) -> String {
115 tools::store_context::execute(self, request).await
116 }
117
118 #[tool(
119 name = "get_context",
120 description = "Primary memory retrieval tool. MUST USE ANYTIME USER ASKS SOMETHING ABOUT REMEMBERING OR MEMORY RELATED INQUIERIES. Returns top resonant memory nodes for the provided AVEC state. Optional context_keywords enables server-side semantic retrieval (with internal embedding generation); keyword fallback is only used when semantic retrieval returns no nodes (or embeddings are unavailable). If session_id is omitted, retrieval is global across sessions. Use list_nodes for inventory when no results comeback after user prompts for memory retrieval."
121 )]
122 async fn get_context(&self, Parameters(request): Parameters<GetContextRequest>) -> String {
123 tools::get_context::execute(self, request).await
124 }
125
126 #[tool(
127 name = "list_nodes",
128 description = "Memory inventory tool. Lists stored nodes newest-first (global when session_id is omitted). Optional context_keywords performs fuzzy and semantic filtering against context_summary for fast discovery. Unlike get_context, list_nodes does not perform AVEC resonance ranking."
129 )]
130 async fn list_nodes(&self, Parameters(request): Parameters<ListNodesRequest>) -> String {
131 tools::list_nodes::execute(self, request).await
132 }
133
134 #[tool(
135 name = "preview_embedding_migration",
136 description = "Preview which nodes would be selected for embedding migration/backfill based on optional filters. Use this before running migration to verify scope and provider availability."
137 )]
138 async fn preview_embedding_migration(
139 &self,
140 Parameters(request): Parameters<PreviewEmbeddingMigrationRequest>,
141 ) -> String {
142 tools::preview_embedding_migration::execute(self, request).await
143 }
144
145 #[tool(
146 name = "run_embedding_migration",
147 description = "Run embedding migration/backfill for selected nodes. Supports dry_run, missing_only mode, and reindex_all mode using the currently configured embedding provider."
148 )]
149 async fn run_embedding_migration(
150 &self,
151 Parameters(request): Parameters<RunEmbeddingMigrationRequest>,
152 ) -> String {
153 tools::run_embedding_migration::execute(self, request).await
154 }
155
156 #[tool(
157 name = "get_moods",
158 description = "Retrieve AVEC mood presets and optional blend preview to intentionally shift reasoning mode (focused, creative, analytical, exploratory, collaborative, defensive, passive) before memory operations. Help maintain coherence and tone. USE WHEN ASKED TO STORE OR RETRIEVE MEMORY WITHOUT INITIAL AVEC CONFIG."
159 )]
160 async fn get_moods(&self, Parameters(request): Parameters<GetMoodsRequest>) -> String {
161 tools::get_moods::execute(self, request).await
162 }
163
164 #[tool(
165 name = "create_monthly_rollup",
166 description = "Aggregate many stored nodes into a compact monthly memory checkpoint. Use this to reduce retrieval noise and preserve high-level memory continuity across long timelines."
167 )]
168 async fn create_monthly_rollup(
169 &self,
170 Parameters(request): Parameters<CreateMonthlyRollupRequest>,
171 ) -> String {
172 tools::create_monthly_rollup::execute(self, request).await
173 }
174}
175
176#[tool_handler]
177impl ServerHandler for SttpMcpServer {}
178
179#[derive(Debug, Deserialize, JsonSchema)]
180pub(crate) struct CalibrateSessionRequest {
181 session_id: String,
182 stability: f32,
183 friction: f32,
184 logic: f32,
185 autonomy: f32,
186 trigger: String,
187}
188
189#[derive(Debug, Deserialize, JsonSchema)]
190pub(crate) struct StoreContextRequest {
191 node: String,
192 session_id: String,
193}
194
195fn default_limit_get_context() -> usize {
196 5
197}
198
199fn default_blend() -> f32 {
200 1.0
201}
202
203#[derive(Debug, Deserialize, JsonSchema)]
204pub(crate) struct GetContextRequest {
205 #[serde(default)]
206 session_id: Option<String>,
207 stability: f32,
208 friction: f32,
209 logic: f32,
210 autonomy: f32,
211 #[serde(default = "default_limit_get_context")]
212 limit: usize,
213 #[serde(default)]
214 from_utc: Option<String>,
215 #[serde(default)]
216 to_utc: Option<String>,
217 #[serde(default)]
218 tiers: Option<Vec<String>>,
219 #[serde(default)]
220 context_keywords: Option<Vec<String>>,
221 #[serde(default)]
222 alpha: Option<f32>,
223 #[serde(default)]
224 beta: Option<f32>,
225}
226
227#[derive(Debug, Deserialize, JsonSchema)]
228pub(crate) struct ListNodesRequest {
229 #[serde(default = "default_limit_list_nodes")]
230 limit: usize,
231 #[serde(default)]
232 session_id: Option<String>,
233 #[serde(default)]
234 context_keywords: Option<Vec<String>>,
235}
236
237fn default_limit_list_nodes() -> usize {
238 50
239}
240
241fn default_sample_limit_preview_migration() -> usize {
242 20
243}
244
245fn default_batch_size_migration() -> usize {
246 100
247}
248
249fn default_max_nodes_migration() -> usize {
250 5000
251}
252
253#[derive(Debug, Deserialize, JsonSchema)]
254pub(crate) struct PreviewEmbeddingMigrationRequest {
255 #[serde(default)]
256 session_id: Option<String>,
257 #[serde(default)]
258 from_utc: Option<String>,
259 #[serde(default)]
260 to_utc: Option<String>,
261 #[serde(default)]
262 tiers: Option<Vec<String>>,
263 #[serde(default)]
264 has_embedding: Option<bool>,
265 #[serde(default)]
266 embedding_model: Option<String>,
267 #[serde(default)]
268 sync_keys: Option<Vec<String>>,
269 #[serde(default = "default_sample_limit_preview_migration")]
270 sample_limit: usize,
271 #[serde(default = "default_max_nodes_migration")]
272 max_nodes: usize,
273}
274
275#[derive(Debug, Deserialize, JsonSchema)]
276pub(crate) struct RunEmbeddingMigrationRequest {
277 #[serde(default)]
278 session_id: Option<String>,
279 #[serde(default)]
280 from_utc: Option<String>,
281 #[serde(default)]
282 to_utc: Option<String>,
283 #[serde(default)]
284 tiers: Option<Vec<String>>,
285 #[serde(default)]
286 has_embedding: Option<bool>,
287 #[serde(default)]
288 embedding_model: Option<String>,
289 #[serde(default)]
290 sync_keys: Option<Vec<String>>,
291 #[serde(default)]
292 mode: Option<String>,
293 #[serde(default = "default_true")]
294 dry_run: bool,
295 #[serde(default = "default_batch_size_migration")]
296 batch_size: usize,
297 #[serde(default = "default_max_nodes_migration")]
298 max_nodes: usize,
299}
300
301#[derive(Debug, Deserialize, JsonSchema)]
302pub(crate) struct GetMoodsRequest {
303 #[serde(default)]
304 target_mood: Option<String>,
305 #[serde(default = "default_blend")]
306 blend: f32,
307 #[serde(default)]
308 current_stability: Option<f32>,
309 #[serde(default)]
310 current_friction: Option<f32>,
311 #[serde(default)]
312 current_logic: Option<f32>,
313 #[serde(default)]
314 current_autonomy: Option<f32>,
315}
316
317#[derive(Debug, Deserialize, JsonSchema)]
318pub(crate) struct CreateMonthlyRollupRequest {
319 session_id: String,
320 start_date_utc: String,
321 end_date_utc: String,
322 #[serde(default)]
323 source_session_id: Option<String>,
324 #[serde(default)]
325 parent_node_id: Option<String>,
326 #[serde(default = "default_true")]
327 persist: bool,
328}
329
330fn default_true() -> bool {
331 true
332}
333
334#[tokio::main]
335async fn main() -> Result<()> {
336 init_logging();
337
338 let args = std::env::args().collect::<Vec<_>>();
339 let use_in_memory = std::env::var("LOCUS_MCP_IN_MEMORY")
340 .map(|value| {
341 let normalized = value.trim().to_ascii_lowercase();
342 normalized == "1" || normalized == "true" || normalized == "yes"
343 })
344 .unwrap_or(false)
345 || std::env::var("LOCUS_MCP_STORAGE")
346 .map(|value| value.eq_ignore_ascii_case("inmemory"))
347 .unwrap_or(false)
348 || args
349 .iter()
350 .any(|arg| arg.eq_ignore_ascii_case("--in-memory"));
351
352 let (store, initializer) = if use_in_memory {
353 let store = Arc::new(InMemoryNodeStore::new());
354 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
355 let node_store: Arc<dyn NodeStore> = store;
356 (node_store, initializer)
357 } else {
358 let settings = load_surreal_settings(&args)?;
359 let runtime_args = runtime_args(&args);
360 let runtime = SurrealDbRuntimeOptions::from_args(&runtime_args, &settings, Some(".locus-mcp"))?;
361
362 let client = Arc::new(
363 RuntimeSurrealDbClient::connect(
364 &runtime,
365 settings.user.as_deref(),
366 settings.password.as_deref(),
367 )
368 .await?,
369 );
370 let store = Arc::new(SurrealDbNodeStore::new(client));
371 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
372 let node_store: Arc<dyn NodeStore> = store;
373
374 tracing::info!(
375 mode = if runtime.use_remote { "remote" } else { "embedded" },
376 endpoint = %runtime.endpoint,
377 namespace = %runtime.namespace,
378 database = %runtime.database,
379 "configured SurrealDB runtime"
380 );
381
382 (node_store, initializer)
383 };
384
385 initializer.initialize_async().await?;
386
387 let validator: Arc<dyn NodeValidator> = Arc::new(TreeSitterValidator::new());
388 let embedding_provider = build_embedding_provider(&args)?;
389 let parse_profile = resolve_parser_profile(&args)?;
390 let parser = SttpNodeParser::with_profile(parse_profile);
391 tracing::info!(parse_profile = ?parse_profile, "configured STTP parser profile for store_context");
392
393 let calibration = Arc::new(CalibrationService::new(store.clone()));
394 let store_context = match embedding_provider.clone() {
395 Some(provider) => Arc::new(StoreContextService::with_embedding_provider(
396 store.clone(),
397 validator.clone(),
398 provider,
399 parser,
400 )),
401 None => Arc::new(StoreContextService::new(store.clone(), validator.clone(), parser)),
402 };
403 let embedding_migration = Arc::new(EmbeddingMigrationService::new(
404 store.clone(),
405 embedding_provider.clone(),
406 ));
407 let moods = Arc::new(MoodCatalogService::new());
408 let monthly_rollup = Arc::new(MonthlyRollupService::new(store.clone(), validator));
409
410 let server = SttpMcpServer::new(
411 store,
412 calibration,
413 store_context,
414 embedding_migration,
415 embedding_provider,
416 moods,
417 monthly_rollup,
418 );
419
420 let running = server
421 .serve((tokio::io::stdin(), tokio::io::stdout()))
422 .await?;
423 running.waiting().await?;
424
425 Ok(())
426}