wm_tools/expansion/
agents.rs

1//! Agent tools — register, list, heartbeat.
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::{Memory, MemoryStore, SearchEngine, content_hash as compute_hash};
11
12pub struct AgentRegisterTool {
13    store: Arc<MemoryStore>,
14    stats: ToolStats,
15    effects: EffectRow,
16}
17
18impl AgentRegisterTool {
19    pub fn new(store: Arc<MemoryStore>) -> Self {
20        Self {
21            store,
22            stats: ToolStats::default(),
23            effects: EffectRow {
24                writes: vec![Resource::Galaxy("substrate".into())],
25                ..Default::default()
26            },
27        }
28    }
29}
30
31#[async_trait]
32impl Tool for AgentRegisterTool {
33    fn input_schema(&self) -> Value {
34        super::common::schema(
35            &json!({
36                "name": super::common::str_prop("Agent name to register (required; recorded in the Substrate galaxy)"),
37                "capabilities": super::common::str_array_prop("Capabilities to record for the agent (optional; default empty)"),
38            }),
39            &["name"],
40        )
41    }
42    fn name(&self) -> &str {
43        "agent.register"
44    }
45    fn gana(&self) -> Gana {
46        Gana::Room
47    }
48    fn effects(&self) -> &EffectRow {
49        &self.effects
50    }
51    fn description(&self) -> &str {
52        "Register a new agent in the Substrate galaxy"
53    }
54    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
55        // Fail closed: the meta-tool table already promises `name` required
56        // for this route, and an "unnamed" registration is not addressable
57        // by any read/update route (they look up by name or id).
58        let name = args
59            .get("name")
60            .and_then(|v| v.as_str())
61            .filter(|s| !s.is_empty())
62            .ok_or_else(|| {
63                wm_core::CoreError::InvalidArgs(
64                    "Missing 'name' parameter — agent.register requires a non-empty name".into(),
65                )
66            })?;
67        let capabilities = args.get("capabilities").cloned().unwrap_or(json!([]));
68        let mut mem = Memory::new(
69            Galaxy::Substrate,
70            json!({
71                "type": "agent_register",
72                "name": name,
73                "capabilities": capabilities,
74            })
75            .to_string(),
76        );
77        mem.metadata.tags = vec!["agent".into(), "registration".into()];
78        mem.metadata.importance = 0.8;
79        self.store.put(Galaxy::Substrate, &mem)?;
80        Ok(json!({
81            "status": "success",
82            "agent_id": mem.metadata.id,
83            "name": name,
84        }))
85    }
86    fn stats(&self) -> &ToolStats {
87        &self.stats
88    }
89}
90
91/// `agent.list` — list registered agents.
92pub struct AgentListTool {
93    store: Arc<MemoryStore>,
94    stats: ToolStats,
95    effects: EffectRow,
96}
97
98impl AgentListTool {
99    pub fn new(store: Arc<MemoryStore>) -> Self {
100        Self {
101            store,
102            stats: ToolStats::default(),
103            effects: EffectRow::read_only(vec![Resource::Galaxy("substrate".into())]),
104        }
105    }
106}
107
108#[async_trait]
109impl Tool for AgentListTool {
110    fn input_schema(&self) -> Value {
111        super::common::schema(
112            &json!({
113                "limit": super::common::int_prop("Maximum Substrate registrations to scan (optional; default 500, clamped 1-500)"),
114            }),
115            &[],
116        )
117    }
118    fn name(&self) -> &str {
119        "agent.list"
120    }
121    fn gana(&self) -> Gana {
122        Gana::Room
123    }
124    fn effects(&self) -> &EffectRow {
125        &self.effects
126    }
127    fn description(&self) -> &str {
128        "List all registered agents"
129    }
130    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
131        let limit = args
132            .get("limit")
133            .and_then(Value::as_u64)
134            .unwrap_or(500)
135            .clamp(1, 500) as usize;
136        let memories = self.store.scan(Galaxy::Substrate, limit)?;
137        let agents: Vec<Value> = memories
138            .iter()
139            .filter(|m| m.metadata.tags.contains(&"agent".to_string()))
140            .map(|m| {
141                json!({
142                    "id": m.metadata.id,
143                    "content": m.content,
144                    "tags": m.metadata.tags,
145                })
146            })
147            .collect();
148        Ok(json!({
149            "status": "success",
150            "count": agents.len(),
151            "agents": agents,
152        }))
153    }
154    fn stats(&self) -> &ToolStats {
155        &self.stats
156    }
157}
158
159/// `agent.heartbeat` — record an agent heartbeat.
160pub struct AgentHeartbeatTool {
161    store: Arc<MemoryStore>,
162    stats: ToolStats,
163    effects: EffectRow,
164}
165
166impl AgentHeartbeatTool {
167    pub fn new(store: Arc<MemoryStore>) -> Self {
168        Self {
169            store,
170            stats: ToolStats::default(),
171            effects: EffectRow {
172                writes: vec![Resource::Galaxy("substrate".into())],
173                ..Default::default()
174            },
175        }
176    }
177}
178
179#[async_trait]
180impl Tool for AgentHeartbeatTool {
181    fn input_schema(&self) -> Value {
182        super::common::schema(
183            &json!({
184                "agent_id": super::common::str_prop("Agent identifier to record the heartbeat for (optional; an empty value records an anonymous heartbeat)"),
185                "status": super::common::str_prop("Status string to record (optional; default 'alive')"),
186            }),
187            &[],
188        )
189    }
190    fn name(&self) -> &str {
191        "agent.heartbeat"
192    }
193    fn gana(&self) -> Gana {
194        Gana::Room
195    }
196    fn effects(&self) -> &EffectRow {
197        &self.effects
198    }
199    fn description(&self) -> &str {
200        "Record an agent heartbeat"
201    }
202    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
203        let agent_id = args.get("agent_id").and_then(|v| v.as_str()).unwrap_or("");
204        let status = args
205            .get("status")
206            .and_then(|v| v.as_str())
207            .unwrap_or("alive");
208        let mem = Memory::new(
209            Galaxy::Substrate,
210            json!({
211                "type": "heartbeat",
212                "agent_id": agent_id,
213                "status": status,
214            })
215            .to_string(),
216        );
217        self.store.put(Galaxy::Substrate, &mem)?;
218        Ok(json!({
219            "status": "success",
220            "agent_id": agent_id,
221            "heartbeat_id": mem.metadata.id,
222        }))
223    }
224    fn stats(&self) -> &ToolStats {
225        &self.stats
226    }
227}
228
229/// `agent.trust` — get or set trust level for an agent.
230///
231/// If `trust_level` is provided, updates the agent's trust level (0.0–1.0).
232/// Otherwise, returns the current trust level.
233pub struct AgentTrustTool {
234    store: Arc<MemoryStore>,
235    stats: ToolStats,
236    effects: EffectRow,
237}
238
239impl AgentTrustTool {
240    pub fn new(store: Arc<MemoryStore>) -> Self {
241        Self {
242            store,
243            stats: ToolStats::default(),
244            effects: EffectRow {
245                writes: vec![Resource::Galaxy("substrate".into())],
246                reads: vec![Resource::Galaxy("substrate".into())],
247                ..Default::default()
248            },
249        }
250    }
251}
252
253#[async_trait]
254impl Tool for AgentTrustTool {
255    fn input_schema(&self) -> Value {
256        super::common::schema(
257            &json!({
258                "agent_id": super::common::str_prop("Agent UUID or registered name to look up (required)"),
259                "trust_level": super::common::num_prop("Trust level to set (0.0-1.0); omit to read the current level"),
260            }),
261            &["agent_id"],
262        )
263    }
264    fn name(&self) -> &str {
265        "agent.trust"
266    }
267    fn gana(&self) -> Gana {
268        Gana::Room
269    }
270    fn effects(&self) -> &EffectRow {
271        &self.effects
272    }
273    fn description(&self) -> &str {
274        "Get or set trust level for an agent (0.0–1.0)"
275    }
276    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
277        let agent_id = args
278            .get("agent_id")
279            .and_then(|v| v.as_str())
280            .ok_or_else(|| {
281                wm_core::CoreError::InvalidArgs("Missing 'agent_id' parameter".into())
282            })?;
283
284        let memories = self.store.scan(Galaxy::Substrate, 500)?;
285        let agent_mem = memories
286            .iter()
287            .filter(|m| m.metadata.tags.contains(&"agent".to_string()))
288            .find(|m| {
289                if let Ok(data) = serde_json::from_str::<serde_json::Value>(&m.content) {
290                    data.get("name").and_then(|v| v.as_str()) == Some(agent_id)
291                        || m.metadata.id.to_string() == agent_id
292                } else {
293                    false
294                }
295            })
296            .ok_or_else(|| wm_core::CoreError::NotFound(format!("Agent '{agent_id}' not found")))?;
297
298        let mut data: serde_json::Value =
299            serde_json::from_str(&agent_mem.content).map_err(|e| {
300                wm_core::CoreError::Memory(format!("Failed to parse agent record: {e}"))
301            })?;
302
303        if let Some(trust) = args.get("trust_level").and_then(serde_json::Value::as_f64) {
304            let trust = trust.clamp(0.0, 1.0);
305            data["trust_level"] = json!(trust);
306            let mut updated = Memory::new(Galaxy::Substrate, data.to_string());
307            updated.metadata.id = agent_mem.metadata.id;
308            updated.metadata.tags.clone_from(&agent_mem.metadata.tags);
309            updated.metadata.importance = agent_mem.metadata.importance;
310            updated.metadata.content_hash = compute_hash(&updated.content);
311            self.store.put(Galaxy::Substrate, &updated)?;
312            Ok(json!({
313                "status": "success",
314                "agent_id": agent_id,
315                "trust_level": trust,
316                "action": "set",
317            }))
318        } else {
319            let trust = data
320                .get("trust_level")
321                .and_then(serde_json::Value::as_f64)
322                .unwrap_or(0.5);
323            Ok(json!({
324                "status": "success",
325                "agent_id": agent_id,
326                "trust_level": trust,
327                "action": "get",
328            }))
329        }
330    }
331    fn stats(&self) -> &ToolStats {
332        &self.stats
333    }
334}
335
336/// `agent.descriptions` — get or set description for an agent.
337pub struct AgentDescriptionsTool {
338    store: Arc<MemoryStore>,
339    stats: ToolStats,
340    effects: EffectRow,
341}
342
343impl AgentDescriptionsTool {
344    pub fn new(store: Arc<MemoryStore>) -> Self {
345        Self {
346            store,
347            stats: ToolStats::default(),
348            effects: EffectRow {
349                writes: vec![Resource::Galaxy("substrate".into())],
350                reads: vec![Resource::Galaxy("substrate".into())],
351                ..Default::default()
352            },
353        }
354    }
355}
356
357#[async_trait]
358impl Tool for AgentDescriptionsTool {
359    fn input_schema(&self) -> Value {
360        super::common::schema(
361            &json!({
362                "agent_id": super::common::str_prop("Agent UUID or registered name to look up (required)"),
363                "description": super::common::str_prop("Description to set; omit to read the current description"),
364            }),
365            &["agent_id"],
366        )
367    }
368    fn name(&self) -> &str {
369        "agent.descriptions"
370    }
371    fn gana(&self) -> Gana {
372        Gana::Room
373    }
374    fn effects(&self) -> &EffectRow {
375        &self.effects
376    }
377    fn description(&self) -> &str {
378        "Get or set description for an agent"
379    }
380    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
381        let agent_id = args
382            .get("agent_id")
383            .and_then(|v| v.as_str())
384            .ok_or_else(|| {
385                wm_core::CoreError::InvalidArgs("Missing 'agent_id' parameter".into())
386            })?;
387
388        let memories = self.store.scan(Galaxy::Substrate, 500)?;
389        let agent_mem = memories
390            .iter()
391            .filter(|m| m.metadata.tags.contains(&"agent".to_string()))
392            .find(|m| {
393                if let Ok(data) = serde_json::from_str::<serde_json::Value>(&m.content) {
394                    data.get("name").and_then(|v| v.as_str()) == Some(agent_id)
395                        || m.metadata.id.to_string() == agent_id
396                } else {
397                    false
398                }
399            })
400            .ok_or_else(|| wm_core::CoreError::NotFound(format!("Agent '{agent_id}' not found")))?;
401
402        let mut data: serde_json::Value =
403            serde_json::from_str(&agent_mem.content).map_err(|e| {
404                wm_core::CoreError::Memory(format!("Failed to parse agent record: {e}"))
405            })?;
406
407        if let Some(desc) = args.get("description").and_then(|v| v.as_str()) {
408            data["description"] = json!(desc);
409            let mut updated = Memory::new(Galaxy::Substrate, data.to_string());
410            updated.metadata.id = agent_mem.metadata.id;
411            updated.metadata.tags.clone_from(&agent_mem.metadata.tags);
412            updated.metadata.importance = agent_mem.metadata.importance;
413            updated.metadata.content_hash = compute_hash(&updated.content);
414            self.store.put(Galaxy::Substrate, &updated)?;
415            Ok(json!({
416                "status": "success",
417                "agent_id": agent_id,
418                "description": desc,
419                "action": "set",
420            }))
421        } else {
422            let desc = data
423                .get("description")
424                .and_then(|v| v.as_str())
425                .unwrap_or("");
426            Ok(json!({
427                "status": "success",
428                "agent_id": agent_id,
429                "description": desc,
430                "action": "get",
431            }))
432        }
433    }
434    fn stats(&self) -> &ToolStats {
435        &self.stats
436    }
437}
438
439/// `agent.capabilities` — get or set capabilities for an agent.
440pub struct AgentCapabilitiesTool {
441    store: Arc<MemoryStore>,
442    stats: ToolStats,
443    effects: EffectRow,
444}
445
446impl AgentCapabilitiesTool {
447    pub fn new(store: Arc<MemoryStore>) -> Self {
448        Self {
449            store,
450            stats: ToolStats::default(),
451            effects: EffectRow {
452                writes: vec![Resource::Galaxy("substrate".into())],
453                reads: vec![Resource::Galaxy("substrate".into())],
454                ..Default::default()
455            },
456        }
457    }
458}
459
460#[async_trait]
461impl Tool for AgentCapabilitiesTool {
462    fn input_schema(&self) -> Value {
463        super::common::schema(
464            &json!({
465                "agent_id": super::common::str_prop("Agent UUID or registered name to look up (required)"),
466                "capabilities": super::common::str_array_prop("Capabilities to set; omit to read the current capabilities"),
467            }),
468            &["agent_id"],
469        )
470    }
471    fn name(&self) -> &str {
472        "agent.capabilities"
473    }
474    fn gana(&self) -> Gana {
475        Gana::Room
476    }
477    fn effects(&self) -> &EffectRow {
478        &self.effects
479    }
480    fn description(&self) -> &str {
481        "Get or set capabilities for an agent"
482    }
483    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
484        let agent_id = args
485            .get("agent_id")
486            .and_then(|v| v.as_str())
487            .ok_or_else(|| {
488                wm_core::CoreError::InvalidArgs("Missing 'agent_id' parameter".into())
489            })?;
490
491        let memories = self.store.scan(Galaxy::Substrate, 500)?;
492        let agent_mem = memories
493            .iter()
494            .filter(|m| m.metadata.tags.contains(&"agent".to_string()))
495            .find(|m| {
496                if let Ok(data) = serde_json::from_str::<serde_json::Value>(&m.content) {
497                    data.get("name").and_then(|v| v.as_str()) == Some(agent_id)
498                        || m.metadata.id.to_string() == agent_id
499                } else {
500                    false
501                }
502            })
503            .ok_or_else(|| wm_core::CoreError::NotFound(format!("Agent '{agent_id}' not found")))?;
504
505        let mut data: serde_json::Value =
506            serde_json::from_str(&agent_mem.content).map_err(|e| {
507                wm_core::CoreError::Memory(format!("Failed to parse agent record: {e}"))
508            })?;
509
510        if let Some(caps) = args.get("capabilities").and_then(|v| v.as_array()) {
511            data["capabilities"] = json!(caps);
512            let mut updated = Memory::new(Galaxy::Substrate, data.to_string());
513            updated.metadata.id = agent_mem.metadata.id;
514            updated.metadata.tags.clone_from(&agent_mem.metadata.tags);
515            updated.metadata.importance = agent_mem.metadata.importance;
516            updated.metadata.content_hash = compute_hash(&updated.content);
517            self.store.put(Galaxy::Substrate, &updated)?;
518            Ok(json!({
519                "status": "success",
520                "agent_id": agent_id,
521                "capabilities": caps,
522                "action": "set",
523            }))
524        } else {
525            let caps = data.get("capabilities").cloned().unwrap_or(json!([]));
526            Ok(json!({
527                "status": "success",
528                "agent_id": agent_id,
529                "capabilities": caps,
530                "action": "get",
531            }))
532        }
533    }
534    fn stats(&self) -> &ToolStats {
535        &self.stats
536    }
537}
538
539/// `agent.heartbeat.history` — retrieve heartbeat history for an agent.
540pub struct AgentHeartbeatHistoryTool {
541    store: Arc<MemoryStore>,
542    stats: ToolStats,
543    effects: EffectRow,
544}
545
546impl AgentHeartbeatHistoryTool {
547    pub fn new(store: Arc<MemoryStore>) -> Self {
548        Self {
549            store,
550            stats: ToolStats::default(),
551            effects: EffectRow::read_only(vec![Resource::Galaxy("substrate".into())]),
552        }
553    }
554}
555
556#[async_trait]
557impl Tool for AgentHeartbeatHistoryTool {
558    fn input_schema(&self) -> Value {
559        super::common::schema(
560            &json!({
561                "agent_id": super::common::str_prop("Agent identifier whose heartbeats to list (required)"),
562                "limit": super::common::int_prop("Maximum heartbeats to return (optional; default 50)"),
563            }),
564            &["agent_id"],
565        )
566    }
567    fn name(&self) -> &str {
568        "agent.heartbeat.history"
569    }
570    fn gana(&self) -> Gana {
571        Gana::Room
572    }
573    fn effects(&self) -> &EffectRow {
574        &self.effects
575    }
576    fn description(&self) -> &str {
577        "Retrieve heartbeat history for an agent"
578    }
579    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
580        let agent_id = args
581            .get("agent_id")
582            .and_then(|v| v.as_str())
583            .ok_or_else(|| {
584                wm_core::CoreError::InvalidArgs("Missing 'agent_id' parameter".into())
585            })?;
586        let limit = args
587            .get("limit")
588            .and_then(serde_json::Value::as_u64)
589            .unwrap_or(50) as usize;
590
591        let memories = self.store.scan(Galaxy::Substrate, 10_000)?;
592        let mut heartbeats: Vec<Value> = memories
593            .iter()
594            .filter(|m| {
595                if let Ok(data) = serde_json::from_str::<serde_json::Value>(&m.content) {
596                    data.get("type").and_then(|v| v.as_str()) == Some("heartbeat")
597                        && data.get("agent_id").and_then(|v| v.as_str()) == Some(agent_id)
598                } else {
599                    false
600                }
601            })
602            .map(|m| {
603                json!({
604                    "id": m.metadata.id,
605                    "timestamp": m.metadata.created_at.to_rfc3339(),
606                    "content": m.content,
607                })
608            })
609            .collect();
610        heartbeats.truncate(limit);
611
612        Ok(json!({
613            "status": "success",
614            "agent_id": agent_id,
615            "count": heartbeats.len(),
616            "heartbeats": heartbeats,
617        }))
618    }
619    fn stats(&self) -> &ToolStats {
620        &self.stats
621    }
622}
623
624/// `agent.deregister` — remove an agent from the registry.
625pub struct AgentDeregisterTool {
626    store: Arc<MemoryStore>,
627    search: Option<Arc<SearchEngine>>,
628    stats: ToolStats,
629    effects: EffectRow,
630}
631
632impl AgentDeregisterTool {
633    pub fn new(store: Arc<MemoryStore>, search: Option<Arc<SearchEngine>>) -> Self {
634        Self {
635            store,
636            search,
637            stats: ToolStats::default(),
638            effects: EffectRow {
639                writes: vec![Resource::Galaxy("substrate".into())],
640                reads: vec![Resource::Galaxy("substrate".into())],
641                ..Default::default()
642            },
643        }
644    }
645}
646
647#[async_trait]
648impl Tool for AgentDeregisterTool {
649    fn input_schema(&self) -> Value {
650        super::common::schema(
651            &json!({
652                "agent_id": super::common::str_prop("Agent UUID or registered name to deregister (required)"),
653            }),
654            &["agent_id"],
655        )
656    }
657    fn name(&self) -> &str {
658        "agent.deregister"
659    }
660    fn gana(&self) -> Gana {
661        Gana::Room
662    }
663    fn effects(&self) -> &EffectRow {
664        &self.effects
665    }
666    fn description(&self) -> &str {
667        "Deregister an agent (removes registration record)"
668    }
669    async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
670        let agent_id = args
671            .get("agent_id")
672            .and_then(|v| v.as_str())
673            .ok_or_else(|| {
674                wm_core::CoreError::InvalidArgs("Missing 'agent_id' parameter".into())
675            })?;
676
677        let memories = self.store.scan(Galaxy::Substrate, 500)?;
678        let agent_mem = memories
679            .iter()
680            .filter(|m| m.metadata.tags.contains(&"agent".to_string()))
681            .find(|m| {
682                if let Ok(data) = serde_json::from_str::<serde_json::Value>(&m.content) {
683                    data.get("name").and_then(|v| v.as_str()) == Some(agent_id)
684                        || m.metadata.id.to_string() == agent_id
685                } else {
686                    false
687                }
688            })
689            .ok_or_else(|| wm_core::CoreError::NotFound(format!("Agent '{agent_id}' not found")))?;
690
691        self.store
692            .delete(Galaxy::Substrate, agent_mem.metadata.id)?;
693        super::common::deindex(self.search.as_deref(), &agent_mem.metadata.id.to_string());
694
695        Ok(json!({
696            "status": "success",
697            "agent_id": agent_id,
698            "deregistered": true,
699        }))
700    }
701    fn stats(&self) -> &ToolStats {
702        &self.stats
703    }
704}
705
706#[cfg(test)]
707mod tests {
708    use super::*;
709    use wm_memory::MemoryStore;
710
711    fn test_store() -> Arc<MemoryStore> {
712        let dir = tempfile::tempdir().unwrap();
713        Arc::new(MemoryStore::open_default(dir.path()).unwrap())
714    }
715
716    #[tokio::test]
717    async fn agent_trust_get_default() {
718        let store = test_store();
719        let reg = AgentRegisterTool::new(store.clone());
720        let mut ctx = Context::default();
721        let _ = reg.call(&mut ctx, json!({"name": "worker-1"})).await;
722
723        let tool = AgentTrustTool::new(store);
724        let result = tool.call(&mut ctx, json!({"agent_id": "worker-1"})).await;
725        assert!(result.is_ok());
726        let v = result.unwrap();
727        assert_eq!(v["action"], "get");
728        assert_eq!(v["trust_level"], 0.5);
729    }
730
731    #[tokio::test]
732    async fn agent_trust_set_and_get() {
733        let store = test_store();
734        let reg = AgentRegisterTool::new(store.clone());
735        let mut ctx = Context::default();
736        let _ = reg.call(&mut ctx, json!({"name": "worker-2"})).await;
737
738        let tool = AgentTrustTool::new(store);
739        let _ = tool
740            .call(
741                &mut ctx,
742                json!({"agent_id": "worker-2", "trust_level": 0.9}),
743            )
744            .await;
745
746        let result = tool.call(&mut ctx, json!({"agent_id": "worker-2"})).await;
747        assert!(result.is_ok());
748        let v = result.unwrap();
749        assert_eq!(v["action"], "get");
750        assert_eq!(v["trust_level"], 0.9);
751    }
752
753    #[tokio::test]
754    async fn agent_descriptions_set_and_get() {
755        let store = test_store();
756        let reg = AgentRegisterTool::new(store.clone());
757        let mut ctx = Context::default();
758        let _ = reg.call(&mut ctx, json!({"name": "worker-3"})).await;
759
760        let tool = AgentDescriptionsTool::new(store);
761        let _ = tool
762            .call(
763                &mut ctx,
764                json!({"agent_id": "worker-3", "description": "A test agent"}),
765            )
766            .await;
767
768        let result = tool.call(&mut ctx, json!({"agent_id": "worker-3"})).await;
769        assert!(result.is_ok());
770        let v = result.unwrap();
771        assert_eq!(v["description"], "A test agent");
772    }
773
774    #[tokio::test]
775    async fn agent_capabilities_set_and_get() {
776        let store = test_store();
777        let reg = AgentRegisterTool::new(store.clone());
778        let mut ctx = Context::default();
779        let _ = reg
780            .call(
781                &mut ctx,
782                json!({"name": "worker-4", "capabilities": ["read"]}),
783            )
784            .await;
785
786        let tool = AgentCapabilitiesTool::new(store);
787        let result = tool.call(&mut ctx, json!({"agent_id": "worker-4"})).await;
788        assert!(result.is_ok());
789        let v = result.unwrap();
790        assert_eq!(v["capabilities"], json!(["read"]));
791
792        let _ = tool
793            .call(
794                &mut ctx,
795                json!({"agent_id": "worker-4", "capabilities": ["read", "write"]}),
796            )
797            .await;
798        let result = tool.call(&mut ctx, json!({"agent_id": "worker-4"})).await;
799        let v = result.unwrap();
800        assert_eq!(v["capabilities"], json!(["read", "write"]));
801    }
802
803    #[tokio::test]
804    async fn agent_heartbeat_history() {
805        let store = test_store();
806        let reg = AgentRegisterTool::new(store.clone());
807        let mut ctx = Context::default();
808        let _ = reg.call(&mut ctx, json!({"name": "worker-5"})).await;
809
810        let hb = AgentHeartbeatTool::new(store.clone());
811        let _ = hb
812            .call(&mut ctx, json!({"agent_id": "worker-5", "status": "alive"}))
813            .await;
814        let _ = hb
815            .call(&mut ctx, json!({"agent_id": "worker-5", "status": "busy"}))
816            .await;
817
818        let tool = AgentHeartbeatHistoryTool::new(store);
819        let result = tool.call(&mut ctx, json!({"agent_id": "worker-5"})).await;
820        assert!(result.is_ok());
821        let v = result.unwrap();
822        assert_eq!(v["count"], 2);
823    }
824
825    #[tokio::test]
826    async fn agent_deregister() {
827        let store = test_store();
828        let reg = AgentRegisterTool::new(store.clone());
829        let mut ctx = Context::default();
830        let _ = reg.call(&mut ctx, json!({"name": "worker-6"})).await;
831
832        let tool = AgentDeregisterTool::new(store.clone(), None);
833        let result = tool.call(&mut ctx, json!({"agent_id": "worker-6"})).await;
834        assert!(result.is_ok());
835        assert_eq!(result.unwrap()["deregistered"], true);
836
837        let list = AgentListTool::new(store);
838        let result = list.call(&mut ctx, json!({})).await;
839        let v = result.unwrap();
840        assert_eq!(v["count"], 0);
841    }
842
843    #[tokio::test]
844    async fn agent_register_requires_name() {
845        let store = test_store();
846        let tool = AgentRegisterTool::new(store);
847        let mut ctx = Context::default();
848
849        let err = tool
850            .call(&mut ctx, json!({}))
851            .await
852            .expect_err("register without a name must fail closed");
853        assert!(matches!(err, wm_core::CoreError::InvalidArgs(_)));
854        assert!(
855            format!("{err}").contains("name"),
856            "the typed error names the missing argument: {err}"
857        );
858
859        let empty = tool.call(&mut ctx, json!({"name": ""})).await;
860        assert!(empty.is_err(), "empty name must fail closed too");
861    }
862
863    #[tokio::test]
864    async fn agent_trust_not_found() {
865        let store = test_store();
866        let tool = AgentTrustTool::new(store);
867        let mut ctx = Context::default();
868        let result = tool
869            .call(&mut ctx, json!({"agent_id": "nonexistent"}))
870            .await;
871        assert!(result.is_err());
872    }
873}
wm_tools/expansion/agents.rs

wm_tools/expansion/
agents.rs
