kcode_k1_chat_boxes/
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 AGENT_RESPONSE_TYPE: &str = "Agent Response";
7pub const USER_ATTACHMENT_TYPE: &str = "User Attachment";
8pub const AGENT_ATTACHMENT_TYPE: &str = "Agent Attachment";
9pub const TOOL_CALL_TYPE: &str = "Tool Call";
10pub const TOOL_MESSAGE_TYPE: &str = "Tool Message";
11pub const TOOL_ATTACHMENT_TYPE: &str = "Tool Attachment";
12pub const TOOL_RESULT_TYPE: &str = "Tool Result";
13pub const ATTACHMENT_TYPE: &str = USER_ATTACHMENT_TYPE;
14
15pub const TOOL_CALL_HIDDEN_TYPE: &str = "k1.tool-call/v1";
16pub const TOOL_MESSAGE_HIDDEN_TYPE: &str = "k1.tool-message/v1";
17pub const TOOL_MESSAGE_V2_HIDDEN_TYPE: &str = "k1.tool-message/v2";
18pub const TOOL_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/v1";
19pub const TOOL_RESULT_V2_HIDDEN_TYPE: &str = "k1.tool-result/v2";
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub struct BoxId(u64);
23
24impl BoxId {
25    pub const fn new(value: u64) -> Self {
26        Self(value)
27    }
28
29    pub const fn get(self) -> u64 {
30        self.0
31    }
32}
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
35pub struct ToolCallId {
36    nonce: [u8; 12],
37    sequence: u64,
38}
39
40impl ToolCallId {
41    pub const fn new(nonce: [u8; 12], sequence: u64) -> Self {
42        Self { nonce, sequence }
43    }
44
45    pub const fn nonce(self) -> [u8; 12] {
46        self.nonce
47    }
48
49    pub const fn sequence(self) -> u64 {
50        self.sequence
51    }
52}
53
54#[derive(Clone, Debug, Eq, PartialEq)]
55pub struct ChatBox {
56    id: BoxId,
57    box_type: String,
58    contents: String,
59    hidden_type: String,
60    hidden_contents: String,
61}
62
63impl ChatBox {
64    pub fn new(
65        id: BoxId,
66        box_type: String,
67        contents: String,
68        hidden_type: String,
69        hidden_contents: String,
70    ) -> Self {
71        Self {
72            id,
73            box_type,
74            contents,
75            hidden_type,
76            hidden_contents,
77        }
78    }
79
80    pub const fn id(&self) -> BoxId {
81        self.id
82    }
83
84    pub fn box_type(&self) -> &str {
85        &self.box_type
86    }
87
88    pub fn contents(&self) -> &str {
89        &self.contents
90    }
91
92    pub fn hidden_type(&self) -> &str {
93        &self.hidden_type
94    }
95
96    pub fn hidden_contents(&self) -> &str {
97        &self.hidden_contents
98    }
99
100    pub fn tool_call_metadata(&self) -> Result<Option<ProviderCall>, MetadataError> {
101        let Some(fields) = self.metadata_fields(TOOL_CALL_TYPE, TOOL_CALL_HIDDEN_TYPE, 4)? else {
102            return Ok(None);
103        };
104        Ok(Some(ProviderCall {
105            tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
106            name: fields[2].to_owned(),
107            arguments: fields[3].to_owned(),
108        }))
109    }
110
111    pub fn tool_message_metadata(&self) -> Result<Option<ToolMessageMetadata>, MetadataError> {
112        let (fields, named) = match self.hidden_type.as_str() {
113            TOOL_MESSAGE_HIDDEN_TYPE => (
114                self.metadata_fields(TOOL_MESSAGE_TYPE, TOOL_MESSAGE_HIDDEN_TYPE, 5)?
115                    .ok_or(MetadataError)?,
116                false,
117            ),
118            TOOL_MESSAGE_V2_HIDDEN_TYPE => {
119                if self.box_type.is_empty() {
120                    return Err(MetadataError);
121                }
122                (
123                    decode_fields(&self.hidden_contents, 5).ok_or(MetadataError)?,
124                    true,
125                )
126            }
127            _ => return Ok(None),
128        };
129        let message_index = parse_u64(fields[3])?;
130        if message_index == 0 || (named && self.contents != fields[4]) {
131            return Err(MetadataError);
132        }
133        Ok(Some(ToolMessageMetadata {
134            tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
135            originating_call: BoxId::new(parse_u64(fields[2])?),
136            message_index,
137            message: fields[4].to_owned(),
138        }))
139    }
140
141    pub fn tool_result_metadata(&self) -> Result<Option<ToolResultMetadata>, MetadataError> {
142        let fields = match self.hidden_type.as_str() {
143            TOOL_RESULT_HIDDEN_TYPE => {
144                self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_HIDDEN_TYPE, 5)?
145            }
146            TOOL_RESULT_V2_HIDDEN_TYPE => {
147                self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_V2_HIDDEN_TYPE, 7)?
148            }
149            _ => None,
150        };
151        fields
152            .map(|fields| parse_result_fields(&fields))
153            .transpose()
154    }
155
156    pub fn tool_result_v2_metadata(&self) -> Result<Option<ToolResultV2Metadata>, MetadataError> {
157        let Some(fields) = self.metadata_fields(TOOL_RESULT_TYPE, TOOL_RESULT_V2_HIDDEN_TYPE, 7)?
158        else {
159            return Ok(None);
160        };
161        let parsed = parse_result_fields(&fields)?;
162        Ok(Some(ToolResultV2Metadata {
163            tool_call_id: parsed.tool_call_id,
164            originating_call: parsed.originating_call,
165            result: parsed.result,
166            metadata_type: fields[5].to_owned(),
167            metadata_contents: fields[6].to_owned(),
168        }))
169    }
170
171    fn metadata_fields<'a>(
172        &'a self,
173        box_type: &str,
174        hidden_type: &str,
175        count: usize,
176    ) -> Result<Option<Vec<&'a str>>, MetadataError> {
177        match (self.hidden_type == hidden_type, self.box_type == box_type) {
178            (false, _) => Ok(None),
179            (true, false) => Err(MetadataError),
180            (true, true) => decode_fields(&self.hidden_contents, count)
181                .map(Some)
182                .ok_or(MetadataError),
183        }
184    }
185}
186
187#[derive(Clone, Copy, Debug, Eq, PartialEq)]
188pub struct MetadataError;
189
190#[derive(Clone, Debug, Eq, PartialEq)]
191pub struct ProviderCall {
192    pub tool_call_id: ToolCallId,
193    pub name: String,
194    pub arguments: String,
195}
196
197#[derive(Clone, Debug, Eq, PartialEq)]
198pub struct ToolMessageMetadata {
199    pub tool_call_id: ToolCallId,
200    pub originating_call: BoxId,
201    pub message_index: u64,
202    pub message: String,
203}
204
205#[derive(Clone, Debug, Eq, PartialEq)]
206pub struct ToolResultMetadata {
207    pub tool_call_id: ToolCallId,
208    pub originating_call: BoxId,
209    pub result: Result<String, String>,
210}
211
212#[derive(Clone, Debug, Eq, PartialEq)]
213pub struct ToolResultV2Metadata {
214    pub tool_call_id: ToolCallId,
215    pub originating_call: BoxId,
216    pub result: Result<String, String>,
217    pub metadata_type: String,
218    pub metadata_contents: String,
219}
220
221pub fn tool_call_box(call: &ProviderCall) -> ChatBox {
222    let call_id = format_tool_call_id(call.tool_call_id);
223    let nonce = encode_nonce(call.tool_call_id.nonce);
224    let sequence = call.tool_call_id.sequence.to_string();
225    let hidden_contents = encode_fields(&[&nonce, &sequence, &call.name, &call.arguments]);
226    ChatBox::new(
227        BoxId::new(0),
228        TOOL_CALL_TYPE.to_owned(),
229        format!(
230            "Call ID: {call_id}\nCall Name: {}\nArguments:\n{}",
231            call.name, call.arguments
232        ),
233        TOOL_CALL_HIDDEN_TYPE.to_owned(),
234        hidden_contents,
235    )
236}
237
238pub fn tool_message_box(metadata: &ToolMessageMetadata) -> Result<ChatBox, MetadataError> {
239    let hidden_contents = tool_message_hidden_contents(metadata)?;
240    let call_id = format_tool_call_id(metadata.tool_call_id);
241    let origin = metadata.originating_call.get();
242    let index = metadata.message_index;
243    Ok(ChatBox::new(
244        BoxId::new(0),
245        TOOL_MESSAGE_TYPE.to_owned(),
246        format!(
247            "Call ID: {call_id}\nOriginating Call Box ID: {origin}\nMessage Index: {index}\nMessage:\n{}",
248            metadata.message
249        ),
250        TOOL_MESSAGE_HIDDEN_TYPE.to_owned(),
251        hidden_contents,
252    ))
253}
254
255pub fn named_tool_message_box(
256    metadata: &ToolMessageMetadata,
257    title: &str,
258) -> Result<ChatBox, MetadataError> {
259    if title.is_empty() {
260        return Err(MetadataError);
261    }
262    let hidden_contents = tool_message_hidden_contents(metadata)?;
263    Ok(ChatBox::new(
264        BoxId::new(0),
265        title.to_owned(),
266        metadata.message.clone(),
267        TOOL_MESSAGE_V2_HIDDEN_TYPE.to_owned(),
268        hidden_contents,
269    ))
270}
271
272fn tool_message_hidden_contents(metadata: &ToolMessageMetadata) -> Result<String, MetadataError> {
273    if metadata.message_index == 0 {
274        return Err(MetadataError);
275    }
276    let nonce = encode_nonce(metadata.tool_call_id.nonce);
277    let sequence = metadata.tool_call_id.sequence.to_string();
278    let origin = metadata.originating_call.get().to_string();
279    let index = metadata.message_index.to_string();
280    Ok(encode_fields(&[
281        &nonce,
282        &sequence,
283        &origin,
284        &index,
285        &metadata.message,
286    ]))
287}
288
289pub fn tool_result_box(
290    tool_call_id: ToolCallId,
291    originating_call: BoxId,
292    result: Result<String, String>,
293) -> ChatBox {
294    tool_result_box_inner(tool_call_id, originating_call, &result, None)
295}
296
297pub fn tool_result_v2_box(metadata: &ToolResultV2Metadata) -> ChatBox {
298    tool_result_box_inner(
299        metadata.tool_call_id,
300        metadata.originating_call,
301        &metadata.result,
302        Some((&metadata.metadata_type, &metadata.metadata_contents)),
303    )
304}
305
306fn tool_result_box_inner(
307    tool_call_id: ToolCallId,
308    originating_call: BoxId,
309    result: &Result<String, String>,
310    metadata: Option<(&str, &str)>,
311) -> ChatBox {
312    let call_id = format_tool_call_id(tool_call_id);
313    let nonce = encode_nonce(tool_call_id.nonce);
314    let sequence = tool_call_id.sequence.to_string();
315    let origin = originating_call.get().to_string();
316    let (hidden_status, visible_status, raw_result) = match result {
317        Ok(contents) => ("ok", "ok", contents),
318        Err(contents) => ("err", "error", contents),
319    };
320    let mut fields = vec![
321        nonce.as_str(),
322        sequence.as_str(),
323        origin.as_str(),
324        hidden_status,
325        raw_result.as_str(),
326    ];
327    if let Some((metadata_type, metadata_contents)) = metadata {
328        fields.extend([metadata_type, metadata_contents]);
329    }
330    let hidden_contents = encode_fields(&fields);
331    let hidden_type = metadata
332        .map(|_| TOOL_RESULT_V2_HIDDEN_TYPE)
333        .unwrap_or(TOOL_RESULT_HIDDEN_TYPE);
334    ChatBox::new(
335        BoxId::new(0),
336        TOOL_RESULT_TYPE.to_owned(),
337        format!(
338            "Call ID: {call_id}\nOriginating Call Box ID: {origin}\nStatus: {visible_status}\nResult:\n{raw_result}"
339        ),
340        hidden_type.to_owned(),
341        hidden_contents,
342    )
343}
344
345fn parse_result_fields(fields: &[&str]) -> Result<ToolResultMetadata, MetadataError> {
346    let result = match fields[3] {
347        "ok" => Ok(fields[4].to_owned()),
348        "err" => Err(fields[4].to_owned()),
349        _ => return Err(MetadataError),
350    };
351    Ok(ToolResultMetadata {
352        tool_call_id: parse_tool_call_id(fields[0], fields[1])?,
353        originating_call: BoxId::new(parse_u64(fields[2])?),
354        result,
355    })
356}
357
358fn parse_tool_call_id(nonce: &str, sequence: &str) -> Result<ToolCallId, MetadataError> {
359    let nonce = decode_nonce(nonce).ok_or(MetadataError)?;
360    Ok(ToolCallId::new(nonce, parse_u64(sequence)?))
361}
362
363fn parse_u64(value: &str) -> Result<u64, MetadataError> {
364    value.parse::<u64>().map_err(|_| MetadataError)
365}
366
367fn format_tool_call_id(tool_call_id: ToolCallId) -> String {
368    format!(
369        "{}/{}",
370        encode_nonce(tool_call_id.nonce),
371        tool_call_id.sequence
372    )
373}
374
375fn encode_fields(fields: &[&str]) -> String {
376    let mut encoded = String::new();
377    for field in fields {
378        encoded.push_str(&field.len().to_string());
379        encoded.push(':');
380        encoded.push_str(field);
381    }
382    encoded
383}
384
385fn decode_fields(input: &str, count: usize) -> Option<Vec<&str>> {
386    let mut fields = Vec::with_capacity(count);
387    let mut cursor = 0;
388    for _ in 0..count {
389        let colon_offset = input
390            .as_bytes()
391            .get(cursor..)?
392            .iter()
393            .position(|byte| *byte == b':')?;
394        let colon = cursor.checked_add(colon_offset)?;
395        let length = input.get(cursor..colon)?.parse::<usize>().ok()?;
396        let start = colon.checked_add(1)?;
397        let end = start.checked_add(length)?;
398        fields.push(input.get(start..end)?);
399        cursor = end;
400    }
401    (cursor == input.len()).then_some(fields)
402}
403
404fn encode_nonce(nonce: [u8; 12]) -> String {
405    const HEX: &[u8; 16] = b"0123456789abcdef";
406    let mut encoded = String::with_capacity(24);
407    for byte in nonce {
408        encoded.push(char::from(HEX[usize::from(byte >> 4)]));
409        encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
410    }
411    encoded
412}
413
414fn decode_nonce(value: &str) -> Option<[u8; 12]> {
415    if value.len() != 24 {
416        return None;
417    }
418    let mut nonce = [0; 12];
419    for (slot, digits) in nonce.iter_mut().zip(value.as_bytes().chunks_exact(2)) {
420        let digits = std::str::from_utf8(digits).ok()?;
421        *slot = u8::from_str_radix(digits, 16).ok()?;
422    }
423    Some(nonce)
424}
425
426#[cfg(test)]
427mod tests;
kcode_k1_chat_boxes/lib.rs

kcode_k1_chat_boxes/
lib.rs
