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| match part {
35                Part::Text { text, .. } => Some(text.clone()),
36                _ => None,
37            })
38            .unwrap_or_else(|| "No text provided".to_string());
39
40        let response = Message::builder()
41            .role(Role::Agent)
42            .parts(vec![Part::text(format!("Echo (from SQLx): {}", text))])
43            .message_id(Uuid::new_v4().to_string())
44            .context_id(message.context_id.clone().unwrap_or_default())
45            .build();
46
47        Ok(Task::builder()
48            .id(task_id.to_string())
49            .context_id(message.context_id.clone().unwrap_or_default())
50            .status(a2a_rs::domain::TaskStatus {
51                state: TaskState::Completed,
52                message: Some(response.clone()),
53                timestamp: Some(chrono::Utc::now()),
54            })
55            .history(vec![message.clone(), response])
56            .build())
57    }
58
59    async fn validate_message(&self, message: &Message) -> Result<(), A2AError> {
60        if message.parts.is_empty() {
61            return Err(A2AError::ValidationError {
62                field: "parts".to_string(),
63                message: "Message must contain at least one part".to_string(),
64            });
65        }
66        Ok(())
67    }
68}
69
70#[tokio::main]
71async fn main() -> Result<(), BuildError> {
72    // Load .env file if present
73    dotenvy::dotenv().ok();
74
75    // Initialize logging
76    tracing_subscriber::fmt()
77        .with_env_filter(
78            tracing_subscriber::EnvFilter::from_default_env()
79                .add_directive(tracing::Level::INFO.into()),
80        )
81        .init();
82
83    println!("🚀 Starting agent with automatic SQLx storage");
84    println!("   Database URL from config: ${{DATABASE_URL}}");
85    println!("   Migrations will be run automatically");
86    println!();
87
88    // Optional: define custom migrations for your agent
89    // These will be run automatically before the agent starts
90    let migrations = &[
91        // Example migration - you can add agent-specific tables here
92        r#"
93        CREATE TABLE IF NOT EXISTS echo_stats (
94            id INTEGER PRIMARY KEY AUTOINCREMENT,
95            message_count INTEGER NOT NULL DEFAULT 0,
96            last_echo TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
97        )
98        "#,
99    ];
100
101    // The builder:
102    // 1. Reads TOML config (including database URL from env)
103    // 2. Creates SQLx connection pool
104    // 3. Runs all migrations (framework + your custom ones)
105    // 4. Wires storage with push notification support
106    // 5. Starts the servers
107    //
108    // All from configuration!
109    AgentBuilder::from_file("examples/auto_storage_sqlx.toml")?
110        .with_handler(PersistentEchoHandler)
111        .build_with_auto_storage_and_migrations(migrations)
112        .await?
113        .run()
114        .await
115        .map_err(|e| BuildError::RuntimeError(e.to_string()))?;
116
117    Ok(())
118}
auto_storage_sqlx/auto_storage_sqlx.rs

auto_storage_sqlx/
auto_storage_sqlx.rs
