sausage/agent/src/main.rs

//! Answer chat calls with an LLM and tools, under one lock for the complete turn.
mod history;
mod llm;
mod tools;

use history::{is_too_large, next_turn_start};
use llm::{Config, Entry, Thinking};
use sausage_plugin::{Ctx, Error, Result, lock, rpc, schema};
use serde_json::json;
use std::io::Write;

/// The instructions sent with every inference request.
const SYSTEM: &str = include_str!("prompt.md");
/// The shared, append-only conversation.
const HISTORY: &str = "/state/agent/history.jsonl";
/// The first retained turn after a context-size rejection.
const SKIP: &str = "/state/agent/history.skip";
/// Each chat process reads the shared history after it acquires the lock.
struct Agent;
/// The person's next contribution to the conversation.
#[schema]
struct ChatParams {
    /// The complete text of the person's request.
    text: String,
    /// The thinking level of this turn only. Without it, the turn uses `thinking` from
    /// `/config/agent.yaml`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    thinking: Option<Thinking>,
}

#[rpc(grace_seconds = 600)]
impl Agent {
    /// Answer a person, with interim text notifications and a final text result.
    ///
    /// The history lock covers inference and tools, so concurrent chats cannot interleave turns.
    async fn chat(&self, ctx: &Ctx, p: ChatParams) -> Result<String> {
        let _guard = lock(HISTORY).await?;
        std::fs::create_dir_all("/state/agent")?;
        let mut cfg = Config::load().map_err(|error| Error::new(-32603, error))?;
        cfg.thinking = p.thinking.or(cfg.thinking);
        let mut entries = load_history(HISTORY)?;
        repair_tools(&mut entries)?;
        let current = entries.len();
        append(&mut entries, Entry::User { text: p.text })?;
        let mut start = std::fs::read_to_string(SKIP)
            .ok()
            .and_then(|s| s.trim().parse().ok())
            .unwrap_or(0);
        start = next_turn_start(&entries, start, current).unwrap_or(current);
        let definitions = tools::definitions();
        loop {
            let response = match llm::complete(&cfg, SYSTEM, &definitions, &entries[start..]).await
            {
                Ok(response) => response,
                Err(error) => {
                    if is_too_large(&error)
                        && let Some(next) = next_turn_start(&entries, start + 1, current)
                    {
                        start = next;
                        continue;
                    }
                    return Err(Error::new(-32603, error));
                }
            };
            std::fs::write(SKIP, start.to_string())?;
            let has_tools = response
                .iter()
                .any(|entry| matches!(entry, Entry::ToolCall { .. }));
            let mut final_text = String::new();
            let mut calls = Vec::new();
            for entry in response {
                match &entry {
                    Entry::Assistant { text, phase }
                        if has_tools || phase.as_deref() == Some("commentary") =>
                    {
                        ctx.notify("text", json!({"text":text}))?
                    }
                    Entry::Assistant { text, .. } => {
                        if !final_text.is_empty() && !text.is_empty() {
                            final_text.push_str("\n\n");
                        }
                        final_text.push_str(text);
                    }
                    Entry::ToolCall {
                        id,
                        name,
                        arguments,
                    } => calls.push((id.clone(), name.clone(), arguments.clone())),
                    _ => {}
                }
                append(&mut entries, entry)?;
            }
            if !has_tools {
                return Ok(final_text);
            }
            for (id, name, arguments) in calls {
                let args = serde_json::from_str::<serde_json::Value>(&arguments);
                let output = match &args {
                    Ok(args) => tools::run(&name, args)
                        .await
                        .unwrap_or_else(|error| json!({"error":error})),
                    Err(error) => json!({"error":format!("arguments are not valid JSON: {error}")}),
                };
                let output = if name == "read_file"
                    && args
                        .as_ref()
                        .ok()
                        .and_then(|a| a["path"].as_str())
                        .is_some_and(|p| p.starts_with("/logs/"))
                {
                    output["content"]
                        .as_str()
                        .map(str::to_owned)
                        .unwrap_or_else(|| output.to_string())
                } else {
                    output.to_string()
                };
                let tail = args
                    .as_ref()
                    .ok()
                    .and_then(|a| a["method"].as_str())
                    .is_some_and(|m| {
                        [".exec", ".log", ".build", ".test"]
                            .iter()
                            .any(|suffix| m.ends_with(suffix))
                    });
                let output = tools::overflow::guard(output, "default", &id, tail);
                append(&mut entries, Entry::ToolResult { id, output })?;
            }
        }
    }
}

/// Read persisted entries in their original order. A new world has no history file.
fn load_history(path: &str) -> Result<Vec<Entry>> {
    match std::fs::read_to_string(path) {
        Ok(text) => Ok(text
            .lines()
            .filter_map(|line| match serde_json::from_str(line) {
                Ok(entry) => Some(entry),
                Err(error) => {
                    eprintln!("skip damaged history line: {error}");
                    None
                }
            })
            .collect::<Vec<_>>()),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Vec::new()),
        Err(error) => Err(error.into()),
    }
}
/// Persist an entry before it enters the current inference context.
fn append(history: &mut Vec<Entry>, entry: Entry) -> Result<()> {
    let mut file = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(HISTORY)?;
    file.write_all(format!("{}\n", serde_json::to_string(&entry)?).as_bytes())?;
    history.push(entry);
    Ok(())
}
/// Complete interrupted tool pairs before a new inference request.
fn repair_tools(history: &mut Vec<Entry>) -> Result<()> {
    let answered: std::collections::HashSet<_> = history
        .iter()
        .filter_map(|entry| match entry {
            Entry::ToolResult { id, .. } => Some(id.clone()),
            _ => None,
        })
        .collect();
    let dangling: Vec<_> = history
        .iter()
        .filter_map(|entry| match entry {
            Entry::ToolCall { id, .. } if !answered.contains(id) => Some(id.clone()),
            _ => None,
        })
        .collect();
    for id in dangling {
        append(
            history,
            Entry::ToolResult {
                id,
                output: json!({"error":"interrupted before the tool finished"}).to_string(),
            },
        )?;
    }
    Ok(())
}
/// Start the plugin's JSON-RPC transport.
fn main() {
    sausage_plugin::run(Agent);
}

#[cfg(test)]
mod tests {
    #[test]
    fn damaged_history_keeps_valid_entries() {
        let path = if cfg!(target_os = "wasi") {
            std::path::PathBuf::from("/history-test.jsonl")
        } else {
            std::env::temp_dir().join(format!("sausage-history-{}", std::process::id()))
        };
        std::fs::write(
            &path,
            "{\"kind\":\"user\",\"text\":\"preserved\"}\n{\"kind\":",
        )
        .unwrap();
        let history = super::load_history(path.to_str().unwrap()).unwrap();
        assert_eq!(history.len(), 1);
        std::fs::remove_file(path).unwrap();
    }
}