use std::collections::{BTreeMap, BTreeSet};
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ToolCallId {
session: [u8; 12],
sequence: u64,
}
impl ToolCallId {
pub const fn new(session: [u8; 12], sequence: u64) -> Self {
Self { session, sequence }
}
pub const fn session(self) -> [u8; 12] {
self.session
}
pub const fn sequence(self) -> u64 {
self.sequence
}
}
#[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 enum BoxContent {
System(String),
User(String),
Kennedy {
text: String,
},
Attachment,
KtoolCall {
tool_call_id: ToolCallId,
name: String,
arguments: String,
},
KtoolReturn {
tool_call_id: ToolCallId,
originating_call: BoxId,
result: Result<String, String>,
},
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ChatBox {
id: BoxId,
content: BoxContent,
}
impl ChatBox {
pub fn new(id: BoxId, content: BoxContent) -> Self {
Self { id, content }
}
pub const fn id(&self) -> BoxId {
self.id
}
pub const fn content(&self) -> &BoxContent {
&self.content
}
}
#[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 DispatchCall {
pub tool_call_id: ToolCallId,
pub call_box_id: BoxId,
pub name: String,
pub arguments: String,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransitionError {
InvalidPhase,
DuplicateToolCall,
UnknownToolCall,
DuplicateReturn,
BoxIdOverflow,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecoveryError {
NonContiguousBoxId,
EmptyKennedy,
DuplicateToolCall,
UnknownToolCall,
WrongOriginatingCall,
DuplicateReturn,
}
pub struct Chatend {
boxes: Vec<ChatBox>,
provider: bool,
queued: Vec<BoxContent>,
tool_calls: BTreeMap<ToolCallId, ToolCallRecord>,
last_id: u64,
}
struct ToolCallRecord {
call_box_id: BoxId,
returned: bool,
}
impl Default for Chatend {
fn default() -> Self {
Self::new()
}
}
impl Chatend {
pub fn new() -> Self {
Self {
boxes: Vec::new(),
provider: false,
queued: Vec::new(),
tool_calls: BTreeMap::new(),
last_id: 0,
}
}
pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
let mut expected_id = 1_u64;
let mut tool_calls: BTreeMap<ToolCallId, ToolCallRecord> = BTreeMap::new();
for chat_box in &boxes {
if chat_box.id.get() != expected_id {
return Err(RecoveryError::NonContiguousBoxId);
}
match &chat_box.content {
BoxContent::Kennedy { text } if text.is_empty() => {
return Err(RecoveryError::EmptyKennedy);
}
BoxContent::KtoolCall { tool_call_id, .. } => {
if tool_calls.contains_key(tool_call_id) {
return Err(RecoveryError::DuplicateToolCall);
}
tool_calls.insert(
*tool_call_id,
ToolCallRecord {
call_box_id: chat_box.id,
returned: false,
},
);
}
BoxContent::KtoolReturn {
tool_call_id,
originating_call,
..
} => {
let record = tool_calls
.get_mut(tool_call_id)
.ok_or(RecoveryError::UnknownToolCall)?;
if record.call_box_id != *originating_call {
return Err(RecoveryError::WrongOriginatingCall);
}
if record.returned {
return Err(RecoveryError::DuplicateReturn);
}
record.returned = true;
}
_ => {}
}
expected_id = expected_id
.checked_add(1)
.ok_or(RecoveryError::NonContiguousBoxId)?;
}
Ok(Self {
last_id: expected_id - 1,
boxes,
provider: false,
queued: Vec::new(),
tool_calls,
})
}
pub fn boxes(&self) -> &[ChatBox] {
&self.boxes
}
pub fn accept_system(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
self.accept_arrival(BoxContent::System(text))
}
pub fn accept_user(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
self.accept_arrival(BoxContent::User(text))
}
pub fn accept_attachment(&mut self) -> Result<Option<BoxId>, TransitionError> {
self.accept_arrival(BoxContent::Attachment)
}
pub fn accept_async_return(
&mut self,
tool_call_id: ToolCallId,
result: Result<String, String>,
) -> Result<Option<BoxId>, TransitionError> {
let record = self
.tool_calls
.get(&tool_call_id)
.ok_or(TransitionError::UnknownToolCall)?;
if record.returned {
return Err(TransitionError::DuplicateReturn);
}
let originating_call = record.call_box_id;
if !self.provider {
self.ensure_capacity(1)?;
}
self.tool_calls.get_mut(&tool_call_id).unwrap().returned = true;
let content = BoxContent::KtoolReturn {
tool_call_id,
originating_call,
result,
};
if self.provider {
self.queued.push(content);
Ok(None)
} else {
Ok(Some(self.append(content)))
}
}
pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
if self.provider {
return Err(TransitionError::InvalidPhase);
}
let frontier = self.boxes.last().map(ChatBox::id);
self.provider = true;
Ok(frontier)
}
pub fn append_stage(
&mut self,
text: String,
calls: Vec<ProviderCall>,
) -> Result<Vec<DispatchCall>, TransitionError> {
if !self.provider {
return Err(TransitionError::InvalidPhase);
}
let mut seen = BTreeSet::new();
for call in &calls {
if !seen.insert(call.tool_call_id) || self.tool_calls.contains_key(&call.tool_call_id) {
return Err(TransitionError::DuplicateToolCall);
}
}
let text_growth = if text.is_empty() { 0 } else { 1 };
let growth = calls
.len()
.checked_add(text_growth)
.ok_or(TransitionError::BoxIdOverflow)?;
self.ensure_capacity(growth)?;
if !text.is_empty() {
self.append(BoxContent::Kennedy { text });
}
let mut dispatches = Vec::with_capacity(calls.len());
for call in calls {
let call_box_id = self.append(BoxContent::KtoolCall {
tool_call_id: call.tool_call_id,
name: call.name.clone(),
arguments: call.arguments.clone(),
});
self.tool_calls.insert(
call.tool_call_id,
ToolCallRecord {
call_box_id,
returned: false,
},
);
dispatches.push(DispatchCall {
tool_call_id: call.tool_call_id,
call_box_id,
name: call.name,
arguments: call.arguments,
});
}
Ok(dispatches)
}
pub fn done(&mut self, text: String) -> Result<(), TransitionError> {
if !self.provider {
return Err(TransitionError::InvalidPhase);
}
let text_growth = if text.is_empty() { 0 } else { 1 };
let growth = self
.queued
.len()
.checked_add(text_growth)
.ok_or(TransitionError::BoxIdOverflow)?;
self.ensure_capacity(growth)?;
if !text.is_empty() {
self.append(BoxContent::Kennedy { text });
}
self.provider = false;
for content in std::mem::take(&mut self.queued) {
self.append(content);
}
Ok(())
}
fn accept_arrival(&mut self, content: BoxContent) -> Result<Option<BoxId>, TransitionError> {
if self.provider {
self.queued.push(content);
Ok(None)
} else {
self.ensure_capacity(1)?;
Ok(Some(self.append(content)))
}
}
fn ensure_capacity(&self, additional: usize) -> Result<(), TransitionError> {
let additional = u64::try_from(additional).map_err(|_| TransitionError::BoxIdOverflow)?;
self.last_id
.checked_add(additional)
.ok_or(TransitionError::BoxIdOverflow)?;
Ok(())
}
fn append(&mut self, content: BoxContent) -> BoxId {
self.last_id += 1;
let id = BoxId(self.last_id);
self.boxes.push(ChatBox { id, content });
id
}
}
#[cfg(test)]
mod tests;