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