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