#![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 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_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/v1";
#[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, 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>, RecoveryError> {
if self.box_type != TOOL_CALL_TYPE || self.hidden_type != TOOL_CALL_HIDDEN_TYPE {
return Ok(None);
}
let fields = decode_fields(&self.hidden_contents, 4)
.ok_or(RecoveryError::MalformedToolConvention)?;
let nonce = decode_nonce(fields[0]).ok_or(RecoveryError::MalformedToolConvention)?;
let sequence = fields[1]
.parse::<u64>()
.map_err(|_| RecoveryError::MalformedToolConvention)?;
Ok(Some(ProviderCall {
tool_call_id: ToolCallId::new(nonce, sequence),
name: fields[2].to_owned(),
arguments: fields[3].to_owned(),
}))
}
pub fn tool_result_metadata(&self) -> Result<Option<ToolResultMetadata>, RecoveryError> {
if self.box_type != TOOL_RESULT_TYPE || self.hidden_type != TOOL_RESULT_HIDDEN_TYPE {
return Ok(None);
}
let fields = decode_fields(&self.hidden_contents, 5)
.ok_or(RecoveryError::MalformedToolConvention)?;
let nonce = decode_nonce(fields[0]).ok_or(RecoveryError::MalformedToolConvention)?;
let sequence = fields[1]
.parse::<u64>()
.map_err(|_| RecoveryError::MalformedToolConvention)?;
let originating_call = fields[2]
.parse::<u64>()
.map_err(|_| RecoveryError::MalformedToolConvention)?;
let result = match fields[3] {
"ok" => Ok(fields[4].to_owned()),
"err" => Err(fields[4].to_owned()),
_ => return Err(RecoveryError::MalformedToolConvention),
};
Ok(Some(ToolResultMetadata {
tool_call_id: ToolCallId::new(nonce, sequence),
originating_call: BoxId::new(originating_call),
result,
}))
}
}
#[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 ProviderCall {
pub tool_call_id: ToolCallId,
pub name: String,
pub arguments: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DispatchedToolCall {
pub tool_call_id: ToolCallId,
pub call_box_id: BoxId,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolResultMetadata {
pub tool_call_id: ToolCallId,
pub originating_call: BoxId,
pub result: Result<String, String>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransitionError {
InvalidPhase,
BoxIdOverflow,
DuplicateToolCall,
UnknownToolCall,
DuplicateReturn,
MalformedToolConvention,
WrongOriginatingCall,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecoveryError {
NonContiguousBoxId,
MalformedToolConvention,
DuplicateToolCall,
UnknownToolCall,
DuplicateReturn,
WrongOriginatingCall,
}
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_result_box(
tool_call_id: ToolCallId,
originating_call: BoxId,
result: Result<String, String>,
) -> 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 originating_call_contents = originating_call.get().to_string();
let (hidden_status, visible_status, result_contents) = match &result {
Ok(contents) => ("ok", "ok", contents.clone()),
Err(contents) => ("err", "error", contents.clone()),
};
let hidden_contents = encode_fields(&[
&nonce,
&sequence,
&originating_call_contents,
hidden_status,
&result_contents,
]);
ChatBox::new(
BoxId::new(0),
TOOL_RESULT_TYPE.to_owned(),
format!(
"Call ID: {call_id}\nOriginating Call Box ID: {originating_call_contents}\nStatus: {visible_status}\nResult:\n{result_contents}"
),
TOOL_RESULT_HIDDEN_TYPE.to_owned(),
hidden_contents,
)
}
pub struct Chatend {
boxes: Vec<ChatBox>,
round_active: bool,
active_arrivals: Vec<ChatBox>,
}
impl Chatend {
pub const fn new() -> Self {
Self {
boxes: Vec::new(),
round_active: false,
active_arrivals: Vec::new(),
}
}
pub fn boxes(&self) -> &[ChatBox] {
&self.boxes
}
pub const fn round_active(&self) -> bool {
self.round_active
}
pub fn accept_box(
&mut self,
box_type: String,
contents: String,
hidden_type: String,
hidden_contents: String,
) -> Result<Option<BoxId>, TransitionError> {
self.accept_arrival(ChatBox::new(
BoxId::new(0),
box_type,
contents,
hidden_type,
hidden_contents,
))
}
pub fn accept_system(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
self.accept_box(
SYSTEM_MESSAGE_TYPE.to_owned(),
contents,
String::new(),
String::new(),
)
}
pub fn accept_user(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
self.accept_box(
USER_MESSAGE_TYPE.to_owned(),
contents,
String::new(),
String::new(),
)
}
pub fn accept_attachment(
&mut self,
contents: String,
hidden_type: String,
hidden_contents: String,
) -> Result<Option<BoxId>, TransitionError> {
self.accept_box(
USER_ATTACHMENT_TYPE.to_owned(),
contents,
hidden_type,
hidden_contents,
)
}
pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
if self.round_active {
return Err(TransitionError::InvalidPhase);
}
self.round_active = true;
Ok(self.boxes.last().map(ChatBox::id))
}
pub fn append_stage(
&mut self,
agent_contents: String,
calls: Vec<ProviderCall>,
) -> Result<Vec<DispatchedToolCall>, TransitionError> {
if !self.round_active {
return Err(TransitionError::InvalidPhase);
}
for (index, call) in calls.iter().enumerate() {
if self.call_box_id(call.tool_call_id)?.is_some()
|| calls[..index]
.iter()
.any(|earlier| earlier.tool_call_id == call.tool_call_id)
{
return Err(TransitionError::DuplicateToolCall);
}
}
let has_agent_contents = !agent_contents.is_empty();
let mut additions = Vec::with_capacity(calls.len() + usize::from(has_agent_contents));
if has_agent_contents {
additions.push(ChatBox::new(
BoxId::new(0),
AGENT_MESSAGE_TYPE.to_owned(),
agent_contents,
String::new(),
String::new(),
));
}
additions.extend(calls.iter().map(tool_call_box));
let appended = self.append_batch(additions)?;
let call_offset = usize::from(has_agent_contents);
Ok(calls
.iter()
.zip(appended[call_offset..].iter())
.map(|(call, value)| DispatchedToolCall {
tool_call_id: call.tool_call_id,
call_box_id: value.id(),
})
.collect())
}
pub fn accept_async_return(
&mut self,
tool_call_id: ToolCallId,
result: Result<String, String>,
) -> Result<Option<BoxId>, TransitionError> {
let originating_call = self
.call_box_id(tool_call_id)?
.ok_or(TransitionError::UnknownToolCall)?;
if self.has_return(tool_call_id)? {
return Err(TransitionError::DuplicateReturn);
}
self.accept_arrival(tool_result_box(tool_call_id, originating_call, result))
}
pub fn flush_active_arrivals(&mut self) -> Result<Vec<ChatBox>, TransitionError> {
if !self.round_active {
return Err(TransitionError::InvalidPhase);
}
let appended = self.append_batch(self.active_arrivals.clone())?;
self.active_arrivals.clear();
Ok(appended)
}
pub fn done(&mut self, final_agent_contents: String) -> Result<Vec<ChatBox>, TransitionError> {
if !self.round_active {
return Err(TransitionError::InvalidPhase);
}
let has_final_contents = !final_agent_contents.is_empty();
let mut additions =
Vec::with_capacity(self.active_arrivals.len() + usize::from(has_final_contents));
if has_final_contents {
additions.push(ChatBox::new(
BoxId::new(0),
AGENT_MESSAGE_TYPE.to_owned(),
final_agent_contents,
String::new(),
String::new(),
));
}
additions.extend(self.active_arrivals.iter().cloned());
let appended = self.append_batch(additions)?;
self.active_arrivals.clear();
self.round_active = false;
Ok(appended)
}
pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
let mut calls = Vec::<(ToolCallId, BoxId)>::new();
let mut returned = Vec::<ToolCallId>::new();
for (index, value) in boxes.iter().enumerate() {
let expected = u64::try_from(index)
.ok()
.and_then(|index| index.checked_add(1))
.ok_or(RecoveryError::NonContiguousBoxId)?;
if value.id().get() != expected {
return Err(RecoveryError::NonContiguousBoxId);
}
if let Some(call) = value.tool_call_metadata()? {
if calls
.iter()
.any(|(tool_call_id, _)| *tool_call_id == call.tool_call_id)
{
return Err(RecoveryError::DuplicateToolCall);
}
calls.push((call.tool_call_id, value.id()));
}
if let Some(result) = value.tool_result_metadata()? {
let Some((_, call_box_id)) = calls
.iter()
.find(|(tool_call_id, _)| *tool_call_id == result.tool_call_id)
else {
return Err(RecoveryError::UnknownToolCall);
};
if *call_box_id != result.originating_call {
return Err(RecoveryError::WrongOriginatingCall);
}
if returned.contains(&result.tool_call_id) {
return Err(RecoveryError::DuplicateReturn);
}
returned.push(result.tool_call_id);
}
}
Ok(Self {
boxes,
round_active: false,
active_arrivals: Vec::new(),
})
}
fn accept_arrival(&mut self, value: ChatBox) -> Result<Option<BoxId>, TransitionError> {
if self.round_active {
self.active_arrivals.push(value);
Ok(None)
} else {
let mut appended = self.append_batch(vec![value])?;
Ok(appended.pop().map(|value| value.id()))
}
}
fn append_batch(
&mut self,
mut additions: Vec<ChatBox>,
) -> Result<Vec<ChatBox>, TransitionError> {
self.ensure_capacity(additions.len())?;
let mut previous = self.boxes.last().map_or(0, |value| value.id().get());
for value in &mut additions {
previous = previous
.checked_add(1)
.ok_or(TransitionError::BoxIdOverflow)?;
value.id = BoxId::new(previous);
}
self.boxes.extend(additions.iter().cloned());
Ok(additions)
}
fn ensure_capacity(&self, additional: usize) -> Result<(), TransitionError> {
let additional = u64::try_from(additional).map_err(|_| TransitionError::BoxIdOverflow)?;
let previous = self.boxes.last().map_or(0, |value| value.id().get());
previous
.checked_add(additional)
.ok_or(TransitionError::BoxIdOverflow)?;
Ok(())
}
fn call_box_id(&self, tool_call_id: ToolCallId) -> Result<Option<BoxId>, TransitionError> {
let mut found = None;
for value in &self.boxes {
let metadata = value
.tool_call_metadata()
.map_err(|_| TransitionError::MalformedToolConvention)?;
if metadata
.as_ref()
.is_some_and(|call| call.tool_call_id == tool_call_id)
{
if found.is_some() {
return Err(TransitionError::DuplicateToolCall);
}
found = Some(value.id());
}
}
Ok(found)
}
fn has_return(&self, tool_call_id: ToolCallId) -> Result<bool, TransitionError> {
for value in self.boxes.iter().chain(&self.active_arrivals) {
let metadata = value
.tool_result_metadata()
.map_err(|_| TransitionError::MalformedToolConvention)?;
if metadata
.as_ref()
.is_some_and(|result| result.tool_call_id == tool_call_id)
{
return Ok(true);
}
}
Ok(false)
}
}
impl Default for Chatend {
fn default() -> Self {
Self::new()
}
}
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 (index, slot) in nonce.iter_mut().enumerate() {
let offset = index.checked_mul(2)?;
let high = decode_hex(*value.as_bytes().get(offset)?)?;
let low = decode_hex(*value.as_bytes().get(offset.checked_add(1)?)?)?;
*slot = (high << 4) | low;
}
Some(nonce)
}
const fn decode_hex(value: u8) -> Option<u8> {
match value {
b'0'..=b'9' => Some(value - b'0'),
b'a'..=b'f' => Some(value - b'a' + 10),
b'A'..=b'F' => Some(value - b'A' + 10),
_ => None,
}
}
#[cfg(test)]
mod tests;