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

kcode_k1_chat_boxes/
lib.rs
