rebranding herald => steep

crates/steep-server/src/sandbox/agent.rs

@@ -0,0 +1,316 @@
+//! Tool-calling loop driving the model against a [`Sandbox`].
+//!
+//! The loop sends the conversation and the available tools to OpenRouter,
+//! executes any requested tool call inside the sandbox and feeds the results
+//! back, until the model produces a final message or the iteration budget is
+//! exhausted.
+
+use std::str::FromStr;
+use std::time::Instant;
+
+use anyhow::Context;
+use serde_json::Value;
+use tracing::{info, warn};
+
+use crate::{
+    open_router::{Message, OpenRouterClient, Role, Tool, ToolCall, ToolChoice},
+    sandbox::{Sandbox, tools},
+    text::{excerpt, truncate_bytes},
+};
+
+/// Number of characters of the arguments of a tool call kept in the logs.
+const MAX_LOGGED_ARGUMENTS: usize = 200;
+
+/// Number of characters of a tool result kept in the logs.
+///
+/// Sizes alone make a broken tool invisible: a command that fails returns a short
+/// `error: …` instead of the expected content, which is exactly what this excerpt
+/// makes obvious.
+const MAX_LOGGED_OUTPUT: usize = 120;
+
+/// Maximum number of bytes of a tool result handed to the model.
+///
+/// A result stays in the conversation and is re-sent on every following
+/// iteration, so an unbounded one (a whole file, a match on every line) inflates
+/// the context and the cost of the whole rest of the run — and can overflow the
+/// model context window outright.
+const MAX_TOOL_OUTPUT_BYTES: usize = 32 * 1024;
+
+/// Asked of the model when its answer does not satisfy the caller's contract.
+const RETRY_PROMPT: &str = "
+    Your answer is not valid: {error}
+
+    Answer again with the requested format only: a raw json object, without any
+    markdown code fence and without any text before or after it.
+";
+
+/// Asked of the model on the last turn, in place of one more exploration turn.
+///
+/// Without it, a model still exploring on its last turn would ask for another
+/// tool, and the run would end on the iteration budget with no answer at all.
+const FINAL_PROMPT: &str = "
+    You are out of turns: this is your last one. Answer now, with the requested
+    format, from what you have already gathered, without calling any tool.
+";
+
+/// Final output of an agent run.
+pub struct AgentResult<A> {
+    /// Answer of the model, parsed into the type the caller asked for.
+    pub answer: A,
+    pub cost: Option<f64>,
+    pub iterations: usize,
+}
+
+/// Runs the tool-calling loop until the model answers or `max_iterations` is
+/// reached.
+///
+/// `tool_definitions` is the set of tools the model may call; it is selected by
+/// the caller based on the webhook action (see [`tools::for_webhook`]).
+///
+/// The answer of the model is parsed into `A`, which is how the caller states the
+/// format it expects: while [`FromStr`] rejects the answer, the error is sent back
+/// to the model so that it can correct itself, which costs an iteration. The
+/// parsing happens here because that is where the conversation and the remaining
+/// budget are at hand.
+pub async fn run<A>(
+    open_router: &OpenRouterClient,
+    sandbox: &Sandbox,
+    tool_definitions: Vec<Tool>,
+    system_prompt: &str,
+    user_prompt: &str,
+    max_iterations: usize,
+) -> anyhow::Result<AgentResult<A>>
+where
+    A: FromStr,
+    A::Err: std::fmt::Display,
+{
+    let mut messages = vec![
+        Message::new(Role::System, system_prompt),
+        Message::new(Role::User, user_prompt),
+    ];
+
+    info!(
+        tools = %tool_definitions
+            .iter()
+            .map(|tool| tool.function.name.as_str())
+            .collect::<Vec<_>>()
+            .join(", "),
+        max_iterations,
+        "Starting tool-calling loop"
+    );
+
+    let mut total_cost = 0.0_f64;
+    let mut has_cost = false;
+    let mut last_rejection: Option<String> = None;
+
+    for iteration in 1..=max_iterations {
+        // Le dernier tour n'est plus un tour d'exploration : le modèle doit rendre
+        // sa réponse, avec ce qu'il a déjà vu.
+        let last = iteration == max_iterations;
+
+        if last {
+            info!(iteration, "Last turn: asking for the final answer");
+            messages.push(Message::new(Role::User, FINAL_PROMPT));
+        }
+
+        let started = Instant::now();
+        let response = open_router
+            .chat_with_tools(
+                &messages,
+                &tool_definitions,
+                if last {
+                    ToolChoice::None
+                } else {
+                    ToolChoice::Auto
+                },
+            )
+            .await?;
+
+        if let Some(cost) = response.cost {
+            total_cost += cost;
+            has_cost = true;
+        }
+
+        info!(
+            iteration,
+            max_iterations,
+            last_turn = last,
+            tool_calls = response.tool_calls.len(),
+            answer_bytes = response.message.as_deref().map_or(0, str::len),
+            elapsed_ms = started.elapsed().as_millis() as u64,
+            "Model answered"
+        );
+
+        // Au dernier tour, les outils sont refusés : même si le fournisseur les
+        // renvoie malgré tout, on ne les exécute pas et on tente la réponse.
+        if last || response.tool_calls.is_empty() {
+            let answer = response.message.unwrap_or_default();
+
+            match answer.parse::<A>() {
+                Ok(parsed) => {
+                    return Ok(AgentResult {
+                        answer: parsed,
+                        cost: has_cost.then_some(total_cost),
+                        iterations: iteration,
+                    });
+                }
+                Err(err) => {
+                    warn!(iteration, %err, "Model answer was rejected, asking for another one");
+
+                    // Keep the rejected answer in the history, so that the model
+                    // sees what it has to fix.
+                    messages.push(Message::new(Role::Assistant, answer));
+                    messages.push(Message::new(
+                        Role::User,
+                        RETRY_PROMPT.replace("{error}", &err.to_string()),
+                    ));
+
+                    last_rejection = Some(err.to_string());
+                    continue;
+                }
+            }
+        }
+
+        messages.push(Message::assistant_with_tool_calls(
+            response.message.unwrap_or_default(),
+            response.tool_calls.clone(),
+        ));
+
+        for call in &response.tool_calls {
+            let started = Instant::now();
+            let (content, truncated) = execute(sandbox, call).await;
+
+            info!(
+                iteration,
+                tool = call.name(),
+                arguments = %excerpt(call.arguments_json(), MAX_LOGGED_ARGUMENTS),
+                output_bytes = content.len(),
+                truncated,
+                // Les sauts de ligne casseraient la lisibilité d'une ligne de log.
+                output = %excerpt(&content, MAX_LOGGED_OUTPUT).replace('\n', " "),
+                elapsed_ms = started.elapsed().as_millis() as u64,
+                "Tool call finished"
+            );
+
+            messages.push(Message::tool_response(call.id(), content));
+        }
+    }
+
+    warn!(
+        max_iterations,
+        "Agent used its last turn without a valid answer"
+    );
+
+    match last_rejection {
+        Some(error) => anyhow::bail!(
+            "agent exceeded the maximum number of iterations ({max_iterations}), \
+             last answer rejected: {error}"
+        ),
+        None => anyhow::bail!("agent exceeded the maximum number of iterations ({max_iterations})"),
+    }
+}
+
+/// Runs a tool call and returns its result for the model, along with whether the
+/// result had to be truncated.
+async fn execute(sandbox: &Sandbox, call: &ToolCall) -> (String, bool) {
+    let args = match parse_args(call) {
+        Ok(args) => args,
+        Err(err) => return cap_output(format!("error: {err}")),
+    };
+
+    match tools::dispatch(sandbox, call.name(), &args).await {
+        Ok(output) => cap_output(output),
+        Err(err) => cap_output(format!("error: {err}")),
+    }
+}
+
+/// Bounds [`MAX_TOOL_OUTPUT_BYTES`] of a tool result, telling the model what was
+/// dropped so that it can narrow its request.
+fn cap_output(mut output: String) -> (String, bool) {
+    let total = output.len();
+
+    if !truncate_bytes(&mut output, MAX_TOOL_OUTPUT_BYTES) {
+        return (output, false);
+    }
+
+    let kept = output.len();
+    output.push_str(&format!(
+        "\n… output truncated: {total} bytes in total, the first {kept} are shown. \
+         Narrow the request (path, pattern or line range) to see the rest."
+    ));
+
+    (output, true)
+}
+
+fn parse_args(call: &ToolCall) -> anyhow::Result<Value> {
+    let raw = call.arguments_json().trim();
+    if raw.is_empty() {
+        return Ok(Value::Object(serde_json::Map::new()));
+    }
+
+    serde_json::from_str(raw)
+        .with_context(|| format!("invalid arguments for tool `{}`", call.name()))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    /// Builds a tool call the way the API returns one, so the fixture also
+    /// covers deserialization.
+    fn tool_call(name: &str, arguments: &str) -> ToolCall {
+        serde_json::from_value(json!({
+            "id": "call_1",
+            "type": "function",
+            "function": { "name": name, "arguments": arguments }
+        }))
+        .unwrap()
+    }
+
+    #[test]
+    fn parse_args_accepts_empty_arguments() {
+        let call = tool_call("ls", "");
+        assert_eq!(
+            parse_args(&call).unwrap(),
+            Value::Object(serde_json::Map::new())
+        );
+    }
+
+    #[test]
+    fn parse_args_parses_json_object() {
+        let call = tool_call("ls", r#"{"path":"src"}"#);
+        assert_eq!(parse_args(&call).unwrap()["path"], "src");
+    }
+
+    #[test]
+    fn parse_args_rejects_invalid_json() {
+        let call = tool_call("ls", "not json");
+        assert!(parse_args(&call).is_err());
+    }
+
+    #[test]
+    fn a_small_tool_output_is_kept_as_is() {
+        let (content, truncated) = cap_output(String::from("src/main.rs"));
+
+        assert_eq!(content, "src/main.rs");
+        assert!(!truncated);
+    }
+
+    #[test]
+    fn a_large_tool_output_is_truncated_and_told_to_the_model() {
+        let (content, truncated) = cap_output("a".repeat(MAX_TOOL_OUTPUT_BYTES + 1));
+
+        assert!(truncated);
+        assert!(content.starts_with(&"a".repeat(MAX_TOOL_OUTPUT_BYTES)));
+        assert!(content.contains("output truncated"));
+    }
+
+    #[test]
+    fn a_truncated_tool_output_stays_valid_utf8() {
+        let (content, truncated) = cap_output("é".repeat(MAX_TOOL_OUTPUT_BYTES));
+
+        assert!(truncated);
+        assert!(content.starts_with('é'));
+    }
+}