1use std::collections::{BTreeMap, BTreeSet};
2
3#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
4pub struct ActionId {
5 session: [u8; 12],
6 sequence: u64,
7}
8
9impl ActionId {
10 pub const fn new(session: [u8; 12], sequence: u64) -> Self {
11 Self { session, sequence }
12 }
13
14 pub const fn session(self) -> [u8; 12] {
15 self.session
16 }
17
18 pub const fn sequence(self) -> u64 {
19 self.sequence
20 }
21}
22
23#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
24pub struct BoxId(u64);
25
26impl BoxId {
27 pub fn get(self) -> u64 {
28 self.0
29 }
30}
31
32#[derive(Clone, Debug, Eq, PartialEq)]
33pub enum BoxContent {
34 System(String),
35 User(String),
36 Kennedy {
37 text: String,
38 complete: bool,
39 },
40 Attachment,
41 KtoolCall {
42 action_id: ActionId,
43 name: String,
44 arguments: String,
45 },
46 KtoolReturn {
47 action_id: ActionId,
48 originating_call: BoxId,
49 result: Result<String, String>,
50 },
51}
52
53#[derive(Clone, Debug, Eq, PartialEq)]
54pub struct ChatBox {
55 id: BoxId,
56 content: BoxContent,
57}
58
59impl ChatBox {
60 pub fn id(&self) -> BoxId {
61 self.id
62 }
63
64 pub fn content(&self) -> &BoxContent {
65 &self.content
66 }
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct DispatchCall {
71 pub action_id: ActionId,
72 pub call_box_id: BoxId,
73 pub name: String,
74 pub arguments: String,
75}
76
77#[derive(Clone, Debug, Eq, PartialEq)]
78pub enum DispatchOutcome {
79 Pending,
80 Terminal(Result<String, String>),
81}
82
83#[derive(Clone, Copy, Debug, Eq, PartialEq)]
84pub enum TransitionError {
85 InvalidPhase,
86 ActionCount,
87 OutcomeCount,
88 DuplicateAction,
89 UnknownAction,
90 DuplicateReturn,
91 BoxIdOverflow,
92}
93
94pub struct Chatend {
95 boxes: Vec<ChatBox>,
96 phase: Phase,
97 queued: Vec<BoxContent>,
98 actions: BTreeMap<ActionId, ActionRecord>,
99 last_id: u64,
100}
101
102enum Phase {
103 Idle,
104 Provider(Vec<StagedCall>),
105 Dispatch(Vec<DispatchCall>),
106}
107
108#[derive(Clone)]
109struct StagedCall {
110 name: String,
111 arguments: String,
112}
113
114struct ActionRecord {
115 call_box_id: BoxId,
116 returned: bool,
117}
118
119impl Default for Chatend {
120 fn default() -> Self {
121 Self::new()
122 }
123}
124
125impl Chatend {
126 pub fn new() -> Self {
127 Self {
128 boxes: Vec::new(),
129 phase: Phase::Idle,
130 queued: Vec::new(),
131 actions: BTreeMap::new(),
132 last_id: 0,
133 }
134 }
135
136 pub fn boxes(&self) -> &[ChatBox] {
137 &self.boxes
138 }
139
140 pub fn accept_system(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
141 self.accept_arrival(BoxContent::System(text))
142 }
143
144 pub fn accept_user(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
145 self.accept_arrival(BoxContent::User(text))
146 }
147
148 pub fn accept_attachment(&mut self) -> Result<Option<BoxId>, TransitionError> {
149 self.accept_arrival(BoxContent::Attachment)
150 }
151
152 pub fn accept_async_return(
153 &mut self,
154 action_id: ActionId,
155 result: Result<String, String>,
156 ) -> Result<Option<BoxId>, TransitionError> {
157 let record = self
158 .actions
159 .get(&action_id)
160 .ok_or(TransitionError::UnknownAction)?;
161 if record.returned {
162 return Err(TransitionError::DuplicateReturn);
163 }
164 let originating_call = record.call_box_id;
165 let idle = matches!(&self.phase, Phase::Idle);
166 if idle {
167 self.ensure_capacity(1)?;
168 }
169 self.actions.get_mut(&action_id).unwrap().returned = true;
170 let content = BoxContent::KtoolReturn {
171 action_id,
172 originating_call,
173 result,
174 };
175 if idle {
176 Ok(Some(self.append(content)?))
177 } else {
178 self.queued.push(content);
179 Ok(None)
180 }
181 }
182
183 pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
184 if !matches!(&self.phase, Phase::Idle) {
185 return Err(TransitionError::InvalidPhase);
186 }
187 let frontier = self.boxes.last().map(ChatBox::id);
188 self.ensure_capacity(1)?;
189 self.append(BoxContent::Kennedy {
190 text: String::new(),
191 complete: false,
192 })?;
193 self.phase = Phase::Provider(Vec::new());
194 Ok(frontier)
195 }
196
197 pub fn append_kennedy_text(&mut self, chunk: &str) -> Result<(), TransitionError> {
198 if !matches!(&self.phase, Phase::Provider(_)) {
199 return Err(TransitionError::InvalidPhase);
200 }
201 let BoxContent::Kennedy { text, .. } = &mut self.boxes.last_mut().unwrap().content else {
202 unreachable!()
203 };
204 text.push_str(chunk);
205 Ok(())
206 }
207
208 pub fn collect_provider_call(
209 &mut self,
210 name: String,
211 arguments: String,
212 ) -> Result<(), TransitionError> {
213 match &mut self.phase {
214 Phase::Provider(calls) => {
215 calls.push(StagedCall { name, arguments });
216 Ok(())
217 }
218 _ => Err(TransitionError::InvalidPhase),
219 }
220 }
221
222 pub fn complete_provider_output(
223 &mut self,
224 action_ids: &[ActionId],
225 ) -> Result<Vec<DispatchCall>, TransitionError> {
226 let staged = match &self.phase {
227 Phase::Provider(calls) => calls.clone(),
228 _ => return Err(TransitionError::InvalidPhase),
229 };
230 if action_ids.len() != staged.len() {
231 return Err(TransitionError::ActionCount);
232 }
233 let mut seen = BTreeSet::new();
234 for action_id in action_ids {
235 if !seen.insert(*action_id) || self.actions.contains_key(action_id) {
236 return Err(TransitionError::DuplicateAction);
237 }
238 }
239 let queued_growth = if staged.is_empty() {
240 self.queued.len()
241 } else {
242 0
243 };
244 let growth = staged
245 .len()
246 .checked_add(queued_growth)
247 .ok_or(TransitionError::BoxIdOverflow)?;
248 self.ensure_capacity(growth)?;
249 let BoxContent::Kennedy { complete, .. } = &mut self.boxes.last_mut().unwrap().content
250 else {
251 unreachable!()
252 };
253 *complete = true;
254 let mut calls = Vec::with_capacity(staged.len());
255 for (staged, action_id) in staged.into_iter().zip(action_ids.iter().copied()) {
256 let call_box_id = self.append(BoxContent::KtoolCall {
257 action_id,
258 name: staged.name.clone(),
259 arguments: staged.arguments.clone(),
260 })?;
261 let record = ActionRecord {
262 call_box_id,
263 returned: false,
264 };
265 self.actions.insert(action_id, record);
266 calls.push(DispatchCall {
267 action_id,
268 call_box_id,
269 name: staged.name,
270 arguments: staged.arguments,
271 });
272 }
273 if calls.is_empty() {
274 self.phase = Phase::Idle;
275 self.drain_queued()?;
276 } else {
277 self.phase = Phase::Dispatch(calls.clone());
278 }
279 Ok(calls)
280 }
281
282 pub fn complete_dispatch(
283 &mut self,
284 outcomes: Vec<DispatchOutcome>,
285 ) -> Result<(), TransitionError> {
286 let calls = match &self.phase {
287 Phase::Dispatch(calls) => calls.clone(),
288 _ => return Err(TransitionError::InvalidPhase),
289 };
290 if outcomes.len() != calls.len() {
291 return Err(TransitionError::OutcomeCount);
292 }
293 for (call, outcome) in calls.iter().zip(&outcomes) {
294 let record = self.actions.get(&call.action_id).unwrap();
295 if matches!(outcome, DispatchOutcome::Terminal(_)) && record.returned {
296 return Err(TransitionError::DuplicateReturn);
297 }
298 }
299 let terminals = outcomes
300 .iter()
301 .filter(|outcome| matches!(outcome, DispatchOutcome::Terminal(_)))
302 .count();
303 let growth = terminals
304 .checked_add(self.queued.len())
305 .ok_or(TransitionError::BoxIdOverflow)?;
306 self.ensure_capacity(growth)?;
307 for (call, outcome) in calls.into_iter().zip(outcomes) {
308 if let DispatchOutcome::Terminal(result) = outcome {
309 let originating_call = {
310 let record = self.actions.get_mut(&call.action_id).unwrap();
311 record.returned = true;
312 record.call_box_id
313 };
314 self.append(BoxContent::KtoolReturn {
315 action_id: call.action_id,
316 originating_call,
317 result,
318 })?;
319 }
320 }
321 self.phase = Phase::Idle;
322 self.drain_queued()
323 }
324
325 fn accept_arrival(&mut self, content: BoxContent) -> Result<Option<BoxId>, TransitionError> {
326 if matches!(&self.phase, Phase::Idle) {
327 self.ensure_capacity(1)?;
328 Ok(Some(self.append(content)?))
329 } else {
330 self.queued.push(content);
331 Ok(None)
332 }
333 }
334
335 fn drain_queued(&mut self) -> Result<(), TransitionError> {
336 let queued = std::mem::take(&mut self.queued);
337 for content in queued {
338 self.append(content)?;
339 }
340 Ok(())
341 }
342
343 fn ensure_capacity(&self, additional: usize) -> Result<(), TransitionError> {
344 let additional = u64::try_from(additional).map_err(|_| TransitionError::BoxIdOverflow)?;
345 self.last_id
346 .checked_add(additional)
347 .ok_or(TransitionError::BoxIdOverflow)?;
348 Ok(())
349 }
350
351 fn append(&mut self, content: BoxContent) -> Result<BoxId, TransitionError> {
352 self.ensure_capacity(1)?;
353 self.last_id += 1;
354 let id = BoxId(self.last_id);
355 self.boxes.push(ChatBox { id, content });
356 Ok(id)
357 }
358}
359
360#[cfg(test)]
361mod tests;