kcode-k1-chat-box 0.1.1

#![forbid(unsafe_code)]

use std::{error::Error, fmt, str::FromStr};

use semver::Version;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

pub const TOOL_CALL_TYPE: &str = "Tool Call";
pub const TOOL_RESULT_TYPE: &str = "Tool Result";
pub const TOOL_CALL_HIDDEN_TYPE: &str = "k1.tool-call/1.0.0";
pub const TOOL_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/1.0.0";

const CALL_PREFIX: &str = "k1.tool-call/";
const RESULT_PREFIX: &str = "k1.tool-result/";

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct BoxId(u64);

impl BoxId {
    pub const fn new(value: u64) -> Self {
        Self(value)
    }

    pub const fn get(self) -> u64 {
        self.0
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ChatBox {
    id: BoxId,
    box_type: String,
    contents: String,
    hidden_type: String,
    hidden_contents: String,
}

impl ChatBox {
    pub fn new(
        id: BoxId,
        box_type: String,
        contents: String,
        hidden_type: String,
        hidden_contents: String,
    ) -> Self {
        Self {
            id,
            box_type,
            contents,
            hidden_type,
            hidden_contents,
        }
    }

    pub const fn id(&self) -> BoxId {
        self.id
    }

    pub fn box_type(&self) -> &str {
        &self.box_type
    }

    pub fn contents(&self) -> &str {
        &self.contents
    }

    pub fn hidden_type(&self) -> &str {
        &self.hidden_type
    }

    pub fn hidden_contents(&self) -> &str {
        &self.hidden_contents
    }

    pub fn tool_call(id: BoxId, call: ToolCall) -> Result<Self, EnvelopeError> {
        Ok(Self::new(
            id,
            TOOL_CALL_TYPE.into(),
            call_contents(&call)?,
            TOOL_CALL_HIDDEN_TYPE.into(),
            compact_json(&CallHidden::from(&call))?,
        ))
    }

    pub fn tool_result(id: BoxId, result: ToolResult) -> Result<Self, EnvelopeError> {
        result.validate()?;
        Ok(Self::new(
            id,
            TOOL_RESULT_TYPE.into(),
            result_contents(&result),
            TOOL_RESULT_HIDDEN_TYPE.into(),
            compact_json(&ResultHidden::from(&result))?,
        ))
    }

    pub fn tool_call_metadata(&self) -> Result<Option<ToolCall>, EnvelopeError> {
        if self.box_type != TOOL_CALL_TYPE || !supported(&self.hidden_type, CALL_PREFIX)? {
            return Ok(None);
        }
        let hidden: CallHidden = decode_json(&self.hidden_contents)?;
        let call = ToolCall::new(
            ToolCallId::from_str(&hidden.call_id)?,
            hidden.tool,
            hidden.tool_version,
            hidden.arguments,
        )?;
        if call_contents(&call)? != self.contents {
            return Err(EnvelopeError::InvalidVisibleContent);
        }
        Ok(Some(call))
    }

    pub fn tool_result_metadata(&self) -> Result<Option<ToolResult>, EnvelopeError> {
        if self.box_type != TOOL_RESULT_TYPE || !supported(&self.hidden_type, RESULT_PREFIX)? {
            return Ok(None);
        }
        let hidden: ResultHidden = decode_json(&self.hidden_contents)?;
        let status = ToolResultStatus::from_str(&hidden.status)?;
        let view = parse_result_view(&self.contents, &hidden.call_id, &hidden.tool, status)?;
        ToolResult::new(
            ToolCallId::from_str(&hidden.call_id)?,
            BoxId::new(hidden.originating_call_box_id),
            hidden.tool,
            hidden.tool_version,
            status,
            hidden.data,
            view,
        )
        .map(Some)
    }
}

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ToolCallId(u64);

impl ToolCallId {
    pub fn new(value: u64) -> Result<Self, EnvelopeError> {
        if value == 0 {
            return Err(EnvelopeError::InvalidCallId);
        }
        Ok(Self(value))
    }

    pub const fn get(self) -> u64 {
        self.0
    }
}

impl fmt::Display for ToolCallId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "c{}", self.0)
    }
}

