1use std::sync::{Arc, atomic::AtomicU8};
2
3pub use kcode_k1_chat_chatend::{
4 AGENT_ATTACHMENT_TYPE, AGENT_MESSAGE_TYPE, ATTACHMENT_TYPE, BoxId, ChatBox, DispatchedToolCall,
5 ProviderCall, RecoveryError, SYSTEM_MESSAGE_TYPE, TOOL_ATTACHMENT_TYPE, TOOL_CALL_TYPE,
6 TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE, ToolCallId, ToolMessageMetadata, ToolResultMetadata,
7 ToolResultV2Metadata, TransitionError, USER_ATTACHMENT_TYPE, USER_MESSAGE_TYPE,
8};
9
10#[derive(Clone, Debug)]
11pub struct InferenceStart {
12 pub job: u64,
13 pub frontier: Option<BoxId>,
14 pub attempt: Arc<AtomicU8>,
15}
16
17#[derive(Debug)]
18pub enum StateError {
19 Transition(TransitionError),
20 Recovery(RecoveryError),
21 WrongInference { expected: Option<u64>, actual: u64 },
22 JobIdExhausted,
23 NotStalled,
24 Busy,
25}
26
27impl From<TransitionError> for StateError {
28 fn from(error: TransitionError) -> Self {
29 Self::Transition(error)
30 }
31}
32
33struct ActiveInference {
34 job: u64,
35 frontier: Option<BoxId>,
36 _attempt: Arc<AtomicU8>,
37}
38
39struct RetryRound {
40 frontier: Option<BoxId>,
41 ready: bool,
42}
43
44pub struct ActorState {
45 chatend: kcode_k1_chat_chatend::Chatend,
46 next_job: u64,
47 active: Option<ActiveInference>,
48 retry: Option<RetryRound>,
49 scheduled: bool,
50 arrival_during_active: bool,
51 halt: Option<String>,
52}
53
54impl ActorState {
55 pub fn new(force: bool) -> Self {
56 Self {
57 chatend: kcode_k1_chat_chatend::Chatend::new(),
58 next_job: 0,
59 active: None,
60 retry: None,
61 scheduled: force,
62 arrival_during_active: false,
63 halt: None,
64 }
65 }
66
67 pub fn recover(boxes: Vec<ChatBox>, force: bool) -> Result<Self, StateError> {
68 Ok(Self {
69 chatend: kcode_k1_chat_chatend::Chatend::recover(boxes)
70 .map_err(StateError::Recovery)?,
71 next_job: 0,
72 active: None,
73 retry: None,
74 scheduled: force,
75 arrival_during_active: false,
76 halt: None,
77 })
78 }
79
80 pub fn boxes(&self) -> &[ChatBox] {
81 self.chatend.boxes()
82 }
83
84 pub fn halted(&self) -> bool {
85 self.halt.is_some() || self.retry.is_some()
86 }
87
88 pub fn halt(&mut self, text: String) -> bool {
89 if self.halted() {
90 false
91 } else {
92 self.halt = Some(text);
93 true
94 }
95 }
96
97 pub fn take_halt(&mut self) -> Option<String> {
98 self.halt.take()
99 }
100
101 pub fn restart(&mut self) -> Result<(), StateError> {
102 if !self.halted() {
103 return Err(StateError::NotStalled);
104 }
105 if self.active.is_some() {
106 return Err(StateError::Busy);
107 }
108 self.halt = None;
109 if let Some(retry) = &mut self.retry {
110 retry.ready = true;
111 } else {
112 self.scheduled = true;
113 }
114 Ok(())
115 }
116
117 pub fn accept_box(
118 &mut self,
119 box_type: String,
120 contents: String,
121 hidden_type: String,
122 hidden_contents: String,
123 ) -> Result<(), StateError> {
124 self.chatend
125 .accept_box(box_type, contents, hidden_type, hidden_contents)?;
126 self.arrival();
127 Ok(())
128 }
129
130 pub fn accept_system(&mut self, contents: String) -> Result<(), StateError> {
131 self.chatend.accept_system(contents)?;
132 self.arrival();
133 Ok(())
134 }
135
136 pub fn accept_user(&mut self, contents: String) -> Result<(), StateError> {
137 self.chatend.accept_user(contents)?;
138 self.arrival();
139 Ok(())
140 }
141
142 pub fn accept_attachment(
143 &mut self,
144 contents: String,
145 hidden_type: String,
146 hidden_contents: String,
147 ) -> Result<(), StateError> {
148 self.chatend
149 .accept_attachment(contents, hidden_type, hidden_contents)?;
150 self.arrival();
151 Ok(())
152 }
153
154 pub fn accept_tool_message(
155 &mut self,
156 tool_call_id: ToolCallId,
157 message: String,
158 ) -> Result<(), StateError> {
159 self.chatend.accept_tool_message(tool_call_id, message)?;
160 Ok(())
161 }
162
163 pub fn accept_async_return(
164 &mut self,
165 tool_call_id: ToolCallId,
166 result: Result<String, String>,
167 ) -> Result<(), StateError> {
168 self.chatend.accept_async_return(tool_call_id, result)?;
169 self.arrival();
170 Ok(())
171 }
172
173 pub fn accept_async_return_v2(
174 &mut self,
175 tool_call_id: ToolCallId,
176 result: Result<String, String>,
177 metadata_type: String,
178 metadata_contents: String,
179 ) -> Result<(), StateError> {
180 self.chatend.accept_async_return_v2(
181 tool_call_id,
182 result,
183 metadata_type,
184 metadata_contents,
185 )?;
186 self.arrival();
187 Ok(())
188 }
189
190 pub fn force_inference(&mut self) {
191 self.scheduled = true;
192 }
193
194 pub fn begin_inference(&mut self) -> Result<Option<InferenceStart>, StateError> {
195 if self.halt.is_some() || self.active.is_some() {
196 return Ok(None);
197 }
198 if let Some(retry) = &self.retry {
199 if !retry.ready {
200 return Ok(None);
201 }
202 let frontier = retry.frontier;
203 let start = self.activate(frontier)?;
204 self.retry = None;
205 return Ok(Some(start));
206 }
207 if !self.scheduled {
208 return Ok(None);
209 }
210 let job = self.next_job()?;
211 self.chatend.start_round()?;
212 let frontier = self.boxes().last().map(ChatBox::id);
213 self.next_job = job;
214 self.scheduled = false;
215 Ok(Some(self.install_active(job, frontier)))
216 }
217
218 pub fn append_stage(
219 &mut self,
220 job: u64,
221 contents: String,
222 calls: Vec<ProviderCall>,
223 ) -> Result<Vec<DispatchedToolCall>, StateError> {
224 self.require_job(job)?;
225 Ok(self.chatend.append_stage(contents, calls)?)
226 }
227
228 pub fn flush_active_arrivals(&mut self, job: u64) -> Result<Vec<ChatBox>, StateError> {
229 self.require_job(job)?;
230 let arrivals = self.chatend.flush_active_arrivals()?;
231 if !arrivals.is_empty() {
232 self.arrival_during_active = false;
233 }
234 Ok(arrivals)
235 }
236
237 pub fn complete_inference(&mut self, job: u64, contents: String) -> Result<(), StateError> {
238 self.require_job(job)?;
239 self.chatend.done(contents)?;
240 self.active = None;
241 if self.arrival_during_active {
242 self.scheduled = true;
243 self.arrival_during_active = false;
244 }
245 Ok(())
246 }
247
248 pub fn stall_inference(&mut self, job: u64, text: String) -> Result<(), StateError> {
249 self.require_job(job)?;
250 let active = self.active.take().expect("validated active inference");
251 self.halt = Some(text);
252 self.retry = Some(RetryRound {
253 frontier: active.frontier,
254 ready: false,
255 });
256 Ok(())
257 }
258
259 pub fn quiet(&self) -> bool {
260 self.active.is_none() && self.retry.is_none() && !self.scheduled
261 }
262
263 fn arrival(&mut self) {
264 if self.active.is_some() || self.retry.is_some() {
265 self.arrival_during_active = true;
266 } else {
267 self.scheduled = true;
268 }
269 }
270
271 fn activate(&mut self, frontier: Option<BoxId>) -> Result<InferenceStart, StateError> {
272 let job = self.next_job()?;
273 self.next_job = job;
274 Ok(self.install_active(job, frontier))
275 }
276
277 fn install_active(&mut self, job: u64, frontier: Option<BoxId>) -> InferenceStart {
278 let attempt = Arc::new(AtomicU8::new(1));
279 self.active = Some(ActiveInference {
280 job,
281 frontier,
282 _attempt: attempt.clone(),
283 });
284 InferenceStart {
285 job,
286 frontier,
287 attempt,
288 }
289 }
290
291 fn next_job(&self) -> Result<u64, StateError> {
292 self.next_job
293 .checked_add(1)
294 .ok_or(StateError::JobIdExhausted)
295 }
296
297 fn require_job(&self, job: u64) -> Result<(), StateError> {
298 let expected = self.active.as_ref().map(|active| active.job);
299 if expected == Some(job) {
300 Ok(())
301 } else {
302 Err(StateError::WrongInference {
303 expected,
304 actual: job,
305 })
306 }
307 }
308}
309
310impl Default for ActorState {
311 fn default() -> Self {
312 Self::new(false)
313 }
314}
315
316#[cfg(test)]
317mod tests;