Skip to main content

codetether_agent/tool/
swarm_share.rs

1//! Swarm Share Tool - Publish and query shared results between sub-agents
2//!
3//! This tool allows sub-agents in a swarm to share intermediate results,
4//! enabling real-time collaboration between concurrent agents.
5
6use super::{Tool, ToolResult};
7use crate::swarm::result_store::ResultStore;
8use anyhow::{Context, Result};
9use async_trait::async_trait;
10use serde::Deserialize;
11use serde_json::{Value, json};
12use std::sync::Arc;
13
14pub struct SwarmShareTool {
15    store: Arc<ResultStore>,
16    /// The subtask ID of the agent using this tool
17    producer_id: String,
18}
19
20impl SwarmShareTool {
21    pub fn new(store: Arc<ResultStore>, producer_id: String) -> Self {
22        Self { store, producer_id }
23    }
24}
25
26#[derive(Deserialize)]
27struct Params {
28    action: String,
29    #[serde(default)]
30    key: Option<String>,
31    #[serde(default)]
32    value: Option<Value>,
33    #[serde(default)]
34    tags: Option<Vec<String>>,
35    #[serde(default)]
36    prefix: Option<String>,
37}
38
39#[async_trait]
40impl Tool for SwarmShareTool {
41    fn id(&self) -> &str {
42        "swarm_share"
43    }
44
45    fn name(&self) -> &str {
46        "Swarm Share"
47    }
48
49    fn description(&self) -> &str {
50        "Share results with other sub-agents in the swarm. Actions: publish (share a result), \
51         get (retrieve a result by key), query_tags (find results by tags), \
52         query_prefix (find results by key prefix), list (show all shared results)."
53    }
54
55    fn parameters(&self) -> Value {
56        json!({
57            "type": "object",
58            "properties": {
59                "action": {
60                    "type": "string",
61                    "enum": ["publish", "get", "query_tags", "query_prefix", "list"],
62                    "description": "Action to perform"
63                },
64                "key": {
65                    "type": "string",
66                    "description": "Result key (for publish/get)"
67                },
68                "value": {
69                    "description": "Result value to publish (any JSON value)"
70                },
71                "tags": {
72                    "type": "array",
73                    "items": {"type": "string"},
74                    "description": "Tags for publish or query_tags"
75                },
76                "prefix": {
77                    "type": "string",
78                    "description": "Key prefix for query_prefix"
79                },
80                "producer": {
81                    "type": "string",
82                    "description": "Producer subtask ID to filter by (for query)"
83                }
84            },
85            "required": ["action"]
86        })
87    }
88
89    async fn execute(&self, params: Value) -> Result<ToolResult> {
90        let p: Params = serde_json::from_value(params).context("Invalid params")?;
91
92        match p.action.as_str() {
93            "publish" => {
94                let key = p
95                    .key
96                    .ok_or_else(|| anyhow::anyhow!("key required for publish"))?;
97                let value = p
98                    .value
99                    .ok_or_else(|| anyhow::anyhow!("value required for publish"))?;
100                let tags = p.tags.unwrap_or_default();
101
102                let result = self
103                    .store
104                    .publish(&key, &self.producer_id, value, tags, None)
105                    .await?;
106
107                Ok(ToolResult::success(format!(
108                    "Published result '{}' (type: {})",
109                    key, result.schema.type_name
110                )))
111            }
112            "get" => {
113                let key = p
114                    .key
115                    .ok_or_else(|| anyhow::anyhow!("key required for get"))?;
116
117                match self.store.get(&key).await {
118                    Some(result) => {
119                        let output = json!({
120                            "key": result.key,
121                            "producer": result.producer_id,
122                            "value": result.value,
123                            "type": result.schema.type_name,
124                            "tags": result.tags,
125                            "published_at": result.published_at.to_rfc3339(),
126                        });
127                        Ok(ToolResult::success(
128                            serde_json::to_string_pretty(&output)
129                                .unwrap_or_else(|_| format!("{:?}", result.value)),
130                        ))
131                    }
132                    None => Ok(ToolResult::error(format!("No result found for key: {key}"))),
133                }
134            }
135            "query_tags" => {
136                let tags = p
137                    .tags
138                    .ok_or_else(|| anyhow::anyhow!("tags required for query_tags"))?;
139
140                let results = self.store.query_by_tags(&tags).await;
141                let output: Vec<Value> = results
142                    .iter()
143                    .map(|r| {
144                        json!({
145                            "key": r.key,
146                            "producer": r.producer_id,
147                            "type": r.schema.type_name,
148                            "tags": r.tags,
149                        })
150                    })
151                    .collect();
152
153                Ok(ToolResult::success(format!(
154                    "Found {} results matching tags {:?}:\n{}",
155                    output.len(),
156                    tags,
157                    serde_json::to_string_pretty(&output).unwrap_or_default()
158                )))
159            }
160            "query_prefix" => {
161                let prefix = p
162                    .prefix
163                    .ok_or_else(|| anyhow::anyhow!("prefix required for query_prefix"))?;
164
165                let results = self.store.query_by_prefix(&prefix).await;
166                let output: Vec<Value> = results
167                    .iter()
168                    .map(|r| {
169                        json!({
170                            "key": r.key,
171                            "producer": r.producer_id,
172                            "type": r.schema.type_name,
173                            "tags": r.tags,
174                        })
175                    })
176                    .collect();
177
178                Ok(ToolResult::success(format!(
179                    "Found {} results with prefix '{}':\n{}",
180                    output.len(),
181                    prefix,
182                    serde_json::to_string_pretty(&output).unwrap_or_default()
183                )))
184            }
185            "list" => {
186                let results = self.store.get_all().await;
187                if results.is_empty() {
188                    return Ok(ToolResult::success(
189                        "No shared results in store".to_string(),
190                    ));
191                }
192
193                let output: Vec<Value> = results
194                    .iter()
195                    .map(|r| {
196                        json!({
197                            "key": r.key,
198                            "producer": r.producer_id,
199                            "type": r.schema.type_name,
200                            "tags": r.tags,
201                        })
202                    })
203                    .collect();
204
205                Ok(ToolResult::success(format!(
206                    "{} shared results:\n{}",
207                    output.len(),
208                    serde_json::to_string_pretty(&output).unwrap_or_default()
209                )))
210            }
211            _ => Ok(ToolResult::error(format!("Unknown action: {}", p.action))),
212        }
213    }
214}