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