kcode_k1_chat_chatend/
lib.rs

1#![forbid(unsafe_code)]
2
3pub const SYSTEM_MESSAGE_TYPE: &str = "System Message";
4pub const USER_MESSAGE_TYPE: &str = "User Message";
5pub const AGENT_MESSAGE_TYPE: &str = "Agent Message";
6pub const ATTACHMENT_TYPE: &str = "Attachment";
7pub const TOOL_CALL_TYPE: &str = "Tool Call";
8pub const TOOL_RESULT_TYPE: &str = "Tool Result";
9
10pub const TOOL_CALL_HIDDEN_TYPE: &str = "k1.tool-call/v1";
11pub const TOOL_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/v1";
12
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14pub struct BoxId(u64);
15
16impl BoxId {
17    pub const fn new(value: u64) -> Self {
18        Self(value)
19    }
20
21    pub const fn get(self) -> u64 {
22        self.0
23    }
24}
25
26#[derive(Clone, Debug, Eq, PartialEq)]
27pub struct ChatBox {
28    id: BoxId,
29    box_type: String,
30    contents: String,
31    hidden_type: String,
32    hidden_contents: String,
33}
34
35impl ChatBox {
36    pub fn new(
37        id: BoxId,
38        box_type: String,
39        contents: String,
40        hidden_type: String,
41        hidden_contents: String,
42    ) -> Self {
43        Self {
44            id,
45            box_type,
46            contents,
47            hidden_type,
48            hidden_contents,
49        }
50    }
51
52    pub const fn id(&self) -> BoxId {
53        self.id
54    }
55
56    pub fn box_type(&self) -> &str {
57        &self.box_type
58    }
59
60    pub fn contents(&self) -> &str {
61        &self.contents
62    }
63
64    pub fn hidden_type(&self) -> &str {
65        &self.hidden_type
66    }
67
68    pub fn hidden_contents(&self) -> &str {
69        &self.hidden_contents
70    }
71
72    pub fn tool_call_metadata(&self) -> Result<Option<ProviderCall>, RecoveryError> {
73        if self.box_type != TOOL_CALL_TYPE || self.hidden_type != TOOL_CALL_HIDDEN_TYPE {
74            return Ok(None);
75        }
76
77        let fields = decode_fields(&self.hidden_contents, 4)
78            .ok_or(RecoveryError::MalformedToolConvention)?;
79        let nonce = decode_nonce(fields[0]).ok_or(RecoveryError::MalformedToolConvention)?;
80        let sequence = fields[1]
81            .parse::<u64>()
82            .map_err(|_| RecoveryError::MalformedToolConvention)?;
83
84        Ok(Some(ProviderCall {
85            tool_call_id: ToolCallId::new(nonce, sequence),
86            name: fields[2].to_owned(),
87            arguments: fields[3].to_owned(),
88        }))
89    }
90
91    pub fn tool_result_metadata(&self) -> Result<Option<ToolResultMetadata>, RecoveryError> {
92        if self.box_type != TOOL_RESULT_TYPE || self.hidden_type != TOOL_RESULT_HIDDEN_TYPE {
93            return Ok(None);
94        }
95
96        let fields = decode_fields(&self.hidden_contents, 5)
97            .ok_or(RecoveryError::MalformedToolConvention)?;
98        let nonce = decode_nonce(fields[0]).ok_or(RecoveryError::MalformedToolConvention)?;
99        let sequence = fields[1]
100            .parse::<u64>()
101            .map_err(|_| RecoveryError::MalformedToolConvention)?;
102        let originating_call = fields[2]
103            .parse::<u64>()
104            .map_err(|_| RecoveryError::MalformedToolConvention)?;
105        let result = match fields[3] {
106            "ok" => Ok(fields[4].to_owned()),
107            "err" => Err(fields[4].to_owned()),
108            _ => return Err(RecoveryError::MalformedToolConvention),
109        };
110
111        Ok(Some(ToolResultMetadata {
112            tool_call_id: ToolCallId::new(nonce, sequence),
113            originating_call: BoxId::new(originating_call),
114            result,
115        }))
116    }
117}
118
119#[derive(Clone, Copy, Debug, Eq, PartialEq)]
120pub struct ToolCallId {
121    nonce: [u8; 12],
122    sequence: u64,
123}
124
125impl ToolCallId {
126    pub const fn new(nonce: [u8; 12], sequence: u64) -> Self {
127        Self { nonce, sequence }
128    }
129
130    pub const fn nonce(self) -> [u8; 12] {
131        self.nonce
132    }
133
134    pub const fn sequence(self) -> u64 {
135        self.sequence
136    }
137}
138
139#[derive(Clone, Debug, Eq, PartialEq)]
140pub struct ProviderCall {
141    pub tool_call_id: ToolCallId,
142    pub name: String,
143    pub arguments: String,
144}
145
146#[derive(Clone, Debug, Eq, PartialEq)]
147pub struct DispatchedToolCall {
148    pub tool_call_id: ToolCallId,
149    pub call_box_id: BoxId,
150}
151
152#[derive(Clone, Debug, Eq, PartialEq)]
153pub struct ToolResultMetadata {
154    pub tool_call_id: ToolCallId,
155    pub originating_call: BoxId,
156    pub result: Result<String, String>,
157}
158
159#[derive(Clone, Copy, Debug, Eq, PartialEq)]
160pub enum TransitionError {
161    InvalidPhase,
162    BoxIdOverflow,
163    DuplicateToolCall,
164    UnknownToolCall,
165    DuplicateReturn,
166    MalformedToolConvention,
167    WrongOriginatingCall,
168}
169
170#[derive(Clone, Copy, Debug, Eq, PartialEq)]
171pub enum RecoveryError {
172    NonContiguousBoxId,
173    MalformedToolConvention,
174    DuplicateToolCall,
175    UnknownToolCall,
176    DuplicateReturn,
177    WrongOriginatingCall,
178}
179
180pub fn tool_call_box(call: &ProviderCall) -> ChatBox {
181    let nonce = encode_nonce(call.tool_call_id.nonce);
182    let sequence = call.tool_call_id.sequence.to_string();
183    let hidden_contents = encode_fields(&[&nonce, &sequence, &call.name, &call.arguments]);
184
185    ChatBox::new(
186        BoxId::new(0),
187        TOOL_CALL_TYPE.to_owned(),
188        format!("Tool: {}\nArguments:\n{}", call.name, call.arguments),
189        TOOL_CALL_HIDDEN_TYPE.to_owned(),
190        hidden_contents,
191    )
192}
193
194pub fn tool_result_box(
195    tool_call_id: ToolCallId,
196    originating_call: BoxId,
197    result: Result<String, String>,
198) -> ChatBox {
199    let nonce = encode_nonce(tool_call_id.nonce);
200    let sequence = tool_call_id.sequence.to_string();
201    let originating_call = originating_call.get().to_string();
202    let (status, contents) = match &result {
203        Ok(contents) => ("ok", contents.clone()),
204        Err(contents) => ("err", contents.clone()),
205    };
206    let hidden_contents = encode_fields(&[&nonce, &sequence, &originating_call, status, &contents]);
207
208    ChatBox::new(
209        BoxId::new(0),
210        TOOL_RESULT_TYPE.to_owned(),
211        contents,
212        TOOL_RESULT_HIDDEN_TYPE.to_owned(),
213        hidden_contents,
214    )
215}
216
217pub struct Chatend {
218    boxes: Vec<ChatBox>,
219    round_active: bool,
220    active_arrivals: Vec<ChatBox>,
221}
222
223impl Chatend {
224    pub const fn new() -> Self {
225        Self {
226            boxes: Vec::new(),
227            round_active: false,
228            active_arrivals: Vec::new(),
229        }
230    }
231
232    pub fn boxes(&self) -> &[ChatBox] {
233        &self.boxes
234    }
235
236    pub const fn round_active(&self) -> bool {
237        self.round_active
238    }
239
240    pub fn accept_box(
241        &mut self,
242        box_type: String,
243        contents: String,
244        hidden_type: String,
245        hidden_contents: String,
246    ) -> Result<Option<BoxId>, TransitionError> {
247        self.accept_arrival(ChatBox::new(
248            BoxId::new(0),
249            box_type,
250            contents,
251            hidden_type,
252            hidden_contents,
253        ))
254    }
255
256    pub fn accept_system(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
257        self.accept_box(
258            SYSTEM_MESSAGE_TYPE.to_owned(),
259            contents,
260            String::new(),
261            String::new(),
262        )
263    }
264
265    pub fn accept_user(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
266        self.accept_box(
267            USER_MESSAGE_TYPE.to_owned(),
268            contents,
269            String::new(),
270            String::new(),
271        )
272    }
273
274    pub fn accept_attachment(
275        &mut self,
276        contents: String,
277        hidden_type: String,
278        hidden_contents: String,
279    ) -> Result<Option<BoxId>, TransitionError> {
280        self.accept_box(
281            ATTACHMENT_TYPE.to_owned(),
282            contents,
283            hidden_type,
284            hidden_contents,
285        )
286    }
287
288    pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
289        if self.round_active {
290            return Err(TransitionError::InvalidPhase);
291        }
292
293        self.round_active = true;
294        Ok(self.boxes.last().map(ChatBox::id))
295    }
296
297    pub fn append_stage(
298        &mut self,
299        agent_contents: String,
300        calls: Vec<ProviderCall>,
301    ) -> Result<Vec<DispatchedToolCall>, TransitionError> {
302        if !self.round_active {
303            return Err(TransitionError::InvalidPhase);
304        }
305
306        for (index, call) in calls.iter().enumerate() {
307            if self.call_box_id(call.tool_call_id)?.is_some()
308                || calls[..index]
309                    .iter()
310                    .any(|earlier| earlier.tool_call_id == call.tool_call_id)
311            {
312                return Err(TransitionError::DuplicateToolCall);
313            }
314        }
315
316        let has_agent_contents = !agent_contents.is_empty();
317        let mut additions = Vec::with_capacity(calls.len() + usize::from(has_agent_contents));
318        if has_agent_contents {
319            additions.push(ChatBox::new(
320                BoxId::new(0),
321                AGENT_MESSAGE_TYPE.to_owned(),
322                agent_contents,
323                String::new(),
324                String::new(),
325            ));
326        }
327        additions.extend(calls.iter().map(tool_call_box));
328
329        let appended = self.append_batch(additions)?;
330        let call_offset = usize::from(has_agent_contents);
331        Ok(calls
332            .iter()
333            .zip(appended[call_offset..].iter())
334            .map(|(call, value)| DispatchedToolCall {
335                tool_call_id: call.tool_call_id,
336                call_box_id: value.id(),
337            })
338            .collect())
339    }
340
341    pub fn accept_async_return(
342        &mut self,
343        tool_call_id: ToolCallId,
344        result: Result<String, String>,
345    ) -> Result<Option<BoxId>, TransitionError> {
346        let originating_call = self
347            .call_box_id(tool_call_id)?
348            .ok_or(TransitionError::UnknownToolCall)?;
349        if self.has_return(tool_call_id)? {
350            return Err(TransitionError::DuplicateReturn);
351        }
352
353        self.accept_arrival(tool_result_box(tool_call_id, originating_call, result))
354    }
355
356    pub fn flush_active_arrivals(&mut self) -> Result<Vec<ChatBox>, TransitionError> {
357        if !self.round_active {
358            return Err(TransitionError::InvalidPhase);
359        }
360
361        let appended = self.append_batch(self.active_arrivals.clone())?;
362        self.active_arrivals.clear();
363        Ok(appended)
364    }
365
366    pub fn done(&mut self, final_agent_contents: String) -> Result<Vec<ChatBox>, TransitionError> {
367        if !self.round_active {
368            return Err(TransitionError::InvalidPhase);
369        }
370
371        let has_final_contents = !final_agent_contents.is_empty();
372        let mut additions =
373            Vec::with_capacity(self.active_arrivals.len() + usize::from(has_final_contents));
374        if has_final_contents {
375            additions.push(ChatBox::new(
376                BoxId::new(0),
377                AGENT_MESSAGE_TYPE.to_owned(),
378                final_agent_contents,
379                String::new(),
380                String::new(),
381            ));
382        }
383        additions.extend(self.active_arrivals.iter().cloned());
384
385        let appended = self.append_batch(additions)?;
386        self.active_arrivals.clear();
387        self.round_active = false;
388        Ok(appended)
389    }
390
391    pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
392        let mut calls = Vec::<(ToolCallId, BoxId)>::new();
393        let mut returned = Vec::<ToolCallId>::new();
394
395        for (index, value) in boxes.iter().enumerate() {
396            let expected = u64::try_from(index)
397                .ok()
398                .and_then(|index| index.checked_add(1))
399                .ok_or(RecoveryError::NonContiguousBoxId)?;
400            if value.id().get() != expected {
401                return Err(RecoveryError::NonContiguousBoxId);
402            }
403
404            if let Some(call) = value.tool_call_metadata()? {
405                if calls
406                    .iter()
407                    .any(|(tool_call_id, _)| *tool_call_id == call.tool_call_id)
408                {
409                    return Err(RecoveryError::DuplicateToolCall);
410                }
411                calls.push((call.tool_call_id, value.id()));
412            }
413
414            if let Some(result) = value.tool_result_metadata()? {
415                let Some((_, call_box_id)) = calls
416                    .iter()
417                    .find(|(tool_call_id, _)| *tool_call_id == result.tool_call_id)
418                else {
419                    return Err(RecoveryError::UnknownToolCall);
420                };
421                if *call_box_id != result.originating_call {
422                    return Err(RecoveryError::WrongOriginatingCall);
423                }
424                if returned.contains(&result.tool_call_id) {
425                    return Err(RecoveryError::DuplicateReturn);
426                }
427                returned.push(result.tool_call_id);
428            }
429        }
430
431        Ok(Self {
432            boxes,
433            round_active: false,
434            active_arrivals: Vec::new(),
435        })
436    }
437
438    fn accept_arrival(&mut self, value: ChatBox) -> Result<Option<BoxId>, TransitionError> {
439        if self.round_active {
440            self.active_arrivals.push(value);
441            Ok(None)
442        } else {
443            let mut appended = self.append_batch(vec![value])?;
444            Ok(appended.pop().map(|value| value.id()))
445        }
446    }
447
448    fn append_batch(
449        &mut self,
450        mut additions: Vec<ChatBox>,
451    ) -> Result<Vec<ChatBox>, TransitionError> {
452        self.ensure_capacity(additions.len())?;
453
454        let mut previous = self.boxes.last().map_or(0, |value| value.id().get());
455        for value in &mut additions {
456            previous = previous
457                .checked_add(1)
458                .ok_or(TransitionError::BoxIdOverflow)?;
459            value.id = BoxId::new(previous);
460        }
461
462        self.boxes.extend(additions.iter().cloned());
463        Ok(additions)
464    }
465
466    fn ensure_capacity(&self, additional: usize) -> Result<(), TransitionError> {
467        let additional = u64::try_from(additional).map_err(|_| TransitionError::BoxIdOverflow)?;
468        let previous = self.boxes.last().map_or(0, |value| value.id().get());
469        previous
470            .checked_add(additional)
471            .ok_or(TransitionError::BoxIdOverflow)?;
472        Ok(())
473    }
474
475    fn call_box_id(&self, tool_call_id: ToolCallId) -> Result<Option<BoxId>, TransitionError> {
476        let mut found = None;
477        for value in &self.boxes {
478            let metadata = value
479                .tool_call_metadata()
480                .map_err(|_| TransitionError::MalformedToolConvention)?;
481            if metadata
482                .as_ref()
483                .is_some_and(|call| call.tool_call_id == tool_call_id)
484            {
485                if found.is_some() {
486                    return Err(TransitionError::DuplicateToolCall);
487                }
488                found = Some(value.id());
489            }
490        }
491        Ok(found)
492    }
493
494    fn has_return(&self, tool_call_id: ToolCallId) -> Result<bool, TransitionError> {
495        for value in self.boxes.iter().chain(&self.active_arrivals) {
496            let metadata = value
497                .tool_result_metadata()
498                .map_err(|_| TransitionError::MalformedToolConvention)?;
499            if metadata
500                .as_ref()
501                .is_some_and(|result| result.tool_call_id == tool_call_id)
502            {
503                return Ok(true);
504            }
505        }
506        Ok(false)
507    }
508}
509
510impl Default for Chatend {
511    fn default() -> Self {
512        Self::new()
513    }
514}
515
516fn encode_fields(fields: &[&str]) -> String {
517    let mut encoded = String::new();
518    for field in fields {
519        encoded.push_str(&field.len().to_string());
520        encoded.push(':');
521        encoded.push_str(field);
522    }
523    encoded
524}
525
526fn decode_fields(input: &str, count: usize) -> Option<Vec<&str>> {
527    let mut fields = Vec::with_capacity(count);
528    let mut cursor = 0;
529
530    for _ in 0..count {
531        let colon_offset = input
532            .as_bytes()
533            .get(cursor..)?
534            .iter()
535            .position(|byte| *byte == b':')?;
536        let colon = cursor.checked_add(colon_offset)?;
537        let length = input.get(cursor..colon)?.parse::<usize>().ok()?;
538        let start = colon.checked_add(1)?;
539        let end = start.checked_add(length)?;
540        fields.push(input.get(start..end)?);
541        cursor = end;
542    }
543
544    (cursor == input.len()).then_some(fields)
545}
546
547fn encode_nonce(nonce: [u8; 12]) -> String {
548    const HEX: &[u8; 16] = b"0123456789abcdef";
549    let mut encoded = String::with_capacity(24);
550    for byte in nonce {
551        encoded.push(char::from(HEX[usize::from(byte >> 4)]));
552        encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
553    }
554    encoded
555}
556
557fn decode_nonce(value: &str) -> Option<[u8; 12]> {
558    if value.len() != 24 {
559        return None;
560    }
561
562    let mut nonce = [0; 12];
563    for (index, slot) in nonce.iter_mut().enumerate() {
564        let offset = index.checked_mul(2)?;
565        let high = decode_hex(*value.as_bytes().get(offset)?)?;
566        let low = decode_hex(*value.as_bytes().get(offset.checked_add(1)?)?)?;
567        *slot = (high << 4) | low;
568    }
569    Some(nonce)
570}
571
572const fn decode_hex(value: u8) -> Option<u8> {
573    match value {
574        b'0'..=b'9' => Some(value - b'0'),
575        b'a'..=b'f' => Some(value - b'a' + 10),
576        b'A'..=b'F' => Some(value - b'A' + 10),
577        _ => None,
578    }
579}
580
581#[cfg(test)]
582mod tests;
kcode_k1_chat_chatend/lib.rs

kcode_k1_chat_chatend/
lib.rs
