1use std::{
2 collections::{BTreeMap, HashSet},
3 sync::{
4 Arc,
5 atomic::{AtomicU8, Ordering},
6 },
7};
8
9use kcode_k1_chat_core::{ChatError, ChatView, PendingAction};
10
11struct Job {
12 action: PendingAction,
13 attempt: Option<Arc<AtomicU8>>,
14 updates: HashSet<u64>,
15}
16
17pub struct ActorState {
18 primary: String,
19 pending: String,
20 history: Vec<String>,
21 sent: usize,
22 next_job: u64,
23 jobs: BTreeMap<u64, Job>,
24 inference: Option<u64>,
25 compaction: Option<u64>,
26 batch_cursor: Option<usize>,
27 force: bool,
28 halt: Option<String>,
29}
30
31impl ActorState {
32 pub fn new(primary: String, force: bool) -> Self {
33 Self {
34 primary,
35 pending: String::new(),
36 history: Vec::new(),
37 sent: 0,
38 next_job: 0,
39 jobs: BTreeMap::new(),
40 inference: None,
41 compaction: None,
42 batch_cursor: None,
43 force,
44 halt: None,
45 }
46 }
47
48 pub fn view(&self) -> ChatView {
49 let actions = self
50 .jobs
51 .values()
52 .map(|job| match &job.attempt {
53 Some(attempt) => PendingAction::Inference {
54 attempt: attempt.load(Ordering::Relaxed),
55 },
56 None => job.action.clone(),
57 })
58 .collect();
59 ChatView {
60 primary: self.primary.clone(),
61 pending: self.pending.clone(),
62 history: self.history.clone(),
63 actions,
64 }
65 }
66
67 pub fn halted(&self) -> bool {
68 self.halt.is_some()
69 }
70
71 pub fn halt(&mut self, text: String) -> bool {
72 if self.halt.is_some() {
73 false
74 } else {
75 self.halt = Some(text);
76 true
77 }
78 }
79
80 pub fn take_halt(&mut self) -> Option<String> {
81 self.halt.take()
82 }
83
84 pub fn append_pending(&mut self, text: String) {
85 self.pending.push_str(&text);
86 }
87
88 pub fn restart(&mut self) -> Result<(), ChatError> {
89 if self.halt.is_none() {
90 return Err(ChatError::NotStalled);
91 }
92 if !self.jobs.is_empty() || self.batch_cursor.is_some() {
93 return Err(ChatError::Busy);
94 }
95 self.primary.push_str(&self.pending);
96 self.pending.clear();
97 self.sent = 0;
98 self.force = true;
99 self.halt = None;
100 Ok(())
101 }
102
103 pub fn begin_inference(&mut self) -> Option<(u64, String, Arc<AtomicU8>)> {
104 if self.halted()
105 || self.inference.is_some()
106 || self.compaction.is_some()
107 || self.batch_cursor.is_some()
108 || (!self.force && self.pending.is_empty())
109 {
110 return None;
111 }
112 self.primary.push_str(&self.pending);
113 self.pending.clear();
114 self.force = false;
115 let delta = self.primary[self.sent..].to_owned();
116 let attempt = Arc::new(AtomicU8::new(1));
117 let job = self.add_job(
118 PendingAction::Inference { attempt: 1 },
119 Some(attempt.clone()),
120 );
121 self.inference = Some(job);
122 Some((job, delta, attempt))
123 }
124
125 pub fn finish_inference(&mut self, job: u64) -> bool {
126 if self.inference != Some(job) {
127 return false;
128 }
129 self.inference = None;
130 self.jobs.remove(&job).is_some()
131 }
132
133 pub fn commit_output(&mut self, text: &str) {
134 self.primary.push_str(text);
135 self.sent = self.primary.len();
136 }
137
138 pub fn force_inference(&mut self) {
139 self.force = true;
140 }
141
142 pub fn begin_tool(&mut self, name: String) -> u64 {
143 self.add_job(PendingAction::Tool { name }, None)
144 }
145
146 pub fn begin_worker(&mut self, llm: String) -> u64 {
147 self.add_job(PendingAction::Worker { llm }, None)
148 }
149
150 pub fn begin_compaction(&mut self) -> (u64, String) {
151 let job = self.add_job(PendingAction::Compaction, None);
152 self.compaction = Some(job);
153 (job, self.primary.clone())
154 }
155
156 pub fn reject_compaction_batch(&mut self) {
157 self.pending
158 .insert_str(0, "compaction must be the sole call");
159 self.force = true;
160 }
161
162 pub fn begin_batch(&mut self) {
163 self.batch_cursor = Some(0);
164 }
165
166 pub fn apply_tool_replies(
167 &mut self,
168 mut entries: Vec<(usize, u64, String, bool)>,
169 finished: bool,
170 ) {
171 let Some(cursor) = self.batch_cursor else {
172 return;
173 };
174 entries.sort_by_key(|entry| entry.0);
175 let mut text = String::new();
176 for (_, job, reply, complete) in entries {
177 text.push_str(&reply);
178 if complete {
179 self.jobs.remove(&job);
180 }
181 }
182 self.pending.insert_str(cursor, &text);
183 self.batch_cursor = if finished {
184 None
185 } else {
186 Some(cursor + text.len())
187 };
188 self.force = true;
189 }
190
191 pub fn complete_action(&mut self, job: u64, text: String) {
192 if self.jobs.remove(&job).is_some() {
193 self.pending.push_str(&text);
194 self.force = true;
195 }
196 }
197
198 pub fn apply_append_update(&mut self, job: u64, identity: u64, text: String) {
199 if self.accept_activity_update(job, identity) {
200 self.pending.push_str(&text);
201 self.force = true;
202 }
203 }
204
205 pub fn accept_activity_update(&mut self, job: u64, identity: u64) -> bool {
206 self.jobs
207 .get_mut(&job)
208 .is_some_and(|job| job.updates.insert(identity))
209 }
210
211 pub fn complete_compaction(
212 &mut self,
213 job: u64,
214 frozen: String,
215 result: Result<String, String>,
216 ) {
217 if self.compaction != Some(job) {
218 return;
219 }
220 self.compaction = None;
221 self.jobs.remove(&job);
222 match result {
223 Ok(replacement) => {
224 self.history.push(frozen);
225 self.primary = replacement;
226 self.primary.push_str(&std::mem::take(&mut self.pending));
227 self.sent = 0;
228 }
229 Err(error) => {
230 self.primary.push_str(&error);
231 self.primary.push_str(&std::mem::take(&mut self.pending));
232 }
233 }
234 self.force = true;
235 }
236
237 pub fn quiet(&self) -> bool {
238 self.jobs.is_empty()
239 && self.pending.is_empty()
240 && !self.force
241 && self.batch_cursor.is_none()
242 }
243
244 fn add_job(&mut self, action: PendingAction, attempt: Option<Arc<AtomicU8>>) -> u64 {
245 let job = self
246 .next_job
247 .checked_add(1)
248 .expect("job ID counter overflowed");
249 self.next_job = job;
250 let replaced = self.jobs.insert(
251 job,
252 Job {
253 action,
254 attempt,
255 updates: HashSet::new(),
256 },
257 );
258 assert!(replaced.is_none(), "job ID was reused");
259 job
260 }
261}
262
263#[cfg(test)]
264mod tests;