Skip to main content

persistence/
persistence.rs

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