kcode-k1-chat-state 0.4.0

use std::sync::{Arc, atomic::AtomicU8};

pub use kcode_k1_chat_chatend::{
    BoxContent, BoxId, ChatBox, DispatchCall, ProviderCall, RecoveryError, ToolCallId,
    TransitionError,
};

#[derive(Clone, Debug)]
pub struct InferenceStart {
    pub job: u64,
    pub frontier: Option<BoxId>,
    pub attempt: Arc<AtomicU8>,
}

#[derive(Debug)]
pub enum StateError {
    Transition(TransitionError),
    Recovery(RecoveryError),
    WrongInference { expected: Option<u64>, actual: u64 },
    JobIdExhausted,
    NotStalled,
    Busy,
}

impl From<TransitionError> for StateError {
    fn from(error: TransitionError) -> Self {
        Self::Transition(error)
    }
}

struct ActiveInference {
    job: u64,
    frontier: Option<BoxId>,
    _attempt: Arc<AtomicU8>,
}

struct RetryRound {
    frontier: Option<BoxId>,
    ready: bool,
}

pub struct ActorState {
    chatend: kcode_k1_chat_chatend::Chatend,
    next_job: u64,
    active: Option<ActiveInference>,
    retry: Option<RetryRound>,
    scheduled: bool,
    arrival_during_active: bool,
    halt: Option<String>,
}

impl ActorState {
    pub fn new(force: bool) -> Self {
        Self {
            chatend: kcode_k1_chat_chatend::Chatend::new(),
            next_job: 0,
            active: None,
            retry: None,
            scheduled: force,
            arrival_during_active: false,
            halt: None,
        }
    }

    pub fn recover(boxes: Vec<ChatBox>, force: bool) -> Result<Self, StateError> {
        Ok(Self {
            chatend: kcode_k1_chat_chatend::Chatend::recover(boxes)
                .map_err(StateError::Recovery)?,
            next_job: 0,
            active: None,
            retry: None,
            scheduled: force,
            arrival_during_active: false,
            halt: None,
        })
    }

    pub fn boxes(&self) -> &[ChatBox] {
        self.chatend.boxes()
    }

    pub fn halted(&self) -> bool {
        self.halt.is_some() || self.retry.is_some()
    }

    pub fn halt(&mut self, text: String) -> bool {
        if self.halted() {
            false
        } else {
            self.halt = Some(text);
            true
        }
    }

    pub fn take_halt(&mut self) -> Option<String> {
        self.halt.take()
    }

    pub fn restart(&mut self) -> Result<(), StateError> {
        if !self.halted() {
            return Err(StateError::NotStalled);
        }
        if self.active.is_some() {
            return Err(StateError::Busy);
        }
        self.halt = None;
        if let Some(retry) = &mut self.retry {
            retry.ready = true;
        } else {
            self.scheduled = true;
        }
        Ok(())
    }

    pub fn accept_system(&mut self, text: String) -> Result<(), StateError> {
        self.arrival();
        self.chatend.accept_system(text)?;
        Ok(())
    }

    pub fn accept_user(&mut self, text: String) -> Result<(), StateError> {
        self.arrival();
        self.chatend.accept_user(text)?;
        Ok(())
    }

    pub fn accept_attachment(&mut self) -> Result<(), StateError> {
        self.arrival();
        self.chatend.accept_attachment()?;
        Ok(())
    }

    pub fn accept_async_return(
        &mut self,
        tool_call_id: ToolCallId,
        result: Result<String, String>,
    ) -> Result<(), StateError> {
        self.arrival();
        self.chatend.accept_async_return(tool_call_id, result)?;
        Ok(())
    }

    pub fn force_inference(&mut self) {
        self.scheduled = true;
    }

    pub fn begin_inference(&mut self) -> Result<Option<InferenceStart>, StateError> {
        if self.halt.is_some() || self.active.is_some() {
            return Ok(None);
        }
        if let Some(retry) = &self.retry {
            if !retry.ready {
                return Ok(None);
            }
            let frontier = retry.frontier;
            let start = self.activate(frontier)?;
            self.retry = None;
            return Ok(Some(start));
        }
        if !self.scheduled {
            return Ok(None);
        }
        let job = self.next_job()?;
        self.chatend.start_round()?;
        let frontier = self.boxes().last().map(ChatBox::id);
        self.next_job = job;
        self.scheduled = false;
        Ok(Some(self.install_active(job, frontier)))
    }

    pub fn append_stage(
        &mut self,
        job: u64,
        text: String,
        calls: Vec<ProviderCall>,
    ) -> Result<Vec<DispatchCall>, StateError> {
        self.require_job(job)?;
        Ok(self.chatend.append_stage(text, calls)?)
    }

    pub fn complete_inference(&mut self, job: u64, final_text: String) -> Result<(), StateError> {
        self.require_job(job)?;
        self.chatend.done(final_text)?;
        self.active = None;
        if self.arrival_during_active {
            self.scheduled = true;
            self.arrival_during_active = false;
        }
        Ok(())
    }

    pub fn stall_inference(&mut self, job: u64, text: String) -> Result<(), StateError> {
        self.require_job(job)?;
        let active = self.active.take().expect("validated active inference");
        self.halt = Some(text);
        self.retry = Some(RetryRound {
            frontier: active.frontier,
            ready: false,
        });
        Ok(())
    }

    pub fn quiet(&self) -> bool {
        self.active.is_none() && self.retry.is_none() && !self.scheduled
    }

    fn arrival(&mut self) {
        if self.active.is_some() || self.retry.is_some() {
            self.arrival_during_active = true;
        } else {
            self.scheduled = true;
        }
    }

    fn activate(&mut self, frontier: Option<BoxId>) -> Result<InferenceStart, StateError> {
        let job = self.next_job()?;
        self.next_job = job;
        Ok(self.install_active(job, frontier))
    }

    fn install_active(&mut self, job: u64, frontier: Option<BoxId>) -> InferenceStart {
        let attempt = Arc::new(AtomicU8::new(1));
        self.active = Some(ActiveInference {
            job,
            frontier,
            _attempt: attempt.clone(),
        });
        InferenceStart {
            job,
            frontier,
            attempt,
        }
    }

    fn next_job(&self) -> Result<u64, StateError> {
        self.next_job
            .checked_add(1)
            .ok_or(StateError::JobIdExhausted)
    }

    fn require_job(&self, job: u64) -> Result<(), StateError> {
        let expected = self.active.as_ref().map(|active| active.job);
        if expected == Some(job) {
            Ok(())
        } else {
            Err(StateError::WrongInference {
                expected,
                actual: job,
            })
        }
    }
}

#[cfg(test)]
mod tests;