lean_ctx_server/tools/registered/
ctx_memory.rs

1// SPDX-License-Identifier: Apache-2.0
2use rmcp::ErrorData;
3use rmcp::model::Tool;
4use serde_json::{Map, Value, json};
5
6use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_str, get_usize};
7use lean_ctx::core::shared_context::SharedContext;
8use lean_ctx::tool_defs::tool_def;
9
10pub struct CtxMemoryTool;
11
12impl McpTool for CtxMemoryTool {
13    fn name(&self) -> &'static str {
14        "ctx_memory"
15    }
16
17    fn tool_def(&self) -> Tool {
18        tool_def(
19            self.name(),
20            "Store and retrieve durable cross-agent facts with BLAKE3 deduplication.",
21            json!({
22                "type": "object",
23                "properties": {
24                    "action": {
25                        "type": "string",
26                        "enum": ["put", "get", "list", "stats", "prune"],
27                        "description": "Memory operation"
28                    },
29                    "content": { "type": "string", "description": "Fact content (required for put)" },
30                    "agent": { "type": "string", "description": "Agent identity for put; defaults to unknown" },
31                    "query": { "type": "string", "description": "Keyword query (required for get)" },
32                    "category": {
33                        "type": "string",
34                        "enum": ["fact", "decision", "blocker", "pattern"],
35                        "description": "Category for put or list filter"
36                    },
37                    "limit": { "type": "integer", "minimum": 0, "default": 10 },
38                    "max_age_seconds": { "type": "integer", "minimum": 0, "description": "Maximum idle age for prune" },
39                    "max_entries": { "type": "integer", "minimum": 0, "description": "Maximum entries retained by prune" }
40                },
41                "required": ["action"]
42            }),
43        )
44    }
45
46    fn handle(
47        &self,
48        args: &Map<String, Value>,
49        ctx: &ToolContext,
50    ) -> Result<ToolOutput, ErrorData> {
51        let action = get_str(args, "action")
52            .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
53        // Facts this path stored under an earlier identity of the project.
54        lean_ctx::core::legacy_adoption::auto_adopt_shared_context(&ctx.project_root);
55        let store = SharedContext::for_project(&ctx.project_root).map_err(memory_error)?;
56        let limit = get_usize(args, "limit").unwrap_or(10);
57        let text = match action.as_str() {
58            "put" => {
59                let content = storable(required(args, "content")?)?;
60                let agent = get_str(args, "agent").unwrap_or_else(|| "unknown".to_string());
61                let category = get_str(args, "category").unwrap_or_else(|| "fact".to_string());
62                let id = store.put(content, agent, category).map_err(memory_error)?;
63                format!("stored {id}")
64            }
65            "get" => {
66                let query = required(args, "query")?;
67                served(&serialize(
68                    &store.get_relevant(query, limit).map_err(memory_error)?,
69                )?)
70            }
71            "list" => {
72                let entries = match get_str(args, "category") {
73                    Some(category) => store
74                        .get_by_category(&category)
75                        .map_err(memory_error)?
76                        .into_iter()
77                        .take(limit)
78                        .collect(),
79                    None => store.get_recent(limit).map_err(memory_error)?,
80                };
81                served(&serialize(&entries)?)
82            }
83            "stats" => serialize(&store.stats().map_err(memory_error)?)?,
84            "prune" => {
85                let max_age_seconds =
86                    get_usize(args, "max_age_seconds").unwrap_or(60 * 60 * 24 * 30) as u64;
87                let max_entries = get_usize(args, "max_entries").unwrap_or(1_000);
88                let removed = store
89                    .prune(max_age_seconds, max_entries)
90                    .map_err(memory_error)?;
91                format!("pruned {removed} entries")
92            }
93            other => {
94                return Err(ErrorData::invalid_params(
95                    format!("unknown action '{other}'"),
96                    None,
97                ));
98            }
99        };
100        Ok(ToolOutput::simple(text))
101    }
102}
103
104fn required<'a>(args: &'a Map<String, Value>, key: &str) -> Result<&'a str, ErrorData> {
105    args.get(key)
106        .and_then(Value::as_str)
107        .filter(|value| !value.trim().is_empty())
108        .ok_or_else(|| ErrorData::invalid_params(format!("{key} is required"), None))
109}
110
111fn serialize<T: serde::Serialize>(value: &T) -> Result<String, ErrorData> {
112    serde_json::to_string_pretty(value).map_err(|error| {
113        ErrorData::internal_error(format!("serialize memory result: {error}"), None)
114    })
115}
116
117fn memory_error(error: impl std::fmt::Display) -> ErrorData {
118    ErrorData::internal_error(format!("shared context: {error}"), None)
119}
120
121/// The text a `put` may persist (E3): caller-supplied content admitted for a
122/// derived store under the current policy. Stored bytes are the admitted
123/// text — a value the gateway masks is stored masked; content that is
124/// withheld or restricted is refused and nothing is written.
125fn storable(content: &str) -> Result<String, ErrorData> {
126    lean_ctx::core::context_admission::carrier::admit_for_store(content)
127        .map_err(|withheld| ErrorData::invalid_params(withheld.message("memory content"), None))
128}
129
130/// Stored entries as they may be served NOW (G04). Shared memory carries no
131/// source binding, so the rendered entries follow the originless rule:
132/// withheld while protection is active in any layer, re-admitted under
133/// today's content rules otherwise. A refusal is the content-free notice.
134fn served(rendered: &str) -> String {
135    crate::tools::admitted_sources::serve_originless(rendered, "shared memory")
136        .unwrap_or_else(|reason| format!("[withheld: {reason}]"))
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    /// E3 / G04: `put` stored caller text raw and `get` / `list` printed the
144    /// store raw. Both ends now pass admission.
145    #[test]
146    fn shared_memory_is_admitted_at_write_and_served_under_the_originless_rule() {
147        const AWS_KEY: &str = concat!("AK", "IAIOSFODNN7EXAMPLE");
148        let _data = lean_ctx::core::data_dir::isolated_data_dir();
149        let entries = "[\n  {\n    \"content\": \"retries live in net/retry.rs\"\n  }\n]";
150        {
151            let _community = lean_ctx::core::policy::runtime::TestPolicyOverride::set(None);
152            let clean = storable("retries live in net/retry.rs");
153            assert!(clean.is_ok_and(|text| text == "retries live in net/retry.rs"));
154            let keyed = storable(&format!("deploy key is {AWS_KEY}"));
155            assert!(keyed.is_ok_and(|text| !text.contains(AWS_KEY)));
156            assert_eq!(served(entries), entries);
157            let legacy = served(&format!("[{{\"content\": \"key {AWS_KEY}\"}}]"));
158            assert!(!legacy.contains(AWS_KEY), "{legacy}");
159        }
160        let policy = lean_ctx::core::policy::load(
161            "name = \"ctx-memory-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
162             [filters]\nclassification = \"block\"\nblocked_labels = [\"CONFIDENTIAL\"]\n",
163        )
164        .unwrap();
165        let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(Some(policy));
166        assert!(storable("CONFIDENTIAL\nroadmap-marker").is_err());
167        let withheld = served(entries);
168        assert!(withheld.starts_with("[withheld: "), "{withheld}");
169        assert!(!withheld.contains("retry.rs"), "{withheld}");
170    }
171
172    #[test]
173    fn memory_tool_schema_has_expected_actions() {
174        let definition = CtxMemoryTool.tool_def();
175        assert_eq!(definition.name, "ctx_memory");
176    }
177}
lean_ctx_server/tools/registered/ctx_memory.rs

lean_ctx_server/tools/registered/
ctx_memory.rs
