kcode_k1_chat_chatend/
lib.rs

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

kcode_k1_chat_chatend/
lib.rs
