procyon 0.1.2

//! The agent's own vocabulary: the shapes a conversation is made of, independent of any provider.
//!
//! Nothing here knows a wire format. That matters most for `ContentPart`: it is what the session
//! log persists, so its serialized shape is a durability commitment, while a provider's request
//! and event shapes follow whatever that API asks for this month. The adapters — `anthropic` and
//! `openai` — translate to and from these types, and each keeps its own protocol details.

pub mod subagent;

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    User,
    Assistant,
}

impl std::fmt::Display for Role {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Role::User => write!(f, "user"),
            Role::Assistant => write!(f, "assistant"),
        }
    }
}

impl std::str::FromStr for Role {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s {
            "user" => Ok(Role::User),
            "assistant" => Ok(Role::Assistant),
            _ => Err(format!("Unknown role: {}", s)),
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct Message {
    pub role: Role,
    pub content: Vec<ContentPart>,
}

impl Message {
    pub fn user(text: &str) -> Self {
        Self {
            role: Role::User,
            content: vec![ContentPart::Text {
                text: text.to_string(),
            }],
        }
    }

    pub fn assistant(blocks: Vec<ContentPart>) -> Self {
        Self {
            role: Role::Assistant,
            content: blocks,
        }
    }

    pub fn tool_results(results: Vec<(String, String)>) -> Self {
        Self {
            role: Role::User,
            content: results
                .into_iter()
                .map(|(tool_use_id, content)| ContentPart::ToolResult {
                    tool_use_id,
                    content,
                })
                .collect(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ContentPart {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "tool_use")]
    ToolUse {
        id: String,
        name: String,
        input: serde_json::Value,
    },
    #[serde(rename = "tool_result")]
    ToolResult {
        tool_use_id: String,
        content: String,
    },
}

// Kept as a suffix rather than a whole replacement prompt so toggling it does not rewrite the
// cached prefix of an ongoing conversation.
pub const EXPLAIN_SYSTEM_PROMPT: &str =
    "Before each tool call, state in one sentence what you are about to do and why. \
     After it returns, say in one sentence what the result means. Keep the narration brief.";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

/// Resolves the arguments of a tool call into the object a `ToolUse` must carry.
///
/// A tool call whose arguments do not parse still has to become a `ToolUse`. Returning a
/// `ToolResult` instead — which is where the error naturally belongs — puts it inside an assistant
/// turn, and that is invalid on both wire formats: the request is rejected outright, and because
/// the block reaches the session log first, a resumed conversation is rejected too. The turn also
/// ends up with no `tool_use` at all, so the loop stops and the model is never told anything.
///
/// So the call is emitted with empty arguments and the tool reports the failure itself, through
/// the one channel that is valid for it — a `tool_result` in the following user turn. The parse
/// error would otherwise be lost, so it is recorded.
pub fn tool_input(name: &str, raw: &str) -> serde_json::Value {
    let empty = serde_json::Value::Object(Default::default());

    if raw.trim().is_empty() {
        // A tool taking no arguments sends no fragments at all.
        return empty;
    }

    match serde_json::from_str::<serde_json::Value>(raw) {
        // Tools read their arguments by key, so anything but an object is unusable. `null` is how
        // some providers spell "no arguments"; the rest is malformed.
        Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
        Ok(serde_json::Value::Null) => empty,
        Ok(other) => {
            crate::diag::warn(format!(
                "tool call {}: arguments are {} rather than an object, treated as empty",
                name,
                kind_of(&other)
            ));
            empty
        }
        Err(e) => {
            crate::diag::warn(format!(
                "tool call {}: arguments did not parse ({}), treated as empty. raw: {}",
                name, e, raw
            ));
            empty
        }
    }
}

/// Turns a tool call the model wrote as prose into a real one.
///
/// Small local models routinely ignore the tool-calling API and emit the call as a JSON object in
/// the reply instead. With no fallback that object went straight to the screen as the agent's
/// answer — the user saw `{"name": "ls", "parameters": {...}}` where a reply should have been, and
/// the tool never ran.
///
/// Recovery is deliberately narrow. It applies only when the model used no real tool call in the
/// same message — if it managed the API, it is not second-guessed — and only when a text block is
/// *entirely* the call. A model explaining a tool call inside a sentence is talking to the user,
/// and executing that would turn documentation into an action.
///
/// An unknown tool name is recovered too, rather than filtered out. The registry answers with
/// "Unknown tool", which reaches the model as a `tool_result` and tells it what it got wrong;
/// dropping the block silently would leave it waiting for a result that never comes. Nothing is
/// waved through by being recovered: a recovered call takes the same path as any other, approval
/// gate included.
pub fn recover_text_tool_calls(blocks: Vec<ContentPart>) -> (Vec<ContentPart>, bool) {
    if blocks
        .iter()
        .any(|b| matches!(b, ContentPart::ToolUse { .. }))
    {
        return (blocks, false);
    }

    let mut recovered = false;
    let out = blocks
        .into_iter()
        .enumerate()
        .map(|(index, block)| match &block {
            ContentPart::Text { text } => match parse_text_tool_call(text) {
                Some((name, input)) => {
                    recovered = true;
                    crate::diag::warn(format!(
                        "recovered a tool call the model wrote as text: {}",
                        name
                    ));
                    ContentPart::ToolUse {
                        // Only has to be unique within the turn and match the `tool_result` that
                        // answers it; no provider prescribes a shape.
                        id: format!("procyon_recovered_{}", index),
                        name,
                        input,
                    }
                }
                None => block,
            },
            _ => block,
        })
        .collect();

    (out, recovered)
}

/// The JSON object in `text`, if the whole block is one and it looks like a tool call.
fn parse_text_tool_call(text: &str) -> Option<(String, serde_json::Value)> {
    let mut body = text.trim();

    // Qwen and the Hermes-style templates wrap the call in a tag. A wrapper says "this is a call"
    // outright, which is worth knowing: without one we have to be much more careful below.
    let mut declared = false;
    for (open, close) in [
        ("<tool_call>", "</tool_call>"),
        ("<tool_use>", "</tool_use>"),
    ] {
        if let Some(inner) = body.strip_prefix(open) {
            body = inner.strip_suffix(close).unwrap_or(inner).trim();
            declared = true;
        }
    }

    // ...and others put it in a fenced code block.
    if let Some(inner) = body.strip_prefix("```") {
        let inner = inner.strip_suffix("```").unwrap_or(inner);
        // Drop an optional language tag on the opening fence.
        body = match inner.split_once('\n') {
            Some((first, rest)) if !first.trim().starts_with('{') => rest.trim(),
            _ => inner.trim(),
        };
    }

    let value: serde_json::Value = serde_json::from_str(body).ok()?;
    let object = value.as_object()?;

    let name = object.get("name")?.as_str()?.trim();
    if name.is_empty() {
        return None;
    }

    // Every spelling of "the arguments" seen in the wild.
    let arguments = ["parameters", "arguments", "input", "args"]
        .iter()
        .find_map(|key| object.get(*key));

    let input = match arguments {
        Some(value) => value.clone(),
        // No arguments field. Inside a `<tool_call>` wrapper that is simply a call that takes
        // none; bare, it is far more likely to be a document the model is showing the user.
        // `{"name": "my-app", "version": "0.1.0"}` is a `package.json`, and in this codebase the
        // model is asked to read those constantly — running a tool called `my-app` would derail
        // the turn and swallow the very content the user asked to see. A missed no-argument call
        // costs only the old behaviour; a false positive costs the answer.
        None if declared => serde_json::Value::Object(Default::default()),
        None => return None,
    };

    // Anything but an object is unusable to a tool.
    if !input.is_object() {
        return None;
    }

    Some((name.to_string(), input))
}

/// How much of one tool result is allowed into the conversation.
///
/// Roughly 8k tokens at the estimator's four-characters-per-token. Chosen to be generous for a
/// source file or a CLI transcript while staying a small fraction of the smallest window this
/// build budgets against, so no single result can dominate the context.
const MAX_TOOL_RESULT_CHARS: usize = 32_000;

/// Clamps a tool result to something a context window can hold.
///
/// `read_file` reads whatever is on disk and the CLI tools return whatever the process printed;
/// neither has an upper bound, and the result went into the history verbatim. One large file was
/// enough to blow the window in a single step — and because the budget is checked *before* a
/// request rather than after a tool returns, the overflow was only discovered on the next turn,
/// when the history already held it.
///
/// The clamp keeps both ends: the head carries the shape of the output and the tail carries the
/// error or summary that a command prints last. The elision is stated in-band, because a model
/// that cannot tell truncated output from complete output will draw conclusions from the gap.
pub fn clamp_tool_result(result: String) -> String {
    if result.len() <= MAX_TOOL_RESULT_CHARS {
        return result;
    }

    // Half the budget each way, split on character boundaries so the result stays valid UTF-8.
    let half = MAX_TOOL_RESULT_CHARS / 2;
    let head_end = floor_boundary(&result, half);
    let tail_start = ceil_boundary(&result, result.len() - half);
    let dropped = tail_start - head_end;

    format!(
        "{}\n\n[... {} bytes elided by Procyon: this tool result was too large for the context \
         window. Narrow the call — a more specific path, a grep, or a smaller range — if the \
         middle matters. ...]\n\n{}",
        &result[..head_end],
        dropped,
        &result[tail_start..]
    )
}

fn floor_boundary(s: &str, mut at: usize) -> usize {
    while at > 0 && !s.is_char_boundary(at) {
        at -= 1;
    }
    at
}

fn ceil_boundary(s: &str, mut at: usize) -> usize {
    while at < s.len() && !s.is_char_boundary(at) {
        at += 1;
    }
    at
}

fn kind_of(value: &serde_json::Value) -> &'static str {
    match value {
        serde_json::Value::Null => "null",
        serde_json::Value::Bool(_) => "a boolean",
        serde_json::Value::Number(_) => "a number",
        serde_json::Value::String(_) => "a string",
        serde_json::Value::Array(_) => "an array",
        serde_json::Value::Object(_) => "an object",
    }
}

/// What the request actually cost, used to anchor the local token estimate.
#[derive(Debug, Default, Clone, Copy, PartialEq)]
pub struct TokenUsage {
    pub input: usize,
    pub cache_read: usize,
    pub cache_write: usize,
    pub output: usize,
}

impl TokenUsage {
    pub fn total(&self) -> usize {
        self.input + self.cache_read + self.cache_write + self.output
    }
}

#[derive(Debug)]
pub struct StreamOutcome {
    pub blocks: Vec<ContentPart>,
    // The loop decides whether to continue from the presence of tool_use blocks, so the reason is
    // carried for diagnostics rather than control flow.
    #[allow(dead_code)]
    pub stop_reason: Option<String>,
    pub usage: Option<TokenUsage>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn usage_total_sums_input_cache_and_output() {
        let usage = TokenUsage {
            input: 1200,
            cache_read: 400,
            cache_write: 30,
            output: 915,
        };
        assert_eq!(usage.total(), 2545);
    }

