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, BoxId, ChatBox, DispatchedToolCall, EnvelopeError,
7 ProviderCall, RecoveryError, ResultView, SYSTEM_MESSAGE_TYPE, TOOL_ATTACHMENT_TYPE,
8 TOOL_CALL_HIDDEN_TYPE, TOOL_CALL_TYPE, TOOL_MESSAGE_TYPE, TOOL_RESULT_HIDDEN_TYPE,
9 TOOL_RESULT_TYPE, ToolCall, ToolCallId, ToolResult, ToolResultStatus, TransitionError,
10 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, Eq, PartialEq)]
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::with_chatend(kcode_k1_chat_chatend::Chatend::new(), force)
60 }
61
62 pub fn recover(boxes: Vec<ChatBox>, force: bool) -> Result<Self, StateError> {
63 Ok(Self::with_chatend(
64 kcode_k1_chat_chatend::Chatend::recover(boxes).map_err(StateError::Recovery)?,
65 force,
66 ))
67 }
68
69 pub fn boxes(&self) -> &[ChatBox] {
70 self.chatend.boxes()
71 }
72
73 pub fn halted(&self) -> bool {
74 self.halt.is_some() || self.retry.is_some()
75 }
76
77 pub fn halt(&mut self, text: String) -> bool {
78 if self.halted() {
79 false
80 } else {
81 self.halt = Some(text);
82 true
83 }
84 }
85
86 pub fn take_halt(&mut self) -> Option<String> {
87 self.halt.take()
88 }
89
90 pub fn restart(&mut self) -> Result<(), StateError> {
91 if !self.halted() {
92 return Err(StateError::NotStalled);
93 }
94 if self.active.is_some() {
95 return Err(StateError::Busy);
96 }
97 self.halt = None;
98 if let Some(retry) = &mut self.retry {
99 retry.ready = true;
100 } else {
101 self.scheduled = true;
102 }
103 Ok(())
104 }
105
106 pub fn accept_box(
107 &mut self,
108 box_type: String,
109 contents: String,
110 hidden_type: String,
111 hidden_contents: String,
112 ) -> Result<(), StateError> {
113 self.chatend
114 .accept_box(box_type, contents, hidden_type, hidden_contents)?;
115 self.arrival();
116 Ok(())
117 }
118
119 pub fn accept_system(&mut self, contents: String) -> Result<(), StateError> {
120 self.accept_known(contents, kcode_k1_chat_chatend::Chatend::accept_system)
121 }
122
123 pub fn accept_user(&mut self, contents: String) -> Result<(), StateError> {
124 self.accept_known(contents, kcode_k1_chat_chatend::Chatend::accept_user)
125 }
126
127 pub fn accept_attachment(
128 &mut self,
129 contents: String,
130 hidden_type: String,
131 hidden_contents: String,
132 ) -> Result<(), StateError> {
133 self.chatend
134 .accept_attachment(contents, hidden_type, hidden_contents)?;
135 self.arrival();
136 Ok(())
137 }
138
139 pub fn accept_async_return(&mut self, result: ToolResult) -> Result<(), StateError> {
140 self.chatend.accept_async_return(result)?;
141 self.arrival();
142 Ok(())
143 }
144
145 pub fn force_inference(&mut self) {
146 self.scheduled = true;
147 }
148
149 pub fn begin_inference(&mut self) -> Result<Option<InferenceStart>, StateError> {
150 if self.halt.is_some() || self.active.is_some() {
151 return Ok(None);
152 }
153 if let Some(retry) = &self.retry {
154 if !retry.ready {
155 return Ok(None);
156 }
157 let frontier = retry.frontier;
158 let start = self.activate(frontier)?;
159 self.retry = None;
160 return Ok(Some(start));
161 }
162 if !self.scheduled {
163 return Ok(None);
164 }
165 let job = self.next_job()?;
166 self.chatend.start_round()?;
167 let frontier = self.boxes().last().map(ChatBox::id);
168 self.next_job = job;
169 self.scheduled = false;
170 Ok(Some(self.install_active(job, frontier)))
171 }
172
173 pub fn append_stage(
174 &mut self,
175 job: u64,
176 contents: String,
177 calls: Vec<ProviderCall>,
178 ) -> Result<Vec<DispatchedToolCall>, StateError> {
179 self.require_job(job)?;
180 Ok(self.chatend.append_stage(contents, calls)?)
181 }
182
183 pub fn flush_active_arrivals(&mut self, job: u64) -> Result<Vec<ChatBox>, StateError> {
184 self.require_job(job)?;
185 let arrivals = self.chatend.flush_active_arrivals()?;
186 if !arrivals.is_empty() {
187 self.arrival_during_active = false;
188 }
189 Ok(arrivals)
190 }
191
192 pub fn complete_inference(&mut self, job: u64, contents: String) -> Result<(), StateError> {
193 self.require_job(job)?;
194 self.chatend.done(contents)?;
195 self.active = None;
196 if self.arrival_during_active {
197 self.scheduled = true;
198 self.arrival_during_active = false;
199 }
200 Ok(())
201 }
202
203 pub fn stall_inference(&mut self, job: u64, text: String) -> Result<(), StateError> {
204 self.require_job(job)?;
205 let active = self.active.take().expect("validated active inference");
206 self.halt = Some(text);
207 self.retry = Some(RetryRound {
208 frontier: active.frontier,
209 ready: false,
210 });
211 Ok(())
212 }
213
214 pub fn quiet(&self) -> bool {
215 self.active.is_none() && self.retry.is_none() && !self.scheduled
216 }
217
218 fn with_chatend(chatend: kcode_k1_chat_chatend::Chatend, force: bool) -> Self {
219 Self {
220 chatend,
221 next_job: 0,
222 active: None,
223 retry: None,
224 scheduled: force,
225 arrival_during_active: false,
226 halt: None,
227 }
228 }
229
230 fn accept_known(
231 &mut self,
232 contents: String,
233 accept: fn(
234 &mut kcode_k1_chat_chatend::Chatend,
235 String,
236 ) -> Result<Option<BoxId>, TransitionError>,
237 ) -> Result<(), StateError> {
238 accept(&mut self.chatend, contents)?;
239 self.arrival();
240 Ok(())
241 }
242
243 fn arrival(&mut self) {
244 if self.active.is_some() || self.retry.is_some() {
245 self.arrival_during_active = true;
246 } else {
247 self.scheduled = true;
248 }
249 }
250
251 fn activate(&mut self, frontier: Option<BoxId>) -> Result<InferenceStart, StateError> {
252 let job = self.next_job()?;
253 self.next_job = job;
254 Ok(self.install_active(job, frontier))
255 }
256
257 fn install_active(&mut self, job: u64, frontier: Option<BoxId>) -> InferenceStart {
258 let attempt = Arc::new(AtomicU8::new(1));
259 self.active = Some(ActiveInference {
260 job,
261 frontier,
262 _attempt: attempt.clone(),
263 });
264 InferenceStart {
265 job,
266 frontier,
267 attempt,
268 }
269 }
270
271 fn next_job(&self) -> Result<u64, StateError> {
272 self.next_job
273 .checked_add(1)
274 .ok_or(StateError::JobIdExhausted)
275 }
276
277 fn require_job(&self, job: u64) -> Result<(), StateError> {
278 let expected = self.active.as_ref().map(|active| active.job);
279 if expected == Some(job) {
280 Ok(())
281 } else {
282 Err(StateError::WrongInference {
283 expected,
284 actual: job,
285 })
286 }
287 }
288}
289
290#[cfg(test)]
291mod tests;