auto_storage_sqlx/
auto_storage_sqlx.rs

1//! Agent with automatic SQLx storage and migrations
2//!
3//! This example shows how to use the builder with SQLx storage.
4//! The builder automatically:
5//! 1. Creates SQLx storage from TOML config
6//! 2. Runs your custom migrations
7//! 3. Wires everything together
8//!
9//! All you provide is your handler and migrations!
10
11use a2a_agents::core::{AgentBuilder, BuildError};
12use a2a_rs::{
13    domain::{A2AError, Message, Part, Role, Task, TaskState},
14    port::AsyncMessageHandler,
15};
16use async_trait::async_trait;
17use uuid::Uuid;
18
19/// Simple echo handler
20#[derive(Clone)]
21struct PersistentEchoHandler;
22
23#[async_trait]
24impl AsyncMessageHandler for PersistentEchoHandler {
25    async fn process_message(
26        &self,
27        task_id: &str,
28        message: &Message,
29        _session_id: Option<&str>,
30    ) -> Result<Task, A2AError> {
31        let text = message
32            .parts
33            .iter()
34            .find_map(|part| part.get_text().map(|s| s.to_string()))
35            .unwrap_or_else(|| "No text provided".to_string());
36
37        let response = Message::builder()
38            .role(Role::Agent)
39            .parts(vec![Part::text(format!("Echo (from SQLx): {}", text))])
40            .message_id(Uuid::new_v4().to_string())
41            .context_id(message.context_id.clone())
42            .build();
43
44        Ok(Task::builder()
45            .id(task_id.to_string())
46            .context_id(message.context_id.clone())
47            .status(a2a_rs::domain::TaskStatus::new(
48                TaskState::Completed,
49                Some(response.clone()),
50            ))
51            .history(vec![message.clone(), response])
52            .build())
53    }
54
55    async fn validate_message(&self, message: &Message) -> Result<(), A2AError> {
56        if message.parts.is_empty() {
57            return Err(A2AError::ValidationError {
58                field: "parts".to_string(),
59                message: "Message must contain at least one part".to_string(),
60            });
61        }
62        Ok(())
63    }
64}
65
66#[tokio::main]
67async fn main() -> Result<(), BuildError> {
68    // Load .env file if present
69    dotenvy::dotenv().ok();
70
71    // Initialize logging
72    tracing_subscriber::fmt()
73        .with_env_filter(
74            tracing_subscriber::EnvFilter::from_default_env()
75                .add_directive(tracing::Level::INFO.into()),
76        )
77        .init();
78
79    println!("🚀 Starting agent with automatic SQLx storage");
80    println!("   Database URL from config: ${{DATABASE_URL}}");
81    println!("   Migrations will be run automatically");
82    println!();
83
84    // Optional: define custom migrations for your agent
85    // These will be run automatically before the agent starts
86    let migrations = &[
87        // Example migration - you can add agent-specific tables here
88        r#"
89        CREATE TABLE IF NOT EXISTS echo_stats (
90            id INTEGER PRIMARY KEY AUTOINCREMENT,
91            message_count INTEGER NOT NULL DEFAULT 0,
92            last_echo TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
93        )
94        "#,
95    ];
96
97    // The builder:
98    // 1. Reads TOML config (including database URL from env)
99    // 2. Creates SQLx connection pool
100    // 3. Runs all migrations (framework + your custom ones)
101    // 4. Wires storage with push notification support
102    // 5. Starts the servers
103    //
104    // All from configuration!
105    AgentBuilder::from_file("examples/auto_storage_sqlx.toml")?
106        .with_handler(PersistentEchoHandler)
107        .build_with_auto_storage_and_migrations(migrations)
108        .await?
109        .run()
110        .await
111        .map_err(|e| BuildError::RuntimeError(e.to_string()))?;
112
113    Ok(())
114}
auto_storage_sqlx/auto_storage_sqlx.rs

auto_storage_sqlx/
auto_storage_sqlx.rs
