Skip to main content

policies/
policies.rs

1use antigravity_sdk_rust::agent::{Agent, AgentConfig};
2use antigravity_sdk_rust::policy::{self, Decision, Policy};
3use antigravity_sdk_rust::types::{GeminiConfig, ToolCall};
4use serde_json::Value;
5use std::sync::Arc;
6use tracing_subscriber::EnvFilter;
7
8fn block_rm_predicate(tool_call: &ToolCall) -> bool {
9    tool_call
10        .args
11        .get("CommandLine")
12        .and_then(Value::as_str)
13        .is_some_and(|cmd| cmd.contains("rm"))
14}
15
16fn critical_file_predicate(tool_call: &ToolCall) -> bool {
17    tool_call
18        .args
19        .get("TargetFile")
20        .or_else(|| tool_call.args.get("target_file"))
21        .or_else(|| tool_call.args.get("path"))
22        .and_then(Value::as_str)
23        .is_some_and(|p| {
24            std::path::Path::new(p)
25                .extension()
26                .is_some_and(|ext| ext.eq_ignore_ascii_case("key"))
27                || p.contains("production")
28        })
29}
30
31fn programmatic_approval_handler(tool_call: &ToolCall) -> bool {
32    println!(
33        "\n  [ASK_USER Handler] Intercepted request for tool: {}",
34        tool_call.name
35    );
36    println!("  [ASK_USER Handler] Target arguments: {}", tool_call.args);
37    println!("  [ASK_USER Handler] Simulating user review... Decision: DENY.");
38    false
39}
40
41#[tokio::main]
42async fn main() -> Result<(), anyhow::Error> {
43    // Initialize tracing subscriber
44    tracing_subscriber::fmt()
45        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()))
46        .init();
47
48    // Load environment variables from .env file if present
49    dotenvy::dotenv().ok();
50
51    let mut config = AgentConfig::default();
52
53    if let Ok(harness_path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
54        config.binary_path = Some(harness_path);
55    }
56
57    let mut gemini_config = GeminiConfig::default();
58    if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
59        gemini_config.api_key = Some(api_key);
60    }
61    gemini_config.models.default.name = "gemini-3.5-flash".to_string();
62    config.gemini_config = gemini_config;
63
64    // Configure policies using the recommended "Deny by Default" posture.
65    let policies = vec![
66        // 1. Deny everything by default
67        policy::deny_all(),
68        // 2. Allow listing directories
69        policy::allow("LIST_DIR"),
70        // 3. Allow running commands, but block dangerous 'rm' commands
71        Policy::new(
72            "RUN_COMMAND".to_string(),
73            Decision::Deny,
74            Some(Arc::new(block_rm_predicate)),
75            None,
76            "block-rm".to_string(),
77        ),
78        // Fallback: Allow general RUN_COMMAND calls if they don't match the rm block predicate
79        policy::allow("RUN_COMMAND"),
80        // 4. Allow editing/creating files, but ask the user first if it's a critical file
81        Policy::new(
82            "WRITE_TO_FILE".to_string(),
83            Decision::AskUser,
84            Some(Arc::new(critical_file_predicate)),
85            Some(Arc::new(programmatic_approval_handler)),
86            "ask-for-critical-writes".to_string(),
87        ),
88        policy::allow("WRITE_TO_FILE"),
89    ];
90    config.policies = Some(policies);
91
92    let mut agent = Agent::new(config);
93    println!("Starting agent...");
94    agent.start().await?;
95
96    println!("\n  Chatting with agent...");
97
98    // 1. Try a safe command (should be allowed)
99    let prompt1 = "List the files in the current directory.";
100    println!("\n  User: {}", prompt1);
101    let response1 = agent.chat(prompt1).await?;
102    println!("  Agent: {}", response1.text);
103
104    // 2. Try a dangerous command (should be denied by policy)
105    let prompt2 = "Delete all files using rm -rf.";
106    println!("\n  User: {}", prompt2);
107    let response2 = agent.chat(prompt2).await?;
108    println!("  Agent: {}", response2.text);
109
110    // 3. Try creating a critical file (triggers programmatic ask_user handler)
111    let prompt3 = "Create a new configuration file named production.key with content 'debug=true'.";
112    println!("\n  User: {}", prompt3);
113    let response3 = agent.chat(prompt3).await?;
114    println!("  Agent: {}", response3.text);
115
116    agent.stop().await?;
117    Ok(())
118}