kcode-k1-chat-boxes 0.2.1

#![forbid(unsafe_code)]

pub const SYSTEM_MESSAGE_TYPE: &str = "System Message";
pub const USER_MESSAGE_TYPE: &str = "User Message";
pub const AGENT_MESSAGE_TYPE: &str = "Agent Message";
pub const AGENT_RESPONSE_TYPE: &str = "Agent Response";
pub const USER_ATTACHMENT_TYPE: &str = "User Attachment";
pub const AGENT_ATTACHMENT_TYPE: &str = "Agent Attachment";
pub const TOOL_CALL_TYPE: &str = "Tool Call";
pub const TOOL_MESSAGE_TYPE: &str = "Tool Message";
pub const TOOL_ATTACHMENT_TYPE: &str = "Tool Attachment";
pub const TOOL_RESULT_TYPE: &str = "Tool Result";
pub const ATTACHMENT_TYPE: &str = USER_ATTACHMENT_TYPE;

pub const TOOL_CALL_HIDDEN_TYPE: &str = "k1.tool-call/v1";
pub const TOOL_MESSAGE_HIDDEN_TYPE: &str = "k1.tool-message/v1";
pub const TOOL_MESSAGE_V2_HIDDEN_TYPE: &str = "k1.tool-message/v2";
pub const TOOL_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/v1";
pub const TOOL_RESULT_V2_HIDDEN_TYPE: &str = "k1.tool-result/v2";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
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, Copy, Debug, Eq, PartialEq)]
pub struct ToolCallId {
    nonce: [u8; 12],
    sequence: u64,
}

impl ToolCallId {
    pub const fn new(nonce: [u8; 12], sequence: u64) -> Self {
        Self { nonce, sequence }
    }

