#![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_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 Some(fields) = self.metadata_fields(TOOL_MESSAGE_TYPE, TOOL_MESSAGE_HIDDEN_TYPE, 5)?
else {
return Ok(None);
};
let message_index = parse_u64(fields[3])?;
if message_index == 0 {
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> {
if metadata.message_index == 0 {
return Err(MetadataError);
}
let call_id = format_tool_call_id(metadata.tool_call_id);
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();
let hidden_contents = encode_fields(&[&nonce, &sequence, &origin, &index, &metadata.message]);
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 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;