policies/
policies.rs

1use antigravity_sdk_rust::agent::Agent;
2use antigravity_sdk_rust::policy::{self, Decision, Policy};
3use antigravity_sdk_rust::types::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 harness_path = std::env::var("ANTIGRAVITY_HARNESS_PATH").ok();
52    let api_key = std::env::var("GEMINI_API_KEY").ok();
53
54    let mut builder = Agent::builder();
55    if let Some(path) = harness_path {
56        builder = builder.binary_path(path);
57    }
58    if let Some(key) = api_key {
59        builder = builder.api_key(key);
60    }
61
62    let agent = builder
63        .default_model("gemini-3.5-flash")
64        // 1. Deny everything by default
65        .policy(policy::deny_all())
66        // 2. Allow listing directories
67        .policy(policy::allow("LIST_DIR"))
68        // 3. Allow running commands, but block dangerous 'rm' commands
69        .policy(Policy::new(
70            "RUN_COMMAND".to_string(),
71            Decision::Deny,
72            Some(Arc::new(block_rm_predicate)),
73            None,
74            "block-rm".to_string(),
75        ))
76        // Fallback: Allow general RUN_COMMAND calls if they don't match the rm block predicate
77        .policy(policy::allow("RUN_COMMAND"))
78        // 4. Allow editing/creating files, but ask the user first if it's a critical file
79        .policy(Policy::new(
80            "WRITE_TO_FILE".to_string(),
81            Decision::AskUser,
82            Some(Arc::new(critical_file_predicate)),
83            Some(Arc::new(programmatic_approval_handler)),
84            "ask-for-critical-writes".to_string(),
85        ))
86        .policy(policy::allow("WRITE_TO_FILE"))
87        .build();
88
89    println!("Starting agent...");
90    let agent = agent.start().await?;
91
92    println!("\n  Chatting with agent...");
93
94    // 1. Try a safe command (should be allowed)
95    let prompt1 = "List the files in the current directory.";
96    println!("\n  User: {}", prompt1);
97    let response1 = agent.chat(prompt1).await?;
98    println!("  Agent: {}", response1.text);
99
100    // 2. Try a dangerous command (should be denied by policy)
101    let prompt2 = "Delete all files using rm -rf.";
102    println!("\n  User: {}", prompt2);
103    let response2 = agent.chat(prompt2).await?;
104    println!("  Agent: {}", response2.text);
105
106    // 3. Try creating a critical file (triggers programmatic ask_user handler)
107    let prompt3 = "Create a new configuration file named production.key with content 'debug=true'.";
108    println!("\n  User: {}", prompt3);
109    let response3 = agent.chat(prompt3).await?;
110    println!("  Agent: {}", response3.text);
111
112    agent.stop().await?;
113    Ok(())
114}
policies/policies.rs

policies/
policies.rs
