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, TransitionError, USER_ATTACHMENT_TYPE,
7 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_async_return(
155 &mut self,
156 tool_call_id: ToolCallId,
157 result: Result<String, String>,
158 ) -> Result<(), StateError> {
159 self.chatend.accept_async_return(tool_call_id, result)?;
160 self.arrival();
161 Ok(())
162 }
163
164 pub fn force_inference(&mut self) {
165 self.scheduled = true;
166 }
167
168 pub fn begin_inference(&mut self) -> Result<Option<InferenceStart>, StateError> {
169 if self.halt.is_some() || self.active.is_some() {
170 return Ok(None);
171 }
172 if let Some(retry) = &self.retry {
173 if !retry.ready {
174 return Ok(None);
175 }
176 let frontier = retry.frontier;
177 let start = self.activate(frontier)?;
178 self.retry = None;
179 return Ok(Some(start));
180 }
181 if !self.scheduled {
182 return Ok(None);
183 }
184 let job = self.next_job()?;
185 self.chatend.start_round()?;
186 let frontier = self.boxes().last().map(ChatBox::id);
187 self.next_job = job;
188 self.scheduled = false;
189 Ok(Some(self.install_active(job, frontier)))
190 }
191
192 pub fn append_stage(
193 &mut self,
194 job: u64,
195 contents: String,
196 calls: Vec<ProviderCall>,
197 ) -> Result<Vec<DispatchedToolCall>, StateError> {
198 self.require_job(job)?;
199 Ok(self.chatend.append_stage(contents, calls)?)
200 }
201
202 pub fn flush_active_arrivals(&mut self, job: u64) -> Result<Vec<ChatBox>, StateError> {
203 self.require_job(job)?;
204 let arrivals = self.chatend.flush_active_arrivals()?;
205 if !arrivals.is_empty() {
206 self.arrival_during_active = false;
207 }
208 Ok(arrivals)
209 }
210
211 pub fn complete_inference(&mut self, job: u64, contents: String) -> Result<(), StateError> {
212 self.require_job(job)?;
213 self.chatend.done(contents)?;
214 self.active = None;
215 if self.arrival_during_active {
216 self.scheduled = true;
217 self.arrival_during_active = false;
218 }
219 Ok(())
220 }
221
222 pub fn stall_inference(&mut self, job: u64, text: String) -> Result<(), StateError> {
223 self.require_job(job)?;
224 let active = self.active.take().expect("validated active inference");
225 self.halt = Some(text);
226 self.retry = Some(RetryRound {
227 frontier: active.frontier,
228 ready: false,
229 });
230 Ok(())
231 }
232
233 pub fn quiet(&self) -> bool {
234 self.active.is_none() && self.retry.is_none() && !self.scheduled
235 }
236
237 fn arrival(&mut self) {
238 if self.active.is_some() || self.retry.is_some() {
239 self.arrival_during_active = true;
240 } else {
241 self.scheduled = true;
242 }
243 }
244
245 fn activate(&mut self, frontier: Option<BoxId>) -> Result<InferenceStart, StateError> {
246 let job = self.next_job()?;
247 self.next_job = job;
248 Ok(self.install_active(job, frontier))
249 }
250
251 fn install_active(&mut self, job: u64, frontier: Option<BoxId>) -> InferenceStart {
252 let attempt = Arc::new(AtomicU8::new(1));
253 self.active = Some(ActiveInference {
254 job,
255 frontier,
256 _attempt: attempt.clone(),
257 });
258 InferenceStart {
259 job,
260 frontier,
261 attempt,
262 }
263 }
264
265 fn next_job(&self) -> Result<u64, StateError> {
266 self.next_job
267 .checked_add(1)
268 .ok_or(StateError::JobIdExhausted)
269 }
270
271 fn require_job(&self, job: u64) -> Result<(), StateError> {
272 let expected = self.active.as_ref().map(|active| active.job);
273 if expected == Some(job) {
274 Ok(())
275 } else {
276 Err(StateError::WrongInference {
277 expected,
278 actual: job,
279 })
280 }
281 }
282}
283
284#[cfg(test)]
285mod tests;