    #[test]
    fn tool_results_share_one_user_message() {
        let msg = Message::tool_results(vec![
            ("id_a".to_string(), "ra".to_string()),
            ("id_b".to_string(), "rb".to_string()),
        ]);
        assert_eq!(msg.role, Role::User);
        assert_eq!(
            msg.content.len(),
            2,
            "the API requires one user message holding every tool_result of a turn"
        );
    }

    // The session log stores these verbatim, so a resumed conversation has to deserialize the
    // shape an earlier run wrote.
    #[test]
    fn content_parts_round_trip_through_serde() {
        let parts = vec![
            ContentPart::Text {
                text: "hi".to_string(),
            },
            ContentPart::ToolUse {
                id: "t1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({"path": "a.rs"}),
            },
            ContentPart::ToolResult {
                tool_use_id: "t1".to_string(),
                content: "ok".to_string(),
            },
        ];

        let written = serde_json::to_string(&parts).unwrap();
        let read: Vec<ContentPart> = serde_json::from_str(&written).unwrap();

        assert_eq!(read, parts);
    }

    #[test]
    fn well_formed_arguments_pass_through() {
        assert_eq!(
            tool_input("read", r#"{"path":"a.rs"}"#),
            serde_json::json!({"path": "a.rs"})
        );
    }

    // A tool taking no arguments sends no fragments at all, and some providers spell the same
    // thing as a literal `null`. Both are the zero-argument case, not a failure.
    #[test]
    fn the_zero_argument_spellings_all_yield_an_object() {
        for raw in ["", "   ", "{}", "null"] {
            assert_eq!(
                tool_input("list", raw),
                serde_json::json!({}),
                "raw was {:?}",
                raw
            );
        }
    }

    // Anything that is not an object is unusable — tools read their arguments by key — but it must
    // not be silent, because the call still goes out as if it had no arguments.
    #[test]
    fn a_non_object_is_emptied_and_recorded() {
        let _guard = crate::diag::test_lock();

        for raw in ["[1,2]", "42", "\"text\"", "{invalid"] {
            crate::diag::drain();
            assert_eq!(tool_input("read", raw), serde_json::json!({}));
            assert!(
                !crate::diag::drain().is_empty(),
                "nothing recorded for {:?}",
                raw
            );
        }
    }

    #[test]
    fn a_result_that_fits_is_returned_untouched() {
        let small = "ok".repeat(100);
        assert_eq!(clamp_tool_result(small.clone()), small);
    }

    // The regression: one unbounded `read_file` used to be able to fill the whole window.
    #[test]
    fn an_oversized_result_is_clamped_and_says_so() {
        let huge = "x".repeat(MAX_TOOL_RESULT_CHARS * 3);
        let clamped = clamp_tool_result(huge);

        assert!(
            clamped.len() < MAX_TOOL_RESULT_CHARS + 500,
            "clamped to {} bytes",
            clamped.len()
        );
        assert!(
            clamped.contains("elided by Procyon"),
            "the model must be able to tell truncated output from complete output"
        );
    }

    // A command prints its error last, so the tail is the half most worth keeping.
    #[test]
    fn both_ends_of_an_oversized_result_survive() {
        let body = format!(
            "FIRST LINE\n{}\nerror: deploy failed",
            "filler ".repeat(MAX_TOOL_RESULT_CHARS)
        );
        let clamped = clamp_tool_result(body);

        assert!(clamped.starts_with("FIRST LINE"));
        assert!(clamped.ends_with("error: deploy failed"));
    }

    // Cutting a multi-byte character in half would panic on the slice.
    #[test]
    fn clamping_never_splits_a_character() {
        for pad in 0..4 {
            let body = format!("{}{}", "a".repeat(pad), "é".repeat(MAX_TOOL_RESULT_CHARS));
            let clamped = clamp_tool_result(body);
            assert!(clamped.contains("elided"), "pad {} was not clamped", pad);
        }
    }

    // --- tool calls written as prose ---------------------------------------------------------

    fn text(body: &str) -> Vec<ContentPart> {
        vec![ContentPart::Text {
            text: body.to_string(),
        }]
    }

    fn recovered_call(blocks: Vec<ContentPart>) -> Option<(String, serde_json::Value)> {
        let (out, flagged) = recover_text_tool_calls(blocks);
        match out.into_iter().next() {
            Some(ContentPart::ToolUse { name, input, .. }) => {
                assert!(flagged, "a recovered call must be reported as one");
                Some((name, input))
            }
            _ => {
                assert!(!flagged, "nothing was recovered but the flag was set");
                None
            }
        }
    }

    // The observed failure: llama3.2 wrote the call into the reply and the JSON went to the screen
    // as the agent's answer.
    #[test]
    fn a_call_written_as_json_becomes_a_real_call() {
        let (name, input) = recovered_call(text(
            r#"{"name": "list_dir", "parameters": {"path": "/w/demo"}}"#,
        ))
        .expect("recovered");

        assert_eq!(name, "list_dir");
        assert_eq!(input, serde_json::json!({"path": "/w/demo"}));
    }

    #[test]
    fn every_spelling_of_the_arguments_is_understood() {
        for key in ["parameters", "arguments", "input", "args"] {
            let body = format!(r#"{{"name": "grep", "{}": {{"q": "x"}}}}"#, key);
            let (_, input) = recovered_call(text(&body)).unwrap_or_else(|| panic!("{}", key));
            assert_eq!(input, serde_json::json!({"q": "x"}), "{}", key);
        }
    }

    #[test]
    fn the_wrappers_models_put_around_a_call_are_stripped() {
        for body in [
            r#"<tool_call>{"name": "glob", "arguments": {}}</tool_call>"#,
            "```json\n{\"name\": \"glob\", \"arguments\": {}}\n```",
            "```\n{\"name\": \"glob\", \"arguments\": {}}\n```",
        ] {
            let (name, _) = recovered_call(text(body)).unwrap_or_else(|| panic!("{}", body));
            assert_eq!(name, "glob", "{}", body);
        }
    }

    // A wrapper declares the intent, so a call with no arguments inside one is unambiguous.
    #[test]
    fn a_declared_call_with_no_arguments_is_still_a_call() {
        let (name, input) =
            recovered_call(text(r#"<tool_call>{"name": "project_info"}</tool_call>"#))
                .expect("recovered");
        assert_eq!(name, "project_info");
        assert_eq!(input, serde_json::json!({}));
    }

    // Bare, it is not. `{"name": ..., "version": ...}` is a `package.json`, and this harness asks
    // the model to read those constantly — see the false-positive tests below. Missing a
    // no-argument call costs the old behaviour; running one costs the user their answer.
    #[test]
    fn a_bare_object_with_only_a_name_is_not_treated_as_a_call() {
        assert!(recovered_call(text(r#"{"name": "project_info"}"#)).is_none());
    }

    // An unknown name is recovered rather than dropped: the registry answers "Unknown tool", and
    // that reaches the model as a tool_result telling it what it got wrong.
    #[test]
    fn an_unknown_tool_name_is_still_recovered_so_the_model_hears_back() {
        let (name, _) = recovered_call(text(r#"{"name": "ls", "parameters": {}}"#)).expect("ls");
        assert_eq!(name, "ls");
    }

    // --- what must NOT be executed -----------------------------------------------------------

    // A model explaining a call is talking to the user. Executing that turns documentation into an
    // action, which is the one way this fix could do real damage.
    #[test]
    fn a_call_described_inside_a_sentence_is_left_as_prose() {
        assert!(recovered_call(text(
            r#"You could call {"name": "write_file", "parameters": {"path": "a"}} to do that."#
        ))
        .is_none());
    }

    #[test]
    fn ordinary_prose_and_ordinary_json_are_left_alone() {
        for body in [
            "Os arquivos no diretório atual são: contracts/, src/.",
            // JSON, but not a tool call.
            r#"{"path": "a.rs", "size": 12}"#,
            // Has a `name`, but nothing that could be arguments.
            r#"{"name": "my-app", "version": "0.1.0"}"#,
            "",
            "{",
        ] {
            assert!(recovered_call(text(body)).is_none(), "recovered {:?}", body);
        }
    }

    // `package.json` has a `name`; showing one to the user must not run anything.
    #[test]
    fn a_file_the_model_is_quoting_is_not_a_tool_call() {
        let body = "```json\n{\"name\": \"my-app\", \"scripts\": {\"dev\": \"vite\"}}\n```";
        // `scripts` is not an arguments key, so there is nothing to call with.
        assert!(recovered_call(text(body)).is_none());
    }

    // If the model managed the tool API, its own call stands and the text beside it is prose.
    #[test]
    fn a_real_tool_call_in_the_same_message_disables_recovery() {
        let blocks = vec![
            ContentPart::Text {
                text: r#"{"name": "list_dir", "parameters": {}}"#.to_string(),
            },
            ContentPart::ToolUse {
                id: "t1".to_string(),
                name: "grep".to_string(),
                input: serde_json::json!({}),
            },
        ];

        let (out, recovered) = recover_text_tool_calls(blocks);
        assert!(!recovered);
        assert!(matches!(out[0], ContentPart::Text { .. }));
    }

    #[test]
    fn a_recovered_call_gets_an_id_a_tool_result_can_answer() {
        let (out, _) = recover_text_tool_calls(text(r#"{"name": "glob", "arguments": {}}"#));
        match &out[0] {
            // Every `tool_use` must be answerable by a matching `tool_result`, or the next
            // request is rejected outright.
            ContentPart::ToolUse { id, .. } => assert!(!id.is_empty()),
            other => panic!("expected a tool call, got {:?}", other),
        }
    }

    #[test]
    fn role_display_and_from_str_agree() {
        for role in [Role::User, Role::Assistant] {
            assert_eq!(role.to_string().parse::<Role>().unwrap(), role);
        }
    }
}