use std::collections::HashMap;
use std::sync::Arc;
use chrono::Utc;
use reqwest::Client;
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::types::*;
pub struct A2aStore {
client: Client,
agents: Arc<RwLock<HashMap<String, RemoteAgent>>>,
tasks: Arc<RwLock<HashMap<String, Task>>>,
push_configs: Arc<RwLock<Vec<PushNotificationConfig>>>,
}
impl A2aStore {
pub fn new() -> Self {
Self {
client: Client::new(),
agents: Arc::new(RwLock::new(HashMap::new())),
tasks: Arc::new(RwLock::new(HashMap::new())),
push_configs: Arc::new(RwLock::new(Vec::new())),
}
}
pub async fn resolve_card(&self, base_url: &str) -> Result<AgentCard, String> {
let url = format!("{}/.well-known/agent.json", base_url.trim_end_matches('/'));
let resp = self.client.get(&url).send().await.map_err(|e| e.to_string())?;
if !resp.status().is_success() {
return Err(format!("HTTP {}", resp.status()));
}
resp.json::<AgentCard>().await.map_err(|e| e.to_string())
}
pub async fn register_agent(&self, base_url: String, card: AgentCard) -> RemoteAgent {
let agent = RemoteAgent {
agent_id: format!("agent_{}", Uuid::new_v4().simple()),
name: card.name.clone(),
base_url,
status: "connected".into(),
card: Some(card),
auth_mode: None,
last_seen_at: Some(Utc::now()),
created_at: Utc::now(),
};
self.agents.write().await.insert(agent.agent_id.clone(), agent.clone());
agent
}
pub async fn list_agents(&self) -> Vec<RemoteAgent> {
self.agents.read().await.values().cloned().collect()
}
pub async fn send_message(&self, agent_id: &str, text: &str) -> Result<Task, String> {
let agents = self.agents.read().await;
let agent = agents.get(agent_id).ok_or_else(|| format!("Agent not found: {}", agent_id))?;
let card = agent.card.as_ref().ok_or("No agent card")?;
let message_id = Uuid::new_v4().to_string();
let request = serde_json::json!({
"jsonrpc": "2.0",
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [{"text": text}],
"messageId": message_id,
}
},
"id": message_id,
});
let resp = self.client.post(&card.url).json(&request).send().await.map_err(|e| e.to_string())?;
let body: serde_json::Value = resp.json().await.map_err(|e| e.to_string())?;
if let Some(error) = body.get("error") {
return Err(format!("RPC error: {}", error));
}
let task: Task = serde_json::from_value(body["result"].clone()).map_err(|e| e.to_string())?;
self.tasks.write().await.insert(task.id.clone(), task.clone());
Ok(task)
}
pub async fn get_task(&self, task_id: &str) -> Option<Task> {
self.tasks.read().await.get(task_id).cloned()
}
pub async fn list_tasks(&self) -> Vec<Task> {
self.tasks.read().await.values().cloned().collect()
}
pub async fn cancel_task(&self, task_id: &str) -> Result<(), String> {
let mut tasks = self.tasks.write().await;
if let Some(task) = tasks.get_mut(task_id) {
task.status = TaskStatus { state: TaskState::Canceled, message: Some("Canceled by operator".into()) };
Ok(())
} else {
Err(format!("Task not found: {}", task_id))
}
}
pub async fn add_push_config(&self, agent_id: String, webhook_url: String, events: Vec<String>) -> PushNotificationConfig {
let config = PushNotificationConfig {
config_id: format!("push_{}", Uuid::new_v4().simple()),
agent_id, webhook_url, events, active: true, created_at: Utc::now(),
};
self.push_configs.write().await.push(config.clone());
config
}
pub async fn list_push_configs(&self) -> Vec<PushNotificationConfig> {
self.push_configs.read().await.clone()
}
}