use std::{
collections::{BTreeMap, HashSet},
sync::{
Arc,
atomic::{AtomicU8, Ordering},
},
};
use kcode_k1_chat_core::{ChatError, ChatView, PendingAction};
struct Job {
action: PendingAction,
attempt: Option<Arc<AtomicU8>>,
updates: HashSet<u64>,
}
pub struct ActorState {
primary: String,
pending: String,
history: Vec<String>,
sent: usize,
next_job: u64,
jobs: BTreeMap<u64, Job>,
inference: Option<u64>,
compaction: Option<u64>,
batch_cursor: Option<usize>,
force: bool,
halt: Option<String>,
}
impl ActorState {
pub fn new(primary: String, force: bool) -> Self {
Self {
primary,
pending: String::new(),
history: Vec::new(),
sent: 0,
next_job: 0,
jobs: BTreeMap::new(),
inference: None,
compaction: None,
batch_cursor: None,
force,
halt: None,
}
}
pub fn view(&self) -> ChatView {
let actions = self
.jobs
.values()
.map(|job| match &job.attempt {
Some(attempt) => PendingAction::Inference {
attempt: attempt.load(Ordering::Relaxed),
},
None => job.action.clone(),
})
.collect();
ChatView {
primary: self.primary.clone(),
pending: self.pending.clone(),
history: self.history.clone(),
actions,
}
}
pub fn halted(&self) -> bool {
self.halt.is_some()
}
pub fn halt(&mut self, text: String) -> bool {
if self.halt.is_some() {
false
} else {
self.halt = Some(text);
true
}
}
pub fn take_halt(&mut self) -> Option<String> {
self.halt.take()
}
pub fn append_pending(&mut self, text: String) {
self.pending.push_str(&text);
}
pub fn restart(&mut self) -> Result<(), ChatError> {
if self.halt.is_none() {
return Err(ChatError::NotStalled);
}
if !self.jobs.is_empty() || self.batch_cursor.is_some() {
return Err(ChatError::Busy);
}
self.primary.push_str(&self.pending);
self.pending.clear();
self.sent = 0;
self.force = true;
self.halt = None;
Ok(())
}
pub fn begin_inference(&mut self) -> Option<(u64, String, Arc<AtomicU8>)> {
if self.halted()
|| self.inference.is_some()
|| self.compaction.is_some()
|| self.batch_cursor.is_some()
|| (!self.force && self.pending.is_empty())
{
return None;
}
self.primary.push_str(&self.pending);
self.pending.clear();
self.force = false;
let delta = self.primary[self.sent..].to_owned();
let attempt = Arc::new(AtomicU8::new(1));
let job = self.add_job(
PendingAction::Inference { attempt: 1 },
Some(attempt.clone()),
);
self.inference = Some(job);
Some((job, delta, attempt))
}
pub fn finish_inference(&mut self, job: u64) -> bool {
if self.inference != Some(job) {
return false;
}
self.inference = None;
self.jobs.remove(&job).is_some()
}
pub fn commit_output(&mut self, text: &str) {
self.primary.push_str(text);
self.sent = self.primary.len();
}
pub fn force_inference(&mut self) {
self.force = true;
}
pub fn begin_tool(&mut self, name: String) -> u64 {
self.add_job(PendingAction::Tool { name }, None)
}
pub fn begin_worker(&mut self, llm: String) -> u64 {
self.add_job(PendingAction::Worker { llm }, None)
}
pub fn begin_compaction(&mut self) -> (u64, String) {
let job = self.add_job(PendingAction::Compaction, None);
self.compaction = Some(job);
(job, self.primary.clone())
}
pub fn reject_compaction_batch(&mut self) {
self.pending
.insert_str(0, "compaction must be the sole call");
self.force = true;
}
pub fn begin_batch(&mut self) {
self.batch_cursor = Some(0);
}
pub fn apply_tool_replies(
&mut self,
mut entries: Vec<(usize, u64, String, bool)>,
finished: bool,
) {
let Some(cursor) = self.batch_cursor else {
return;
};
entries.sort_by_key(|entry| entry.0);
let mut text = String::new();
for (_, job, reply, complete) in entries {
text.push_str(&reply);
if complete {
self.jobs.remove(&job);
}
}
self.pending.insert_str(cursor, &text);
self.batch_cursor = if finished {
None
} else {
Some(cursor + text.len())
};
self.force = true;
}
pub fn complete_action(&mut self, job: u64, text: String) {
if self.jobs.remove(&job).is_some() {
self.pending.push_str(&text);
self.force = true;
}
}
pub fn apply_append_update(&mut self, job: u64, identity: u64, text: String) {
if self.accept_activity_update(job, identity) {
self.pending.push_str(&text);
self.force = true;
}
}
pub fn accept_activity_update(&mut self, job: u64, identity: u64) -> bool {
self.jobs
.get_mut(&job)
.is_some_and(|job| job.updates.insert(identity))
}
pub fn complete_compaction(
&mut self,
job: u64,
frozen: String,
result: Result<String, String>,
) {
if self.compaction != Some(job) {
return;
}
self.compaction = None;
self.jobs.remove(&job);
match result {
Ok(replacement) => {
self.history.push(frozen);
self.primary = replacement;
self.primary.push_str(&std::mem::take(&mut self.pending));
self.sent = 0;
}
Err(error) => {
self.primary.push_str(&error);
self.primary.push_str(&std::mem::take(&mut self.pending));
}
}
self.force = true;
}
pub fn quiet(&self) -> bool {
self.jobs.is_empty()
&& self.pending.is_empty()
&& !self.force
&& self.batch_cursor.is_none()
}
fn add_job(&mut self, action: PendingAction, attempt: Option<Arc<AtomicU8>>) -> u64 {
let job = self
.next_job
.checked_add(1)
.expect("job ID counter overflowed");
self.next_job = job;
let replaced = self.jobs.insert(
job,
Job {
action,
attempt,
updates: HashSet::new(),
},
);
assert!(replaced.is_none(), "job ID was reused");
job
}
}
#[cfg(test)]
mod tests;