Skip to main content

wm_tools/expansion/
additional.rs

1//! Additional tools — count, tags, session_list, citta_coherence, dharma_profiles, nearby.
2
3#![forbid(unsafe_code)]
4
5use async_trait::async_trait;
6
7use serde_json::{Value, json};
8use std::sync::Arc;
9use wm_core::{Context, EffectRow, Galaxy, Gana, Resource, Tool, ToolStats};
10use wm_memory::MemoryStore;
11
12use super::common::{galaxy_name, parse_galaxy, parse_galaxy_or};
13
14pub struct MemoryCountTool {
15    store: Arc<MemoryStore>,
16    stats: ToolStats,
17    effects: EffectRow,
18}
19
20impl MemoryCountTool {
21    pub fn new(store: Arc<MemoryStore>) -> Self {
22        Self {
23            store,
24            stats: ToolStats::default(),
25            effects: EffectRow::read_only(vec![Resource::Galaxy("codex".into())]),
26        }
27    }
28}
29
30#[async_trait]
31impl Tool for MemoryCountTool {
32    fn input_schema(&self) -> Value {
33        super::common::schema(
34            &json!({
35                "galaxy": super::common::str_prop("Galaxy to count (optional; all memory galaxies when absent)"),
36            }),
37            &[],
38        )
39    }
40    fn name(&self) -> &str {
41        "memory.count"
42    }
43    fn gana(&self) -> Gana {
44        Gana::WinnowingBasket
45    }
46    fn effects(&self) -> &EffectRow {
47        &self.effects
48    }
49    fn description(&self) -> &str {
50        "Count memories in a galaxy"
51    }
52    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
53        let galaxy = parse_galaxy_or(args.get("galaxy").and_then(|v| v.as_str()), Galaxy::Codex)?;
54        let count = self.store.count(galaxy)?;
55        Ok(json!({ "status": "success", "galaxy": galaxy_name(galaxy), "count": count }))
56    }
57    fn stats(&self) -> &ToolStats {
58        &self.stats
59    }
60}
61
62/// `memory.tags` — list all unique tags in a galaxy.
63pub struct MemoryTagsTool {
64    store: Arc<MemoryStore>,
65    stats: ToolStats,
66    effects: EffectRow,
67}
68
69impl MemoryTagsTool {
70    pub fn new(store: Arc<MemoryStore>) -> Self {
71        Self {
72            store,
73            stats: ToolStats::default(),
74            effects: EffectRow::read_only(vec![Resource::Galaxy("codex".into())]),
75        }
76    }
77}
78
79#[async_trait]
80impl Tool for MemoryTagsTool {
81    fn input_schema(&self) -> Value {
82        super::common::schema(
83            &json!({
84                "galaxy": super::common::str_prop("Galaxy whose tags to list (optional; default codex)"),
85            }),
86            &[],
87        )
88    }
89    fn name(&self) -> &str {
90        "memory.tags"
91    }
92    fn gana(&self) -> Gana {
93        Gana::Net
94    }
95    fn effects(&self) -> &EffectRow {
96        &self.effects
97    }
98    fn description(&self) -> &str {
99        "List all unique tags in a galaxy"
100    }
101    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
102        let galaxy = parse_galaxy_or(args.get("galaxy").and_then(|v| v.as_str()), Galaxy::Codex)?;
103        let memories = self.store.scan(galaxy, 10_000)?;
104        let tags: std::collections::HashSet<String> = memories
105            .iter()
106            .flat_map(|m| m.metadata.tags.iter().cloned())
107            .collect();
108        let mut tag_list: Vec<String> = tags.into_iter().collect();
109        tag_list.sort_unstable();
110        Ok(
111            json!({ "status": "success", "galaxy": galaxy_name(galaxy), "unique_tags": tag_list.len(), "tags": tag_list }),
112        )
113    }
114    fn stats(&self) -> &ToolStats {
115        &self.stats
116    }
117}
118
119/// `session.list` — list all sessions.
120pub struct SessionListTool {
121    store: Arc<MemoryStore>,
122    stats: ToolStats,
123    effects: EffectRow,
124}
125
126impl SessionListTool {
127    pub fn new(store: Arc<MemoryStore>) -> Self {
128        Self {
129            store,
130            stats: ToolStats::default(),
131            effects: EffectRow::read_only(vec![Resource::Galaxy("sessions".into())]),
132        }
133    }
134}
135
136#[async_trait]
137impl Tool for SessionListTool {
138    fn input_schema(&self) -> Value {
139        super::common::schema(
140            &json!({
141                "session_id": super::common::str_prop("Filter by session UUID"),
142                "title": super::common::str_prop("Filter by title substring"),
143                "sequence": super::common::int_prop("Filter by start sequence number"),
144                "type": super::common::str_prop("Filter by session type"),
145            }),
146            &[],
147        )
148    }
149    fn name(&self) -> &str {
150        "session.list"
151    }
152    fn gana(&self) -> Gana {
153        Gana::StraddlingLegs
154    }
155    fn effects(&self) -> &EffectRow {
156        &self.effects
157    }
158    fn description(&self) -> &str {
159        "List session summaries in the Sessions galaxy (turns grouped by session)"
160    }
161    async fn call(&self, _ctx: &mut Context, _args: Value) -> wm_core::Result<Value> {
162        let memories = self.store.scan_all(Galaxy::Sessions)?;
163
164        // Group turns and start markers into session summaries.
165        #[derive(Default)]
166        struct Summary {
167            title: Option<String>,
168            turns: u64,
169            earliest: Option<chrono::DateTime<chrono::Utc>>,
170            latest: Option<chrono::DateTime<chrono::Utc>>,
171        }
172        let mut summaries: std::collections::HashMap<String, Summary> =
173            std::collections::HashMap::new();
174
175        for m in &memories {
176            if let Ok(v) = serde_json::from_str::<Value>(&m.content) {
177                // session_start: session id lives in the tag.
178                if v.get("type").and_then(Value::as_str) == Some("session_start") {
179                    if let Some(sid) = m
180                        .metadata
181                        .tags
182                        .iter()
183                        .find_map(|t| t.strip_prefix("session:"))
184                    {
185                        let entry = summaries.entry(sid.to_string()).or_default();
186                        entry.title = v.get("title").and_then(Value::as_str).map(String::from);
187                    }
188                    continue;
189                }
190                // session_turn (and other session memories): session id in content.
191                if let Some(sid) = v.get("session_id").and_then(Value::as_str) {
192                    let entry = summaries.entry(sid.to_string()).or_default();
193                    let ts = m.metadata.created_at;
194                    entry.earliest = Some(entry.earliest.map_or(ts, |e| e.min(ts)));
195                    entry.latest = Some(entry.latest.map_or(ts, |l| l.max(ts)));
196                    if v.get("sequence").and_then(Value::as_u64).is_some() {
197                        entry.turns += 1;
198                    }
199                }
200            }
201        }
202
203        let mut sessions: Vec<(String, Summary)> = summaries.into_iter().collect();
204        // Most recently active sessions first.
205        sessions.sort_by_key(|(_, s)| {
206            std::cmp::Reverse(
207                s.latest
208                    .unwrap_or(chrono::DateTime::<chrono::Utc>::UNIX_EPOCH),
209            )
210        });
211        let total = sessions.len();
212        sessions.truncate(100);
213
214        let sessions: Vec<Value> = sessions
215            .into_iter()
216            .map(|(sid, s)| {
217                json!({
218                    "session_id": sid,
219                    "title": s.title.unwrap_or_else(|| format!("Session {}", &sid[..sid.len().min(8)])),
220                    "turns": s.turns,
221                    "first_activity": s.earliest.map(|t| t.to_rfc3339()),
222                    "last_activity": s.latest.map(|t| t.to_rfc3339()),
223                })
224            })
225            .collect();
226
227        Ok(json!({
228            "status": "success",
229            "count": total,
230            "sessions": sessions,
231        }))
232    }
233    fn stats(&self) -> &ToolStats {
234        &self.stats
235    }
236}
237
238/// `citta.coherence` — check coherence threshold.
239pub struct CittaCoherenceTool {
240    stats: ToolStats,
241    effects: EffectRow,
242}
243
244impl CittaCoherenceTool {
245    #[must_use]
246    pub fn new() -> Self {
247        Self {
248            stats: ToolStats::default(),
249            effects: EffectRow::pure(),
250        }
251    }
252}
253
254impl Default for CittaCoherenceTool {
255    fn default() -> Self {
256        Self::new()
257    }
258}
259
260#[async_trait]
261impl Tool for CittaCoherenceTool {
262    fn name(&self) -> &str {
263        "citta.coherence"
264    }
265    fn gana(&self) -> Gana {
266        Gana::Ghost
267    }
268    fn effects(&self) -> &EffectRow {
269        &self.effects
270    }
271    fn description(&self) -> &str {
272        "Check citta coherence level and whether writes are permitted"
273    }
274    async fn call(&self, ctx: &mut Context, _args: Value) -> wm_core::Result<Value> {
275        let threshold = 0.3f32;
276        let can_write = ctx.citta_coherence >= threshold;
277        Ok(json!({
278            "status": "success",
279            "coherence": ctx.citta_coherence,
280            "valence": ctx.citta_valence,
281            "write_threshold": threshold,
282            "can_write": can_write,
283        }))
284    }
285    fn stats(&self) -> &ToolStats {
286        &self.stats
287    }
288}
289
290/// `dharma.profiles` — list available dharma profiles.
291pub struct DharmaProfilesTool {
292    stats: ToolStats,
293    effects: EffectRow,
294}
295
296impl DharmaProfilesTool {
297    #[must_use]
298    pub fn new() -> Self {
299        Self {
300            stats: ToolStats::default(),
301            effects: EffectRow::pure(),
302        }
303    }
304}
305
306impl Default for DharmaProfilesTool {
307    fn default() -> Self {
308        Self::new()
309    }
310}
311
312#[async_trait]
313impl Tool for DharmaProfilesTool {
314    fn name(&self) -> &str {
315        "dharma.profiles"
316    }
317    fn gana(&self) -> Gana {
318        Gana::ExtendedNet
319    }
320    fn effects(&self) -> &EffectRow {
321        &self.effects
322    }
323    fn description(&self) -> &str {
324        "List available dharma governance profiles"
325    }
326    async fn call(&self, _ctx: &mut Context, _args: Value) -> wm_core::Result<Value> {
327        Ok(json!({
328            "status": "success",
329            "profiles": [
330                { "name": "default", "description": "Standard governance — observe and advise" },
331                { "name": "strict", "description": "Strict governance — intervene on writes in low coherence" },
332                { "name": "research", "description": "Lenient governance — allow experimental tools" },
333                { "name": "production", "description": "Hardened governance — panic on dharma violations" },
334            ],
335        }))
336    }
337    fn stats(&self) -> &ToolStats {
338        &self.stats
339    }
340}
341
342/// `memory.nearby` — find memories spatially near a query using 5D coordinates.
343pub struct MemoryNearbyTool {
344    store: Arc<MemoryStore>,
345    stats: ToolStats,
346    effects: EffectRow,
347}
348
349impl MemoryNearbyTool {
350    pub fn new(store: Arc<MemoryStore>) -> Self {
351        Self {
352            store,
353            stats: ToolStats::default(),
354            effects: EffectRow::read_only(vec![Resource::Galaxy("codex".into())]),
355        }
356    }
357}
358
359#[async_trait]
360impl Tool for MemoryNearbyTool {
361    fn name(&self) -> &str {
362        "memory.nearby"
363    }
364    fn gana(&self) -> Gana {
365        Gana::Star
366    }
367    fn effects(&self) -> &EffectRow {
368        &self.effects
369    }
370    fn description(&self) -> &str {
371        "Find memories spatially near a query text using 5D holographic coordinates"
372    }
373    fn input_schema(&self) -> Value {
374        super::common::schema(
375            &json!({
376                "query": super::common::str_prop("Query text to locate in 5D space"),
377                "limit": super::common::int_prop("Maximum results (default 10)"),
378                "radius": super::common::num_prop("Spatial radius cutoff (optional)"),
379                "galaxy": super::common::str_prop("Galaxy (default: codex)"),
380            }),
381            &["query"],
382        )
383    }
384    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
385        let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
386        if query.is_empty() {
387            return Err(wm_core::CoreError::InvalidArgs(
388                "Missing 'query' parameter".into(),
389            ));
390        }
391        let galaxy_name_str = args
392            .get("galaxy")
393            .and_then(|v| v.as_str())
394            .unwrap_or("codex");
395        let galaxy = parse_galaxy(galaxy_name_str)?;
396        let radius = args
397            .get("radius")
398            .and_then(serde_json::Value::as_f64)
399            .unwrap_or(0.5) as f32;
400        let limit = args
401            .get("limit")
402            .and_then(serde_json::Value::as_u64)
403            .unwrap_or(20) as usize;
404
405        let center = wm_core::Coordinate5D::encode(query);
406        let memories = self.store.scan(galaxy, 1000)?;
407
408        let candidates: Vec<(usize, wm_core::Coordinate5D)> = memories
409            .iter()
410            .enumerate()
411            .map(|(i, m)| (i, m.metadata.coord5d.clone()))
412            .collect();
413
414        let nearby = wm_core::find_nearby(&center, &candidates, radius);
415
416        let results: Vec<Value> = nearby
417            .iter()
418            .filter(|(idx, _)| crate::expansion::common::mcp_visible(&memories[*idx]))
419            .filter(|(idx, _)| crate::expansion::common::validity_visible(&memories[*idx]))
420            .take(limit)
421            .map(|(idx, dist)| {
422                let mem = &memories[*idx];
423                json!({
424                    "id": mem.metadata.id,
425                    "content": mem.content.chars().take(100).collect::<String>(),
426                    "distance": dist,
427                    "zone": mem.metadata.coord5d.zone().name(),
428                    "importance": mem.metadata.importance,
429                    "tags": mem.metadata.tags,
430                })
431            })
432            .collect();
433
434        Ok(json!({
435            "status": "success",
436            "query": query,
437            "galaxy": galaxy_name_str,
438            "radius": radius,
439            "center": {
440                "x": center.x,
441                "y": center.y,
442                "z": center.z,
443                "w": center.w,
444                "v": center.v,
445            },
446            "found": results.len(),
447            "scanned": memories.len(),
448            "nearby": results,
449        }))
450    }
451    fn stats(&self) -> &ToolStats {
452        &self.stats
453    }
454}
455
456#[cfg(test)]
457mod tests {
458    use super::*;
459    use wm_memory::Memory;
460
461    fn open_store() -> (tempfile::TempDir, Arc<MemoryStore>) {
462        let tmp = tempfile::tempdir().unwrap();
463        let store = Arc::new(MemoryStore::open_default(tmp.path()).unwrap());
464        (tmp, store)
465    }
466
467    #[tokio::test]
468    async fn memory_tags_returns_unique_sorted_tags() {
469        let (_tmp, store) = open_store();
470        for tags in [["zeta", "alpha"], ["beta", "alpha"]] {
471            let mut memory = Memory::new(Galaxy::Codex, "invented tag fixture".into());
472            memory.metadata.tags = tags.into_iter().map(String::from).collect();
473            store.put(Galaxy::Codex, &memory).unwrap();
474        }
475        let result = MemoryTagsTool::new(store)
476            .call(&mut Context::default(), json!({"galaxy": "codex"}))
477            .await
478            .unwrap();
479        assert_eq!(result["unique_tags"], 3);
480        assert_eq!(result["tags"], json!(["alpha", "beta", "zeta"]));
481    }
482
483    #[tokio::test]
484    async fn citta_coherence_reflects_context_without_claiming_global_state() {
485        let mut context = Context {
486            citta_coherence: 0.29,
487            citta_valence: -0.2,
488            ..Context::default()
489        };
490        let result = CittaCoherenceTool::new()
491            .call(&mut context, json!({}))
492            .await
493            .unwrap();
494        assert!((result["coherence"].as_f64().unwrap() - 0.29).abs() < 1e-6);
495        assert_eq!(result["can_write"], false);
496        assert!((result["write_threshold"].as_f64().unwrap() - 0.3).abs() < 1e-6);
497    }
498
499    #[tokio::test]
500    async fn dharma_profiles_reports_static_descriptive_catalog() {
501        let result = DharmaProfilesTool::new()
502            .call(&mut Context::default(), json!({}))
503            .await
504            .unwrap();
505        assert_eq!(result["profiles"].as_array().unwrap().len(), 4);
506        assert_eq!(result["profiles"][0]["name"], "default");
507    }
508}