Skip to main content

atman_runtime/tools/
memory_stubs.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use tokio::sync::RwLock;
5
6use crate::error::RuntimeError;
7use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
8use crate::value::Value;
9
10use crate::migration::MigratedRule;
11
12#[derive(Default, Clone)]
13pub struct RuleFetch {
14    entries: Arc<RwLock<HashMap<String, String>>>,
15    migrated: Arc<RwLock<Vec<MigratedRule>>>,
16}
17
18impl RuleFetch {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub async fn insert(&self, name: impl Into<String>, content: impl Into<String>) {
24        self.entries
25            .write()
26            .await
27            .insert(name.into(), content.into());
28    }
29
30    pub async fn set_migrated(&self, rules: Vec<MigratedRule>) {
31        *self.migrated.write().await = rules;
32    }
33
34    pub async fn migrated_count(&self) -> usize {
35        self.migrated.read().await.len()
36    }
37}
38
39impl Tool for RuleFetch {
40    fn name(&self) -> &str {
41        "rule.fetch"
42    }
43
44    fn tier(&self) -> Tier {
45        Tier::Zero
46    }
47
48    fn description(&self) -> Option<&str> {
49        Some(
50            "Load a skill/rule's full content by name, or search the rule index by keyword. \
51             Sources: CLAUDE.md, AGENTS.md, ~/.claude/skills/*/SKILL.md, .cursorrules, \
52             .kiro/steering/*.md, and aider conventions.",
53        )
54    }
55
56    fn input_schema(&self) -> serde_json::Value {
57        serde_json::json!({
58            "type": "object",
59            "properties": {
60                "name": {
61                    "type": "string",
62                    "description": "Exact rule name to load full content (e.g. 'skill:code-review::references/rules.md')."
63                },
64                "query": {
65                    "type": "string",
66                    "description": "Keyword to search across rule names and descriptions. Returns a list of {name, description, scope, source}. Omit both name and query to list the full index."
67                }
68            }
69        })
70    }
71
72    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
73        Box::pin(async move {
74            // Search mode: query keyword → list of matching rules (index).
75            if let Some(Value::Str(query)) = args.named("query") {
76                let migrated = self.migrated.read().await;
77                let lower = query.to_lowercase();
78                let results: Vec<Value> = migrated
79                    .iter()
80                    .filter(|r| {
81                        r.name.to_lowercase().contains(&lower)
82                            || r.description
83                                .as_deref()
84                                .map(|d| d.to_lowercase().contains(&lower))
85                                .unwrap_or(false)
86                    })
87                    .map(rule_to_index_value)
88                    .collect();
89                return Ok(Value::List(results));
90            }
91
92            // No args → return the full compact index (name + description + scope + source).
93            if args.named("name").is_none() && args.positional(0).is_err() {
94                let migrated = self.migrated.read().await;
95                let index: Vec<Value> = migrated.iter().map(rule_to_index_value).collect();
96                return Ok(Value::List(index));
97            }
98
99            let name = extract_string(&args, "name", 0)?;
100            let entries = self.entries.read().await;
101            if let Some(content) = entries.get(&name) {
102                return Ok(Value::Str(content.clone()));
103            }
104            drop(entries);
105            let migrated = self.migrated.read().await;
106            if let Some(rule) = crate::migration::resolve_by_name(&migrated, &name) {
107                return Ok(Value::Str(rule.content.clone()));
108            }
109            Ok(Value::Str(String::new()))
110        })
111    }
112}
113
114/// Serialize a migrated rule into a compact index entry `{name, description, scope, source}`.
115fn rule_to_index_value(rule: &MigratedRule) -> Value {
116    use crate::value::Value as V;
117    V::Struct(vec![
118        ("name".into(), V::Str(rule.name.clone())),
119        (
120            "description".into(),
121            V::Str(rule.description.clone().unwrap_or_default()),
122        ),
123        ("scope".into(), V::Str(rule.scope.as_str().to_string())),
124        ("source".into(), V::Str(rule.source_tool.clone())),
125    ])
126}
127
128fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
129    let value = match args.named(name) {
130        Some(v) => v,
131        None => args.positional(pos)?,
132    };
133    match value {
134        Value::Str(s) => Ok(s.clone()),
135        other => Err(RuntimeError::TypeMismatch {
136            expected: "string".into(),
137            actual: other.kind_name().into(),
138        }),
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[tokio::test]
147    async fn rule_fetch_returns_stored_content() {
148        let tool = RuleFetch::new();
149        tool.insert("code-review", "review carefully").await;
150        let out = tool
151            .call(
152                ToolArgs {
153                    positional: vec![Value::Str("code-review".into())],
154                    named: vec![],
155                },
156                &ToolCtx::new(),
157            )
158            .await
159            .unwrap();
160        assert!(matches!(out, Value::Str(s) if s == "review carefully"));
161    }
162
163    #[tokio::test]
164    async fn rule_fetch_missing_returns_empty_string() {
165        let tool = RuleFetch::new();
166        let out = tool
167            .call(
168                ToolArgs {
169                    positional: vec![Value::Str("missing".into())],
170                    named: vec![],
171                },
172                &ToolCtx::new(),
173            )
174            .await
175            .unwrap();
176        assert!(matches!(out, Value::Str(s) if s.is_empty()));
177    }
178
179    fn migrated_rule(name: &str, description: &str) -> MigratedRule {
180        MigratedRule {
181            name: name.to_string(),
182            source_tool: "skill".to_string(),
183            source_path: "/tmp/x".into(),
184            scope: crate::migration::RuleScope::Global,
185            content: "content".to_string(),
186            description: Some(description.to_string()),
187        }
188    }
189
190    #[tokio::test]
191    async fn rule_fetch_query_searches_name_and_description() {
192        let tool = RuleFetch::new();
193        tool.set_migrated(vec![
194            migrated_rule("skill:code-review::a.md", "structured code review"),
195            migrated_rule("skill:deploy::b.md", "deployment checklist"),
196        ])
197        .await;
198
199        let out = tool
200            .call(
201                ToolArgs {
202                    positional: vec![],
203                    named: vec![("query".into(), Value::Str("review".into()))],
204                },
205                &ToolCtx::new(),
206            )
207            .await
208            .unwrap();
209        let Value::List(items) = out else {
210            panic!("expected list, got {out:?}")
211        };
212        assert_eq!(items.len(), 1, "only the review rule should match");
213        let Value::Struct(fields) = &items[0] else {
214            panic!("expected struct entry")
215        };
216        assert_eq!(fields[0].0, "name");
217        assert!(matches!(&fields[0].1, Value::Str(s) if s == "skill:code-review::a.md"));
218    }
219
220    #[tokio::test]
221    async fn rule_fetch_no_args_returns_full_index() {
222        let tool = RuleFetch::new();
223        tool.set_migrated(vec![
224            migrated_rule("a", "first"),
225            migrated_rule("b", "second"),
226        ])
227        .await;
228
229        let out = tool
230            .call(ToolArgs::default(), &ToolCtx::new())
231            .await
232            .unwrap();
233        let Value::List(items) = out else {
234            panic!("expected list, got {out:?}")
235        };
236        assert_eq!(items.len(), 2);
237    }
238}