Skip to main content

start_command/
formal_ai.rs

1//! The task loop of the experimental `--use-formal-ai` mode (issue #255).
2//! Mirrors js/src/lib/formal-ai.js.
3//!
4//! The natural-language request goes to Formal AI verbatim. Formal AI decides
5//! what to observe and which tool calls to make; start checks every call
6//! against the effect policy, runs approved calls, reports each receipt back
7//! and keeps evidence of what actually happened. The run ends when Formal AI
8//! gives its final answer, says it did not understand, or a bound, refusal,
9//! protocol error or cancellation stops it, and the exit status says which.
10
11use std::path::Path;
12use std::time::{Duration, Instant};
13
14use serde_json::{json, Value};
15
16use crate::formal_ai_client::{CompletionClient, ToolCall};
17use crate::formal_ai_policy::{
18    self as policy, Bounds, EffectDiff, PolicyContext, EXIT_BUDGET, EXIT_PROTOCOL, EXIT_REFUSED,
19    EXIT_UNAVAILABLE, EXIT_UNSUPPORTED,
20};
21use crate::formal_ai_tool::{format_tool_result, run_tool, ToolOptions};
22
23const PREVIEW_CHARS: usize = 2000;
24
25/// Exit status of a cancellation by signal name (128 + signal number).
26pub fn signal_exit(name: &str) -> i32 {
27    match name {
28        "SIGHUP" => 129,
29        "SIGTERM" => 143,
30        _ => 130,
31    }
32}
33
34fn preview(text: &str) -> String {
35    let count = text.chars().count();
36    if count > PREVIEW_CHARS {
37        format!(
38            "{}… [{} more chars]",
39            text.chars().take(PREVIEW_CHARS).collect::<String>(),
40            count - PREVIEW_CHARS
41        )
42    } else {
43        text.to_string()
44    }
45}
46
47/// Observed working-directory changes, in words.
48pub fn describe_effects(effects: &EffectDiff) -> String {
49    let parts: Vec<String> = [
50        ("added", &effects.added),
51        ("removed", &effects.removed),
52        ("changed", &effects.changed),
53    ]
54    .iter()
55    .filter(|(_, names)| !names.is_empty())
56    .map(|(key, names)| {
57        let more = if names.len() > 10 {
58            format!(" and {} more", names.len() - 10)
59        } else {
60            String::new()
61        };
62        let shown: Vec<&str> = names.iter().take(10).map(String::as_str).collect();
63        format!("{} {}{}", key, shown.join(", "), more)
64    })
65    .collect();
66    let text = if parts.is_empty() {
67        "no file changed".to_string()
68    } else {
69        parts.join("; ")
70    };
71    if effects.complete {
72        text
73    } else {
74        format!("{} (fingerprint incomplete)", text)
75    }
76}
77
78/// What one task run needs besides the Formal AI client.
79pub struct TaskOptions<'a> {
80    /// Natural-language request, never executed.
81    pub request: &'a str,
82    pub cwd: &'a Path,
83    /// START_* variables added to the tools' environment.
84    pub env: &'a [(String, String)],
85    pub bounds: &'a Bounds,
86    pub deadline: Instant,
87    /// Mirror tool output to this process's stdout/stderr.
88    pub mirror: bool,
89    pub note: &'a dyn Fn(&str),
90    /// Exit status of a received cancellation signal, if any.
91    pub cancel: &'a dyn Fn() -> Option<i32>,
92}
93
94/// The final state of a run.
95#[derive(Debug, Clone, PartialEq)]
96pub struct Outcome {
97    pub status: String,
98    pub exit_code: i32,
99    pub message: String,
100    pub remaining: Vec<String>,
101}
102
103impl Outcome {
104    pub fn new(
105        status: &str,
106        exit_code: i32,
107        message: impl Into<String>,
108        remaining: &[&str],
109    ) -> Self {
110        Self {
111            status: status.to_string(),
112            exit_code,
113            message: message.into(),
114            remaining: remaining.iter().map(|s| s.to_string()).collect(),
115        }
116    }
117
118    pub fn to_json(&self) -> Value {
119        json!({
120            "status": self.status,
121            "exitCode": self.exit_code,
122            "message": self.message,
123            "remainingObligations": self.remaining,
124        })
125    }
126}
127
128fn cancelled(code: i32) -> Outcome {
129    let name = crate::exit_reason::signal_name_for_exit_code(Some(code))
130        .unwrap_or_else(|| "SIGINT".to_string());
131    Outcome::new(
132        "cancelled",
133        code,
134        format!("cancelled by {}", name),
135        &["the task was interrupted; later steps did not run"],
136    )
137}
138
139fn steps_mut(evidence: &mut Value) -> &mut Vec<Value> {
140    if !evidence["steps"].is_array() {
141        evidence["steps"] = json!([]);
142    }
143    evidence["steps"].as_array_mut().expect("steps is an array")
144}
145
146/// Run one request to completion. Fills `evidence` (a JSON object) as it
147/// goes, sets `evidence.outcome` and returns the outcome.
148pub fn run_formal_ai_task(
149    options: &TaskOptions<'_>,
150    client: &mut dyn CompletionClient,
151    evidence: &mut Value,
152) -> Outcome {
153    if !evidence.is_object() {
154        *evidence = json!({});
155    }
156    evidence["request"] = json!(options.request);
157    evidence["cwd"] = json!(options.cwd.to_string_lossy());
158    evidence["effectsMode"] = json!(options.bounds.effects);
159    evidence["bounds"] = options.bounds.to_json();
160    evidence["steps"] = json!([]);
161    evidence["effects"] = EffectDiff::default().to_json();
162    evidence["outcome"] = Value::Null;
163    let outcome = task_loop(options, client, evidence);
164    evidence["outcome"] = outcome.to_json();
165    outcome
166}
167
168fn task_loop(
169    options: &TaskOptions<'_>,
170    client: &mut dyn CompletionClient,
171    evidence: &mut Value,
172) -> Outcome {
173    let bounds = options.bounds;
174    let mut context = PolicyContext::new(options.cwd, &bounds.effects);
175    let mut total = EffectDiff::default();
176    let mut messages = vec![json!({ "role": "user", "content": options.request })];
177    let mut last_exit: Option<i32> = None;
178    let mut tool_calls: u64 = 0;
179    let is_cancelled = || (options.cancel)().is_some();
180
181    for turn in 1..=bounds.max_steps {
182        if let Some(code) = (options.cancel)() {
183            return cancelled(code);
184        }
185        if Instant::now() >= options.deadline {
186            return Outcome::new(
187                "budget-exhausted",
188                EXIT_BUDGET,
189                "time budget exhausted",
190                &["the task did not finish within START_FORMAL_AI_TIMEOUT_MS"],
191            );
192        }
193        let response = match client.complete(&messages, &is_cancelled) {
194            Ok(response) => response,
195            Err(error) => {
196                if let Some(code) = (options.cancel)() {
197                    return cancelled(code);
198                }
199                let status = match error.exit_code {
200                    EXIT_BUDGET => "budget-exhausted",
201                    EXIT_UNAVAILABLE => "unavailable",
202                    _ => "protocol-error",
203                };
204                return Outcome::new(
205                    status,
206                    error.exit_code,
207                    error.message,
208                    &["Formal AI did not complete the task"],
209                );
210            }
211        };
212        steps_mut(evidence).push(json!({
213            "kind": "response",
214            "turn": turn,
215            "finishReason": response.finish_reason,
216            "handler": response.handler,
217            "content": preview(&response.content),
218            "toolCalls": response.tool_calls.iter().map(|c| c.command.clone()).collect::<Vec<_>>(),
219        }));
220        if response.tool_calls.is_empty() {
221            if !response.content.is_empty() {
222                (options.note)(&response.content);
223            }
224            if response.handler.as_deref() == Some("unknown") && tool_calls == 0 {
225                return Outcome::new(
226                    "unsupported",
227                    EXIT_UNSUPPORTED,
228                    "Formal AI did not understand the request; no tool ran",
229                    &["rephrase the request or perform it manually"],
230                );
231            }
232            let Some(last) = last_exit else {
233                return Outcome::new("answered", 0, "Formal AI answered without tools", &[]);
234            };
235            // The status reports what the tools did, not what Formal AI
236            // claims: a run of reads is never presented as a completed change.
237            let status = if last == 0 {
238                format!("finished after {} tool call(s)", tool_calls)
239            } else {
240                format!("finished; the last tool call exited with {}", last)
241            };
242            let remaining: &[&str] = if last == 0 {
243                &[]
244            } else {
245                &["the last tool call failed"]
246            };
247            return Outcome::new(
248                "completed",
249                last,
250                format!("{}; {}", status, describe_effects(&total)),
251                remaining,
252            );
253        }
254        let ids: Vec<String> = response
255            .tool_calls
256            .iter()
257            .enumerate()
258            .map(|(index, call)| {
259                call.id
260                    .clone()
261                    .unwrap_or_else(|| format!("call_{}_{}", turn, index))
262            })
263            .collect();
264        messages.push(json!({
265            "role": "assistant",
266            "content": response.content,
267            "tool_calls": response.tool_calls.iter().zip(&ids).map(|(call, id)| json!({
268                "id": id,
269                "type": "function",
270                "function": {
271                    "name": call.name,
272                    "arguments": json!({ "command": call.command }).to_string(),
273                },
274            })).collect::<Vec<_>>(),
275        }));
276        for (call, id) in response.tool_calls.iter().zip(&ids) {
277            tool_calls += 1;
278            if tool_calls > bounds.max_steps {
279                return Outcome::new(
280                    "budget-exhausted",
281                    EXIT_BUDGET,
282                    format!("tool call budget of {} exhausted", bounds.max_steps),
283                    &["the task did not finish within START_FORMAL_AI_MAX_STEPS"],
284                );
285            }
286            let result = run_call(options, call, &mut context, &mut total, evidence);
287            if let Some(code) = (options.cancel)() {
288                return cancelled(code);
289            }
290            let (exit_code, content) = match result {
291                Ok(done) => done,
292                Err(outcome) => return outcome,
293            };
294            last_exit = Some(exit_code);
295            messages.push(json!({ "role": "tool", "tool_call_id": id, "content": content }));
296        }
297    }
298    Outcome::new(
299        "budget-exhausted",
300        EXIT_BUDGET,
301        format!("step budget of {} exhausted", bounds.max_steps),
302        &["the task did not finish within START_FORMAL_AI_MAX_STEPS"],
303    )
304}
305
306/// Check and run one call: Ok((exit code, tool message)) or the outcome that
307/// ends the run.
308fn run_call(
309    options: &TaskOptions<'_>,
310    call: &ToolCall,
311    context: &mut PolicyContext,
312    total: &mut EffectDiff,
313    evidence: &mut Value,
314) -> Result<(i32, String), Outcome> {
315    let bounds = options.bounds;
316    let note = options.note;
317    let decision = policy::evaluate_command(&call.command, context);
318    let mut step = json!({ "kind": "tool", "command": call.command, "argv": decision.tokens });
319    if !decision.allowed {
320        let reason = decision.reason.clone().unwrap_or_default();
321        step["decision"] = json!("refused");
322        step["reason"] = json!(reason);
323        steps_mut(evidence).push(step);
324        let message = format!("refused `{}`: {}", call.command, reason);
325        note(&message);
326        return Err(Outcome::new(
327            "refused",
328            EXIT_REFUSED,
329            message,
330            &["the refused call did not run; nothing after it ran either"],
331        ));
332    }
333    step["decision"] = json!("allowed");
334    step["effect"] = json!(decision.kind);
335    note(&format!("$ {}", call.command));
336    let before = policy::fingerprint(options.cwd, bounds.max_fingerprint_entries);
337    let remaining = options.deadline.saturating_duration_since(Instant::now());
338    let limit = Duration::from_millis(bounds.tool_timeout_ms)
339        .min(remaining)
340        .max(Duration::from_millis(1));
341    let cancel = || (options.cancel)().is_some();
342    let receipt = run_tool(&ToolOptions {
343        argv: &decision.tokens,
344        cwd: options.cwd,
345        env: options.env,
346        timeout: limit,
347        max_output_bytes: bounds.max_output_bytes,
348        mirror: options.mirror,
349        cancel: &cancel,
350    });
351    let diff = policy::diff_fingerprints(
352        &before,
353        &policy::fingerprint(options.cwd, bounds.max_fingerprint_entries),
354    );
355    step["exitCode"] = json!(receipt.exit_code);
356    step["signal"] = json!(receipt.signal);
357    step["timedOut"] = json!(receipt.timed_out);
358    step["durationMs"] = json!(receipt.duration_ms);
359    step["truncated"] = json!(receipt.truncated);
360    step["spawnError"] = json!(receipt.spawn_error);
361    step["stdout"] = json!(preview(&receipt.stdout));
362    step["stderr"] = json!(preview(&receipt.stderr));
363    step["effects"] = diff.to_json();
364    total.added.extend(diff.added.iter().cloned());
365    total.removed.extend(diff.removed.iter().cloned());
366    total.changed.extend(diff.changed.iter().cloned());
367    total.complete = total.complete && diff.complete;
368    evidence["effects"] = total.to_json();
369    if receipt.truncated {
370        note(&format!(
371            "output truncated at START_FORMAL_AI_MAX_OUTPUT_BYTES={}",
372            bounds.max_output_bytes
373        ));
374    }
375    let cwd_text = options.cwd.to_string_lossy().to_string();
376    let result = if receipt.timed_out {
377        Err(Outcome::new(
378            "budget-exhausted",
379            EXIT_BUDGET,
380            format!(
381                "`{}` exceeded the tool time limit and was stopped",
382                call.command
383            ),
384            &["the stopped call may have done part of its work"],
385        ))
386    } else if decision.kind == Some("read") {
387        policy::record_observation(&decision, context);
388        if diff.has_effects() {
389            step["unexpectedEffects"] = json!(true);
390            let names: Vec<&str> = diff
391                .added
392                .iter()
393                .chain(&diff.removed)
394                .chain(&diff.changed)
395                .take(5)
396                .map(String::as_str)
397                .collect();
398            note(&format!(
399                "warning: the working directory changed during a read-only call ({})",
400                names.join(", ")
401            ));
402        }
403        Ok((
404            receipt.exit_code,
405            format_tool_result(&receipt, &call.command, &cwd_text),
406        ))
407    } else if receipt.exit_code == 0 {
408        let verification = policy::verify_change(&decision.tokens, options.cwd);
409        step["verification"] = verification.to_json();
410        if verification.verified {
411            for target in &decision.targets {
412                context.observed.insert(target.clone());
413            }
414            Ok((0, format_tool_result(&receipt, &call.command, &cwd_text)))
415        } else {
416            Err(Outcome::new(
417                "unverified",
418                EXIT_PROTOCOL,
419                format!("`{}` exited 0 but {}", call.command, verification.detail),
420                &["the change could not be confirmed by readback"],
421            ))
422        }
423    } else {
424        Ok((
425            receipt.exit_code,
426            format_tool_result(&receipt, &call.command, &cwd_text),
427        ))
428    };
429    steps_mut(evidence).push(step);
430    result
431}
432
433#[cfg(test)]
434#[path = "formal_ai_cases.rs"]
435mod formal_ai_cases;