impl FromStr for ToolCallId {
    type Err = EnvelopeError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        let Some(number) = value.strip_prefix('c') else {
            return Err(EnvelopeError::InvalidCallId);
        };
        if number.is_empty() || number.starts_with('0') {
            return Err(EnvelopeError::InvalidCallId);
        }
        let value = number
            .parse::<u64>()
            .map_err(|_| EnvelopeError::InvalidCallId)?;
        Self::new(value)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct ToolCall {
    call_id: ToolCallId,
    tool: String,
    tool_version: String,
    arguments: Value,
}

impl ToolCall {
    pub fn new(
        call_id: ToolCallId,
        tool: String,
        tool_version: String,
        arguments: Value,
    ) -> Result<Self, EnvelopeError> {
        validate_tool(&tool)?;
        validate_version(&tool_version)?;
        let Some(properties) = arguments.as_object() else {
            return Err(EnvelopeError::InvalidArguments);
        };
        if properties.contains_key("tool") {
            return Err(EnvelopeError::InvalidArguments);
        }
        Ok(Self {
            call_id,
            tool,
            tool_version,
            arguments,
        })
    }

    pub const fn call_id(&self) -> ToolCallId {
        self.call_id
    }

    pub fn tool(&self) -> &str {
        &self.tool
    }

    pub fn tool_version(&self) -> &str {
        &self.tool_version
    }

    pub fn arguments(&self) -> &Value {
        &self.arguments
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ToolResultStatus {
    Ok,
    Error,
}

impl ToolResultStatus {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Ok => "ok",
            Self::Error => "error",
        }
    }
}

impl FromStr for ToolResultStatus {
    type Err = EnvelopeError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "ok" => Ok(Self::Ok),
            "error" => Ok(Self::Error),
            _ => Err(EnvelopeError::InvalidStatus),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum ResultView {
    OneLine(String),
    Multiline(String),
    Error(String),
}

#[derive(Clone, Debug, PartialEq)]
pub struct ToolResult {
    call_id: ToolCallId,
    originating_call_box_id: BoxId,
    tool: String,
    tool_version: String,
    status: ToolResultStatus,
    data: Value,
    view: ResultView,
}

impl ToolResult {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        call_id: ToolCallId,
        originating_call_box_id: BoxId,
        tool: String,
        tool_version: String,
        status: ToolResultStatus,
        data: Value,
        view: ResultView,
    ) -> Result<Self, EnvelopeError> {
        let result = Self {
            call_id,
            originating_call_box_id,
            tool,
            tool_version,
            status,
            data,
            view,
        };
        result.validate()?;
        Ok(result)
    }

    pub const fn call_id(&self) -> ToolCallId {
        self.call_id
    }

    pub const fn originating_call_box_id(&self) -> BoxId {
        self.originating_call_box_id
    }

    pub fn tool(&self) -> &str {
        &self.tool
    }

    pub fn tool_version(&self) -> &str {
        &self.tool_version
    }

    pub const fn status(&self) -> ToolResultStatus {
        self.status
    }

    pub fn data(&self) -> &Value {
        &self.data
    }

    pub fn view(&self) -> &ResultView {
        &self.view
    }

