start_command/
formal_ai_client.rs

1//! Minimal client of the OpenAI-compatible agent-mode endpoint of Formal AI
2//! (`POST /api/openai/v1/chat/completions`), used by `--use-formal-ai`
3//! (issue #255). Mirrors js/src/lib/formal-ai-client.js. Only the `bash`
4//! function tool is offered: Formal AI picks the command, start decides
5//! whether and how it runs.
6//!
7//! The runtime is always a loopback `http://` server, so a small blocking
8//! HTTP/1.1 client over `std::net` is enough; it polls for cancellation and
9//! the time budget while it waits.
10
11use std::io::{Read, Write};
12use std::net::{TcpStream, ToSocketAddrs};
13use std::time::{Duration, Instant};
14
15use serde_json::{json, Value};
16
17use crate::formal_ai_policy::{EXIT_BUDGET, EXIT_PROTOCOL, EXIT_UNAVAILABLE};
18
19pub const API_PREFIX: &str = "/api/openai/v1";
20pub const TOOL_NAMES: [&str; 3] = ["bash", "shell", "run_command"];
21const MAX_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
22const POLL: Duration = Duration::from_millis(100);
23
24/// The only tool offered to Formal AI.
25pub fn bash_tool() -> Value {
26    json!({
27        "type": "function",
28        "function": {
29            "name": "bash",
30            "description": "Run one command in the working directory. Pipes, redirects, globs and variables are not supported.",
31            "parameters": {
32                "type": "object",
33                "properties": { "command": { "type": "string" } },
34                "required": ["command"],
35            },
36        },
37    })
38}
39
40/// A failed exchange with Formal AI and the exit status it maps to.
41#[derive(Debug, Clone, PartialEq)]
42pub struct ClientError {
43    pub message: String,
44    pub exit_code: i32,
45    pub cancelled: bool,
46}
47
48impl ClientError {
49    pub fn protocol(message: impl Into<String>) -> Self {
50        Self::with_code(message, EXIT_PROTOCOL)
51    }
52
53    pub fn with_code(message: impl Into<String>, exit_code: i32) -> Self {
54        Self {
55            message: message.into(),
56            exit_code,
57            cancelled: false,
58        }
59    }
60}
61
62/// One tool call Formal AI asked for.
63#[derive(Debug, Clone, PartialEq)]
64pub struct ToolCall {
65    pub id: Option<String>,
66    pub name: String,
67    pub command: String,
68}
69
70/// One chat completion, normalized for the task loop.
71#[derive(Debug, Clone, PartialEq, Default)]
72pub struct Completion {
73    pub content: String,
74    pub finish_reason: Option<String>,
75    pub tool_calls: Vec<ToolCall>,
76    pub handler: Option<String>,
77}
78
79fn value_text(value: &Value) -> String {
80    match value {
81        Value::String(text) => text.clone(),
82        other => other.to_string(),
83    }
84}
85
86/// Normalize one chat completion into what the loop needs.
87pub fn parse_completion(body: &Value) -> Result<Completion, ClientError> {
88    let choice = body
89        .get("choices")
90        .and_then(Value::as_array)
91        .and_then(|choices| choices.first())
92        .filter(|choice| choice.get("message").is_some_and(Value::is_object))
93        .ok_or_else(|| ClientError::protocol("Formal AI response has no choices[0].message"))?;
94    let message = &choice["message"];
95    let handler = message
96        .get("thinking_steps")
97        .and_then(Value::as_array)
98        .and_then(|steps| {
99            steps
100                .iter()
101                .find(|s| s.get("step").and_then(Value::as_str) == Some("dispatch_handler"))
102        })
103        .map(|step| value_text(step.get("detail").unwrap_or(&Value::Null)));
104    let mut tool_calls = Vec::new();
105    if let Some(calls) = message.get("tool_calls").and_then(Value::as_array) {
106        for call in calls {
107            let function = call.get("function");
108            let name = function
109                .and_then(|f| f.get("name"))
110                .and_then(Value::as_str)
111                .unwrap_or("");
112            if !TOOL_NAMES.contains(&name) {
113                let shown = if name.is_empty() { "undefined" } else { name };
114                return Err(ClientError::protocol(format!(
115                    "Formal AI called an unsupported tool \"{}\"",
116                    shown
117                )));
118            }
119            let raw = function
120                .and_then(|f| f.get("arguments"))
121                .and_then(Value::as_str)
122                .filter(|text| !text.is_empty())
123                .unwrap_or("{}");
124            let args: Value = serde_json::from_str(raw).map_err(|_| {
125                ClientError::protocol(format!("Formal AI sent non-JSON arguments for {}", name))
126            })?;
127            let command = args
128                .get("command")
129                .filter(|v| !v.is_null())
130                .or_else(|| args.get("cmd"))
131                .and_then(Value::as_str)
132                .filter(|text| !text.trim().is_empty())
133                .ok_or_else(|| {
134                    ClientError::protocol(format!("Formal AI called {} without a command", name))
135                })?;
136            tool_calls.push(ToolCall {
137                id: call
138                    .get("id")
139                    .and_then(Value::as_str)
140                    .filter(|id| !id.is_empty())
141                    .map(str::to_string),
142                name: name.to_string(),
143                command: command.to_string(),
144            });
145        }
146    }
147    Ok(Completion {
148        content: message
149            .get("content")
150            .and_then(Value::as_str)
151            .unwrap_or("")
152            .to_string(),
153        finish_reason: choice
154            .get("finish_reason")
155            .and_then(Value::as_str)
156            .filter(|reason| !reason.is_empty())
157            .map(str::to_string),
158        tool_calls,
159        handler,
160    })
161}
162
163/// Something that answers chat completions (HTTP, or a replay in tests).
164pub trait CompletionClient {
165    /// Send the conversation; `cancel` is polled while waiting.
166    fn complete(
167        &mut self,
168        messages: &[Value],
169        cancel: &dyn Fn() -> bool,
170    ) -> Result<Completion, ClientError>;
171    /// Requests sent so far (retries included).
172    fn requests(&self) -> u64;
173}
174
175/// Why an HTTP exchange failed.
176#[derive(Debug, Clone, PartialEq)]
177pub enum HttpError {
178    Timeout,
179    Cancelled,
180    Io(String),
181}
182
183/// Host, port and path prefix of an `http://host:port/prefix` URL.
184pub fn parse_http_url(url: &str) -> Result<(String, u16, String), String> {
185    let rest = url
186        .strip_prefix("http://")
187        .ok_or_else(|| format!("not an http:// URL: {}", url))?;
188    let (authority, prefix) = match rest.find('/') {
189        Some(index) => (&rest[..index], rest[index..].trim_end_matches('/')),
190        None => (rest, ""),
191    };
192    let (host, port) = if let Some(stripped) = authority.strip_prefix('[') {
193        let end = stripped
194            .find(']')
195            .ok_or_else(|| format!("invalid host in {}", url))?;
196        let port = stripped[end + 1..].strip_prefix(':').unwrap_or("80");
197        (stripped[..end].to_string(), port)
198    } else {
199        match authority.rsplit_once(':') {
200            Some((host, port)) => (host.to_string(), port),
201            None => (authority.to_string(), "80"),
202        }
203    };
204    if host.is_empty() {
205        return Err(format!("missing host in {}", url));
206    }
207    let port = port
208        .parse::<u16>()
209        .map_err(|_| format!("invalid port in {}", url))?;
210    Ok((host, port, prefix.to_string()))
211}
212
213fn header_end(buf: &[u8]) -> Option<usize> {
214    buf.windows(4).position(|w| w == b"\r\n\r\n").map(|i| i + 4)
215}
216
217/// Decode a chunked body; None while it is incomplete.
218fn decode_chunked(mut body: &[u8]) -> Option<Vec<u8>> {
219    let mut out = Vec::new();
220    loop {
221        let line_end = body.windows(2).position(|w| w == b"\r\n")?;
222        let size_text = String::from_utf8_lossy(&body[..line_end]);
223        let size_text = size_text.split(';').next().unwrap_or("").trim();
224        let size = usize::from_str_radix(size_text, 16).ok()?;
225        body = &body[line_end + 2..];
226        if size == 0 {
227            return Some(out);
228        }
229        if body.len() < size + 2 {
230            return None;
231        }
232        out.extend_from_slice(&body[..size]);
233        body = &body[size + 2..];
234    }
235}
236
237struct Head {
238    status: u16,
239    content_length: Option<usize>,
240    chunked: bool,
241}
242
243fn parse_head(head: &[u8]) -> Result<Head, HttpError> {
244    let text = String::from_utf8_lossy(head);
245    let mut lines = text.split("\r\n");
246    let status = lines
247        .next()
248        .and_then(|line| line.split_whitespace().nth(1))
249        .and_then(|code| code.parse::<u16>().ok())
250        .ok_or_else(|| HttpError::Io("malformed HTTP response".to_string()))?;
251    let mut content_length = None;
252    let mut chunked = false;
253    for line in lines {
254        if let Some((name, value)) = line.split_once(':') {
255            let name = name.trim().to_ascii_lowercase();
256            let value = value.trim();
257            if name == "content-length" {
258                content_length = value.parse::<usize>().ok();
259            } else if name == "transfer-encoding" && value.to_ascii_lowercase().contains("chunked")
260            {
261                chunked = true;
262            }
263        }
264    }
265    Ok(Head {
266        status,
267        content_length,
268        chunked,
269    })
270}
271
272/// Complete body of a buffered response, once it has fully arrived.
273fn complete_body(buf: &[u8], closed: bool) -> Result<Option<(u16, Vec<u8>)>, HttpError> {
274    let Some(end) = header_end(buf) else {
275        return if closed {
276            Err(HttpError::Io(
277                "connection closed before a response".to_string(),
278            ))
279        } else {
280            Ok(None)
281        };
282    };
283    let head = parse_head(&buf[..end])?;
284    let body = &buf[end..];
285    if head.chunked {
286        return match decode_chunked(body) {
287            Some(decoded) => Ok(Some((head.status, decoded))),
288            None if closed => Err(HttpError::Io("truncated chunked response".to_string())),
289            None => Ok(None),
290        };
291    }
292    match head.content_length {
293        Some(length) if body.len() >= length => Ok(Some((head.status, body[..length].to_vec()))),
294        Some(_) if closed => Err(HttpError::Io("truncated response".to_string())),
295        Some(_) => Ok(None),
296        None if closed => Ok(Some((head.status, body.to_vec()))),
297        None => Ok(None),
298    }
299}
300
301fn io_error(error: std::io::Error) -> HttpError {
302    HttpError::Io(error.to_string())
303}
304
305/// One blocking HTTP/1.1 exchange with `Connection: close`.
306pub fn http_request(
307    base_url: &str,
308    method: &str,
309    path: &str,
310    body: Option<&str>,
311    timeout: Duration,
312    cancel: &dyn Fn() -> bool,
313) -> Result<(u16, String), HttpError> {
314    let (host, port, prefix) = parse_http_url(base_url).map_err(HttpError::Io)?;
315    let deadline = Instant::now() + timeout;
316    let remaining = || deadline.saturating_duration_since(Instant::now());
317    let mut last = HttpError::Io(format!("cannot resolve {}", host));
318    let mut stream = None;
319    for addr in (host.as_str(), port).to_socket_addrs().map_err(io_error)? {
320        if remaining().is_zero() {
321            return Err(HttpError::Timeout);
322        }
323        match TcpStream::connect_timeout(&addr, remaining()) {
324            Ok(connected) => {
325                stream = Some(connected);
326                break;
327            }
328            Err(error) if error.kind() == std::io::ErrorKind::TimedOut => {
329                last = HttpError::Timeout;
330            }
331            Err(error) => last = io_error(error),
332        }
333    }
334    let mut stream = stream.ok_or(last)?;
335    let mut request = format!(
336        "{} {}{} HTTP/1.1\r\nHost: {}:{}\r\nConnection: close\r\nAccept: application/json\r\n",
337        method, prefix, path, host, port
338    );
339    if let Some(body) = body {
340        request.push_str(&format!(
341            "Content-Type: application/json\r\nContent-Length: {}\r\n",
342            body.len()
343        ));
344    }
345    request.push_str("\r\n");
346    request.push_str(body.unwrap_or(""));
347    stream
348        .set_write_timeout(Some(remaining().max(Duration::from_millis(1))))
349        .map_err(io_error)?;
350    stream.write_all(request.as_bytes()).map_err(io_error)?;
351    stream.set_read_timeout(Some(POLL)).map_err(io_error)?;
352    let mut buf = Vec::new();
353    let mut chunk = [0u8; 16 * 1024];
354    loop {
355        if cancel() {
356            return Err(HttpError::Cancelled);
357        }
358        if remaining().is_zero() {
359            return Err(HttpError::Timeout);
360        }
361        match stream.read(&mut chunk) {
362            Ok(0) => {
363                let (status, body) = complete_body(&buf, true)?.unwrap_or_default();
364                return Ok((status, String::from_utf8_lossy(&body).to_string()));
365            }
366            Ok(n) => {
367                buf.extend_from_slice(&chunk[..n]);
368                if buf.len() > MAX_RESPONSE_BYTES {
369                    return Err(HttpError::Io("response too large".to_string()));
370                }
371                if let Some((status, body)) = complete_body(&buf, false)? {
372                    return Ok((status, String::from_utf8_lossy(&body).to_string()));
373                }
374            }
375            Err(error)
376                if matches!(
377                    error.kind(),
378                    std::io::ErrorKind::WouldBlock
379                        | std::io::ErrorKind::TimedOut
380                        | std::io::ErrorKind::Interrupted
381                ) => {}
382            Err(error) => return Err(io_error(error)),
383        }
384    }
385}
386
387/// Client bound to one runtime over HTTP.
388pub struct HttpClient {
389    pub base_url: String,
390    pub max_retries: u64,
391    pub deadline: Instant,
392    pub requests: u64,
393}
394
395impl HttpClient {
396    pub fn new(base_url: &str, max_retries: u64, deadline: Instant) -> Self {
397        Self {
398            base_url: base_url.to_string(),
399            max_retries,
400            deadline,
401            requests: 0,
402        }
403    }
404}
405
406fn budget_error() -> ClientError {
407    ClientError::with_code(
408        "time budget exhausted while waiting for Formal AI",
409        EXIT_BUDGET,
410    )
411}
412
413impl CompletionClient for HttpClient {
414    fn complete(
415        &mut self,
416        messages: &[Value],
417        cancel: &dyn Fn() -> bool,
418    ) -> Result<Completion, ClientError> {
419        let body = json!({
420            "model": "formal-ai",
421            "messages": messages,
422            "tools": [bash_tool()],
423        })
424        .to_string();
425        let mut last_error = ClientError::protocol("Formal AI was not asked");
426        for _ in 0..=self.max_retries {
427            let remaining = self.deadline.saturating_duration_since(Instant::now());
428            if remaining.is_zero() {
429                return Err(budget_error());
430            }
431            self.requests += 1;
432            let path = "/chat/completions";
433            let full = format!("{}{}", API_PREFIX, path);
434            match http_request(
435                &self.base_url,
436                "POST",
437                &full,
438                Some(&body),
439                remaining,
440                cancel,
441            ) {
442                Ok((status, text)) => {
443                    if status >= 500 {
444                        last_error =
445                            ClientError::protocol(format!("Formal AI returned HTTP {}", status));
446                        continue;
447                    }
448                    if !(200..300).contains(&status) {
449                        let shown: String = text.chars().take(300).collect();
450                        return Err(ClientError::protocol(format!(
451                            "Formal AI returned HTTP {}: {}",
452                            status, shown
453                        )));
454                    }
455                    let parsed: Value = serde_json::from_str(&text).map_err(|_| {
456                        ClientError::protocol("Formal AI returned a non-JSON response")
457                    })?;
458                    return parse_completion(&parsed);
459                }
460                Err(HttpError::Cancelled) => {
461                    let mut error = ClientError::protocol("cancelled");
462                    error.cancelled = true;
463                    return Err(error);
464                }
465                Err(HttpError::Timeout) => return Err(budget_error()),
466                Err(HttpError::Io(message)) => {
467                    last_error = ClientError::with_code(
468                        format!("Formal AI request failed: {}", message),
469                        EXIT_UNAVAILABLE,
470                    );
471                }
472            }
473        }
474        if self.max_retries > 0 {
475            let plural = if self.max_retries == 1 { "y" } else { "ies" };
476            last_error.message = format!(
477                "{} (after {} retr{})",
478                last_error.message, self.max_retries, plural
479            );
480        }
481        Err(last_error)
482    }
483
484    fn requests(&self) -> u64 {
485        self.requests
486    }
487}
488
489#[cfg(test)]
490mod tests {
491    use super::*;
492
493    #[test]
494    fn parses_tool_calls_and_handler() {
495        let body = json!({"choices": [{"finish_reason": "tool_calls", "message": {
496            "content": "",
497            "tool_calls": [{"id": "c1", "function": {"name": "bash", "arguments": "{\"command\":\"ls\"}"}}],
498            "thinking_steps": [{"step": "dispatch_handler", "detail": "list_files"}]
499        }}]});
500        let completion = parse_completion(&body).unwrap();
501        assert_eq!(completion.handler.as_deref(), Some("list_files"));
502        assert_eq!(completion.finish_reason.as_deref(), Some("tool_calls"));
503        assert_eq!(
504            completion.tool_calls,
505            vec![ToolCall {
506                id: Some("c1".to_string()),
507                name: "bash".to_string(),
508                command: "ls".to_string()
509            }]
510        );
511    }
512
513    #[test]
514    fn rejects_malformed_completions() {
515        let cases = [
516            (json!({}), "Formal AI response has no choices[0].message"),
517            (
518                json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "python", "arguments": "{}"}}]}}]}),
519                "Formal AI called an unsupported tool \"python\"",
520            ),
521            (
522                json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "bash", "arguments": "ls"}}]}}]}),
523                "Formal AI sent non-JSON arguments for bash",
524            ),
525            (
526                json!({"choices": [{"message": {"tool_calls": [{"function": {"name": "shell", "arguments": "{\"cmd\":\" \"}"}}]}}]}),
527                "Formal AI called shell without a command",
528            ),
529        ];
530        for (body, message) in cases {
531            let error = parse_completion(&body).unwrap_err();
532            assert_eq!(error.message, message);
533            assert_eq!(error.exit_code, EXIT_PROTOCOL);
534        }
535    }
536
537    #[test]
538    fn accepts_cmd_arguments() {
539        let body = json!({"choices": [{"message": {"content": "ok",
540            "tool_calls": [{"function": {"name": "run_command", "arguments": "{\"cmd\":\"pwd\"}"}}]}}]});
541        let completion = parse_completion(&body).unwrap();
542        assert_eq!(completion.tool_calls[0].command, "pwd");
543        assert_eq!(completion.tool_calls[0].id, None);
544        assert_eq!(completion.content, "ok");
545    }
546
547    #[test]
548    fn parses_http_urls() {
549        assert_eq!(
550            parse_http_url("http://127.0.0.1:8080").unwrap(),
551            ("127.0.0.1".to_string(), 8080, String::new())
552        );
553        assert_eq!(
554            parse_http_url("http://localhost/base/").unwrap(),
555            ("localhost".to_string(), 80, "/base".to_string())
556        );
557        assert_eq!(
558            parse_http_url("http://[::1]:9/").unwrap(),
559            ("::1".to_string(), 9, String::new())
560        );
561        assert!(parse_http_url("https://example.com").is_err());
562        assert!(parse_http_url("http://:80").is_err());
563    }
564
565    #[test]
566    fn decodes_content_length_and_chunked_bodies() {
567        let fixed = b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nokEXTRA";
568        assert_eq!(
569            complete_body(fixed, false).unwrap(),
570            Some((200, b"ok".to_vec()))
571        );
572        let partial = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nok";
573        assert_eq!(complete_body(partial, false).unwrap(), None);
574        assert!(complete_body(partial, true).is_err());
575        let chunked = b"HTTP/1.1 503 Busy\r\nTransfer-Encoding: chunked\r\n\r\n3\r\nabc\r\n2\r\nde\r\n0\r\n\r\n";
576        assert_eq!(
577            complete_body(chunked, false).unwrap(),
578            Some((503, b"abcde".to_vec()))
579        );
580        let open = b"HTTP/1.1 200 OK\r\n\r\nbody";
581        assert_eq!(complete_body(open, false).unwrap(), None);
582        assert_eq!(
583            complete_body(open, true).unwrap(),
584            Some((200, b"body".to_vec()))
585        );
586    }
587
588    #[test]
589    fn unreachable_server_is_unavailable_after_retries() {
590        // Port 9 (discard) on loopback refuses connections on CI runners.
591        let mut client = HttpClient::new(
592            "http://127.0.0.1:9",
593            1,
594            Instant::now() + Duration::from_secs(10),
595        );
596        let error = client.complete(&[], &|| false).unwrap_err();
597        assert_eq!(client.requests(), 2);
598        if error.exit_code == EXIT_UNAVAILABLE {
599            assert!(error.message.starts_with("Formal AI request failed: "));
600            assert!(error.message.ends_with(" (after 1 retry)"));
601        } else {
602            assert_eq!(error.exit_code, EXIT_BUDGET);
603        }
604    }
605}
start_command/formal_ai_client.rs

start_command/
formal_ai_client.rs
