1use std::sync::Arc;
7
8use anyhow::Result;
9use locus_core_rs::domain::contracts::EmbeddingProvider;
10use locus_core_rs::{
11 CalibrationService, EmbeddingMigrationService, InMemoryNodeStore, InMemorySemanticIndexStore,
12 MonthlyRollupService, MoodCatalogService, NodeStore, NodeStoreInitializer, NodeValidator,
13 SemanticIndexStore, SemanticIndexStoreInitializer, StoreContextService, SttpNodeParser,
14 SurrealDbNodeStore, SurrealDbRuntimeOptions, SurrealDbSemanticIndexStore, TreeSitterValidator,
15};
16use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters};
17use rmcp::{ServerHandler, ServiceExt, tool, tool_handler, tool_router};
18use schemars::JsonSchema;
19use serde::Deserialize;
20
21mod composition;
22mod shared;
23mod tools;
24
25use composition::{
26 RuntimeSurrealDbClient, build_embedding_provider, init_logging, load_surreal_settings,
27 resolve_parser_profile, runtime_args,
28};
29
30pub(crate) use shared::{
31 avec_to_json, expanded_limit, filter_nodes_by_context_keywords, infer_store_error_code,
32 mode_to_string, normalize_context_keywords, normalize_tiers, parse_migration_mode,
33 parse_utc_optional, parse_utc_required, schema_first_guidance_payload,
34 strict_typed_ir_profile_name, sttp_node_to_json, to_json_string, tool_error,
35 validate_batch_size, validate_limit, validate_max_nodes,
36};
37
38#[derive(Clone)]
39pub(crate) struct SttpMcpServer {
40 pub(crate) node_store: Arc<dyn NodeStore>,
41 pub(crate) semantic_index: Arc<dyn SemanticIndexStore>,
42 pub(crate) calibration: Arc<CalibrationService>,
43 pub(crate) store_context: Arc<StoreContextService>,
44 pub(crate) embedding_migration: Arc<EmbeddingMigrationService>,
45 pub(crate) embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
46 pub(crate) moods: Arc<MoodCatalogService>,
47 pub(crate) monthly_rollup: Arc<MonthlyRollupService>,
48 #[allow(dead_code)]
49 tool_router: ToolRouter<Self>,
50}
51
52impl SttpMcpServer {
53 fn new(
54 node_store: Arc<dyn NodeStore>,
55 semantic_index: Arc<dyn SemanticIndexStore>,
56 calibration: Arc<CalibrationService>,
57 store_context: Arc<StoreContextService>,
58 embedding_migration: Arc<EmbeddingMigrationService>,
59 embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
60 moods: Arc<MoodCatalogService>,
61 monthly_rollup: Arc<MonthlyRollupService>,
62 ) -> Self {
63 Self {
64 node_store,
65 semantic_index,
66 calibration,
67 store_context,
68 embedding_migration,
69 embedding_provider,
70 moods,
71 monthly_rollup,
72 tool_router: Self::tool_router(),
73 }
74 }
75
76 pub(crate) async fn embed_context_keywords(&self, keywords: &[String]) -> Option<Vec<f32>> {
77 let provider = self.embedding_provider.as_ref()?;
78 let prompt = keywords.join(" ");
79 let prompt = prompt.trim();
80
81 if prompt.is_empty() {
82 return None;
83 }
84
85 provider
86 .embed_async(prompt)
87 .await
88 .ok()
89 .filter(|vector| !vector.is_empty())
90 }
91}
92
93#[tool_router]
94impl SttpMcpServer {
95 #[tool(
96 name = "get_schema",
97 description = "Get a canonical example of what an STTP node should look like before storage."
98 )]
99 async fn get_schema(&self) -> String {
100 tools::get_schema::execute().await
101 }
102
103 #[tool(
104 name = "calibrate_session",
105 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."
106 )]
107 async fn calibrate_session(
108 &self,
109 Parameters(request): Parameters<CalibrateSessionRequest>,
110 ) -> String {
111 tools::calibrate_session::execute(self, request).await
112 }
113
114 #[tool(
115 name = "store_context",
116 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."
117 )]
118 async fn store_context(&self, Parameters(request): Parameters<StoreContextRequest>) -> String {
119 tools::store_context::execute(self, request).await
120 }
121
122 #[tool(
123 name = "get_context",
124 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."
125 )]
126 async fn get_context(&self, Parameters(request): Parameters<GetContextRequest>) -> String {
127 tools::get_context::execute(self, request).await
128 }
129
130 #[tool(
131 name = "list_nodes",
132 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."
133 )]
134 async fn list_nodes(&self, Parameters(request): Parameters<ListNodesRequest>) -> String {
135 tools::list_nodes::execute(self, request).await
136 }
137
138 #[tool(
139 name = "get_graph",
140 description = "Return a memory graph with session topology, lineage, and semantic link edges materialized from stored nodes. Supports optional session scope plus semantic tag and link filters."
141 )]
142 async fn get_graph(&self, Parameters(request): Parameters<GetGraphRequest>) -> String {
143 tools::get_graph::execute(self, request).await
144 }
145
146 #[tool(
147 name = "evict_nodes",
148 description = "Explicitly delete memory nodes by sync key, node id, semantic filter, or full session purge. Supports dry_run preview (recommended first), force to bypass inbound reference blocking, and optional calibration/checkpoint cleanup on session purge."
149 )]
150 async fn evict_nodes(&self, Parameters(request): Parameters<EvictNodesRequest>) -> String {
151 tools::evict_nodes::execute(self, request).await
152 }
153
154 #[tool(
155 name = "preview_embedding_migration",
156 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."
157 )]
158 async fn preview_embedding_migration(
159 &self,
160 Parameters(request): Parameters<PreviewEmbeddingMigrationRequest>,
161 ) -> String {
162 tools::preview_embedding_migration::execute(self, request).await
163 }
164
165 #[tool(
166 name = "run_embedding_migration",
167 description = "Run embedding migration/backfill for selected nodes. Supports dry_run, missing_only mode, and reindex_all mode using the currently configured embedding provider."
168 )]
169 async fn run_embedding_migration(
170 &self,
171 Parameters(request): Parameters<RunEmbeddingMigrationRequest>,
172 ) -> String {
173 tools::run_embedding_migration::execute(self, request).await
174 }
175
176 #[tool(
177 name = "get_moods",
178 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."
179 )]
180 async fn get_moods(&self, Parameters(request): Parameters<GetMoodsRequest>) -> String {
181 tools::get_moods::execute(self, request).await
182 }
183
184 #[tool(
185 name = "create_monthly_rollup",
186 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."
187 )]
188 async fn create_monthly_rollup(
189 &self,
190 Parameters(request): Parameters<CreateMonthlyRollupRequest>,
191 ) -> String {
192 tools::create_monthly_rollup::execute(self, request).await
193 }
194}
195
196#[tool_handler]
197impl ServerHandler for SttpMcpServer {}
198
199#[derive(Debug, Deserialize, JsonSchema)]
200pub(crate) struct CalibrateSessionRequest {
201 session_id: String,
202 stability: f32,
203 friction: f32,
204 logic: f32,
205 autonomy: f32,
206 trigger: String,
207}
208
209#[derive(Debug, Deserialize, JsonSchema)]
210pub(crate) struct StoreContextRequest {
211 node: String,
212 session_id: String,
213}
214
215fn default_limit_get_context() -> usize {
216 5
217}
218
219fn default_blend() -> f32 {
220 1.0
221}
222
223#[derive(Debug, Deserialize, JsonSchema)]
224pub(crate) struct GetContextRequest {
225 #[serde(default)]
226 session_id: Option<String>,
227 stability: f32,
228 friction: f32,
229 logic: f32,
230 autonomy: f32,
231 #[serde(default = "default_limit_get_context")]
232 limit: usize,
233 #[serde(default)]
234 from_utc: Option<String>,
235 #[serde(default)]
236 to_utc: Option<String>,
237 #[serde(default)]
238 tiers: Option<Vec<String>>,
239 #[serde(default)]
240 context_keywords: Option<Vec<String>>,
241 #[serde(default)]
242 alpha: Option<f32>,
243 #[serde(default)]
244 beta: Option<f32>,
245 #[serde(default)]
246 gamma: Option<f32>,
247 #[serde(default)]
248 semantic_tags: Option<Vec<String>>,
249 #[serde(default)]
250 link_rel: Option<String>,
251 #[serde(default)]
252 link_target: Option<String>,
253}
254
255#[derive(Debug, Deserialize, JsonSchema)]
256pub(crate) struct GetGraphRequest {
257 #[serde(default)]
258 session_id: Option<String>,
259 #[serde(default = "default_limit_get_graph")]
260 limit: usize,
261 #[serde(default)]
262 rel: Option<String>,
263 #[serde(default)]
264 target_prefix: Option<String>,
265 #[serde(default)]
266 semantic_tags: Option<Vec<String>>,
267 #[serde(default)]
268 link_rel: Option<String>,
269 #[serde(default)]
270 link_target: Option<String>,
271 #[serde(default)]
272 links_to_ref: Option<String>,
273 #[serde(default)]
274 tag_prefix: Option<String>,
275 #[serde(default)]
276 has_semantic_links: Option<bool>,
277 #[serde(default)]
278 include_lineage: Option<bool>,
279 #[serde(default)]
280 include_semantic: Option<bool>,
281 #[serde(default)]
282 include_session_topology: Option<bool>,
283}
284
285fn default_limit_get_graph() -> usize {
286 1000
287}
288
289#[derive(Debug, Deserialize, JsonSchema)]
290pub(crate) struct EvictNodesRequest {
291 session_id: String,
292 #[serde(default)]
293 sync_keys: Option<Vec<String>>,
294 #[serde(default)]
295 node_ids: Option<Vec<String>>,
296 #[serde(default)]
297 semantic_tags: Option<Vec<String>>,
298 #[serde(default)]
299 link_rel: Option<String>,
300 #[serde(default)]
301 link_target: Option<String>,
302 #[serde(default)]
303 links_to_ref: Option<String>,
304 #[serde(default)]
305 tag_prefix: Option<String>,
306 #[serde(default)]
307 has_semantic_links: Option<bool>,
308 #[serde(default)]
309 purge_session: Option<bool>,
310 #[serde(default)]
311 dry_run: Option<bool>,
312 #[serde(default)]
313 force: Option<bool>,
314 #[serde(default)]
315 max_nodes: Option<usize>,
316 #[serde(default)]
317 include_calibration: Option<bool>,
318 #[serde(default)]
319 include_checkpoints: Option<bool>,
320}
321
322#[derive(Debug, Deserialize, JsonSchema)]
323pub(crate) struct ListNodesRequest {
324 #[serde(default = "default_limit_list_nodes")]
325 limit: usize,
326 #[serde(default)]
327 session_id: Option<String>,
328 #[serde(default)]
329 context_keywords: Option<Vec<String>>,
330 #[serde(default)]
331 semantic_tags: Option<Vec<String>>,
332 #[serde(default)]
333 link_rel: Option<String>,
334 #[serde(default)]
335 link_target: Option<String>,
336}
337
338fn default_limit_list_nodes() -> usize {
339 50
340}
341
342fn default_sample_limit_preview_migration() -> usize {
343 20
344}
345
346fn default_batch_size_migration() -> usize {
347 100
348}
349
350fn default_max_nodes_migration() -> usize {
351 5000
352}
353
354#[derive(Debug, Deserialize, JsonSchema)]
355pub(crate) struct PreviewEmbeddingMigrationRequest {
356 #[serde(default)]
357 session_id: Option<String>,
358 #[serde(default)]
359 from_utc: Option<String>,
360 #[serde(default)]
361 to_utc: Option<String>,
362 #[serde(default)]
363 tiers: Option<Vec<String>>,
364 #[serde(default)]
365 has_embedding: Option<bool>,
366 #[serde(default)]
367 embedding_model: Option<String>,
368 #[serde(default)]
369 sync_keys: Option<Vec<String>>,
370 #[serde(default = "default_sample_limit_preview_migration")]
371 sample_limit: usize,
372 #[serde(default = "default_max_nodes_migration")]
373 max_nodes: usize,
374}
375
376#[derive(Debug, Deserialize, JsonSchema)]
377pub(crate) struct RunEmbeddingMigrationRequest {
378 #[serde(default)]
379 session_id: Option<String>,
380 #[serde(default)]
381 from_utc: Option<String>,
382 #[serde(default)]
383 to_utc: Option<String>,
384 #[serde(default)]
385 tiers: Option<Vec<String>>,
386 #[serde(default)]
387 has_embedding: Option<bool>,
388 #[serde(default)]
389 embedding_model: Option<String>,
390 #[serde(default)]
391 sync_keys: Option<Vec<String>>,
392 #[serde(default)]
393 mode: Option<String>,
394 #[serde(default = "default_true")]
395 dry_run: bool,
396 #[serde(default = "default_batch_size_migration")]
397 batch_size: usize,
398 #[serde(default = "default_max_nodes_migration")]
399 max_nodes: usize,
400}
401
402#[derive(Debug, Deserialize, JsonSchema)]
403pub(crate) struct GetMoodsRequest {
404 #[serde(default)]
405 target_mood: Option<String>,
406 #[serde(default = "default_blend")]
407 blend: f32,
408 #[serde(default)]
409 current_stability: Option<f32>,
410 #[serde(default)]
411 current_friction: Option<f32>,
412 #[serde(default)]
413 current_logic: Option<f32>,
414 #[serde(default)]
415 current_autonomy: Option<f32>,
416}
417
418#[derive(Debug, Deserialize, JsonSchema)]
419pub(crate) struct CreateMonthlyRollupRequest {
420 session_id: String,
421 start_date_utc: String,
422 end_date_utc: String,
423 #[serde(default)]
424 source_session_id: Option<String>,
425 #[serde(default)]
426 parent_node_id: Option<String>,
427 #[serde(default = "default_true")]
428 persist: bool,
429}
430
431fn default_true() -> bool {
432 true
433}
434
435#[tokio::main]
436async fn main() -> Result<()> {
437 init_logging();
438
439 let args = std::env::args().collect::<Vec<_>>();
440 let use_in_memory = std::env::var("LOCUS_MCP_IN_MEMORY")
441 .map(|value| {
442 let normalized = value.trim().to_ascii_lowercase();
443 normalized == "1" || normalized == "true" || normalized == "yes"
444 })
445 .unwrap_or(false)
446 || std::env::var("LOCUS_MCP_STORAGE")
447 .map(|value| value.eq_ignore_ascii_case("inmemory"))
448 .unwrap_or(false)
449 || args
450 .iter()
451 .any(|arg| arg.eq_ignore_ascii_case("--in-memory"));
452
453 let (store, semantic_index, initializer, semantic_initializer) = if use_in_memory {
454 let store = Arc::new(InMemoryNodeStore::new());
455 let semantic_index = Arc::new(InMemorySemanticIndexStore::new());
456 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
457 let semantic_initializer: Arc<dyn SemanticIndexStoreInitializer> = semantic_index.clone();
458 let node_store: Arc<dyn NodeStore> = store;
459 let semantic_trait: Arc<dyn SemanticIndexStore> = semantic_index;
460 (
461 node_store,
462 semantic_trait,
463 initializer,
464 semantic_initializer,
465 )
466 } else {
467 let settings = load_surreal_settings(&args)?;
468 let runtime_args = runtime_args(&args);
469 let runtime = SurrealDbRuntimeOptions::from_args(&runtime_args, &settings, Some(".locus-mcp"))?;
470
471 let client = Arc::new(
472 RuntimeSurrealDbClient::connect(
473 &runtime,
474 settings.user.as_deref(),
475 settings.password.as_deref(),
476 )
477 .await?,
478 );
479 let semantic_index = Arc::new(SurrealDbSemanticIndexStore::new(client.clone()));
480 let store = Arc::new(SurrealDbNodeStore::new(client));
481 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
482 let semantic_initializer: Arc<dyn SemanticIndexStoreInitializer> = semantic_index.clone();
483 let node_store: Arc<dyn NodeStore> = store;
484 let semantic_trait: Arc<dyn SemanticIndexStore> = semantic_index;
485
486 tracing::info!(
487 mode = if runtime.use_remote { "remote" } else { "embedded" },
488 endpoint = %runtime.endpoint,
489 namespace = %runtime.namespace,
490 database = %runtime.database,
491 "configured SurrealDB runtime"
492 );
493
494 (
495 node_store,
496 semantic_trait,
497 initializer,
498 semantic_initializer,
499 )
500 };
501
502 initializer.initialize_async().await?;
503 semantic_initializer.initialize_async().await?;
504
505 let validator: Arc<dyn NodeValidator> = Arc::new(TreeSitterValidator::new());
506 let embedding_provider = build_embedding_provider(&args)?;
507 let parse_profile = resolve_parser_profile(&args)?;
508 let parser = SttpNodeParser::with_profile(parse_profile);
509 tracing::info!(parse_profile = ?parse_profile, "configured STTP parser profile for store_context");
510
511 let calibration = Arc::new(CalibrationService::new(store.clone()));
512 let store_context = match embedding_provider.clone() {
513 Some(provider) => Arc::new(
514 StoreContextService::with_embedding_provider(
515 store.clone(),
516 validator.clone(),
517 provider,
518 parser,
519 )
520 .with_semantic_index(semantic_index.clone()),
521 ),
522 None => Arc::new(
523 StoreContextService::new(store.clone(), validator.clone(), parser)
524 .with_semantic_index(semantic_index.clone()),
525 ),
526 };
527 let embedding_migration = Arc::new(EmbeddingMigrationService::new(
528 store.clone(),
529 embedding_provider.clone(),
530 ));
531 let mut monthly_rollup =
532 MonthlyRollupService::new(store.clone(), validator.clone())
533 .with_semantic_index(semantic_index.clone());
534 if let Some(provider) = embedding_provider.clone() {
535 monthly_rollup = monthly_rollup.with_embedding_provider(provider);
536 }
537 let moods = Arc::new(MoodCatalogService::new());
538
539 let server = SttpMcpServer::new(
540 store,
541 semantic_index,
542 calibration,
543 store_context,
544 embedding_migration,
545 embedding_provider,
546 moods,
547 Arc::new(monthly_rollup),
548 );
549
550 let running = server
551 .serve((tokio::io::stdin(), tokio::io::stdout()))
552 .await?;
553 running.waiting().await?;
554
555 Ok(())
556}