kcode_k1_chat_chatend/

lib.rs

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;