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

kcode_k1_chat_boxes/
lib.rs
