Skip to main content

origin_mcp/
tools.rs

1use crate::client::{OriginClient, OriginError};
2use crate::types::*;
3use rmcp::{
4    handler::server::router::tool::ToolRouter,
5    handler::server::wrapper::Parameters,
6    model::{CallToolResult, Content, Implementation, InitializeResult, ServerCapabilities},
7    service::{NotificationContext, RoleServer},
8    tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler,
9};
10use serde::{Deserialize, Deserializer};
11
12/// Deserialize an `Option<usize>` that also accepts stringified numbers (e.g. `"10"`).
13/// MCP clients like Claude Desktop sometimes send numeric params as strings.
14fn deserialize_optional_usize_lenient<'de, D>(deserializer: D) -> Result<Option<usize>, D::Error>
15where
16    D: Deserializer<'de>,
17{
18    #[derive(Deserialize)]
19    #[serde(untagged)]
20    enum StringOrNumber {
21        Number(usize),
22        Str(String),
23    }
24
25    match Option::<StringOrNumber>::deserialize(deserializer)? {
26        None => Ok(None),
27        Some(StringOrNumber::Number(n)) => Ok(Some(n)),
28        Some(StringOrNumber::Str(s)) => s
29            .parse::<usize>()
30            .map(Some)
31            .map_err(serde::de::Error::custom),
32    }
33}
34
35/// Controls which operations are allowed based on transport.
36#[derive(Clone, Debug, PartialEq)]
37pub enum TransportMode {
38    /// Local stdio — full access, all tools
39    Stdio,
40    /// Remote HTTP — block deletes, inject source_agent
41    Http,
42}
43
44#[derive(Clone)]
45pub struct OriginMcpServer {
46    #[allow(dead_code)]
47    tool_router: ToolRouter<Self>,
48    client: OriginClient,
49    transport: TransportMode,
50    agent_name: String,
51    /// Client name from MCP initialize handshake (e.g., "Claude Code", "Claude Desktop")
52    client_name: std::sync::Arc<std::sync::Mutex<Option<String>>>,
53    user_id: Option<String>,
54}
55
56// ===== Parameter Structs =====
57
58// --- Primary tool params ---
59
60#[derive(Debug, Deserialize, schemars::JsonSchema)]
61pub struct CaptureParams {
62    #[schemars(
63        description = "The memory content. Write as a complete statement with context and reasoning, not shorthand. One idea per memory."
64    )]
65    pub content: String,
66    #[schemars(
67        description = "\"profile\" (about the user) or \"knowledge\" (about the world) — or precise: \"identity\", \"preference\", \"goal\", \"fact\", \"decision\" — auto-classified if omitted"
68    )]
69    pub memory_type: Option<String>,
70    #[schemars(
71        description = "Topic scope (e.g. 'rust', 'work', 'health', 'origin'). Auto-detected if omitted."
72    )]
73    pub domain: Option<String>,
74    #[schemars(
75        description = "Person, project, or tool name to anchor to (e.g. 'Alice', 'Origin', 'PostgreSQL'). Helps build the knowledge graph."
76    )]
77    pub entity: Option<String>,
78    #[schemars(
79        description = "0.0-1.0. Leave unset for auto-calculation based on type and trust level. Set low (0.3-0.5) for uncertain info, high (0.8-1.0) for user-stated facts."
80    )]
81    pub confidence: Option<f32>,
82    #[schemars(
83        description = "source_id of a memory this replaces. Use when correcting or updating an existing memory — get the ID from recall first."
84    )]
85    pub supersedes: Option<String>,
86    #[schemars(
87        description = "Pre-extracted structured fields as a JSON object. Auto-extracted by backend; only supply if you have high-quality structured data already."
88    )]
89    pub structured_fields: Option<serde_json::Map<String, serde_json::Value>>,
90    #[schemars(
91        description = "A question this memory answers, for search matching. Auto-generated by backend; only supply to override."
92    )]
93    pub retrieval_cue: Option<String>,
94}
95
96#[derive(Debug, Deserialize, schemars::JsonSchema)]
97pub struct RecallParams {
98    #[schemars(
99        description = "Natural language search. Be specific: 'Alice database preference' finds more than 'database stuff'."
100    )]
101    pub query: String,
102    #[schemars(
103        description = "Max results, default 10. Use 3-5 for quick lookups, 10-20 for exploration."
104    )]
105    #[serde(default, deserialize_with = "deserialize_optional_usize_lenient")]
106    pub limit: Option<usize>,
107    #[schemars(
108        description = "Filter by type. Two-level filter: \"profile\" (user-facing) or \"knowledge\" (world-facing), or precise: identity, preference, goal, fact, decision."
109    )]
110    pub memory_type: Option<String>,
111    #[schemars(description = "Filter by topic scope.")]
112    pub domain: Option<String>,
113}
114
115#[derive(Debug, Deserialize, schemars::JsonSchema)]
116pub struct ContextParams {
117    #[schemars(
118        description = "Topic or conversation summary to focus context retrieval. Omit at session start for general orientation; provide when shifting topics."
119    )]
120    pub topic: Option<String>,
121    #[schemars(
122        description = "Max context chunks, default 20. Increase for complex topics, decrease for quick check-ins."
123    )]
124    #[serde(default, deserialize_with = "deserialize_optional_usize_lenient")]
125    pub limit: Option<usize>,
126    #[schemars(
127        description = "Scope context to a domain/space (e.g. 'work', 'personal'). Auto-detected from conversation if omitted."
128    )]
129    pub domain: Option<String>,
130}
131
132#[derive(Debug, Deserialize, schemars::JsonSchema)]
133pub struct ForgetParams {
134    #[schemars(
135        description = "The source_id of the memory to delete. Get this from recall results first."
136    )]
137    pub memory_id: String,
138}
139
140#[derive(Debug, Deserialize, schemars::JsonSchema)]
141pub struct DistillParams {
142    #[schemars(
143        description = "Optional page ID. If provided, re-distills only that page from its current sources. If omitted, runs a full distillation pass over any clusters with new sources."
144    )]
145    pub page_id: Option<String>,
146}
147
148#[derive(Debug, Deserialize, schemars::JsonSchema)]
149pub struct ListPendingParams {
150    #[schemars(
151        description = "Max results, default 20. Increase for full audit, decrease for quick check-in."
152    )]
153    #[serde(default, deserialize_with = "deserialize_optional_usize_lenient")]
154    pub limit: Option<usize>,
155}
156
157#[derive(Debug, Deserialize, schemars::JsonSchema)]
158pub struct ConfirmMemoryParams {
159    #[schemars(
160        description = "The source_id of the memory to confirm. Get this from list_pending or recall results."
161    )]
162    pub memory_id: String,
163}
164
165// ===== Internal Implementations =====
166
167fn format_capture_success(resp: &StoreMemoryResponse) -> String {
168    let mut msg = format!("Stored {}", resp.source_id);
169    if !resp.warnings.is_empty() {
170        msg.push_str("\nWarnings:");
171        for warning in &resp.warnings {
172            msg.push_str(&format!("\n  - {}", warning));
173        }
174    }
175    msg
176}
177
178fn daemon_setup_hint() -> &'static str {
179    "Install the Origin desktop app, or install the headless runtime and run `origin setup`.
180
181Setup choices:
182- Basic Memory: store, search, and recall now. No model download or API key.
183- On-device Model: private local extraction and background refinement after model download.
184- Anthropic Key: richer extraction and background refinement using your API key.
185
186Desktop: https://github.com/7xuanlu/origin/releases/latest
187Headless:
188  curl -fsSL https://raw.githubusercontent.com/7xuanlu/origin/main/install.sh | bash
189  export PATH=\"$HOME/.origin/bin:$PATH\"
190  origin setup
191  origin install
192  origin status"
193}
194
195/// Convert a backend error into a tool-level error result (isError: true)
196/// with an actionable message. This keeps the MCP transport healthy
197/// (no protocol-level McpError) while telling the caller what happened.
198fn tool_error(e: OriginError, verb: &str) -> CallToolResult {
199    let msg = match &e {
200        OriginError::Unreachable(_) => format!(
201            "Origin daemon is not reachable (retried 3x over ~6s). \
202             The {verb} was NOT completed.\n\n{}",
203            daemon_setup_hint()
204        ),
205        OriginError::Api { status, body } => format!(
206            "Origin daemon returned HTTP {status}: {body}. The {verb} may not have completed."
207        ),
208        OriginError::Deserialize(detail) => format!(
209            "Failed to parse daemon response: {detail}. \
210             This may indicate a version mismatch between origin-mcp and the daemon."
211        ),
212    };
213    CallToolResult::error(vec![Content::text(msg)])
214}
215
216fn format_doctor_message(status: &serde_json::Value) -> String {
217    let mode = status
218        .get("mode")
219        .and_then(|v| v.as_str())
220        .unwrap_or("unknown");
221    let setup_completed = status
222        .get("setup_completed")
223        .and_then(|v| v.as_bool())
224        .unwrap_or(false);
225    let anthropic_key_configured = status
226        .get("anthropic_key_configured")
227        .and_then(|v| v.as_bool())
228        .unwrap_or(false);
229    let local_model_selected = status.get("local_model_selected").and_then(|v| v.as_str());
230    let local_model_loaded = status.get("local_model_loaded").and_then(|v| v.as_str());
231    let local_model_cached = status
232        .get("local_model_cached")
233        .and_then(|v| v.as_bool())
234        .unwrap_or(false);
235
236    let mode_label = match mode {
237        "basic-memory" => "Basic Memory",
238        "local-model" => "On-device Model",
239        "anthropic-key" => "Anthropic Key",
240        other => other,
241    };
242    let local_model_line = match local_model_selected {
243        Some(id) => {
244            let cache_status = if local_model_cached {
245                "downloaded"
246            } else {
247                "not downloaded"
248            };
249            let loaded_status = if Some(id) == local_model_loaded {
250                ", loaded"
251            } else {
252                ""
253            };
254            format!("{id} ({cache_status}{loaded_status})")
255        }
256        None => "not selected".to_string(),
257    };
258    let refinement_line = if anthropic_key_configured || local_model_loaded.is_some() {
259        "enabled (richer extraction and background refinement are active)"
260    } else if setup_completed {
261        "paused (Basic Memory stores, searches, and recalls now. Choose an on-device model or Anthropic key for richer extraction.)"
262    } else {
263        "not configured"
264    };
265
266    let mut msg = format!(
267        "Origin daemon: running\n\
268         Setup: {}\n\
269         Mode: {mode_label}\n\
270         Anthropic key: {}\n\
271         On-device model: {local_model_line}\n\
272         Background refinement: {refinement_line}",
273        if setup_completed {
274            "completed"
275        } else {
276            "not completed"
277        },
278        if anthropic_key_configured {
279            "configured"
280        } else {
281            "not configured"
282        }
283    );
284
285    if !setup_completed {
286        msg.push_str(
287            "\n\nRun `origin setup` to choose Basic Memory, On-device Model, or Anthropic Key.",
288        );
289    } else if !anthropic_key_configured && local_model_loaded.is_none() {
290        msg.push_str(
291            "\n\nBasic Memory works now: capture, recall, and context are available. \
292             To enable richer extraction and background refinement, run `origin model install` \
293             or `origin key set anthropic`.",
294        );
295    }
296
297    msg
298}
299
300impl OriginMcpServer {
301    /// Resolve the source_agent for a write operation.
302    /// Priority: explicit param > MCP client name (from initialize) > configured agent_name.
303    fn resolve_source_agent(&self, param_agent: Option<String>) -> Option<String> {
304        // 1. Explicit param from tool call
305        if let Some(ref agent) = param_agent {
306            if !agent.is_empty() {
307                return param_agent;
308            }
309        }
310        // 2. Client name captured from MCP initialize handshake
311        if let Ok(guard) = self.client_name.lock() {
312            if let Some(ref name) = *guard {
313                return Some(name.clone());
314            }
315        }
316        // 3. Configured --agent-name flag
317        Some(self.agent_name.clone())
318    }
319
320    /// Resolve a local user_id for logging or future use.
321    /// This value is intentionally not sent on the wire (D4).
322    fn resolve_user_id(&self, param_user_id: Option<String>) -> Option<String> {
323        if self.transport == TransportMode::Http {
324            self.user_id.clone().or(param_user_id)
325        } else {
326            param_user_id
327        }
328    }
329
330    pub async fn capture_impl(&self, params: CaptureParams) -> Result<CallToolResult, McpError> {
331        // Tool was renamed `remember -> capture` in v0.4. The HTTP request
332        // body shape (StoreMemoryRequest) is unchanged; only the MCP-facing
333        // tool name shifted.
334        let source_agent = self.resolve_source_agent(None);
335        if let Some(uid) = self.resolve_user_id(None) {
336            tracing::debug!(user_id = %uid, "capture invoked");
337        }
338
339        let req = StoreMemoryRequest {
340            content: params.content,
341            memory_type: params.memory_type,
342            domain: params.domain,
343            source_agent,
344            title: None,
345            confidence: params.confidence,
346            supersedes: params.supersedes,
347            entity: params.entity,
348            entity_id: None,
349            structured_fields: params.structured_fields.map(serde_json::Value::Object),
350            retrieval_cue: params.retrieval_cue,
351        };
352
353        let resp: StoreMemoryResponse = match self.client.post("/api/memory/store", &req).await {
354            Ok(r) => r,
355            Err(e) => return Ok(tool_error(e, "memory store")),
356        };
357
358        Ok(CallToolResult::success(vec![Content::text(
359            format_capture_success(&resp),
360        )]))
361    }
362
363    pub async fn recall_impl(&self, params: RecallParams) -> Result<CallToolResult, McpError> {
364        let req = SearchMemoryRequest {
365            query: params.query,
366            limit: params.limit.unwrap_or(10),
367            memory_type: params.memory_type,
368            domain: params.domain,
369            source_agent: self.resolve_source_agent(None),
370        };
371
372        let resp: SearchMemoryResponse = match self.client.post("/api/memory/search", &req).await {
373            Ok(r) => r,
374            Err(e) => return Ok(tool_error(e, "search")),
375        };
376
377        let json = serde_json::to_string_pretty(&resp.results)
378            .map_err(|e| McpError::internal_error(e.to_string(), None))?;
379
380        Ok(CallToolResult::success(vec![Content::text(format!(
381            "{} results ({:.1}ms)\n{}",
382            resp.results.len(),
383            resp.took_ms,
384            json
385        ))]))
386    }
387
388    pub async fn context_impl(&self, params: ContextParams) -> Result<CallToolResult, McpError> {
389        let req = ChatContextRequest {
390            query: None,
391            conversation_id: params.topic,
392            max_chunks: params.limit.unwrap_or(20),
393            relevance_threshold: None,
394            include_goals: true,
395            domain: params.domain,
396        };
397
398        // Extract only the `context` string field from the response.
399        //
400        // The full ChatContextResponse embeds Vec<SearchResult> which may
401        // contain fields added after the published origin-types version.
402        // Since context_impl only uses `resp.context`, we parse the raw
403        // JSON and pull that field directly — this makes the tool forward-
404        // compatible with any new fields the daemon might add.
405        let raw: serde_json::Value = match self.client.post("/api/chat-context", &req).await {
406            Ok(r) => r,
407            Err(e) => return Ok(tool_error(e, "context load")),
408        };
409
410        let context = raw
411            .get("context")
412            .and_then(|v| v.as_str())
413            .unwrap_or_default()
414            .to_string();
415
416        if context.is_empty() {
417            Ok(CallToolResult::success(vec![Content::text(
418                "No relevant context found".to_string(),
419            )]))
420        } else {
421            Ok(CallToolResult::success(vec![Content::text(context)]))
422        }
423    }
424
425    pub async fn doctor_impl(&self) -> Result<CallToolResult, McpError> {
426        let status: serde_json::Value = match self.client.get("/api/setup/status").await {
427            Ok(r) => r,
428            Err(OriginError::Api { status: 404, .. }) => {
429                return Ok(CallToolResult::error(vec![Content::text(
430                    "Origin daemon is running, but it does not expose /api/setup/status. \
431                     Update Origin, then run `origin doctor`."
432                        .to_string(),
433                )]));
434            }
435            Err(e) => return Ok(tool_error(e, "status check")),
436        };
437
438        Ok(CallToolResult::success(vec![Content::text(
439            format_doctor_message(&status),
440        )]))
441    }
442
443    pub async fn forget_impl(&self, memory_id: &str) -> Result<CallToolResult, McpError> {
444        if self.transport == TransportMode::Http {
445            return Ok(CallToolResult::error(vec![Content::text(
446                "Delete operations are not available over remote connections. \
447                 Use the Origin desktop app to delete memories."
448                    .to_string(),
449            )]));
450        }
451
452        let resp: DeleteResponse = match self
453            .client
454            .delete(&format!("/api/memory/delete/{}", memory_id))
455            .await
456        {
457            Ok(r) => r,
458            Err(e) => return Ok(tool_error(e, "delete")),
459        };
460
461        Ok(CallToolResult::success(vec![Content::text(
462            if resp.deleted {
463                "Memory deleted"
464            } else {
465                "Memory not found"
466            }
467            .to_string(),
468        )]))
469    }
470
471    pub async fn distill_impl(&self, params: DistillParams) -> Result<CallToolResult, McpError> {
472        let path = match params.page_id.as_deref() {
473            Some(id) if !id.is_empty() => format!("/api/distill/{}", id),
474            _ => "/api/distill".to_string(),
475        };
476        match self
477            .client
478            .post::<serde_json::Value, serde_json::Value>(&path, &serde_json::json!({}))
479            .await
480        {
481            Ok(_) => Ok(CallToolResult::success(vec![Content::text(
482                match params.page_id {
483                    Some(id) => format!("Re-distilled page {}.", id),
484                    None => "Distillation pass triggered.".to_string(),
485                },
486            )])),
487            Err(e) => Ok(tool_error(e, "distill")),
488        }
489    }
490
491    pub async fn list_pending_impl(
492        &self,
493        params: ListPendingParams,
494    ) -> Result<CallToolResult, McpError> {
495        let limit = params.limit.unwrap_or(20).min(100);
496        let path = format!("/api/memory/list?confirmed=false&limit={}", limit);
497        let value: serde_json::Value = match self.client.get(&path).await {
498            Ok(v) => v,
499            Err(e) => return Ok(tool_error(e, "list_pending")),
500        };
501        let body = serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string());
502        Ok(CallToolResult::success(vec![Content::text(body)]))
503    }
504
505    pub async fn confirm_memory_impl(&self, memory_id: &str) -> Result<CallToolResult, McpError> {
506        if self.transport == TransportMode::Http {
507            return Ok(CallToolResult::error(vec![Content::text(
508                "Confirm operations are not available over remote connections. \
509                 Use the Origin desktop app or local MCP for review."
510                    .to_string(),
511            )]));
512        }
513        let path = format!("/api/memory/confirm/{}", memory_id);
514        match self
515            .client
516            .post::<serde_json::Value, serde_json::Value>(&path, &serde_json::json!({}))
517            .await
518        {
519            Ok(_) => Ok(CallToolResult::success(vec![Content::text(format!(
520                "Memory {} confirmed.",
521                memory_id
522            ))])),
523            Err(e) => Ok(tool_error(e, "confirm_memory")),
524        }
525    }
526}
527
528// ===== Tool Registrations =====
529
530#[tool_router]
531impl OriginMcpServer {
532    pub fn new(
533        client: OriginClient,
534        transport: TransportMode,
535        agent_name: String,
536        user_id: Option<String>,
537    ) -> Self {
538        Self {
539            tool_router: Self::tool_router(),
540            client,
541            transport,
542            agent_name,
543            client_name: std::sync::Arc::new(std::sync::Mutex::new(None)),
544            user_id,
545        }
546    }
547
548    // --- Primary Tools ---
549
550    #[tool(
551        description = "Capture a memory. Call PROACTIVELY when you learn something durable about the user — preferences, decisions, corrections, or facts about people/projects/tools they care about. Don't wait for the user to say 'remember this' or 'capture that' — that phrasing is a floor, not a trigger.\n\nWrite content as a complete, self-contained statement — someone reading it months later with no conversation context should understand it. Include the WHY, not just the WHAT. Name people, projects, and tools explicitly.\n\nThe backend auto-classifies type, extracts structured fields, detects entities, and links to the knowledge graph. You don't need to set memory_type or structured_fields unless you're confident — omitting them gets better results than guessing wrong.\n\nDo NOT store: system prompts, boot logs, heartbeat/health checks, transient task state ('currently working on...'), tool output/responses, architecture dumps, single-word acknowledgments, or content you have already stored. Focus on durable facts, preferences, decisions, lessons, gotchas, and identity information. Each call is one atomic idea — \"prefers TDD\" and \"uses pytest\" are two calls, not one.",
552        annotations(
553            title = "Capture",
554            read_only_hint = false,
555            destructive_hint = false,
556            idempotent_hint = false,
557            open_world_hint = false
558        )
559    )]
560    async fn capture(
561        &self,
562        Parameters(params): Parameters<CaptureParams>,
563    ) -> Result<CallToolResult, McpError> {
564        self.capture_impl(params).await
565    }
566
567    #[tool(
568        description = "Search memories by query. Use when the user asks 'do you remember', 'what do you know about', 'look up', or when you need a specific fact before acting.\n\nWrite queries as natural language — the search engine handles semantic matching. For precision, use filters (memory_type, domain) to narrow results. If you get too many results, add filters rather than making the query longer.\n\nThis is for targeted lookups. For broad session orientation, use context instead.",
569        annotations(title = "Recall", read_only_hint = true, open_world_hint = false)
570    )]
571    async fn recall(
572        &self,
573        Parameters(params): Parameters<RecallParams>,
574    ) -> Result<CallToolResult, McpError> {
575        self.recall_impl(params).await
576    }
577
578    #[tool(
579        description = "Load session context — identity, preferences, goals, and topic-relevant memories. Call this FIRST at the start of every session before doing anything else. Also call on major topic shifts or when the user says 'catch me up' or 'what's the background on'.\n\nThis returns a curated blend of who the user is and what's relevant. For specific factual lookups, use recall instead. Use the result to model how the user thinks, not just to look things up — their preferences and corrections tell you how they want to be helped.",
580        annotations(title = "Context", read_only_hint = true, open_world_hint = false)
581    )]
582    async fn context(
583        &self,
584        Parameters(params): Parameters<ContextParams>,
585    ) -> Result<CallToolResult, McpError> {
586        self.context_impl(params).await
587    }
588
589    #[tool(
590        description = "Diagnose the local Origin runtime. This is not part of the memory loop. Use only when Origin tools fail, when onboarding a new MCP client, or when the user asks why setup, extraction, or background refinement is paused. Reports daemon reachability, setup mode, Basic Memory, On-device Model, Anthropic key state, and on-device model state.",
591        annotations(title = "Doctor", read_only_hint = true, open_world_hint = false)
592    )]
593    async fn doctor(&self) -> Result<CallToolResult, McpError> {
594        self.doctor_impl().await
595    }
596
597    #[tool(
598        description = "Delete a memory by ID. Use when the user says 'forget this', 'delete that', 'that's wrong and should be removed'. Requires the source_id — get it from recall first.\n\nThis is destructive and cannot be undone. For corrections, prefer storing a new memory with the supersedes param pointing to the old one — this preserves history.",
599        annotations(
600            title = "Forget",
601            read_only_hint = false,
602            destructive_hint = true,
603            idempotent_hint = true,
604            open_world_hint = false
605        )
606    )]
607    async fn forget(
608        &self,
609        Parameters(params): Parameters<ForgetParams>,
610    ) -> Result<CallToolResult, McpError> {
611        self.forget_impl(&params.memory_id).await
612    }
613
614    #[tool(
615        description = "Trigger Origin's distillation pass. With no `page_id`, runs a full pass that clusters new memories into pages and refreshes the wiki view. With a `page_id`, re-distills that single page from its current sources. Use when the user explicitly asks to synthesize, distill, or rebuild a page. The daemon also runs distillation periodically in the background, so don't trigger redundantly during normal flow.",
616        annotations(
617            title = "Distill",
618            read_only_hint = false,
619            destructive_hint = false,
620            idempotent_hint = true,
621            open_world_hint = false
622        )
623    )]
624    async fn distill(
625        &self,
626        Parameters(params): Parameters<DistillParams>,
627    ) -> Result<CallToolResult, McpError> {
628        self.distill_impl(params).await
629    }
630
631    #[tool(
632        description = "List unconfirmed memories pending review. Use when the user wants to audit what got captured before it becomes authoritative — typical phrases: 'review pending', 'show unconfirmed', 'what got captured'. Pair with `confirm_memory` to accept and `forget` to reject.",
633        annotations(title = "List pending", read_only_hint = true, open_world_hint = false)
634    )]
635    async fn list_pending(
636        &self,
637        Parameters(params): Parameters<ListPendingParams>,
638    ) -> Result<CallToolResult, McpError> {
639        self.list_pending_impl(params).await
640    }
641
642    #[tool(
643        description = "Confirm a pending memory by source_id. Use during review to accept a memory the agent captured. The user typically picks from a `list_pending` result. To reject instead, call `forget` with the same `memory_id`.",
644        annotations(
645            title = "Confirm memory",
646            read_only_hint = false,
647            destructive_hint = false,
648            idempotent_hint = true,
649            open_world_hint = false
650        )
651    )]
652    async fn confirm_memory(
653        &self,
654        Parameters(params): Parameters<ConfirmMemoryParams>,
655    ) -> Result<CallToolResult, McpError> {
656        self.confirm_memory_impl(&params.memory_id).await
657    }
658}
659
660// ===== ServerHandler =====
661
662#[tool_handler]
663impl ServerHandler for OriginMcpServer {
664    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
665        // Capture client name from MCP initialize handshake
666        if let Some(client_info) = context.peer.peer_info() {
667            let name = &client_info.client_info.name;
668            if !name.is_empty() {
669                if let Ok(mut guard) = self.client_name.lock() {
670                    tracing::info!("MCP client identified: {}", name);
671                    *guard = Some(name.clone());
672                }
673            }
674        }
675    }
676
677    fn get_info(&self) -> InitializeResult {
678        InitializeResult::new(
679            ServerCapabilities::builder()
680                .enable_tools()
681                .build(),
682        )
683        .with_server_info(
684            Implementation::new("origin-mcp", env!("CARGO_PKG_VERSION"))
685        )
686        .with_instructions(
687            "Origin is your personal memory layer — a local knowledge base that persists across sessions and tools.\n\
688             Think of yourself as a curator, not a logger. Store insights, not conversation artifacts.\n\n\
689             Origin is self-evolving — each memory you store contributes to a knowledge structure that grows over time. \
690             It's also shared across all the user's tools: what you write, other agents (Claude Desktop, Claude Code, \
691             ChatGPT, Cursor, etc.) will read later. Write for any future reader, not just this conversation.\n\n\
692             FIRST THING EVERY SESSION: Call context to load the user's identity, preferences, goals, and\n\
693             topic-relevant memories. This is how you know who you're talking to. Use the result to model how the \
694             user thinks — their preferences, corrections, and past decisions tell you how they want to be helped, \
695             not just what they already know.\n\n\
696             STORE PROACTIVELY — don't wait for the user to ask.\n\
697             - The user states a preference (\"I use X because...\", \"I prefer Y over Z\")\n\
698             - The user makes a decision (\"going with approach A\", \"switching to B\")\n\
699             - The user corrects you or prior info (\"actually, it's C, not D\") — store the correction so it sticks\n\
700             - The user shares a durable fact about themselves, their work, or people/projects/tools they care about — \
701               anchor it to the entity\n\n\
702             If the user asks explicitly (\"remember this\", \"save this\", \"don't forget\"), that's a floor — you \
703             should have already stored it.\n\n\
704             WHEN NOT TO STORE:\n\
705             - Conversation filler (\"ok\", \"thanks\", \"let's move on\")\n\
706             - Things the user can trivially re-derive (file paths, recent git history)\n\
707             - Anything already stored — recall first if unsure\n\
708             - Tool output or command results (file contents, git history, build logs) — these are derivable\n\
709             - General world facts or documentation that aren't personal to this user (e.g., \"Rust has a borrow \
710               checker\", \"PostgreSQL supports JSONB\") — those are not memory material.\n\
711             - Your own inferences about the user that they didn't express. Store what they said; infer from that \
712               when responding.\n\n\
713             CONTENT QUALITY — this is where you make the biggest difference:\n\
714             - Specific beats vague: \"prefers Rust for CLI tools because of compile-time safety\" > \"likes Rust\"\n\
715             - Include the WHY: the backend can classify \"dark mode\" as a preference, but only you know\n\
716               \"switched to dark mode because of migraines from bright screens\"\n\
717             - Name the entities: mention people, projects, tools by name — this powers the knowledge graph\n\
718             - Atomic: one idea per memory — \"prefers TDD\" and \"uses pytest\" should be two memories, not one\n\
719             - Declarative, not narrative: \"User prefers X because Y\" — not \"User said today they prefer X\". \
720               Memories outlive the conversation that produced them.\n\n\
721             MEMORY TYPES — omit and trust the backend.\n\n\
722             By default, do NOT set memory_type. The backend auto-classifies into identity / preference / goal / \
723             fact / decision with more context than you have. Agents that over-specify types tend to pick wrong.\n\n\
724             Opt-in specification:\n\
725             - \"profile\"   — you're sure it's about the user (identity/preference/goal)\n\
726             - \"knowledge\" — you're sure it's about the world (fact/decision)\n\
727             - Precise type — only if you're confident and the distinction matters.\n\n\
728             EXCEPTION — decisions carry structured fields (alternatives considered, reversibility, domain) \
729             that power the Decision Log view. Set memory_type=\"decision\" explicitly ONLY when the user \
730             articulated alternatives weighed AND the reasoning for the choice. A bare \"I'm switching to Cursor\" \
731             is just a preference change — omit the type. \"Switching to Cursor over VSCode because of better \
732             Claude integration, and we can always go back\" — that's a decision.\n\n\
733             RECALL vs CONTEXT:\n\
734             - context: broad orientation, session start, topic shifts, \"catch me up\"\n\
735             - recall: specific lookup (\"what's Alice's role?\", \"database preferences\", \"our auth decision\")\n\n\
736             The backend handles classification, entity extraction, structured fields, quality scoring,\n\
737             and dedup — you don't need to replicate that logic. Focus on what only you know:\n\
738             the conversational context, why something matters, and what the user actually cares about."
739        )
740    }
741}
742
743#[cfg(test)]
744mod tests {
745    use super::*;
746    use crate::client::OriginClient;
747    use crate::types::{
748        ChatContextRequest, ChatContextResponse, SearchMemoryRequest, SearchResult,
749        StoreMemoryRequest, StoreMemoryResponse,
750    };
751
752    fn make_server(
753        transport: TransportMode,
754        agent_name: &str,
755        user_id: Option<&str>,
756    ) -> OriginMcpServer {
757        let client = OriginClient::new("http://127.0.0.1:19999".into());
758        OriginMcpServer::new(
759            client,
760            transport,
761            agent_name.into(),
762            user_id.map(String::from),
763        )
764    }
765
766    // ===== Transport resolution (existing) =====
767
768    #[test]
769    fn test_http_mode_prefers_param_over_agent_name() {
770        let server = make_server(TransportMode::Http, "claude.ai", None);
771        // Explicit param has highest priority
772        let result = server.resolve_source_agent(Some("user-provided".into()));
773        assert_eq!(result, Some("user-provided".into()));
774    }
775
776    #[test]
777    fn test_http_mode_sets_source_agent_when_none() {
778        let server = make_server(TransportMode::Http, "chatgpt", None);
779        let result = server.resolve_source_agent(None);
780        assert_eq!(result, Some("chatgpt".into()));
781    }
782
783    #[test]
784    fn test_stdio_mode_passes_through_source_agent() {
785        let server = make_server(TransportMode::Stdio, "ignored", None);
786        let result = server.resolve_source_agent(Some("user-provided".into()));
787        assert_eq!(result, Some("user-provided".into()));
788    }
789
790    #[test]
791    fn test_stdio_mode_falls_back_to_agent_name() {
792        let server = make_server(TransportMode::Stdio, "fallback", None);
793        // No param, no client_name → falls back to configured agent_name
794        let result = server.resolve_source_agent(None);
795        assert_eq!(result, Some("fallback".into()));
796    }
797
798    #[test]
799    fn test_http_mode_resolves_configured_user_id_for_local_use() {
800        let server = make_server(TransportMode::Http, "agent", Some("lucian"));
801        let result = server.resolve_user_id(None);
802        assert_eq!(result, Some("lucian".into()));
803    }
804
805    #[test]
806    fn test_transport_mode_equality() {
807        assert_eq!(TransportMode::Stdio, TransportMode::Stdio);
808        assert_eq!(TransportMode::Http, TransportMode::Http);
809        assert_ne!(TransportMode::Stdio, TransportMode::Http);
810    }
811
812    // ===== Param deserialization: CaptureParams =====
813
814    #[test]
815    fn test_capture_params_minimal() {
816        let json = r#"{"content": "Lucian prefers dark mode"}"#;
817        let params: CaptureParams = serde_json::from_str(json).unwrap();
818        assert_eq!(params.content, "Lucian prefers dark mode");
819        assert!(params.memory_type.is_none());
820        assert!(params.domain.is_none());
821        assert!(params.entity.is_none());
822        assert!(params.confidence.is_none());
823        assert!(params.supersedes.is_none());
824    }
825
826    #[test]
827    fn test_capture_params_full() {
828        let json = r#"{
829            "content": "We chose PostgreSQL over MongoDB",
830            "memory_type": "decision",
831            "domain": "origin",
832            "entity": "PostgreSQL",
833            "confidence": 0.95,
834            "supersedes": "mem_abc123"
835        }"#;
836        let params: CaptureParams = serde_json::from_str(json).unwrap();
837        assert_eq!(params.content, "We chose PostgreSQL over MongoDB");
838        assert_eq!(params.memory_type.as_deref(), Some("decision"));
839        assert_eq!(params.domain.as_deref(), Some("origin"));
840        assert_eq!(params.entity.as_deref(), Some("PostgreSQL"));
841        assert_eq!(params.confidence, Some(0.95));
842        assert_eq!(params.supersedes.as_deref(), Some("mem_abc123"));
843    }
844
845    #[test]
846    fn test_capture_params_missing_content_fails() {
847        let json = r#"{"memory_type": "fact"}"#;
848        let result = serde_json::from_str::<CaptureParams>(json);
849        assert!(result.is_err());
850    }
851
852    // ===== Param deserialization: RecallParams =====
853
854    #[test]
855    fn test_recall_params_minimal() {
856        let json = r#"{"query": "what does Alice work on?"}"#;
857        let params: RecallParams = serde_json::from_str(json).unwrap();
858        assert_eq!(params.query, "what does Alice work on?");
859        assert!(params.limit.is_none());
860    }
861
862    #[test]
863    fn test_recall_params_full() {
864        let json = r#"{
865            "query": "database preferences",
866            "limit": 5,
867            "memory_type": "decision",
868            "domain": "origin"
869        }"#;
870        let params: RecallParams = serde_json::from_str(json).unwrap();
871        assert_eq!(params.query, "database preferences");
872        assert_eq!(params.limit, Some(5));
873        assert_eq!(params.memory_type.as_deref(), Some("decision"));
874        assert_eq!(params.domain.as_deref(), Some("origin"));
875    }
876
877    #[test]
878    fn test_recall_params_limit_as_string() {
879        let json = r#"{"query": "test", "limit": "10"}"#;
880        let params: RecallParams = serde_json::from_str(json).unwrap();
881        assert_eq!(params.limit, Some(10));
882    }
883
884    #[test]
885    fn test_recall_params_missing_query_fails() {
886        let json = r#"{"limit": 5}"#;
887        let result = serde_json::from_str::<RecallParams>(json);
888        assert!(result.is_err());
889    }
890
891    // ===== Param deserialization: ContextParams =====
892
893    #[test]
894    fn test_context_params_empty() {
895        let json = r#"{}"#;
896        let params: ContextParams = serde_json::from_str(json).unwrap();
897        assert!(params.topic.is_none());
898        assert!(params.limit.is_none());
899        assert!(params.domain.is_none());
900    }
901
902    #[test]
903    fn test_context_params_full() {
904        let json = r#"{"topic": "project Origin architecture", "limit": 30, "domain": "work"}"#;
905        let params: ContextParams = serde_json::from_str(json).unwrap();
906        assert_eq!(params.topic.as_deref(), Some("project Origin architecture"));
907        assert_eq!(params.limit, Some(30));
908        assert_eq!(params.domain.as_deref(), Some("work"));
909    }
910
911    #[test]
912    fn test_context_params_limit_as_string() {
913        let json = r#"{"limit": "20"}"#;
914        let params: ContextParams = serde_json::from_str(json).unwrap();
915        assert_eq!(params.limit, Some(20));
916    }
917
918    #[test]
919    fn store_memory_request_serialization_excludes_user_id() {
920        let req = StoreMemoryRequest {
921            content: "test content".into(),
922            memory_type: None,
923            domain: None,
924            source_agent: Some("test-agent".into()),
925            title: None,
926            confidence: None,
927            supersedes: None,
928            entity: None,
929            entity_id: None,
930            structured_fields: None,
931            retrieval_cue: None,
932        };
933        let json = serde_json::to_value(&req).unwrap();
934        let obj = json.as_object().unwrap();
935        assert!(
936            !obj.contains_key("user_id"),
937            "user_id must not be on the wire; got: {:?}",
938            obj.keys().collect::<Vec<_>>()
939        );
940    }
941
942    #[test]
943    fn capture_success_message_is_terse() {
944        let resp = StoreMemoryResponse {
945            source_id: "mem_abc".into(),
946            chunks_created: 3,
947            memory_type: "fact".into(),
948            entity_id: Some("ent_xyz".into()),
949            quality: Some("high".into()),
950            warnings: vec![],
951            extraction_method: "llm".into(),
952            enrichment: String::new(),
953            hint: String::new(),
954        };
955        let msg = format_capture_success(&resp);
956        assert_eq!(msg, "Stored mem_abc");
957        assert!(!msg.contains("chunks"));
958        assert!(!msg.contains("quality"));
959        assert!(!msg.contains("entity"));
960    }
961
962    #[test]
963    fn capture_success_message_surfaces_warnings() {
964        let resp = StoreMemoryResponse {
965            source_id: "mem_abc".into(),
966            chunks_created: 1,
967            memory_type: "decision".into(),
968            entity_id: None,
969            quality: None,
970            warnings: vec!["decision memory missing required 'claim' field".into()],
971            extraction_method: "agent".into(),
972            enrichment: String::new(),
973            hint: String::new(),
974        };
975        let msg = format_capture_success(&resp);
976        assert!(msg.starts_with("Stored mem_abc"));
977        assert!(msg.contains("Warnings:"));
978        assert!(msg.contains("decision memory missing required 'claim' field"));
979    }
980
981    #[test]
982    fn doctor_basic_memory_message_sets_expectations() {
983        let msg = format_doctor_message(&serde_json::json!({
984            "setup_completed": true,
985            "mode": "basic-memory",
986            "anthropic_key_configured": false,
987            "local_model_selected": null,
988            "local_model_loaded": null,
989            "local_model_cached": false
990        }));
991
992        assert!(msg.contains("Mode: Basic Memory"));
993        assert!(msg.contains("On-device model: not selected"));
994        assert!(msg.contains("Background refinement: paused"));
995        assert!(msg.contains("Basic Memory works now: capture, recall, and context are available"));
996        assert!(msg.contains("origin model install"));
997        assert!(msg.contains("origin key set anthropic"));
998    }
999
1000    #[test]
1001    fn doctor_on_device_model_message_shows_loaded_model() {
1002        let msg = format_doctor_message(&serde_json::json!({
1003            "setup_completed": true,
1004            "mode": "local-model",
1005            "anthropic_key_configured": false,
1006            "local_model_selected": "qwen3-1.7b",
1007            "local_model_loaded": "qwen3-1.7b",
1008            "local_model_cached": true
1009        }));
1010
1011        assert!(msg.contains("Mode: On-device Model"), "{msg}");
1012        assert!(
1013            msg.contains("On-device model: qwen3-1.7b (downloaded, loaded)"),
1014            "{msg}"
1015        );
1016        assert!(msg.contains("Background refinement: enabled"), "{msg}");
1017        assert!(!msg.contains("Basic Memory works now"));
1018    }
1019
1020    #[test]
1021    fn doctor_unconfigured_message_names_three_setup_paths() {
1022        let msg = format_doctor_message(&serde_json::json!({
1023            "setup_completed": false,
1024            "mode": "unknown",
1025            "anthropic_key_configured": false,
1026            "local_model_selected": null,
1027            "local_model_loaded": null,
1028            "local_model_cached": false
1029        }));
1030
1031        assert!(msg.contains("Setup: not completed"));
1032        assert!(msg.contains("Run `origin setup`"));
1033        assert!(msg.contains("Basic Memory, On-device Model, or Anthropic Key"));
1034    }
1035
1036    #[test]
1037    fn search_memory_request_serialization_excludes_entity() {
1038        let req = SearchMemoryRequest {
1039            query: "test".into(),
1040            limit: 10,
1041            memory_type: None,
1042            domain: None,
1043            source_agent: None,
1044        };
1045        let json = serde_json::to_value(&req).unwrap();
1046        let obj = json.as_object().unwrap();
1047        assert!(
1048            !obj.contains_key("entity"),
1049            "entity must not be on the wire; got keys: {:?}",
1050            obj.keys().collect::<Vec<_>>()
1051        );
1052    }
1053
1054    #[test]
1055    fn chat_context_request_serialization_includes_domain() {
1056        let req = ChatContextRequest {
1057            query: None,
1058            conversation_id: Some("topic".into()),
1059            max_chunks: 20,
1060            relevance_threshold: None,
1061            include_goals: true,
1062            domain: Some("work".into()),
1063        };
1064        let json = serde_json::to_value(&req).unwrap();
1065        assert_eq!(json["domain"], serde_json::json!("work"));
1066        assert_eq!(json["conversation_id"], serde_json::json!("topic"));
1067    }
1068
1069    #[test]
1070    fn chat_context_response_deserializes_with_profile_and_knowledge() {
1071        let json = r#"{
1072            "context": "user is Lucian, prefers Rust",
1073            "profile": {
1074                "narrative": "n",
1075                "identity": ["rust"],
1076                "preferences": [],
1077                "goals": []
1078            },
1079            "knowledge": {
1080                "pages": [],
1081                "decisions": [],
1082                "relevant_memories": [],
1083                "graph_context": []
1084            },
1085            "took_ms": 42.0,
1086            "token_estimates": {
1087                "tier1_identity": 10,
1088                "tier2_project": 20,
1089                "tier3_relevant": 30,
1090                "total": 60
1091            }
1092        }"#;
1093        let parsed: ChatContextResponse = serde_json::from_str(json).unwrap();
1094        assert_eq!(parsed.context, "user is Lucian, prefers Rust");
1095        assert_eq!(parsed.profile.identity, vec!["rust"]);
1096        assert_eq!(parsed.token_estimates.total, 60);
1097    }
1098
1099    #[test]
1100    fn capture_params_structured_fields_schema_is_object() {
1101        use schemars::schema_for;
1102
1103        let schema = schema_for!(CaptureParams);
1104        let json = serde_json::to_value(&schema).unwrap();
1105        let sf_schema = json
1106            .pointer("/properties/structured_fields")
1107            .expect("structured_fields property in schema");
1108        let type_val = sf_schema
1109            .pointer("/type")
1110            .unwrap_or(&serde_json::Value::Null);
1111        let type_str = match type_val {
1112            serde_json::Value::String(s) => s.clone(),
1113            serde_json::Value::Array(arr) => arr
1114                .iter()
1115                .filter_map(|v| v.as_str())
1116                .collect::<Vec<_>>()
1117                .join(","),
1118            other => panic!(
1119                "structured_fields schema lacks type constraint; got: {:?}",
1120                other
1121            ),
1122        };
1123        assert!(
1124            type_str.contains("object"),
1125            "expected object type, got: {}",
1126            type_str
1127        );
1128    }
1129
1130    // ===== Param deserialization: ForgetParams =====
1131
1132    #[test]
1133    fn test_forget_params() {
1134        let json = r#"{"memory_id": "mem_abc123"}"#;
1135        let params: ForgetParams = serde_json::from_str(json).unwrap();
1136        assert_eq!(params.memory_id, "mem_abc123");
1137    }
1138
1139    #[test]
1140    fn test_forget_params_missing_id_fails() {
1141        let json = r#"{}"#;
1142        let result = serde_json::from_str::<ForgetParams>(json);
1143        assert!(result.is_err());
1144    }
1145
1146    // ===== Request serialization: StoreMemoryRequest =====
1147
1148    #[test]
1149    fn test_store_request_includes_new_fields() {
1150        let req = StoreMemoryRequest {
1151            content: "test".into(),
1152            memory_type: Some("decision".into()),
1153            domain: None,
1154            source_agent: Some("claude".into()),
1155            title: None,
1156            confidence: Some(0.9),
1157            supersedes: Some("old_id".into()),
1158            entity: Some("PostgreSQL".into()),
1159            entity_id: None,
1160            structured_fields: None,
1161            retrieval_cue: None,
1162        };
1163        let json = serde_json::to_value(&req).unwrap();
1164        assert_eq!(json["entity"], "PostgreSQL");
1165        assert_eq!(json["supersedes"], "old_id");
1166        assert!(json["confidence"].as_f64().unwrap() > 0.89);
1167        assert_eq!(json["source_agent"], "claude");
1168        assert!(json.get("user_id").is_none());
1169    }
1170
1171    #[test]
1172    fn test_store_request_minimal() {
1173        let req = StoreMemoryRequest {
1174            content: "hello".into(),
1175            memory_type: Some("fact".into()),
1176            domain: None,
1177            source_agent: None,
1178            title: None,
1179            confidence: None,
1180            supersedes: None,
1181            entity: None,
1182            entity_id: None,
1183            structured_fields: None,
1184            retrieval_cue: None,
1185        };
1186        let json = serde_json::to_value(&req).unwrap();
1187        assert_eq!(json["content"], "hello");
1188        assert_eq!(json["memory_type"], "fact");
1189        assert!(json.get("user_id").is_none());
1190    }
1191
1192    // ===== Response deserialization: StoreMemoryResponse =====
1193
1194    #[test]
1195    fn test_store_response_with_new_fields() {
1196        let json = r#"{
1197            "source_id": "mem_xyz",
1198            "chunks_created": 2,
1199            "memory_type": "fact",
1200            "entity_id": "ent_abc",
1201            "quality": "high",
1202            "warnings": ["decision memory missing claim"],
1203            "extraction_method": "agent"
1204        }"#;
1205        let resp: StoreMemoryResponse = serde_json::from_str(json).unwrap();
1206        assert_eq!(resp.source_id, "mem_xyz");
1207        assert_eq!(resp.chunks_created, 2);
1208        assert_eq!(resp.memory_type, "fact");
1209        assert_eq!(resp.entity_id.as_deref(), Some("ent_abc"));
1210        assert_eq!(resp.quality.as_deref(), Some("high"));
1211        assert_eq!(resp.warnings, vec!["decision memory missing claim"]);
1212        assert_eq!(resp.extraction_method, "agent");
1213    }
1214
1215    #[test]
1216    fn test_store_response_backward_compat_no_new_fields() {
1217        // Old backend response without warnings/extraction_method
1218        let json = r#"{
1219            "source_id": "mem_old",
1220            "chunks_created": 1,
1221            "memory_type": "fact"
1222        }"#;
1223        let resp: StoreMemoryResponse = serde_json::from_str(json).unwrap();
1224        assert_eq!(resp.source_id, "mem_old");
1225        assert_eq!(resp.chunks_created, 1);
1226        assert_eq!(resp.memory_type, "fact");
1227        assert!(resp.entity_id.is_none());
1228        assert!(resp.quality.is_none());
1229        assert!(resp.warnings.is_empty());
1230        assert_eq!(resp.extraction_method, "unknown");
1231    }
1232
1233    #[test]
1234    fn test_store_response_with_warnings_and_extraction_method() {
1235        let json = r#"{
1236            "source_id": "mem_xyz",
1237            "chunks_created": 1,
1238            "memory_type": "decision",
1239            "warnings": ["decision memory missing required 'claim' field"],
1240            "extraction_method": "llm"
1241        }"#;
1242        let resp: StoreMemoryResponse = serde_json::from_str(json).unwrap();
1243        assert_eq!(resp.memory_type, "decision");
1244        assert_eq!(
1245            resp.warnings,
1246            vec!["decision memory missing required 'claim' field"]
1247        );
1248        assert_eq!(resp.extraction_method, "llm");
1249    }
1250
1251    // ===== Response deserialization: SearchResult =====
1252
1253    #[test]
1254    fn test_search_result_with_new_fields() {
1255        let json = r#"{
1256            "id": "1",
1257            "content": "We chose Postgres",
1258            "source": "memory",
1259            "source_id": "mem_1",
1260            "title": "DB decision",
1261            "url": null,
1262            "chunk_index": 0,
1263            "last_modified": 1711000000,
1264            "score": 0.95,
1265            "chunk_type": "memory",
1266            "language": "en",
1267            "semantic_unit": "sentence",
1268            "memory_type": "decision",
1269            "domain": "origin",
1270            "source_agent": "claude",
1271            "confidence": 0.9,
1272            "confirmed": true,
1273            "stability": "standard",
1274            "supersedes": "mem_0",
1275            "summary": "DB choice",
1276            "entity_id": "ent_pg",
1277            "entity_name": "PostgreSQL",
1278            "quality": "high",
1279            "is_archived": false,
1280            "is_recap": false,
1281            "source_text": "We chose Postgres",
1282            "raw_score": 0.42
1283        }"#;
1284        let result: SearchResult = serde_json::from_str(json).unwrap();
1285        assert_eq!(result.chunk_type.as_deref(), Some("memory"));
1286        assert_eq!(result.language.as_deref(), Some("en"));
1287        assert_eq!(result.semantic_unit.as_deref(), Some("sentence"));
1288        assert_eq!(result.stability.as_deref(), Some("standard"));
1289        assert_eq!(result.supersedes.as_deref(), Some("mem_0"));
1290        assert_eq!(result.summary.as_deref(), Some("DB choice"));
1291        assert_eq!(result.entity_id.as_deref(), Some("ent_pg"));
1292        assert_eq!(result.entity_name.as_deref(), Some("PostgreSQL"));
1293        assert_eq!(result.quality.as_deref(), Some("high"));
1294        assert!(!result.is_archived);
1295        assert!(!result.is_recap);
1296        assert_eq!(result.source_text.as_deref(), Some("We chose Postgres"));
1297        assert!((result.raw_score - 0.42).abs() < f32::EPSILON);
1298    }
1299
1300    #[test]
1301    fn test_search_result_backward_compat_no_new_fields() {
1302        // Old backend response without entity/quality/archive/recap
1303        let json = r#"{
1304            "id": "1",
1305            "content": "test",
1306            "source": "memory",
1307            "source_id": "mem_1",
1308            "title": "test",
1309            "url": null,
1310            "chunk_index": 0,
1311            "last_modified": 1711000000,
1312            "score": 0.8,
1313            "memory_type": "fact",
1314            "domain": null,
1315            "source_agent": null,
1316            "confidence": null,
1317            "confirmed": null
1318        }"#;
1319        let result: SearchResult = serde_json::from_str(json).unwrap();
1320        assert!(result.entity_id.is_none());
1321        assert!(result.entity_name.is_none());
1322        assert!(result.quality.is_none());
1323        assert!(!result.is_archived);
1324        assert!(!result.is_recap);
1325        assert!(result.structured_fields.is_none());
1326        assert!(result.retrieval_cue.is_none());
1327        assert_eq!(result.raw_score, 0.0);
1328    }
1329
1330    #[test]
1331    fn test_search_result_with_structured_fields_and_retrieval_cue() {
1332        let json = r#"{
1333            "id": "1",
1334            "content": "Lucian prefers dark mode",
1335            "source": "memory",
1336            "source_id": "mem_1",
1337            "title": "Dark mode preference",
1338            "url": null,
1339            "chunk_index": 0,
1340            "last_modified": 1711000000,
1341            "score": 0.92,
1342            "memory_type": "preference",
1343            "domain": null,
1344            "source_agent": null,
1345            "confidence": null,
1346            "confirmed": null,
1347            "structured_fields": "{\"theme\":\"dark\",\"applies_to\":\"all_apps\"}",
1348            "retrieval_cue": "What UI theme does Lucian prefer?"
1349        }"#;
1350        let result: SearchResult = serde_json::from_str(json).unwrap();
1351        assert_eq!(
1352            result.structured_fields.as_deref(),
1353            Some("{\"theme\":\"dark\",\"applies_to\":\"all_apps\"}")
1354        );
1355        assert_eq!(
1356            result.retrieval_cue.as_deref(),
1357            Some("What UI theme does Lucian prefer?")
1358        );
1359        assert!(!result.is_archived);
1360        assert!(!result.is_recap);
1361        assert_eq!(result.raw_score, 0.0);
1362    }
1363
1364    #[test]
1365    fn test_search_result_knowledge_graph_source() {
1366        // Entity-boosted observation results from knowledge graph
1367        let json = r#"{
1368            "id": "obs_1",
1369            "content": "Prefers Rust over Go",
1370            "source": "knowledge_graph",
1371            "source_id": "ent_lucian",
1372            "title": "Lucian",
1373            "url": null,
1374            "chunk_index": 0,
1375            "last_modified": 1711000000,
1376            "score": 1.14,
1377            "memory_type": null,
1378            "domain": null,
1379            "source_agent": null,
1380            "confidence": null,
1381            "confirmed": null,
1382            "entity_id": "ent_lucian",
1383            "entity_name": "Lucian"
1384        }"#;
1385        let result: SearchResult = serde_json::from_str(json).unwrap();
1386        assert_eq!(result.source, "knowledge_graph");
1387        assert_eq!(result.entity_id.as_deref(), Some("ent_lucian"));
1388        assert_eq!(result.entity_name.as_deref(), Some("Lucian"));
1389        assert!(!result.is_archived);
1390        assert!(!result.is_recap);
1391        assert_eq!(result.raw_score, 0.0);
1392    }
1393
1394    // ===== Transport security: forget blocks on HTTP =====
1395
1396    #[tokio::test]
1397    async fn test_forget_blocked_on_http_transport() {
1398        let server = make_server(TransportMode::Http, "agent", None);
1399        let result = server.forget_impl("mem_123").await.unwrap();
1400        // Should return error content, not an Err
1401        let content = &result.content[0];
1402        match content.raw {
1403            rmcp::model::RawContent::Text(ref tc) => {
1404                assert!(tc.text.contains("not available over remote connections"));
1405            }
1406            _ => panic!("expected text content"),
1407        }
1408    }
1409
1410    #[tokio::test]
1411    async fn test_forget_allowed_on_stdio_transport() {
1412        // This will fail with connection error (no server), which proves
1413        // the transport check passed and it tried to make the HTTP call.
1414        // The error comes back as CallToolResult with is_error: true
1415        // (tool-level failure), not McpError (protocol-level).
1416        let server = make_server(TransportMode::Stdio, "agent", None);
1417        let result = server.forget_impl("mem_123").await.unwrap();
1418        assert!(
1419            result.is_error.unwrap_or(false),
1420            "should fail with connection error, not transport block"
1421        );
1422    }
1423
1424    // ===== Context default limit =====
1425
1426    #[test]
1427    fn test_context_request_default_limit() {
1428        let params = ContextParams {
1429            topic: Some("test".into()),
1430            limit: None,
1431            domain: None,
1432        };
1433        let req = ChatContextRequest {
1434            query: None,
1435            conversation_id: params.topic,
1436            max_chunks: params.limit.unwrap_or(20),
1437            relevance_threshold: None,
1438            include_goals: true,
1439            domain: params.domain,
1440        };
1441        assert_eq!(req.max_chunks, 20);
1442    }
1443
1444    #[test]
1445    fn test_context_request_custom_limit() {
1446        let params = ContextParams {
1447            topic: None,
1448            limit: Some(5),
1449            domain: Some("work".into()),
1450        };
1451        let req = ChatContextRequest {
1452            query: None,
1453            conversation_id: params.topic,
1454            max_chunks: params.limit.unwrap_or(20),
1455            relevance_threshold: None,
1456            include_goals: true,
1457            domain: params.domain,
1458        };
1459        assert_eq!(req.max_chunks, 5);
1460        assert_eq!(req.domain.as_deref(), Some("work"));
1461    }
1462
1463    #[test]
1464    fn test_context_maps_topic_to_conversation_id() {
1465        let params = ContextParams {
1466            topic: Some("project Origin".into()),
1467            limit: None,
1468            domain: None,
1469        };
1470        let req = ChatContextRequest {
1471            query: None,
1472            conversation_id: params.topic.clone(),
1473            max_chunks: params.limit.unwrap_or(20),
1474            relevance_threshold: None,
1475            include_goals: true,
1476            domain: params.domain,
1477        };
1478        assert_eq!(req.conversation_id.as_deref(), Some("project Origin"));
1479    }
1480
1481    // ===== Remember request construction =====
1482
1483    #[test]
1484    fn test_capture_constructs_store_request_with_entity() {
1485        let server = make_server(TransportMode::Stdio, "claude", None);
1486        let params = CaptureParams {
1487            content: "Alice manages the frontend team".into(),
1488            memory_type: Some("fact".into()),
1489            domain: Some("work".into()),
1490            entity: Some("Alice".into()),
1491            confidence: Some(0.9),
1492            supersedes: None,
1493            structured_fields: None,
1494            retrieval_cue: None,
1495        };
1496
1497        // Replicate capture_impl's request construction
1498        let source_agent = server.resolve_source_agent(None);
1499
1500        let req = StoreMemoryRequest {
1501            content: params.content,
1502            memory_type: params.memory_type,
1503            domain: params.domain,
1504            source_agent,
1505            title: None,
1506            confidence: params.confidence,
1507            supersedes: params.supersedes,
1508            entity: params.entity,
1509            entity_id: None,
1510            structured_fields: params.structured_fields.map(serde_json::Value::Object),
1511            retrieval_cue: params.retrieval_cue,
1512        };
1513
1514        let json = serde_json::to_value(&req).unwrap();
1515        assert_eq!(json["content"], "Alice manages the frontend team");
1516        assert_eq!(json["memory_type"], "fact");
1517        assert_eq!(json["domain"], "work");
1518        assert_eq!(json["entity"], "Alice");
1519        assert!(json["confidence"].as_f64().unwrap() > 0.89);
1520        // stdio mode: no param, no client_name → falls back to agent_name "claude"
1521        assert_eq!(json["source_agent"], "claude");
1522    }
1523
1524    #[test]
1525    fn test_remember_http_mode_injects_agent() {
1526        let server = make_server(TransportMode::Http, "claude.ai", Some("lucian"));
1527        let source_agent = server.resolve_source_agent(None);
1528
1529        assert_eq!(source_agent, Some("claude.ai".into()));
1530    }
1531
1532    // ===== Recall request construction =====
1533
1534    #[test]
1535    fn test_recall_constructs_search_request() {
1536        let params = RecallParams {
1537            query: "database choices".into(),
1538            limit: Some(5),
1539            memory_type: Some("decision".into()),
1540            domain: None,
1541        };
1542
1543        let req = SearchMemoryRequest {
1544            query: params.query,
1545            limit: params.limit.unwrap_or(10),
1546            memory_type: params.memory_type,
1547            domain: params.domain,
1548            source_agent: None,
1549        };
1550
1551        let json = serde_json::to_value(&req).unwrap();
1552        assert_eq!(json["query"], "database choices");
1553        assert_eq!(json["limit"], 5);
1554        assert_eq!(json["memory_type"], "decision");
1555        assert!(json.get("entity").is_none());
1556        assert!(json["domain"].is_null());
1557        assert!(json["source_agent"].is_null());
1558    }
1559
1560    // ===== Memory type backward compat =====
1561
1562    #[test]
1563    fn test_remember_passes_through_all_5_types() {
1564        for t in &["identity", "preference", "fact", "decision", "goal"] {
1565            let params = CaptureParams {
1566                content: "test".into(),
1567                memory_type: Some(t.to_string()),
1568                domain: None,
1569                entity: None,
1570                confidence: None,
1571                supersedes: None,
1572                structured_fields: None,
1573                retrieval_cue: None,
1574            };
1575            assert_eq!(params.memory_type.as_deref(), Some(*t));
1576        }
1577    }
1578
1579    // ===== Structured fields in remember params =====
1580
1581    #[test]
1582    fn test_capture_params_with_structured_fields_and_cue() {
1583        let json = r#"{
1584            "content": "Lucian prefers dark mode",
1585            "structured_fields": {"theme":"dark"},
1586            "retrieval_cue": "What theme does Lucian prefer?"
1587        }"#;
1588        let params: CaptureParams = serde_json::from_str(json).unwrap();
1589        let structured_fields = params.structured_fields.expect("structured_fields");
1590        assert_eq!(
1591            structured_fields.get("theme"),
1592            Some(&serde_json::Value::String("dark".into()))
1593        );
1594        assert_eq!(
1595            params.retrieval_cue.as_deref(),
1596            Some("What theme does Lucian prefer?")
1597        );
1598    }
1599
1600    #[test]
1601    fn test_store_request_with_structured_fields() {
1602        let req = StoreMemoryRequest {
1603            content: "test".into(),
1604            memory_type: Some("fact".into()),
1605            domain: None,
1606            source_agent: None,
1607            title: None,
1608            confidence: None,
1609            supersedes: None,
1610            entity: None,
1611            entity_id: None,
1612            structured_fields: Some(serde_json::json!({"key":"val"})),
1613            retrieval_cue: Some("What is the key?".into()),
1614        };
1615        let json = serde_json::to_value(&req).unwrap();
1616        assert_eq!(json["structured_fields"], serde_json::json!({"key":"val"}));
1617        assert_eq!(json["retrieval_cue"], "What is the key?");
1618    }
1619
1620    // ===== ChatContextResponse deserialization =====
1621
1622    #[test]
1623    fn test_chat_context_response() {
1624        let json = r#"{
1625            "context": "User prefers dark mode. Works on Origin project.",
1626            "profile": {
1627                "narrative": "narrative",
1628                "identity": [],
1629                "preferences": [],
1630                "goals": []
1631            },
1632            "knowledge": {
1633                "pages": [],
1634                "decisions": [],
1635                "relevant_memories": [],
1636                "graph_context": []
1637            },
1638            "took_ms": 12.5,
1639            "token_estimates": {
1640                "tier1_identity": 1,
1641                "tier2_project": 2,
1642                "tier3_relevant": 3,
1643                "total": 6
1644            }
1645        }"#;
1646        let resp: ChatContextResponse = serde_json::from_str(json).unwrap();
1647        assert!(!resp.context.is_empty());
1648        assert!(resp.profile.identity.is_empty());
1649        assert_eq!(resp.took_ms, 12.5);
1650        assert_eq!(resp.token_estimates.total, 6);
1651    }
1652
1653    #[test]
1654    fn test_chat_context_response_empty() {
1655        let json = r#"{
1656            "context": "",
1657            "profile": {
1658                "narrative": "",
1659                "identity": [],
1660                "preferences": [],
1661                "goals": []
1662            },
1663            "knowledge": {
1664                "pages": [],
1665                "decisions": [],
1666                "relevant_memories": [],
1667                "graph_context": []
1668            },
1669            "took_ms": 1.0,
1670            "token_estimates": {
1671                "tier1_identity": 0,
1672                "tier2_project": 0,
1673                "tier3_relevant": 0,
1674                "total": 0
1675            }
1676        }"#;
1677        let resp: ChatContextResponse = serde_json::from_str(json).unwrap();
1678        assert!(resp.context.is_empty());
1679    }
1680
1681    // ===== with_instructions content assertions =====
1682    // These tests lock in the refined agent-facing guidance. If any
1683    // assertion fails, either the rule was intentionally changed
1684    // (update the test) or the refinement was accidentally dropped
1685    // (restore the rule).
1686
1687    fn server_instructions() -> String {
1688        let s = make_server(TransportMode::Stdio, "test", None);
1689        s.get_info()
1690            .instructions
1691            .expect("server must ship with_instructions")
1692    }
1693
1694    #[test]
1695    fn instructions_mention_self_evolving_knowledge() {
1696        assert!(
1697            server_instructions().contains("self-evolving"),
1698            "with_instructions must describe Origin as self-evolving"
1699        );
1700    }
1701
1702    #[test]
1703    fn instructions_mention_shared_across_tools() {
1704        assert!(
1705            server_instructions().contains("shared across all"),
1706            "with_instructions must tell agents the store is shared across tools"
1707        );
1708    }
1709
1710    #[test]
1711    fn instructions_mention_how_user_thinks() {
1712        assert!(
1713            server_instructions().contains("how the user thinks"),
1714            "with_instructions must frame context as modeling how the user thinks"
1715        );
1716    }
1717
1718    #[test]
1719    fn instructions_use_proactive_framing() {
1720        assert!(
1721            server_instructions().contains("STORE PROACTIVELY"),
1722            "with_instructions must use STORE PROACTIVELY framing (not passive WHEN TO STORE)"
1723        );
1724    }
1725
1726    #[test]
1727    fn instructions_ban_tool_output_storage() {
1728        assert!(
1729            server_instructions().contains("Tool output or command results"),
1730            "with_instructions must explicitly rule out tool output as storage material"
1731        );
1732    }
1733
1734    #[test]
1735    fn instructions_ban_ghost_inferences() {
1736        assert!(
1737            server_instructions().contains("Your own inferences"),
1738            "with_instructions must rule out storing agent's own inferences user didn't express"
1739        );
1740    }
1741
1742    #[test]
1743    fn instructions_call_out_atomic_memory() {
1744        assert!(
1745            server_instructions().contains("Atomic: one idea per memory"),
1746            "with_instructions must call out the atomic-memory rule explicitly by name"
1747        );
1748    }
1749
1750    #[test]
1751    fn instructions_specify_declarative_writing() {
1752        assert!(
1753            server_instructions().contains("Declarative, not narrative"),
1754            "with_instructions must require declarative (not narrative) writing style"
1755        );
1756    }
1757
1758    #[test]
1759    fn instructions_default_to_omit_memory_type() {
1760        let i = server_instructions();
1761        assert!(
1762            i.contains("omit and trust the backend"),
1763            "with_instructions must default agents to omitting memory_type"
1764        );
1765        assert!(
1766            i.contains("do NOT set memory_type"),
1767            "with_instructions must explicitly say do NOT set memory_type by default"
1768        );
1769    }
1770
1771    #[test]
1772    fn instructions_carve_out_decisions_for_decision_log() {
1773        let i = server_instructions();
1774        assert!(
1775            i.contains("Decision Log"),
1776            "with_instructions must name the Decision Log as the reason for explicit decision typing"
1777        );
1778        assert!(
1779            i.contains("memory_type=\"decision\""),
1780            "with_instructions must tell agents to set memory_type=\"decision\" explicitly for decisions"
1781        );
1782    }
1783
1784    // ===== tool-level and param-level description assertions =====
1785
1786    fn tool_descriptions() -> std::collections::HashMap<String, String> {
1787        let server = make_server(TransportMode::Stdio, "test", None);
1788        server
1789            .tool_router
1790            .list_all()
1791            .into_iter()
1792            .filter_map(|t| {
1793                let desc = t.description.as_ref()?.to_string();
1794                Some((t.name.to_string(), desc))
1795            })
1796            .collect()
1797    }
1798
1799    #[test]
1800    fn capture_description_calls_out_atomic() {
1801        let descriptions = tool_descriptions();
1802        let capture = descriptions.get("capture").expect("capture tool exists");
1803        assert!(
1804            capture.contains("Each call is one atomic idea"),
1805            "capture description must call out atomic-per-call explicitly, got: {capture}"
1806        );
1807    }
1808
1809    #[test]
1810    fn context_description_frames_modeling_user() {
1811        let descriptions = tool_descriptions();
1812        let ctx = descriptions.get("context").expect("context tool exists");
1813        assert!(
1814            ctx.contains("how the user thinks"),
1815            "context description must frame the result as modeling how the user thinks, got: {ctx}"
1816        );
1817    }
1818
1819    #[test]
1820    fn doctor_description_mentions_setup_mode() {
1821        let descriptions = tool_descriptions();
1822        let status = descriptions.get("doctor").expect("doctor tool exists");
1823        assert!(
1824            status.contains("Basic Memory"),
1825            "doctor description must mention setup modes, got: {status}"
1826        );
1827        assert!(
1828            status.contains("On-device Model"),
1829            "doctor description must mention on-device setup, got: {status}"
1830        );
1831        assert!(
1832            status.contains("not part of the memory loop"),
1833            "doctor description must frame itself as diagnostic-only, got: {status}"
1834        );
1835    }
1836
1837    #[test]
1838    fn recall_memory_type_param_lists_two_level_filter() {
1839        let params_schema = serde_json::to_string(&schemars::schema_for!(RecallParams))
1840            .expect("RecallParams schema serializes");
1841        assert!(
1842            params_schema.contains("Two-level filter"),
1843            "RecallParams.memory_type must advertise the two-level filter, got schema: {params_schema}"
1844        );
1845        assert!(
1846            params_schema.contains("profile"),
1847            "RecallParams.memory_type must mention profile alias"
1848        );
1849        assert!(
1850            params_schema.contains("knowledge"),
1851            "RecallParams.memory_type must mention knowledge alias"
1852        );
1853    }
1854}