    fn validate(&self) -> Result<(), EnvelopeError> {
        validate_tool(&self.tool)?;
        validate_version(&self.tool_version)?;
        if self.originating_call_box_id.get() == 0 {
            return Err(EnvelopeError::InvalidOriginatingCallBoxId);
        }
        match (self.status, &self.view) {
            (ToolResultStatus::Ok, ResultView::OneLine(text)) if is_line(text) => Ok(()),
            (ToolResultStatus::Ok, ResultView::Multiline(_)) => Ok(()),
            (ToolResultStatus::Error, ResultView::Error(text)) if is_line(text) => Ok(()),
            _ => Err(EnvelopeError::InvalidVisibleContent),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EnvelopeError {
    InvalidCallId,
    InvalidTool,
    InvalidToolVersion,
    InvalidArguments,
    InvalidOriginatingCallBoxId,
    InvalidStatus,
    MalformedEnvelope,
    UnsupportedEnvelopeVersion,
    InvalidVisibleContent,
}

impl fmt::Display for EnvelopeError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let message = match self {
            Self::InvalidCallId => "invalid tool call ID",
            Self::InvalidTool => "invalid tool name",
            Self::InvalidToolVersion => "invalid tool version",
            Self::InvalidArguments => "invalid tool arguments",
            Self::InvalidOriginatingCallBoxId => "invalid originating call box ID",
            Self::InvalidStatus => "invalid tool result status",
            Self::MalformedEnvelope => "malformed tool envelope",
            Self::UnsupportedEnvelopeVersion => "unsupported tool envelope version",
            Self::InvalidVisibleContent => "invalid visible tool content",
        };
        formatter.write_str(message)
    }
}

impl Error for EnvelopeError {}

#[derive(Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
struct CallHidden {
    call_id: String,
    tool: String,
    tool_version: String,
    arguments: Value,
}

impl From<&ToolCall> for CallHidden {
    fn from(value: &ToolCall) -> Self {
        Self {
            call_id: value.call_id.to_string(),
            tool: value.tool.clone(),
            tool_version: value.tool_version.clone(),
            arguments: value.arguments.clone(),
        }
    }
}

#[derive(Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
struct ResultHidden {
    call_id: String,
    originating_call_box_id: u64,
    tool: String,
    tool_version: String,
    status: String,
    data: Value,
}

impl From<&ToolResult> for ResultHidden {
    fn from(value: &ToolResult) -> Self {
        Self {
            call_id: value.call_id.to_string(),
            originating_call_box_id: value.originating_call_box_id.get(),
            tool: value.tool.clone(),
            tool_version: value.tool_version.clone(),
            status: value.status.as_str().into(),
            data: value.data.clone(),
        }
    }
}

fn validate_tool(value: &str) -> Result<(), EnvelopeError> {
    if value.is_empty() || !is_line(value) {
        return Err(EnvelopeError::InvalidTool);
    }
    Ok(())
}

fn validate_version(value: &str) -> Result<(), EnvelopeError> {
    Version::parse(value)
        .map(|_| ())
        .map_err(|_| EnvelopeError::InvalidToolVersion)
}

fn is_line(value: &str) -> bool {
    !value.contains(['\n', '\r'])
}

fn supported(hidden_type: &str, prefix: &str) -> Result<bool, EnvelopeError> {
    let Some(version) = hidden_type.strip_prefix(prefix) else {
        return Ok(false);
    };
    let version = Version::parse(version).map_err(|_| EnvelopeError::MalformedEnvelope)?;
    Ok(version.major == 1)
}

fn compact_json<T: Serialize>(value: &T) -> Result<String, EnvelopeError> {
    serde_json::to_string(value).map_err(|_| EnvelopeError::MalformedEnvelope)
}

fn decode_json<T: for<'a> Deserialize<'a>>(value: &str) -> Result<T, EnvelopeError> {
    serde_json::from_str(value).map_err(|_| EnvelopeError::MalformedEnvelope)
}

fn readable_json(value: &Value) -> Result<String, EnvelopeError> {
    match value {
        Value::Array(values) => {
            let values = values
                .iter()
                .map(readable_json)
                .collect::<Result<Vec<_>, _>>()?;
            Ok(format!("[{}]", values.join(", ")))
        }
        Value::Object(properties) => {
            let properties = properties
                .iter()
                .map(|(key, value)| {
                    Ok(format!("{}: {}", compact_json(key)?, readable_json(value)?))
                })
                .collect::<Result<Vec<_>, EnvelopeError>>()?;
            Ok(format!("{{{}}}", properties.join(", ")))
        }
        _ => compact_json(value),
    }
}

fn call_contents(call: &ToolCall) -> Result<String, EnvelopeError> {
    let Some(arguments) = call.arguments.as_object() else {
        return Err(EnvelopeError::InvalidArguments);
    };
    let mut visible = Map::new();
    visible.insert("tool".into(), Value::String(call.tool.clone()));
    visible.extend(arguments.clone());
    Ok(format!(
        "Call ID: {}\nArgs: {}",
        call.call_id,
        readable_json(&Value::Object(visible))?
    ))
}

fn result_contents(result: &ToolResult) -> String {
    let prefix = format!("Call ID: {}\n{} call", result.call_id, result.tool);
    match &result.view {
        ResultView::OneLine(message) => format!("{prefix} result: {message}"),
        ResultView::Multiline(body) => format!("{prefix} result:\n\n{body}"),
        ResultView::Error(message) => format!("{prefix} error: {message}"),
    }
}

fn parse_result_view(
    contents: &str,
    call_id: &str,
    tool: &str,
    status: ToolResultStatus,
) -> Result<ResultView, EnvelopeError> {
    let call_id = ToolCallId::from_str(call_id)?;
    validate_tool(tool)?;
    let prefix = format!("Call ID: {call_id}\n{tool} call");
    match status {
        ToolResultStatus::Ok => {
            if let Some(message) = contents.strip_prefix(&format!("{prefix} result: ")) {
                if !is_line(message) {
                    return Err(EnvelopeError::InvalidVisibleContent);
                }
                return Ok(ResultView::OneLine(message.into()));
            }
            contents
                .strip_prefix(&format!("{prefix} result:\n\n"))
                .map(|body| ResultView::Multiline(body.into()))
                .ok_or(EnvelopeError::InvalidVisibleContent)
        }
        ToolResultStatus::Error => contents
            .strip_prefix(&format!("{prefix} error: "))
            .filter(|message| is_line(message))
            .map(|message| ResultView::Error(message.into()))
            .ok_or(EnvelopeError::InvalidVisibleContent),
    }
}

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

