haematite-cli 0.21.3

use super::handshake::Version;
use agent_client_protocol_schema::{v1, v2};
use serde_json::{Value, json};

pub(super) fn project(
    version: Version,
    params: Value,
    session: &str,
    bound: usize,
) -> anyhow::Result<Option<Value>> {
    let (id, update) = match version {
        Version::One => {
            let note: v1::SessionNotification = super::handshake::decode(params, "session_update")?;
            (
                note.session_id.to_string(),
                serde_json::to_value(note.update)?,
            )
        }
        Version::Two => {
            let note: v2::UpdateSessionNotification =
                super::handshake::decode(params, "session_update")?;
            (
                note.session_id.to_string(),
                serde_json::to_value(note.update)?,
            )
        }
    };
    anyhow::ensure!(id == session, "stream_session_differs");
    let kind = update["sessionUpdate"]
        .as_str()
        .ok_or_else(|| anyhow::anyhow!("stream_update_kind_missing"))?;
    let mut item = match kind {
        "agent_message_chunk" | "user_message_chunk" | "agent_thought_chunk"
            if update["content"]["type"] == "text" =>
        {
            json!({"kind":if kind=="agent_thought_chunk" {"thought"}else{"text"},
                "role":if kind=="user_message_chunk" {"user"}else{"assistant"},
                "text":update["content"]["text"],"message_id":update["messageId"],"complete":false})
        }
        "agent_message" | "user_message" => {
            let mut item = json!({"kind":"text","role":if kind=="user_message" {"user"}else{"assistant"},"message_id":update["messageId"],"complete":false,"replace":false});
            // An omitted patch leaves the previous content intact; null explicitly clears it.
            if let Some(content) = update.get("content") {
                let mut text = String::new();
                if let Some(blocks) = content.as_array() {
                    for block in blocks.iter().filter(|block| block["type"] == "text") {
                        if let Some(part) = block["text"].as_str() {
                            text.push_str(part);
                        }
                    }
                }
                item["text"] = json!(text);
                item["replace"] = json!(true);
                item["complete"] = json!(true);
            }
            item
        }
        "tool_call" | "tool_call_update" => {
            let mut item = json!({"kind":"tool_call","tool_call_id":update["toolCallId"]});
            for (source, target) in [
                ("title", "title"),
                ("status", "status"),
                ("kind", "tool_kind"),
                ("rawInput", "input"),
                ("rawOutput", "output"),
            ] {
                if let Some(value) = update.get(source) {
                    item[target] = value.clone();
                }
            }
            item
        }
        "plan" | "plan_update" => {
            json!({"kind":"plan","plan_id":if kind=="plan" {update["planId"].clone()}else{update["plan"]["planId"].clone()},"entries":if kind=="plan" {update["entries"].clone()}else{update["plan"]["entries"].clone()}})
        }
        "available_commands_update" => {
            json!({"kind":"available_commands","commands":update["availableCommands"]})
        }
        "config_option_update" => {
            json!({"kind":"config_options","options":update["configOptions"]})
        }
        "state_update" => {
            json!({"kind":"state","state":update["state"],"stop_reason":update["stopReason"]})
        }
        _ => json!({"kind":"omitted","reason":"content outside the supported stream projection"}),
    };
    item["truncated"] = json!(false);
    fit(&mut item, bound)?;
    Ok(Some(item))
}

pub(super) fn fit(item: &mut Value, bound: usize) -> anyhow::Result<()> {
    if serde_json::to_vec(item)?.len() <= bound {
        return Ok(());
    }
    item["truncated"] = json!(true);
    if let Some(fields) = item.as_object_mut() {
        fields.remove("input");
        fields.remove("output");
    }
    for field in ["entries", "commands", "options", "sessions"] {
        while serde_json::to_vec(item)?.len() > bound {
            let Some(entries) = item.get_mut(field).and_then(Value::as_array_mut) else {
                break;
            };
            if entries.pop().is_none() {
                break;
            }
        }
    }
    for field in ["text", "title", "message", "reason"] {
        if serde_json::to_vec(item)?.len() <= bound {
            return Ok(());
        }
        let Some(text) = item.get(field).and_then(Value::as_str).map(str::to_owned) else {
            continue;
        };
        item[field] = json!("");
        if serde_json::to_vec(item)?.len() > bound {
            continue;
        }
        let mut low = 0;
        let mut high = text.chars().count();
        while low < high {
            let middle = low + (high - low).div_ceil(2);
            item[field] = json!(text.chars().take(middle).collect::<String>());
            if serde_json::to_vec(item)?.len() <= bound {
                low = middle;
            } else {
                high = middle - 1;
            }
        }
        item[field] = json!(text.chars().take(low).collect::<String>());
    }
    anyhow::ensure!(
        serde_json::to_vec(item)?.len() <= bound,
        "stream_envelope_exceeds_content_bytes"
    );
    Ok(())
}