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