    fn call() -> ToolCall {
        ToolCall::new(
            ToolCallId::new(2).unwrap(),
            "Search".into(),
            "1.0.0".into(),
            json!({"query":{"tags":["rust", {"stable":true}]},"limit":2}),
        )
        .unwrap()
    }

    fn result(status: ToolResultStatus, view: ResultView) -> ToolResult {
        ToolResult::new(
            ToolCallId::new(2).unwrap(),
            BoxId::new(9),
            "Search".into(),
            "1.0.0".into(),
            status,
            json!({"matches":[1, 2]}),
            view,
        )
        .unwrap()
    }

    #[test]
    fn call_ids_are_canonical() {
        assert_eq!(ToolCallId::from_str("c1").unwrap().to_string(), "c1");
        assert_eq!(
            ToolCallId::from_str(&format!("c{}", u64::MAX))
                .unwrap()
                .get(),
            u64::MAX
        );
        for invalid in ["", "1", "c", "c0", "c01", "c-1", "c18446744073709551616"] {
            assert_eq!(
                ToolCallId::from_str(invalid),
                Err(EnvelopeError::InvalidCallId)
            );
        }
    }

    #[test]
    fn call_has_exact_readable_and_hidden_json() {
        let value = ChatBox::tool_call(BoxId::new(9), call()).unwrap();
        assert_eq!(
            value.contents(),
            "Call ID: c2\nArgs: {\"tool\": \"Search\", \"query\": {\"tags\": [\"rust\", {\"stable\": true}]}, \"limit\": 2}"
        );
        assert_eq!(
            value.hidden_contents(),
            "{\"callId\":\"c2\",\"tool\":\"Search\",\"toolVersion\":\"1.0.0\",\"arguments\":{\"query\":{\"tags\":[\"rust\",{\"stable\":true}]},\"limit\":2}}"
        );
        assert_eq!(value.tool_call_metadata().unwrap(), Some(call()));
    }

    #[test]
    fn all_result_views_are_exact_and_round_trip() {
        let cases = [
            (
                ToolResultStatus::Ok,
                ResultView::OneLine("success".into()),
                "Call ID: c2\nSearch call result: success",
            ),
            (
                ToolResultStatus::Ok,
                ResultView::Multiline("first\nsecond".into()),
                "Call ID: c2\nSearch call result:\n\nfirst\nsecond",
            ),
            (
                ToolResultStatus::Error,
                ResultView::Error("node unavailable".into()),
                "Call ID: c2\nSearch call error: node unavailable",
            ),
        ];
        for (status, view, expected) in cases {
            let result = result(status, view);
            let value = ChatBox::tool_result(BoxId::new(10), result.clone()).unwrap();
            assert_eq!(value.contents(), expected);
            assert_eq!(value.tool_result_metadata().unwrap(), Some(result));
        }
    }

