sausage/agent/src/main.rs

//! The Sausage agent actor.

mod history;
mod llm;
mod tools;

use history::{is_too_large, next_turn_start};
use llm::{Config, Entry};
use sausage_actor::{Inbox, Message, Value, json, publish, send};
use std::collections::{HashSet, VecDeque};
use std::io::Write;
use wstd::future::FutureExt;
use wstd::time::{Duration, Instant};

const SYSTEM: &str = include_str!("prompt.md");
const CONFIG: &str = "/data/config.json";
const HISTORY: &str = "/data/history.jsonl";
const HISTORY_SKIP: &str = "/data/history.skip";
const TURN: &str = "/data/turn.json";

fn main() {
    wstd::runtime::block_on(async { Agent::new().run().await });
}

struct Agent {
    inbox: Inbox,
    history: Vec<Entry>,
    history_skip: usize,
    active: bool,
    queue: VecDeque<Message>,
    notes: Vec<Message>,
}

impl Agent {
    fn new() -> Self {
        let history = std::fs::read_to_string(HISTORY)
            .unwrap_or_default()
            .lines()
            .filter_map(|line| serde_json::from_str(line).ok())
            .collect();
        let history_skip = std::fs::read_to_string(HISTORY_SKIP)
            .ok()
            .and_then(|skip| skip.trim().parse::<usize>().ok())
            .unwrap_or(0);
        let active = std::fs::read_to_string(TURN)
            .ok()
            .and_then(|contents| serde_json::from_str::<bool>(&contents).ok())
            .unwrap_or(false);
        Self {
            inbox: Inbox::start(),
            history,
            history_skip,
            active,
            queue: VecDeque::new(),
            notes: Vec::new(),
        }
    }

    fn add_history(&mut self, entry: Entry) {
        let line = serde_json::to_string(&entry).unwrap();
        match std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(HISTORY)
        {
            Ok(mut file) => {
                if let Err(e) = writeln!(file, "{line}") {
                    eprintln!("agent: cannot append history: {e}");
                }
            }
            Err(e) => eprintln!("agent: cannot append history: {e}"),
        }
        self.history.push(entry);
    }

    fn finish_turn(&mut self) {
        self.active = false;
        let _ = std::fs::remove_file(TURN);
    }

    fn closing(&self) -> bool {
        self.inbox.is_closed()
    }

    fn repair_dangling_tools(&mut self) {
        let answered: HashSet<String> = self
            .history
            .iter()
            .filter_map(|entry| match entry {
                Entry::ToolResult { id, .. } => Some(id.clone()),
                _ => None,
            })
            .collect();
        let dangling: Vec<String> = self
            .history
            .iter()
            .filter_map(|entry| match entry {
                Entry::ToolCall { id, .. } if !answered.contains(id) => Some(id.clone()),
                _ => None,
            })
            .collect();
        for id in dangling {
            self.add_history(Entry::ToolResult {
                id,
                output:
                    json!({"error": "interrupted: the agent restarted before this tool finished"})
                        .to_string(),
            });
        }
    }

    fn turn_text(&self, message: &Message) -> String {
        let mut text = format!(
            "Message {} from {}:\n{}",
            message.id,
            message.from,
            message.body["text"].as_str().unwrap_or("")
        );
        if !self.notes.is_empty() {
            let notes = self
                .notes
                .iter()
                .map(|note| serde_json::to_string(note).unwrap())
                .collect::<Vec<_>>()
                .join("\n");
            text = format!("Messages that arrived while you were idle:\n{notes}\n\n{text}");
        }
        text
    }

    fn not_configured(&mut self, error: String) {
        let _ = publish(json!({
            "type": "text",
            "text": format!("I am not configured: cannot read {CONFIG}: {error}")
        }));
        self.finish_turn();
    }

    async fn run(mut self) {
        self.repair_dangling_tools();
        if self.active {
            self.turn(None).await;
            if self.closing() && self.active {
                return;
            }
        }
        loop {
            let message = match self.queue.pop_front() {
                Some(message) => message,
                None => match self.inbox.recv().await {
                    Some(message) => message,
                    None => break,
                },
            };
            if starts_turn(&message) {
                self.turn(Some(message)).await;
                if self.closing() && self.active {
                    break;
                }
            } else {
                self.notes.push(message);
            }
        }
    }

