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 commit_output_with_deferred(&mut self, text: &str, deferred: &str) {
139 self.primary.push_str(text);
140 self.sent = self.primary.len();
141 self.primary.push_str(deferred);
142 }
143
144 pub fn force_inference(&mut self) {
145 self.force = true;
146 }
147
148 pub fn begin_tool(&mut self, name: String) -> u64 {
149 self.add_job(PendingAction::Tool { name }, None)
150 }
151
152 pub fn begin_worker(&mut self, llm: String) -> u64 {
153 self.add_job(PendingAction::Worker { llm }, None)
154 }
155
156 pub fn begin_compaction(&mut self) -> (u64, String) {
157 let job = self.add_job(PendingAction::Compaction, None);
158 self.compaction = Some(job);
159 (job, self.primary.clone())
160 }
161
162 pub fn reject_compaction_batch(&mut self) {
163 self.pending
164 .insert_str(0, "compaction must be the sole call");
165 self.force = true;
166 }
167
168 pub fn begin_batch(&mut self) {
169 self.batch_cursor = Some(0);
170 }
171
172 pub fn apply_tool_replies(
173 &mut self,
174 mut entries: Vec<(usize, u64, String, bool)>,
175 finished: bool,
176 ) {
177 let Some(cursor) = self.batch_cursor else {
178 return;
179 };
180 entries.sort_by_key(|entry| entry.0);
181 let mut text = String::new();
182 for (_, job, reply, complete) in entries {
183 text.push_str(&reply);
184 if complete {
185 self.jobs.remove(&job);
186 }
187 }
188 self.pending.insert_str(cursor, &text);
189 self.batch_cursor = if finished {
190 None
191 } else {
192 Some(cursor + text.len())
193 };
194 self.force = true;
195 }
196
197 pub fn complete_action(&mut self, job: u64, text: String) {
198 if self.jobs.remove(&job).is_some() {
199 self.pending.push_str(&text);
200 self.force = true;
201 }
202 }
203
204 pub fn apply_append_update(&mut self, job: u64, identity: u64, text: String) {
205 if self.accept_activity_update(job, identity) {
206 self.pending.push_str(&text);
207 self.force = true;
208 }
209 }
210
211 pub fn accept_activity_update(&mut self, job: u64, identity: u64) -> bool {
212 self.jobs
213 .get_mut(&job)
214 .is_some_and(|job| job.updates.insert(identity))
215 }
216
217 pub fn complete_compaction(
218 &mut self,
219 job: u64,
220 frozen: String,
221 result: Result<String, String>,
222 ) {
223 if self.compaction != Some(job) {
224 return;
225 }
226 self.compaction = None;
227 self.jobs.remove(&job);
228 match result {
229 Ok(replacement) => {
230 self.history.push(frozen);
231 self.primary = replacement;
232 self.primary.push_str(&std::mem::take(&mut self.pending));
233 self.sent = 0;
234 }
235 Err(error) => {
236 self.primary.push_str(&error);
237 self.primary.push_str(&std::mem::take(&mut self.pending));
238 }
239 }
240 self.force = true;
241 }
242
243 pub fn quiet(&self) -> bool {
244 self.jobs.is_empty()
245 && self.pending.is_empty()
246 && !self.force
247 && self.batch_cursor.is_none()
248 }
249
250 fn add_job(&mut self, action: PendingAction, attempt: Option<Arc<AtomicU8>>) -> u64 {
251 let job = self
252 .next_job
253 .checked_add(1)
254 .expect("job ID counter overflowed");
255 self.next_job = job;
256 let replaced = self.jobs.insert(
257 job,
258 Job {
259 action,
260 attempt,
261 updates: HashSet::new(),
262 },
263 );
264 assert!(replaced.is_none(), "job ID was reused");
265 job
266 }
267}
268
269#[cfg(test)]
270mod tests;