    #[test]
    fn result_hidden_json_is_canonical() {
        let value = ChatBox::tool_result(
            BoxId::new(10),
            result(ToolResultStatus::Ok, ResultView::OneLine("success".into())),
        )
        .unwrap();
        assert_eq!(
            value.hidden_contents(),
            "{\"callId\":\"c2\",\"originatingCallBoxId\":9,\"tool\":\"Search\",\"toolVersion\":\"1.0.0\",\"status\":\"ok\",\"data\":{\"matches\":[1,2]}}"
        );
    }

    #[test]
    fn compatible_minor_accepts_options_and_unsupported_major_is_opaque() {
        let good = ChatBox::tool_call(BoxId::new(1), call()).unwrap();
        let future = ChatBox::new(
            BoxId::new(1),
            TOOL_CALL_TYPE.into(),
            good.contents().into(),
            "k1.tool-call/1.1.0".into(),
            format!("{{\"optional\":true,{}", &good.hidden_contents()[1..]),
        );
        let unsupported = ChatBox::new(
            BoxId::new(1),
            TOOL_CALL_TYPE.into(),
            "anything".into(),
            "k1.tool-call/2.0.0".into(),
            "not json".into(),
        );
        assert_eq!(future.tool_call_metadata().unwrap(), Some(call()));
        assert_eq!(unsupported.tool_call_metadata().unwrap(), None);
    }

    #[test]
    fn malformed_owned_envelopes_and_mismatched_visible_text_fail() {
        let malformed_version = ChatBox::new(
            BoxId::new(1),
            TOOL_CALL_TYPE.into(),
            "anything".into(),
            "k1.tool-call/not-semver".into(),
            "opaque".into(),
        );
        assert_eq!(
            malformed_version.tool_call_metadata(),
            Err(EnvelopeError::MalformedEnvelope)
        );

        let good = ChatBox::tool_call(BoxId::new(1), call()).unwrap();
        let mismatch = ChatBox::new(
            good.id(),
            good.box_type().into(),
            "different".into(),
            good.hidden_type().into(),
            good.hidden_contents().into(),
        );
        assert_eq!(
            mismatch.tool_call_metadata(),
            Err(EnvelopeError::InvalidVisibleContent)
        );
    }

    #[test]
    fn unknown_boxes_and_hidden_types_stay_opaque() {
        let unknown = ChatBox::new(
            BoxId::new(7),
            "Future Box".into(),
            "visible".into(),
            "future.hidden/not-semver".into(),
            "opaque".into(),
        );
        assert_eq!(unknown.tool_call_metadata().unwrap(), None);
        assert_eq!(unknown.tool_result_metadata().unwrap(), None);

        let unrelated = ChatBox::new(
            BoxId::new(8),
            TOOL_CALL_TYPE.into(),
            "visible".into(),
            "other.tool/1.0.0".into(),
            "opaque".into(),
        );
        assert_eq!(unrelated.tool_call_metadata().unwrap(), None);
    }

    #[test]
    fn constructors_reject_invalid_structures() {
        assert_eq!(
            ToolCall::new(
                ToolCallId::new(1).unwrap(),
                "Tool".into(),
                "1.0.0".into(),
                json!({"tool":"duplicate"}),
            ),
            Err(EnvelopeError::InvalidArguments)
        );
        assert_eq!(
            ToolResult::new(
                ToolCallId::new(1).unwrap(),
                BoxId::new(0),
                "Tool".into(),
                "1.0.0".into(),
                ToolResultStatus::Ok,
                Value::Null,
                ResultView::OneLine("success".into()),
            ),
            Err(EnvelopeError::InvalidOriginatingCallBoxId)
        );
    }
}