    async fn turn(&mut self, message: Option<Message>) {
        let current_start = if message.is_some() {
            self.active = true;
            let _ = std::fs::write(TURN, "true");
            self.history.len()
        } else {
            let current_start = self
                .history
                .iter()
                .rposition(|entry| matches!(entry, Entry::User { text } if !text.starts_with("[note from your runtime]")))
                .unwrap_or(self.history.len());
            let version = std::env::var("SAUSAGE_VERSION").unwrap_or_default();
            self.add_history(Entry::User {
                text: format!(
                    "[note from your runtime] you restarted and now run version {version}. Continue the active turn."
                ),
            });
            current_start
        };
        let cfg = match load_config() {
            Ok(cfg) => cfg,
            Err(error) => {
                self.not_configured(error);
                return;
            }
        };
        if let Some(message) = message {
            let text = self.turn_text(&message);
            self.notes.clear();
            self.add_history(Entry::User { text });
        }
        self.steps(&cfg, current_start).await;
        if !self.closing() {
            self.finish_turn();
        }
    }

    async fn steps(&mut self, cfg: &Config, current_start: usize) {
        let tools = tools::definitions();
        loop {
            if self.closing() {
                return;
            }
            let start = next_turn_start(&self.history, self.history_skip, current_start)
                .unwrap_or(current_start);
            if start != self.history_skip {
                self.history_skip = start;
                let _ = std::fs::write(HISTORY_SKIP, start.to_string());
            }
            let response = llm::complete(cfg, SYSTEM, &tools, &self.history[start..]).await;
            let entries = match response {
                Ok(entries) => entries,
                Err(e) => {
                    if is_too_large(&e) {
                        if let Some(next) =
                            next_turn_start(&self.history, start.saturating_add(1), current_start)
                        {
                            self.history_skip = next;
                            let _ = std::fs::write(HISTORY_SKIP, next.to_string());
                            continue;
                        }
                    }
                    let text = format!("Inference failed: {e}");
                    self.add_history(Entry::Assistant {
                        text: text.clone(),
                        phase: None,
                    });
                    let _ = publish(json!({"type": "text", "text": text}));
                    return;
                }
            };
            let calls: Vec<(String, String, String)> = entries
                .iter()
                .filter_map(|entry| match entry {
                    Entry::ToolCall {
                        id,
                        name,
                        arguments,
                    } => Some((id.clone(), name.clone(), arguments.clone())),
                    _ => None,
                })
                .collect();
            for entry in entries {
                if let Entry::Assistant { text, .. } = &entry {
                    if !text.is_empty() {
                        let _ = publish(json!({"type": "text", "text": text}));
                    }
                }
                self.add_history(entry);
            }
            if self.closing() {
                return;
            }
            if calls.is_empty() {
                return;
            }
            for (id, name, args) in calls {
                if self.closing() {
                    return;
                }
                let output = self.tool(&name, &args).await;
                self.add_history(Entry::ToolResult {
                    id,
                    output: output.to_string(),
                });
            }
        }
    }

    fn stash(&mut self, message: Message, others: &mut Vec<Value>) {
        if starts_turn(&message) {
            others.push(json!({"queued_for_a_later_turn": true, "message": message}));
            self.queue.push_back(message);
        } else {
            others.push(serde_json::to_value(message).unwrap());
        }
    }

    fn drain_inbox(&mut self, others: &mut Vec<Value>) {
        while let Some(message) = self.inbox.try_recv() {
            self.stash(message, others);
        }
    }

    async fn tool(&mut self, name: &str, args: &str) -> Value {
        let args: Value = match serde_json::from_str(args) {
            Ok(args) => args,
            Err(error) => return json!({"error": format!("arguments are not valid JSON: {error}")}),
        };
        if name == "send_message" {
            return self.send_message(&args).await;
        }
        tools::run(name, &args).unwrap_or_else(|error| json!({"error": error.to_string()}))
    }

    async fn send_message(&mut self, args: &Value) -> Value {
        let Some(to) = args["to"].as_str() else {
            return json!({"error": "missing \"to\""});
        };
        if !args["body"].is_object() {
            return json!({"error": "\"body\" must be a JSON object"});
        }
        let id = send(to, args["body"].clone());
        let wait = args["wait_seconds"].as_f64().unwrap_or(0.0).max(0.0);
        let mut reply = Value::Null;
        let mut others = Vec::new();
        if wait > 0.0 {
            let deadline = Instant::now() + Duration::from_secs_f64(wait);
            loop {
                match self.inbox.recv().timeout(deadline).await {
                    Ok(Some(message)) if message.reply_to.as_deref() == Some(&id) => {
                        reply = serde_json::to_value(message).unwrap();
                        break;
                    }
                    Ok(Some(message)) => self.stash(message, &mut others),
                    Ok(None) => break,
                    Err(_) => break,
                }
            }
        }
        self.drain_inbox(&mut others);
        json!({"id": id, "reply": reply, "waited": wait > 0.0, "other_messages": others})
    }
}

fn load_config() -> Result<Config, String> {
    std::fs::read_to_string(CONFIG)
        .map_err(|e| e.to_string())
        .and_then(|contents| serde_json::from_str(&contents).map_err(|e| e.to_string()))
}

fn starts_turn(message: &Message) -> bool {
    message.kind() == "text" && message.reply_to.is_none()
}