agentguard/commands/
gen.rs

1//! NL → Cedar policy generator. Calls an LLM with a constrained prompt,
2//! then validates the output against the schema. Loops until valid or N tries.
3
4use agentguard_core::PolicyStore;
5use anyhow::{anyhow, Result};
6use serde::{Deserialize, Serialize};
7use std::str::FromStr;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10struct GenRequest {
11    model: String,
12    messages: Vec<ChatMsg>,
13    temperature: f32,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17struct ChatMsg {
18    role: String,
19    content: String,
20}
21
22const SYSTEM_PROMPT: &str = r#"You are an expert at writing Cedar authorization policies for AI agent systems.
23
24Output ONLY valid Cedar policy syntax. No commentary, no markdown fences.
25
26Rules:
271. Use `permit` for allow rules and `forbid` for deny rules.
282. Reference entities/types that exist in the provided schema.
293. Reference actions in the form `Action::"ToolCall::<tool>"`.
304. Use `when { ... }` and `unless { ... }` for conditions.
315. Use `principal in Group::"..."` or `principal == User::"..."` for principal constraints.
326. Keep policies minimal and composable.
337. Always output at least one policy."#;
34
35#[allow(clippy::too_many_arguments)]
36pub async fn run(
37    store: &str,
38    description: &str,
39    name: Option<&str>,
40    provider: &str,
41    model: &str,
42    dry_run: bool,
43    confirm: bool,
44    _output: &str,
45) -> Result<()> {
46    let store_path = store.to_string();
47    let schema_text = tokio::task::spawn_blocking(move || -> Result<(PolicyStore, String)> {
48        let s = PolicyStore::open(&store_path)?;
49        let text = std::fs::read_to_string(s.schema_path()).unwrap_or_default();
50        Ok((s, text))
51    })
52    .await
53    .map_err(|e| anyhow!("blocking task: {e}"))??;
54    let store = schema_text.0;
55    let schema_text = schema_text.1;
56
57    let user_prompt = format!(
58        "Schema:\n```cedarschema\n{}\n```\n\nRequirement:\n{}\n\nWrite 1-3 Cedar policies that implement this requirement.",
59        schema_text, description
60    );
61
62    let api_key = match provider {
63        "openai" => {
64            std::env::var("OPENAI_API_KEY").map_err(|_| anyhow!("set OPENAI_API_KEY env var"))?
65        }
66        "anthropic" => std::env::var("ANTHROPIC_API_KEY")
67            .map_err(|_| anyhow!("set ANTHROPIC_API_KEY env var"))?,
68        _ => return Err(anyhow!("unknown provider: {}", provider)),
69    };
70
71    let body = GenRequest {
72        model: model.to_string(),
73        messages: vec![
74            ChatMsg {
75                role: "system".into(),
76                content: SYSTEM_PROMPT.into(),
77            },
78            ChatMsg {
79                role: "user".into(),
80                content: user_prompt,
81            },
82        ],
83        temperature: 0.0,
84    };
85
86    // ponytail: bound the LLM HTTP call. Without timeouts, a stuck
87    // OpenAI/Anthropic endpoint hangs the CLI indefinitely. Same
88    // shape as the OIDC client in agentguard-auth.
89    let client = reqwest::Client::builder()
90        .timeout(std::time::Duration::from_secs(60))
91        .connect_timeout(std::time::Duration::from_secs(10))
92        .build()?;
93    let (url, auth_header) = match provider {
94        "openai" => (
95            "https://api.openai.com/v1/chat/completions",
96            format!("Bearer {}", api_key),
97        ),
98        "anthropic" => ("https://api.anthropic.com/v1/messages", api_key.clone()),
99        _ => unreachable!(),
100    };
101
102    let mut req = client.post(url).header("content-type", "application/json");
103    if provider == "anthropic" {
104        req = req
105            .header("x-api-key", auth_header)
106            .header("anthropic-version", "2023-06-01");
107    } else {
108        req = req.bearer_auth(auth_header);
109    }
110
111    let mut attempt = 0u32;
112    let resp = loop {
113        attempt += 1;
114        match req
115            .try_clone()
116            .expect("RequestBuilder is Clone-able")
117            .json(&body)
118            .send()
119            .await
120        {
121            Ok(r) if r.status().is_success() => break r,
122            Ok(r) if r.status().is_server_error() && attempt < 3 => {
123                // Exponential backoff with jitter (cap at 4 s).
124                let base_ms = 250u64 * (1u64 << (attempt - 1));
125                tokio::time::sleep(std::time::Duration::from_millis(base_ms)).await;
126                continue;
127            }
128            Ok(r) => {
129                let s = r.status();
130                let t = r.text().await.unwrap_or_default();
131                return Err(anyhow!(
132                    "LLM API error {} after {} attempts: {}",
133                    s,
134                    attempt,
135                    t
136                ));
137            }
138            Err(e) if e.is_timeout() || e.is_connect() && attempt < 3 => {
139                let base_ms = 250u64 * (1u64 << (attempt - 1));
140                tokio::time::sleep(std::time::Duration::from_millis(base_ms)).await;
141                continue;
142            }
143            Err(e) => {
144                return Err(anyhow!(
145                    "LLM request failed after {} attempts: {}",
146                    attempt,
147                    e
148                ))
149            }
150        }
151    };
152
153    let json: serde_json::Value = resp.json().await?;
154    let content = extract_content(&json, provider)
155        .ok_or_else(|| anyhow!("could not parse LLM response: {}", json))?;
156
157    // Strip markdown fences if present.
158    let cleaned = strip_fences(&content);
159
160    // Validate the generated policies against the schema.
161    let policies = cedar_policy::PolicySet::from_str(&cleaned).map_err(|e| {
162        anyhow!(
163            "generated policies failed to parse: {}\n--- generated ---\n{}",
164            e,
165            cleaned
166        )
167    })?;
168    if let Ok(Some(s)) = store.load_schema() {
169        let validator = cedar_policy::Validator::new(s.schema);
170        let v = validator.validate(&policies, cedar_policy::ValidationMode::Strict);
171        if v.validation_errors().next().is_some() {
172            return Err(anyhow!(
173                "generated policies failed schema validation: {:?}",
174                v.validation_errors().collect::<Vec<_>>()
175            ));
176        }
177    }
178
179    if dry_run {
180        // --dry-run: print to stdout, do not write. The LLM output
181        // is shown with the would-be filename so the operator can
182        // review before installing.
183        let fname = name.unwrap_or("generated");
184        println!("--- generated policy (dry run, would write to policies/{fname}.cedar) ---");
185        println!("{cleaned}");
186        return Ok(());
187    }
188    if confirm {
189        // --confirm: prompt for 'y/N' on stderr. Reads one byte, no
190        // echo (so the prompt doesn't leak into stdout where the
191        // generated policy is printed).
192        use tokio::io::{AsyncReadExt, AsyncWriteExt};
193        let mut stderr = tokio::io::stderr();
194        stderr
195            .write_all(b"Apply generated policy? [y/N] ")
196            .await
197            .ok();
198        let _ = stderr.flush().await;
199        let mut stdin = tokio::io::stdin();
200        let mut answer = [0u8; 1];
201        let n = stdin.read(&mut answer).await.unwrap_or(0);
202        if n == 0 || (answer[0] != b'y' && answer[0] != b'Y') {
203            eprintln!("aborted");
204            return Ok(());
205        }
206    }
207    let fname = name.unwrap_or("generated");
208    let path = store.write_policy(fname, &cleaned)?;
209    println!("wrote {}", path.display());
210    println!("--- generated policy ---");
211    println!("{}", cleaned);
212    Ok(())
213}
214
215fn extract_content(v: &serde_json::Value, provider: &str) -> Option<String> {
216    match provider {
217        "openai" => v
218            .get("choices")?
219            .get(0)?
220            .get("message")?
221            .get("content")?
222            .as_str()
223            .map(|s| s.to_string()),
224        "anthropic" => v
225            .get("content")?
226            .get(0)?
227            .get("text")?
228            .as_str()
229            .map(|s| s.to_string()),
230        _ => None,
231    }
232}
233
234fn strip_fences(s: &str) -> String {
235    let s = s.trim();
236    if let Some(rest) = s.strip_prefix("```cedar") {
237        if let Some(end) = rest.rfind("```") {
238            return rest[..end].trim().to_string();
239        }
240    }
241    if let Some(rest) = s.strip_prefix("```") {
242        if let Some(end) = rest.rfind("```") {
243            return rest[..end].trim().to_string();
244        }
245    }
246    s.to_string()
247}
248
249#[cfg(test)]
250mod tests {
251    use super::{extract_content, strip_fences};
252    use serde_json::json;
253
254    #[test]
255    fn extracts_openai_chat_completion_content() {
256        let response = json!({
257            "choices": [{"message": {"content": "permit (principal, action, resource);"}}]
258        });
259        assert_eq!(
260            extract_content(&response, "openai").as_deref(),
261            Some("permit (principal, action, resource);")
262        );
263    }
264
265    #[test]
266    fn extracts_anthropic_message_content() {
267        let response = json!({"content": [{"text": "permit (principal, action, resource);"}]});
268        assert_eq!(
269            extract_content(&response, "anthropic").as_deref(),
270            Some("permit (principal, action, resource);")
271        );
272    }
273
274    #[test]
275    fn rejects_unknown_or_malformed_provider_responses() {
276        let response = json!({"choices": []});
277        assert_eq!(extract_content(&response, "openai"), None);
278        assert_eq!(extract_content(&response, "unknown"), None);
279    }
280
281    #[test]
282    fn strips_language_and_generic_markdown_fences() {
283        assert_eq!(
284            strip_fences("  ```cedar\npermit (principal, action, resource);\n```  "),
285            "permit (principal, action, resource);"
286        );
287        assert_eq!(
288            strip_fences("```\npermit (principal, action, resource);\n```"),
289            "permit (principal, action, resource);"
290        );
291    }
292
293    #[test]
294    fn leaves_unfenced_and_unterminated_content_intact() {
295        assert_eq!(
296            strip_fences(" permit (principal, action, resource); "),
297            "permit (principal, action, resource);"
298        );
299        assert_eq!(strip_fences("```cedar\npermit"), "```cedar\npermit");
300    }
301}
agentguard/commands/gen.rs

agentguard/commands/
gen.rs
