#![forbid(unsafe_code)]
use serde_json::Value;
pub use kcode_k1_chat_box::{
BoxId, ChatBox, EnvelopeError, ResultView, TOOL_CALL_HIDDEN_TYPE, TOOL_CALL_TYPE,
TOOL_RESULT_HIDDEN_TYPE, TOOL_RESULT_TYPE, ToolCall, ToolCallId, ToolResult, ToolResultStatus,
};
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_MESSAGE_TYPE: &str = "Tool Message";
pub const TOOL_ATTACHMENT_TYPE: &str = "Tool Attachment";
pub const ATTACHMENT_TYPE: &str = USER_ATTACHMENT_TYPE;
#[derive(Clone, Debug, PartialEq)]
pub struct ProviderCall {
pub tool: String,
pub tool_version: String,
pub arguments: Value,
}
#[derive(Clone, Debug, PartialEq)]
pub struct DispatchedToolCall {
pub call: ToolCall,
pub call_box_id: BoxId,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransitionError {
InvalidPhase,
BoxIdOverflow,
CallIdOverflow,
InvalidToolEnvelope,
DuplicateToolCall,
UnknownToolCall,
DuplicateReturn,
WrongOriginatingCall,
MismatchedTool,
MalformedToolConvention,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecoveryError {
NonContiguousBoxId,
InvalidCallSequence,
UnknownToolCall,
DuplicateReturn,
WrongOriginatingCall,
MismatchedTool,
MalformedToolConvention,
}
pub struct Chatend {
boxes: Vec<ChatBox>,
round_active: bool,
active_arrivals: Vec<ChatBox>,
next_call: Option<u64>,
}
impl Chatend {
pub const fn new() -> Self {
Self {
boxes: Vec::new(),
round_active: false,
active_arrivals: Vec::new(),
next_call: Some(1),
}
}
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.into(),
contents,
String::new(),
String::new(),
)
}
pub fn accept_user(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
self.accept_box(
USER_MESSAGE_TYPE.into(),
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.into(),
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);
}
let (calls, next_call) = self.assign_calls(calls)?;
let includes_agent = !agent_contents.is_empty();
self.ensure_capacity(calls.len() + usize::from(includes_agent))?;
let mut id = self.last_id();
let mut additions = Vec::with_capacity(calls.len() + usize::from(includes_agent));
if includes_agent {
id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
additions.push(ChatBox::new(
BoxId::new(id),
AGENT_MESSAGE_TYPE.into(),
agent_contents,
String::new(),
String::new(),
));
}
let mut dispatched = Vec::with_capacity(calls.len());
for call in calls {
id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
let call_box_id = BoxId::new(id);
let box_value = ChatBox::tool_call(call_box_id, call.clone())
.map_err(|_| TransitionError::InvalidToolEnvelope)?;
additions.push(box_value);
dispatched.push(DispatchedToolCall { call, call_box_id });
}
self.boxes.extend(additions);
self.next_call = next_call;
Ok(dispatched)
}
pub fn accept_async_return(
&mut self,
result: ToolResult,
) -> Result<Option<BoxId>, TransitionError> {
let (call, call_box_id) = self
.call_for(result.call_id())?
.ok_or(TransitionError::UnknownToolCall)?;
if self.has_return(result.call_id())? {
return Err(TransitionError::DuplicateReturn);
}
if result.originating_call_box_id() != call_box_id {
return Err(TransitionError::WrongOriginatingCall);
}
if result.tool() != call.tool() || result.tool_version() != call.tool_version() {
return Err(TransitionError::MismatchedTool);
}
let box_value = ChatBox::tool_result(BoxId::new(0), result)
.map_err(|_| TransitionError::InvalidToolEnvelope)?;
self.accept_arrival(box_value)
}
pub fn flush_active_arrivals(&mut self) -> Result<Vec<ChatBox>, TransitionError> {
if !self.round_active {
return Err(TransitionError::InvalidPhase);
}
let arrivals = self.active_arrivals.clone();
let appended = self.append_batch(&arrivals)?;
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 mut arrivals = self.active_arrivals.clone();
if !final_agent_contents.is_empty() {
arrivals.insert(
0,
ChatBox::new(
BoxId::new(0),
AGENT_MESSAGE_TYPE.into(),
final_agent_contents,
String::new(),
String::new(),
),
);
}
let appended = self.append_batch(&arrivals)?;
self.active_arrivals.clear();
self.round_active = false;
Ok(appended)
}
pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
let mut calls = Vec::<(ToolCall, BoxId)>::new();
let mut returned = Vec::<ToolCallId>::new();
let mut next_call = Some(1);
for (index, value) in boxes.iter().enumerate() {
let expected_box_id = u64::try_from(index)
.ok()
.and_then(|value| value.checked_add(1))
.ok_or(RecoveryError::NonContiguousBoxId)?;
if value.id() != BoxId::new(expected_box_id) {
return Err(RecoveryError::NonContiguousBoxId);
}
if let Some(call) = value
.tool_call_metadata()
.map_err(|_| RecoveryError::MalformedToolConvention)?
{
let expected_call = next_call.ok_or(RecoveryError::InvalidCallSequence)?;
if call.call_id().get() != expected_call {
return Err(RecoveryError::InvalidCallSequence);
}
next_call = expected_call.checked_add(1);
calls.push((call, value.id()));
}
if let Some(result) = value
.tool_result_metadata()
.map_err(|_| RecoveryError::MalformedToolConvention)?
{
let Some((call, call_box_id)) = calls
.iter()
.find(|(call, _)| call.call_id() == result.call_id())
else {
return Err(RecoveryError::UnknownToolCall);
};
if result.originating_call_box_id() != *call_box_id {
return Err(RecoveryError::WrongOriginatingCall);
}
if result.tool() != call.tool() || result.tool_version() != call.tool_version() {
return Err(RecoveryError::MismatchedTool);
}
if returned.contains(&result.call_id()) {
return Err(RecoveryError::DuplicateReturn);
}
returned.push(result.call_id());
}
}
Ok(Self {
boxes,
round_active: false,
active_arrivals: Vec::new(),
next_call,
})
}
fn assign_calls(
&self,
calls: Vec<ProviderCall>,
) -> Result<(Vec<ToolCall>, Option<u64>), TransitionError> {
let mut next = self.next_call;
let mut assigned = Vec::with_capacity(calls.len());
for call in calls {
let id = next.ok_or(TransitionError::CallIdOverflow)?;
let call_id = ToolCallId::new(id).map_err(|_| TransitionError::CallIdOverflow)?;
let call = ToolCall::new(call_id, call.tool, call.tool_version, call.arguments)
.map_err(|_| TransitionError::InvalidToolEnvelope)?;
assigned.push(call);
next = id.checked_add(1);
}
Ok((assigned, next))
}
fn accept_arrival(&mut self, value: ChatBox) -> Result<Option<BoxId>, TransitionError> {
if self.round_active {
self.active_arrivals.push(value);
return Ok(None);
}
Ok(self.append_batch(&[value])?.pop().map(|value| value.id()))
}
fn append_batch(&mut self, additions: &[ChatBox]) -> Result<Vec<ChatBox>, TransitionError> {
self.ensure_capacity(additions.len())?;
let mut id = self.last_id();
let appended = additions
.iter()
.map(|value| {
id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
Ok(ChatBox::new(
BoxId::new(id),
value.box_type().into(),
value.contents().into(),
value.hidden_type().into(),
value.hidden_contents().into(),
))
})
.collect::<Result<Vec<_>, _>>()?;
self.boxes.extend(appended.iter().cloned());
Ok(appended)
}
fn last_id(&self) -> u64 {
self.boxes.last().map_or(0, |value| value.id().get())
}
fn ensure_capacity(&self, count: usize) -> Result<(), TransitionError> {
let count = u64::try_from(count).map_err(|_| TransitionError::BoxIdOverflow)?;
self.last_id()
.checked_add(count)
.ok_or(TransitionError::BoxIdOverflow)
.map(|_| ())
}
fn call_for(&self, id: ToolCallId) -> Result<Option<(ToolCall, BoxId)>, TransitionError> {
let mut found = None;
for value in &self.boxes {
if let Some(call) = value
.tool_call_metadata()
.map_err(|_| TransitionError::MalformedToolConvention)?
&& call.call_id() == id
{
if found.is_some() {
return Err(TransitionError::DuplicateToolCall);
}
found = Some((call, value.id()));
}
}
Ok(found)
}
fn has_return(&self, id: ToolCallId) -> Result<bool, TransitionError> {
for value in self.boxes.iter().chain(&self.active_arrivals) {
let result = value
.tool_result_metadata()
.map_err(|_| TransitionError::MalformedToolConvention)?;
if result.is_some_and(|result| result.call_id() == id) {
return Ok(true);
}
}
Ok(false)
}
}
impl Default for Chatend {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;