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