Skip to main content

kcode_k1_chat_box/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::{error::Error, fmt, str::FromStr};
4
5use semver::Version;
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8
9pub const TOOL_CALL_TYPE: &str = "Tool Call";
10pub const TOOL_RESULT_TYPE: &str = "Tool Result";
11pub const TOOL_CALL_HIDDEN_TYPE: &str = "k1.tool-call/1.0.0";
12pub const TOOL_RESULT_HIDDEN_TYPE: &str = "k1.tool-result/1.0.0";
13
14const CALL_PREFIX: &str = "k1.tool-call/";
15const RESULT_PREFIX: &str = "k1.tool-result/";
16
17#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
18pub struct BoxId(u64);
19
20impl BoxId {
21    pub const fn new(value: u64) -> Self {
22        Self(value)
23    }
24
25    pub const fn get(self) -> u64 {
26        self.0
27    }
28}
29
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct ChatBox {
32    id: BoxId,
33    box_type: String,
34    contents: String,
35    hidden_type: String,
36    hidden_contents: String,
37}
38
39impl ChatBox {
40    pub fn new(
41        id: BoxId,
42        box_type: String,
43        contents: String,
44        hidden_type: String,
45        hidden_contents: String,
46    ) -> Self {
47        Self {
48            id,
49            box_type,
50            contents,
51            hidden_type,
52            hidden_contents,
53        }
54    }
55
56    pub const fn id(&self) -> BoxId {
57        self.id
58    }
59
60    pub fn box_type(&self) -> &str {
61        &self.box_type
62    }
63
64    pub fn contents(&self) -> &str {
65        &self.contents
66    }
67
68    pub fn hidden_type(&self) -> &str {
69        &self.hidden_type
70    }
71
72    pub fn hidden_contents(&self) -> &str {
73        &self.hidden_contents
74    }
75
76    pub fn tool_call(id: BoxId, call: ToolCall) -> Result<Self, EnvelopeError> {
77        Ok(Self::new(
78            id,
79            TOOL_CALL_TYPE.into(),
80            call_contents(&call)?,
81            TOOL_CALL_HIDDEN_TYPE.into(),
82            compact_json(&CallHidden::from(&call))?,
83        ))
84    }
85
86    pub fn tool_result(id: BoxId, result: ToolResult) -> Result<Self, EnvelopeError> {
87        result.validate()?;
88        Ok(Self::new(
89            id,
90            TOOL_RESULT_TYPE.into(),
91            result_contents(&result),
92            TOOL_RESULT_HIDDEN_TYPE.into(),
93            compact_json(&ResultHidden::from(&result))?,
94        ))
95    }
96
97    pub fn tool_call_metadata(&self) -> Result<Option<ToolCall>, EnvelopeError> {
98        if self.box_type != TOOL_CALL_TYPE || !supported(&self.hidden_type, CALL_PREFIX)? {
99            return Ok(None);
100        }
101        let hidden: CallHidden = decode_json(&self.hidden_contents)?;
102        let call = ToolCall::new(
103            ToolCallId::from_str(&hidden.call_id)?,
104            hidden.tool,
105            hidden.tool_version,
106            hidden.arguments,
107        )?;
108        if call_contents(&call)? != self.contents {
109            return Err(EnvelopeError::InvalidVisibleContent);
110        }
111        Ok(Some(call))
112    }
113
114    pub fn tool_result_metadata(&self) -> Result<Option<ToolResult>, EnvelopeError> {
115        if self.box_type != TOOL_RESULT_TYPE || !supported(&self.hidden_type, RESULT_PREFIX)? {
116            return Ok(None);
117        }
118        let hidden: ResultHidden = decode_json(&self.hidden_contents)?;
119        let status = ToolResultStatus::from_str(&hidden.status)?;
120        let view = parse_result_view(&self.contents, &hidden.call_id, &hidden.tool, status)?;
121        ToolResult::new(
122            ToolCallId::from_str(&hidden.call_id)?,
123            BoxId::new(hidden.originating_call_box_id),
124            hidden.tool,
125            hidden.tool_version,
126            status,
127            hidden.data,
128            view,
129        )
130        .map(Some)
131    }
132}
133
134#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
135pub struct ToolCallId(u64);
136
137impl ToolCallId {
138    pub fn new(value: u64) -> Result<Self, EnvelopeError> {
139        if value == 0 {
140            return Err(EnvelopeError::InvalidCallId);
141        }
142        Ok(Self(value))
143    }
144
145    pub const fn get(self) -> u64 {
146        self.0
147    }
148}
149
150impl fmt::Display for ToolCallId {
151    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
152        write!(formatter, "c{}", self.0)
153    }
154}
155
156impl FromStr for ToolCallId {
157    type Err = EnvelopeError;
158
159    fn from_str(value: &str) -> Result<Self, Self::Err> {
160        let Some(number) = value.strip_prefix('c') else {
161            return Err(EnvelopeError::InvalidCallId);
162        };
163        if number.is_empty() || number.starts_with('0') {
164            return Err(EnvelopeError::InvalidCallId);
165        }
166        let value = number
167            .parse::<u64>()
168            .map_err(|_| EnvelopeError::InvalidCallId)?;
169        Self::new(value)
170    }
171}
172
173#[derive(Clone, Debug, PartialEq)]
174pub struct ToolCall {
175    call_id: ToolCallId,
176    tool: String,
177    tool_version: String,
178    arguments: Value,
179}
180
181impl ToolCall {
182    pub fn new(
183        call_id: ToolCallId,
184        tool: String,
185        tool_version: String,
186        arguments: Value,
187    ) -> Result<Self, EnvelopeError> {
188        validate_tool(&tool)?;
189        validate_version(&tool_version)?;
190        let Some(properties) = arguments.as_object() else {
191            return Err(EnvelopeError::InvalidArguments);
192        };
193        if properties.contains_key("tool") {
194            return Err(EnvelopeError::InvalidArguments);
195        }
196        Ok(Self {
197            call_id,
198            tool,
199            tool_version,
200            arguments,
201        })
202    }
203
204    pub const fn call_id(&self) -> ToolCallId {
205        self.call_id
206    }
207
208    pub fn tool(&self) -> &str {
209        &self.tool
210    }
211
212    pub fn tool_version(&self) -> &str {
213        &self.tool_version
214    }
215
216    pub fn arguments(&self) -> &Value {
217        &self.arguments
218    }
219}
220
221#[derive(Clone, Copy, Debug, Eq, PartialEq)]
222pub enum ToolResultStatus {
223    Ok,
224    Error,
225}
226
227impl ToolResultStatus {
228    pub const fn as_str(self) -> &'static str {
229        match self {
230            Self::Ok => "ok",
231            Self::Error => "error",
232        }
233    }
234}
235
236impl FromStr for ToolResultStatus {
237    type Err = EnvelopeError;
238
239    fn from_str(value: &str) -> Result<Self, Self::Err> {
240        match value {
241            "ok" => Ok(Self::Ok),
242            "error" => Ok(Self::Error),
243            _ => Err(EnvelopeError::InvalidStatus),
244        }
245    }
246}
247
248#[derive(Clone, Debug, PartialEq)]
249pub enum ResultView {
250    OneLine(String),
251    Multiline(String),
252    Error(String),
253}
254
255#[derive(Clone, Debug, PartialEq)]
256pub struct ToolResult {
257    call_id: ToolCallId,
258    originating_call_box_id: BoxId,
259    tool: String,
260    tool_version: String,
261    status: ToolResultStatus,
262    data: Value,
263    view: ResultView,
264}
265
266impl ToolResult {
267    #[allow(clippy::too_many_arguments)]
268    pub fn new(
269        call_id: ToolCallId,
270        originating_call_box_id: BoxId,
271        tool: String,
272        tool_version: String,
273        status: ToolResultStatus,
274        data: Value,
275        view: ResultView,
276    ) -> Result<Self, EnvelopeError> {
277        let result = Self {
278            call_id,
279            originating_call_box_id,
280            tool,
281            tool_version,
282            status,
283            data,
284            view,
285        };
286        result.validate()?;
287        Ok(result)
288    }
289
290    pub const fn call_id(&self) -> ToolCallId {
291        self.call_id
292    }
293
294    pub const fn originating_call_box_id(&self) -> BoxId {
295        self.originating_call_box_id
296    }
297
298    pub fn tool(&self) -> &str {
299        &self.tool
300    }
301
302    pub fn tool_version(&self) -> &str {
303        &self.tool_version
304    }
305
306    pub const fn status(&self) -> ToolResultStatus {
307        self.status
308    }
309
310    pub fn data(&self) -> &Value {
311        &self.data
312    }
313
314    pub fn view(&self) -> &ResultView {
315        &self.view
316    }
317
318    fn validate(&self) -> Result<(), EnvelopeError> {
319        validate_tool(&self.tool)?;
320        validate_version(&self.tool_version)?;
321        if self.originating_call_box_id.get() == 0 {
322            return Err(EnvelopeError::InvalidOriginatingCallBoxId);
323        }
324        match (self.status, &self.view) {
325            (ToolResultStatus::Ok, ResultView::OneLine(text)) if is_line(text) => Ok(()),
326            (ToolResultStatus::Ok, ResultView::Multiline(_)) => Ok(()),
327            (ToolResultStatus::Error, ResultView::Error(text)) if is_line(text) => Ok(()),
328            _ => Err(EnvelopeError::InvalidVisibleContent),
329        }
330    }
331}
332
333#[derive(Clone, Copy, Debug, Eq, PartialEq)]
334pub enum EnvelopeError {
335    InvalidCallId,
336    InvalidTool,
337    InvalidToolVersion,
338    InvalidArguments,
339    InvalidOriginatingCallBoxId,
340    InvalidStatus,
341    MalformedEnvelope,
342    UnsupportedEnvelopeVersion,
343    InvalidVisibleContent,
344}
345
346impl fmt::Display for EnvelopeError {
347    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
348        let message = match self {
349            Self::InvalidCallId => "invalid tool call ID",
350            Self::InvalidTool => "invalid tool name",
351            Self::InvalidToolVersion => "invalid tool version",
352            Self::InvalidArguments => "invalid tool arguments",
353            Self::InvalidOriginatingCallBoxId => "invalid originating call box ID",
354            Self::InvalidStatus => "invalid tool result status",
355            Self::MalformedEnvelope => "malformed tool envelope",
356            Self::UnsupportedEnvelopeVersion => "unsupported tool envelope version",
357            Self::InvalidVisibleContent => "invalid visible tool content",
358        };
359        formatter.write_str(message)
360    }
361}
362
363impl Error for EnvelopeError {}
364
365#[derive(Deserialize, Serialize)]
366#[serde(rename_all = "camelCase")]
367struct CallHidden {
368    call_id: String,
369    tool: String,
370    tool_version: String,
371    arguments: Value,
372}
373
374impl From<&ToolCall> for CallHidden {
375    fn from(value: &ToolCall) -> Self {
376        Self {
377            call_id: value.call_id.to_string(),
378            tool: value.tool.clone(),
379            tool_version: value.tool_version.clone(),
380            arguments: value.arguments.clone(),
381        }
382    }
383}
384
385#[derive(Deserialize, Serialize)]
386#[serde(rename_all = "camelCase")]
387struct ResultHidden {
388    call_id: String,
389    originating_call_box_id: u64,
390    tool: String,
391    tool_version: String,
392    status: String,
393    data: Value,
394}
395
396impl From<&ToolResult> for ResultHidden {
397    fn from(value: &ToolResult) -> Self {
398        Self {
399            call_id: value.call_id.to_string(),
400            originating_call_box_id: value.originating_call_box_id.get(),
401            tool: value.tool.clone(),
402            tool_version: value.tool_version.clone(),
403            status: value.status.as_str().into(),
404            data: value.data.clone(),
405        }
406    }
407}
408
409fn validate_tool(value: &str) -> Result<(), EnvelopeError> {
410    if value.is_empty() || !is_line(value) {
411        return Err(EnvelopeError::InvalidTool);
412    }
413    Ok(())
414}
415
416fn validate_version(value: &str) -> Result<(), EnvelopeError> {
417    Version::parse(value)
418        .map(|_| ())
419        .map_err(|_| EnvelopeError::InvalidToolVersion)
420}
421
422fn is_line(value: &str) -> bool {
423    !value.contains(['\n', '\r'])
424}
425
426fn supported(hidden_type: &str, prefix: &str) -> Result<bool, EnvelopeError> {
427    let Some(version) = hidden_type.strip_prefix(prefix) else {
428        return Ok(false);
429    };
430    let version = Version::parse(version).map_err(|_| EnvelopeError::MalformedEnvelope)?;
431    Ok(version.major == 1)
432}
433
434fn compact_json<T: Serialize>(value: &T) -> Result<String, EnvelopeError> {
435    serde_json::to_string(value).map_err(|_| EnvelopeError::MalformedEnvelope)
436}
437
438fn decode_json<T: for<'a> Deserialize<'a>>(value: &str) -> Result<T, EnvelopeError> {
439    serde_json::from_str(value).map_err(|_| EnvelopeError::MalformedEnvelope)
440}
441
442fn readable_json(value: &Value) -> Result<String, EnvelopeError> {
443    match value {
444        Value::Array(values) => {
445            let values = values
446                .iter()
447                .map(readable_json)
448                .collect::<Result<Vec<_>, _>>()?;
449            Ok(format!("[{}]", values.join(", ")))
450        }
451        Value::Object(properties) => {
452            let properties = properties
453                .iter()
454                .map(|(key, value)| {
455                    Ok(format!("{}: {}", compact_json(key)?, readable_json(value)?))
456                })
457                .collect::<Result<Vec<_>, EnvelopeError>>()?;
458            Ok(format!("{{{}}}", properties.join(", ")))
459        }
460        _ => compact_json(value),
461    }
462}
463
464fn call_contents(call: &ToolCall) -> Result<String, EnvelopeError> {
465    let Some(arguments) = call.arguments.as_object() else {
466        return Err(EnvelopeError::InvalidArguments);
467    };
468    let mut visible = Map::new();
469    visible.insert("tool".into(), Value::String(call.tool.clone()));
470    visible.extend(arguments.clone());
471    Ok(format!(
472        "Call ID: {}\nArgs: {}",
473        call.call_id,
474        readable_json(&Value::Object(visible))?
475    ))
476}
477
478fn result_contents(result: &ToolResult) -> String {
479    let prefix = format!("Call ID: {}\n{} call", result.call_id, result.tool);
480    match &result.view {
481        ResultView::OneLine(message) => format!("{prefix} result: {message}"),
482        ResultView::Multiline(body) => format!("{prefix} result:\n\n{body}"),
483        ResultView::Error(message) => format!("{prefix} error: {message}"),
484    }
485}
486
487fn parse_result_view(
488    contents: &str,
489    call_id: &str,
490    tool: &str,
491    status: ToolResultStatus,
492) -> Result<ResultView, EnvelopeError> {
493    let call_id = ToolCallId::from_str(call_id)?;
494    validate_tool(tool)?;
495    let prefix = format!("Call ID: {call_id}\n{tool} call");
496    match status {
497        ToolResultStatus::Ok => {
498            if let Some(message) = contents.strip_prefix(&format!("{prefix} result: ")) {
499                if !is_line(message) {
500                    return Err(EnvelopeError::InvalidVisibleContent);
501                }
502                return Ok(ResultView::OneLine(message.into()));
503            }
504            contents
505                .strip_prefix(&format!("{prefix} result:\n\n"))
506                .map(|body| ResultView::Multiline(body.into()))
507                .ok_or(EnvelopeError::InvalidVisibleContent)
508        }
509        ToolResultStatus::Error => contents
510            .strip_prefix(&format!("{prefix} error: "))
511            .filter(|message| is_line(message))
512            .map(|message| ResultView::Error(message.into()))
513            .ok_or(EnvelopeError::InvalidVisibleContent),
514    }
515}
516
517#[cfg(test)]
518mod tests {
519    use super::*;
520    use serde_json::json;
521
522    fn call() -> ToolCall {
523        ToolCall::new(
524            ToolCallId::new(2).unwrap(),
525            "Search".into(),
526            "1.0.0".into(),
527            json!({"query":{"tags":["rust", {"stable":true}]},"limit":2}),
528        )
529        .unwrap()
530    }
531
532    fn result(status: ToolResultStatus, view: ResultView) -> ToolResult {
533        ToolResult::new(
534            ToolCallId::new(2).unwrap(),
535            BoxId::new(9),
536            "Search".into(),
537            "1.0.0".into(),
538            status,
539            json!({"matches":[1, 2]}),
540            view,
541        )
542        .unwrap()
543    }
544
545    #[test]
546    fn call_ids_are_canonical() {
547        assert_eq!(ToolCallId::from_str("c1").unwrap().to_string(), "c1");
548        assert_eq!(
549            ToolCallId::from_str(&format!("c{}", u64::MAX))
550                .unwrap()
551                .get(),
552            u64::MAX
553        );
554        for invalid in ["", "1", "c", "c0", "c01", "c-1", "c18446744073709551616"] {
555            assert_eq!(
556                ToolCallId::from_str(invalid),
557                Err(EnvelopeError::InvalidCallId)
558            );
559        }
560    }
561
562    #[test]
563    fn call_has_exact_readable_and_hidden_json() {
564        let value = ChatBox::tool_call(BoxId::new(9), call()).unwrap();
565        assert_eq!(
566            value.contents(),
567            "Call ID: c2\nArgs: {\"tool\": \"Search\", \"query\": {\"tags\": [\"rust\", {\"stable\": true}]}, \"limit\": 2}"
568        );
569        assert_eq!(
570            value.hidden_contents(),
571            "{\"callId\":\"c2\",\"tool\":\"Search\",\"toolVersion\":\"1.0.0\",\"arguments\":{\"query\":{\"tags\":[\"rust\",{\"stable\":true}]},\"limit\":2}}"
572        );
573        assert_eq!(value.tool_call_metadata().unwrap(), Some(call()));
574    }
575
576    #[test]
577    fn all_result_views_are_exact_and_round_trip() {
578        let cases = [
579            (
580                ToolResultStatus::Ok,
581                ResultView::OneLine("success".into()),
582                "Call ID: c2\nSearch call result: success",
583            ),
584            (
585                ToolResultStatus::Ok,
586                ResultView::Multiline("first\nsecond".into()),
587                "Call ID: c2\nSearch call result:\n\nfirst\nsecond",
588            ),
589            (
590                ToolResultStatus::Error,
591                ResultView::Error("node unavailable".into()),
592                "Call ID: c2\nSearch call error: node unavailable",
593            ),
594        ];
595        for (status, view, expected) in cases {
596            let result = result(status, view);
597            let value = ChatBox::tool_result(BoxId::new(10), result.clone()).unwrap();
598            assert_eq!(value.contents(), expected);
599            assert_eq!(value.tool_result_metadata().unwrap(), Some(result));
600        }
601    }
602
603    #[test]
604    fn result_hidden_json_is_canonical() {
605        let value = ChatBox::tool_result(
606            BoxId::new(10),
607            result(ToolResultStatus::Ok, ResultView::OneLine("success".into())),
608        )
609        .unwrap();
610        assert_eq!(
611            value.hidden_contents(),
612            "{\"callId\":\"c2\",\"originatingCallBoxId\":9,\"tool\":\"Search\",\"toolVersion\":\"1.0.0\",\"status\":\"ok\",\"data\":{\"matches\":[1,2]}}"
613        );
614    }
615
616    #[test]
617    fn compatible_minor_accepts_options_and_unsupported_major_is_opaque() {
618        let good = ChatBox::tool_call(BoxId::new(1), call()).unwrap();
619        let future = ChatBox::new(
620            BoxId::new(1),
621            TOOL_CALL_TYPE.into(),
622            good.contents().into(),
623            "k1.tool-call/1.1.0".into(),
624            format!("{{\"optional\":true,{}", &good.hidden_contents()[1..]),
625        );
626        let unsupported = ChatBox::new(
627            BoxId::new(1),
628            TOOL_CALL_TYPE.into(),
629            "anything".into(),
630            "k1.tool-call/2.0.0".into(),
631            "not json".into(),
632        );
633        assert_eq!(future.tool_call_metadata().unwrap(), Some(call()));
634        assert_eq!(unsupported.tool_call_metadata().unwrap(), None);
635    }
636
637    #[test]
638    fn malformed_owned_envelopes_and_mismatched_visible_text_fail() {
639        let malformed_version = ChatBox::new(
640            BoxId::new(1),
641            TOOL_CALL_TYPE.into(),
642            "anything".into(),
643            "k1.tool-call/not-semver".into(),
644            "opaque".into(),
645        );
646        assert_eq!(
647            malformed_version.tool_call_metadata(),
648            Err(EnvelopeError::MalformedEnvelope)
649        );
650
651        let good = ChatBox::tool_call(BoxId::new(1), call()).unwrap();
652        let mismatch = ChatBox::new(
653            good.id(),
654            good.box_type().into(),
655            "different".into(),
656            good.hidden_type().into(),
657            good.hidden_contents().into(),
658        );
659        assert_eq!(
660            mismatch.tool_call_metadata(),
661            Err(EnvelopeError::InvalidVisibleContent)
662        );
663    }
664
665    #[test]
666    fn unknown_boxes_and_hidden_types_stay_opaque() {
667        let unknown = ChatBox::new(
668            BoxId::new(7),
669            "Future Box".into(),
670            "visible".into(),
671            "future.hidden/not-semver".into(),
672            "opaque".into(),
673        );
674        assert_eq!(unknown.tool_call_metadata().unwrap(), None);
675        assert_eq!(unknown.tool_result_metadata().unwrap(), None);
676
677        let unrelated = ChatBox::new(
678            BoxId::new(8),
679            TOOL_CALL_TYPE.into(),
680            "visible".into(),
681            "other.tool/1.0.0".into(),
682            "opaque".into(),
683        );
684        assert_eq!(unrelated.tool_call_metadata().unwrap(), None);
685    }
686
687    #[test]
688    fn constructors_reject_invalid_structures() {
689        assert_eq!(
690            ToolCall::new(
691                ToolCallId::new(1).unwrap(),
692                "Tool".into(),
693                "1.0.0".into(),
694                json!({"tool":"duplicate"}),
695            ),
696            Err(EnvelopeError::InvalidArguments)
697        );
698        assert_eq!(
699            ToolResult::new(
700                ToolCallId::new(1).unwrap(),
701                BoxId::new(0),
702                "Tool".into(),
703                "1.0.0".into(),
704                ToolResultStatus::Ok,
705                Value::Null,
706                ResultView::OneLine("success".into()),
707            ),
708            Err(EnvelopeError::InvalidOriginatingCallBoxId)
709        );
710    }
711}