Skip to main content

persistence/
persistence.rs

1use antigravity_sdk_rust::agent::{Agent, AgentConfig};
2use antigravity_sdk_rust::policy;
3use antigravity_sdk_rust::types::GeminiConfig;
4use std::time::SystemTime;
5use tracing_subscriber::EnvFilter;
6
7#[tokio::main]
8async fn main() -> Result<(), anyhow::Error> {
9    // Initialize tracing subscriber
10    tracing_subscriber::fmt()
11        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()))
12        .init();
13
14    // Load environment variables from .env file if present
15    dotenvy::dotenv().ok();
16
17    // Create a temporary path for save_dir
18    let mut save_dir = std::env::temp_dir();
19    let epoch = SystemTime::now()
20        .duration_since(SystemTime::UNIX_EPOCH)
21        .unwrap_or_default()
22        .as_secs();
23    save_dir.push(format!("agent_session_{}", epoch));
24    let save_dir_str = save_dir.to_string_lossy().to_string();
25    println!("  Save directory: {}", save_dir_str);
26
27    // Ensure the save_dir exists
28    std::fs::create_dir_all(&save_dir)?;
29
30    println!("\n  === Session 1: establishing context ===");
31
32    let mut config1 = AgentConfig::default();
33    if let Ok(harness_path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
34        config1.binary_path = Some(harness_path);
35    }
36    let mut gemini_config1 = GeminiConfig::default();
37    if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
38        gemini_config1.api_key = Some(api_key);
39    }
40    gemini_config1.models.default.name = "gemini-3.5-flash".to_string();
41    config1.gemini_config = gemini_config1;
42    config1.policies = Some(vec![policy::allow_all()]);
43    config1.save_dir = Some(save_dir_str.clone());
44
45    let conversation_id = {
46        let mut agent1 = Agent::new(config1);
47        agent1.start().await?;
48
49        let prompt1 = "Remember this: my favorite color is blue.";
50        println!("  User: {}", prompt1);
51        let response1 = agent1.chat(prompt1).await?;
52        println!("  Agent: {}", response1.text);
53
54        let conv_id = agent1
55            .conversation_id()
56            .ok_or_else(|| anyhow::anyhow!("Failed to retrieve conversation ID"))?;
57        println!("  Assigned conversation ID: {}", conv_id);
58
59        agent1.stop().await?;
60        println!("  Session 1 ended.\n");
61        conv_id
62    };
63
64    println!("  === Session 2: resuming and verifying recall ===");
65
66    let mut config2 = AgentConfig::default();
67    if let Ok(harness_path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
68        config2.binary_path = Some(harness_path);
69    }
70    let mut gemini_config2 = GeminiConfig::default();
71    if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
72        gemini_config2.api_key = Some(api_key);
73    }
74    gemini_config2.models.default.name = "gemini-3.5-flash".to_string();
75    config2.gemini_config = gemini_config2;
76    config2.policies = Some(vec![policy::allow_all()]);
77    config2.save_dir = Some(save_dir_str.clone());
78    config2.conversation_id = Some(conversation_id);
79
80    let mut agent2 = Agent::new(config2);
81    agent2.start().await?;
82
83    let prompt2 = "What is my favorite color?";
84    println!("  User: {}", prompt2);
85    let response2 = agent2.chat(prompt2).await?;
86    println!("  Agent: {}", response2.text);
87
88    agent2.stop().await?;
89    println!("  Session 2 ended.");
90
91    // Clean up temporary session directory
92    let _ = std::fs::remove_dir_all(&save_dir);
93
94    Ok(())
95}