pub use kcode_k1_chat_core::ActionId;
use std::collections::{BTreeMap, BTreeSet};
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct BoxId(u64);
impl BoxId {
pub fn get(self) -> u64 {
self.0
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum BoxContent {
System(String),
User(String),
Kennedy {
text: String,
complete: bool,
},
Attachment,
KtoolCall {
action_id: ActionId,
name: String,
arguments: String,
},
KtoolReturn {
action_id: ActionId,
originating_call: BoxId,
result: Result<String, String>,
},
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ChatBox {
id: BoxId,
content: BoxContent,
}
impl ChatBox {
pub fn id(&self) -> BoxId {
self.id
}
pub fn content(&self) -> &BoxContent {
&self.content
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DispatchCall {
pub action_id: ActionId,
pub call_box_id: BoxId,
pub name: String,
pub arguments: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DispatchOutcome {
Pending,
Terminal(Result<String, String>),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransitionError {
InvalidPhase,
ActionCount,
OutcomeCount,
DuplicateAction,
UnknownAction,
DuplicateReturn,
BoxIdOverflow,
}
pub struct Chatend {
boxes: Vec<ChatBox>,
phase: Phase,
queued: Vec<BoxContent>,
actions: BTreeMap<ActionId, ActionRecord>,
last_id: u64,
}
enum Phase {
Idle,
Provider(Vec<StagedCall>),
Dispatch(Vec<DispatchCall>),
}
#[derive(Clone)]
struct StagedCall {
name: String,
arguments: String,
}
struct ActionRecord {
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(),
phase: Phase::Idle,
queued: Vec::new(),
actions: BTreeMap::new(),
last_id: 0,
}
}
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,
action_id: ActionId,
result: Result<String, String>,
) -> Result<Option<BoxId>, TransitionError> {
let record = self
.actions
.get(&action_id)
.ok_or(TransitionError::UnknownAction)?;
if record.returned {
return Err(TransitionError::DuplicateReturn);
}
let originating_call = record.call_box_id;
let idle = matches!(&self.phase, Phase::Idle);
if idle {
self.ensure_capacity(1)?;
}
self.actions.get_mut(&action_id).unwrap().returned = true;
let content = BoxContent::KtoolReturn {
action_id,
originating_call,
result,
};
if idle {
Ok(Some(self.append(content)?))
} else {
self.queued.push(content);
Ok(None)
}
}
pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
if !matches!(&self.phase, Phase::Idle) {
return Err(TransitionError::InvalidPhase);
}
let frontier = self.boxes.last().map(ChatBox::id);
self.ensure_capacity(1)?;
self.append(BoxContent::Kennedy {
text: String::new(),
complete: false,
})?;
self.phase = Phase::Provider(Vec::new());
Ok(frontier)
}
pub fn append_kennedy_text(&mut self, chunk: &str) -> Result<(), TransitionError> {
if !matches!(&self.phase, Phase::Provider(_)) {
return Err(TransitionError::InvalidPhase);
}
let BoxContent::Kennedy { text, .. } = &mut self.boxes.last_mut().unwrap().content else {
unreachable!()
};
text.push_str(chunk);
Ok(())
}
pub fn collect_provider_call(
&mut self,
name: String,
arguments: String,
) -> Result<(), TransitionError> {
match &mut self.phase {
Phase::Provider(calls) => {
calls.push(StagedCall { name, arguments });
Ok(())
}
_ => Err(TransitionError::InvalidPhase),
}
}
pub fn complete_provider_output(
&mut self,
action_ids: &[ActionId],
) -> Result<Vec<DispatchCall>, TransitionError> {
let staged = match &self.phase {
Phase::Provider(calls) => calls.clone(),
_ => return Err(TransitionError::InvalidPhase),
};
if action_ids.len() != staged.len() {
return Err(TransitionError::ActionCount);
}
let mut seen = BTreeSet::new();
for action_id in action_ids {
if !seen.insert(*action_id) || self.actions.contains_key(action_id) {
return Err(TransitionError::DuplicateAction);
}
}
let queued_growth = if staged.is_empty() {
self.queued.len()
} else {
0
};
let growth = staged
.len()
.checked_add(queued_growth)
.ok_or(TransitionError::BoxIdOverflow)?;
self.ensure_capacity(growth)?;
let BoxContent::Kennedy { complete, .. } = &mut self.boxes.last_mut().unwrap().content
else {
unreachable!()
};
*complete = true;
let mut calls = Vec::with_capacity(staged.len());
for (staged, action_id) in staged.into_iter().zip(action_ids.iter().copied()) {
let call_box_id = self.append(BoxContent::KtoolCall {
action_id,
name: staged.name.clone(),
arguments: staged.arguments.clone(),
})?;
let record = ActionRecord {
call_box_id,
returned: false,
};
self.actions.insert(action_id, record);
calls.push(DispatchCall {
action_id,
call_box_id,
name: staged.name,
arguments: staged.arguments,
});
}
if calls.is_empty() {
self.phase = Phase::Idle;
self.drain_queued()?;
} else {
self.phase = Phase::Dispatch(calls.clone());
}
Ok(calls)
}
pub fn complete_dispatch(
&mut self,
outcomes: Vec<DispatchOutcome>,
) -> Result<(), TransitionError> {
let calls = match &self.phase {
Phase::Dispatch(calls) => calls.clone(),
_ => return Err(TransitionError::InvalidPhase),
};
if outcomes.len() != calls.len() {
return Err(TransitionError::OutcomeCount);
}
for (call, outcome) in calls.iter().zip(&outcomes) {
let record = self.actions.get(&call.action_id).unwrap();
if matches!(outcome, DispatchOutcome::Terminal(_)) && record.returned {
return Err(TransitionError::DuplicateReturn);
}
}
let terminals = outcomes
.iter()
.filter(|outcome| matches!(outcome, DispatchOutcome::Terminal(_)))
.count();
let growth = terminals
.checked_add(self.queued.len())
.ok_or(TransitionError::BoxIdOverflow)?;
self.ensure_capacity(growth)?;
for (call, outcome) in calls.into_iter().zip(outcomes) {
if let DispatchOutcome::Terminal(result) = outcome {
let originating_call = {
let record = self.actions.get_mut(&call.action_id).unwrap();
record.returned = true;
record.call_box_id
};
self.append(BoxContent::KtoolReturn {
action_id: call.action_id,
originating_call,
result,
})?;
}
}
self.phase = Phase::Idle;
self.drain_queued()
}
fn accept_arrival(&mut self, content: BoxContent) -> Result<Option<BoxId>, TransitionError> {
if matches!(&self.phase, Phase::Idle) {
self.ensure_capacity(1)?;
Ok(Some(self.append(content)?))
} else {
self.queued.push(content);
Ok(None)
}
}
fn drain_queued(&mut self) -> Result<(), TransitionError> {
let queued = std::mem::take(&mut self.queued);
for content in queued {
self.append(content)?;
}
Ok(())
}
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) -> Result<BoxId, TransitionError> {
self.ensure_capacity(1)?;
self.last_id += 1;
let id = BoxId(self.last_id);
self.boxes.push(ChatBox { id, content });
Ok(id)
}
}
#[cfg(test)]
mod tests;