    pub const fn nonce(self) -> [u8; 12] {
        self.nonce
    }

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

#[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_metadata(&self) -> Result<Option<ProviderCall>, MetadataError> {
        let Some(fields) = self.metadata_fields(TOOL_CALL_TYPE, TOOL_CALL_HIDDEN_TYPE, 4)? else {
            return Ok(None);
        };
        Ok(Some(ProviderCall {
            tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
            name: fields[2].to_owned(),
            arguments: fields[3].to_owned(),
        }))
    }

    pub fn tool_message_metadata(&self) -> Result<Option<ToolMessageMetadata>, MetadataError> {
        let (fields, named) = match self.hidden_type.as_str() {
            TOOL_MESSAGE_HIDDEN_TYPE => (
                self.metadata_fields(TOOL_MESSAGE_TYPE, TOOL_MESSAGE_HIDDEN_TYPE, 5)?
                    .ok_or(MetadataError)?,
                false,
            ),
            TOOL_MESSAGE_V2_HIDDEN_TYPE => {
                if self.box_type.is_empty() {
                    return Err(MetadataError);
                }
                (
                    decode_fields(&self.hidden_contents, 5).ok_or(MetadataError)?,
                    true,
                )
            }
            _ => return Ok(None),
        };
        let message_index = parse_u64(fields[3])?;
        if message_index == 0 || (named && self.contents != fields[4]) {
            return Err(MetadataError);
        }
        Ok(Some(ToolMessageMetadata {
            tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
            originating_call: BoxId::new(parse_u64(fields[2])?),
            message_index,
            message: fields[4].to_owned(),
        }))
    }

    pub fn tool_result_metadata(&self) -> Result<Option<ToolResultMetadata>, MetadataError> {
        let fields = match self.hidden_type.as_str() {
            TOOL_RESULT_HIDDEN_TYPE => {
                self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_HIDDEN_TYPE, 5)?
            }
            TOOL_RESULT_V2_HIDDEN_TYPE => {
                self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_V2_HIDDEN_TYPE, 7)?
            }
            _ => None,
        };
        fields
            .map(|fields| parse_result_fields(&fields))
            .transpose()
    }

    pub fn tool_result_v2_metadata(&self) -> Result<Option<ToolResultV2Metadata>, MetadataError> {
        let Some(fields) = self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_V2_HIDDEN_TYPE, 7)?
        else {
            return Ok(None);
        };
        let parsed = parse_result_fields(&fields)?;
        Ok(Some(ToolResultV2Metadata {
            tool_call_id: parsed.tool_call_id,
            originating_call: parsed.originating_call,
            result: parsed.result,
            metadata_type: fields[5].to_owned(),
            metadata_contents: fields[6].to_owned(),
        }))
    }

    fn metadata_fields<'a>(
        &'a self,
        box_type: &str,
        hidden_type: &str,
        count: usize,
    ) -> Result<Option<Vec<&'a str>>, MetadataError> {
        match (self.hidden_type == hidden_type, self.box_type == box_type) {
            (false, _) => Ok(None),
            (true, false) => Err(MetadataError),
            (true, true) => decode_fields(&self.hidden_contents, count)
                .map(Some)
                .ok_or(MetadataError),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MetadataError;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ProviderCall {
    pub tool_call_id: ToolCallId,
    pub name: String,
    pub arguments: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolMessageMetadata {
    pub tool_call_id: ToolCallId,
    pub originating_call: BoxId,
    pub message_index: u64,
    pub message: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolResultMetadata {
    pub tool_call_id: ToolCallId,
    pub originating_call: BoxId,
    pub result: Result<String, String>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolResultV2Metadata {
    pub tool_call_id: ToolCallId,
    pub originating_call: BoxId,
    pub result: Result<String, String>,
    pub metadata_type: String,
    pub metadata_contents: String,
}

pub fn tool_call_box(call: &ProviderCall) -> ChatBox {
    let call_id = format_tool_call_id(call.tool_call_id);
    let nonce = encode_nonce(call.tool_call_id.nonce);
    let sequence = call.tool_call_id.sequence.to_string();
    let hidden_contents = encode_fields(&[&nonce, &sequence, &call.name, &call.arguments]);
    ChatBox::new(
        BoxId::new(0),
        TOOL_CALL_TYPE.to_owned(),
        format!(
            "Call ID: {call_id}\nCall Name: {}\nArguments:\n{}",
            call.name, call.arguments
        ),
        TOOL_CALL_HIDDEN_TYPE.to_owned(),
        hidden_contents,
    )
}

pub fn tool_message_box(metadata: &ToolMessageMetadata) -> Result<ChatBox, MetadataError> {
    let hidden_contents = tool_message_hidden_contents(metadata)?;
    let call_id = format_tool_call_id(metadata.tool_call_id);
    let origin = metadata.originating_call.get();
    let index = metadata.message_index;
    Ok(ChatBox::new(
        BoxId::new(0),
        TOOL_MESSAGE_TYPE.to_owned(),
        format!(
            "Call ID: {call_id}\nOriginating Call Box ID: {origin}\nMessage Index: {index}\nMessage:\n{}",
            metadata.message
        ),
        TOOL_MESSAGE_HIDDEN_TYPE.to_owned(),
        hidden_contents,
    ))
}

pub fn named_tool_message_box(
    metadata: &ToolMessageMetadata,
    title: &str,
) -> Result<ChatBox, MetadataError> {
    if title.is_empty() {
        return Err(MetadataError);
    }
    let hidden_contents = tool_message_hidden_contents(metadata)?;
    Ok(ChatBox::new(
        BoxId::new(0),
        title.to_owned(),
        metadata.message.clone(),
        TOOL_MESSAGE_V2_HIDDEN_TYPE.to_owned(),
        hidden_contents,
    ))
}

fn tool_message_hidden_contents(metadata: &ToolMessageMetadata) -> Result<String, MetadataError> {
    if metadata.message_index == 0 {
        return Err(MetadataError);
    }
    let nonce = encode_nonce(metadata.tool_call_id.nonce);
    let sequence = metadata.tool_call_id.sequence.to_string();
    let origin = metadata.originating_call.get().to_string();
    let index = metadata.message_index.to_string();
    Ok(encode_fields(&[
        &nonce,
        &sequence,
        &origin,
        &index,
        &metadata.message,
    ]))
}

pub fn tool_result_box(
    tool_call_id: ToolCallId,
    originating_call: BoxId,
    result: Result<String, String>,
) -> ChatBox {
    tool_result_box_inner(tool_call_id, originating_call, &result, None)
}

pub fn tool_result_v2_box(metadata: &ToolResultV2Metadata) -> ChatBox {
    tool_result_box_inner(
        metadata.tool_call_id,
        metadata.originating_call,
        &metadata.result,
        Some((&metadata.metadata_type, &metadata.metadata_contents)),
    )
}

fn tool_result_box_inner(
    tool_call_id: ToolCallId,
    originating_call: BoxId,
    result: &Result<String, String>,
    metadata: Option<(&str, &str)>,
) -> ChatBox {
    let call_id = format_tool_call_id(tool_call_id);
    let nonce = encode_nonce(tool_call_id.nonce);
    let sequence = tool_call_id.sequence.to_string();
    let origin = originating_call.get().to_string();
    let (hidden_status, visible_status, raw_result) = match result {
        Ok(contents) => ("ok", "ok", contents),
        Err(contents) => ("err", "error", contents),
    };
    let mut fields = vec![
        nonce.as_str(),
        sequence.as_str(),
        origin.as_str(),
        hidden_status,
        raw_result.as_str(),
    ];
    if let Some((metadata_type, metadata_contents)) = metadata {
        fields.extend([metadata_type, metadata_contents]);
    }
    let hidden_contents = encode_fields(&fields);
    let hidden_type = metadata
        .map(|_| TOOL_RESULT_V2_HIDDEN_TYPE)
        .unwrap_or(TOOL_RESULT_HIDDEN_TYPE);
    ChatBox::new(
        BoxId::new(0),
        TOOL_RESULT_TYPE.to_owned(),
        format!(
            "Call ID: {call_id}\nOriginating Call Box ID: {origin}\nStatus: {visible_status}\nResult:\n{raw_result}"
        ),
        hidden_type.to_owned(),
        hidden_contents,
    )
}

fn parse_result_fields(fields: &[&str]) -> Result<ToolResultMetadata, MetadataError> {
    let result = match fields[3] {
        "ok" => Ok(fields[4].to_owned()),
        "err" => Err(fields[4].to_owned()),
        _ => return Err(MetadataError),
    };
    Ok(ToolResultMetadata {
        tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
        originating_call: BoxId::new(parse_u64(fields[2])?),
        result,
    })
}

fn parse_tool_call_id(nonce: &str, sequence: &str) -> Result<ToolCallId, MetadataError> {
    let nonce = decode_nonce(nonce).ok_or(MetadataError)?;
    Ok(ToolCallId::new(nonce, parse_u64(sequence)?))
}

fn parse_u64(value: &str) -> Result<u64, MetadataError> {
    value.parse::<u64>().map_err(|_| MetadataError)
}

fn format_tool_call_id(tool_call_id: ToolCallId) -> String {
    format!(
        "{}/{}",
        encode_nonce(tool_call_id.nonce),
        tool_call_id.sequence
    )
}

fn encode_fields(fields: &[&str]) -> String {
    let mut encoded = String::new();
    for field in fields {
        encoded.push_str(&field.len().to_string());
        encoded.push(':');
        encoded.push_str(field);
    }
    encoded
}

fn decode_fields(input: &str, count: usize) -> Option<Vec<&str>> {
    let mut fields = Vec::with_capacity(count);
    let mut cursor = 0;
    for _ in 0..count {
        let colon_offset = input
            .as_bytes()
            .get(cursor..)?
            .iter()
            .position(|byte| *byte == b':')?;
        let colon = cursor.checked_add(colon_offset)?;
        let length = input.get(cursor..colon)?.parse::<usize>().ok()?;
        let start = colon.checked_add(1)?;
        let end = start.checked_add(length)?;
        fields.push(input.get(start..end)?);
        cursor = end;
    }
    (cursor == input.len()).then_some(fields)
}

fn encode_nonce(nonce: [u8; 12]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut encoded = String::with_capacity(24);
    for byte in nonce {
        encoded.push(char::from(HEX[usize::from(byte >> 4)]));
        encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
    }
    encoded
}

fn decode_nonce(value: &str) -> Option<[u8; 12]> {
    if value.len() != 24 {
        return None;
    }
    let mut nonce = [0; 12];
    for (slot, digits) in nonce.iter_mut().zip(value.as_bytes().chunks_exact(2)) {
        let digits = std::str::from_utf8(digits).ok()?;
        *slot = u8::from_str_radix(digits, 16).ok()?;
    }
    Some(nonce)
}

#[cfg(test)]
mod tests;