1#![forbid(unsafe_code)]
2
3use serde_json::Value;
4
5pub use kcode_k1_chat_box::{
6 BoxId, ChatBox, EnvelopeError, ResultView, TOOL_CALL_HIDDEN_TYPE, TOOL_CALL_TYPE,
7 TOOL_RESULT_HIDDEN_TYPE, TOOL_RESULT_TYPE, ToolCall, ToolCallId, ToolResult, ToolResultStatus,
8};
9
10pub const SYSTEM_MESSAGE_TYPE: &str = "System Message";
11pub const USER_MESSAGE_TYPE: &str = "User Message";
12pub const AGENT_MESSAGE_TYPE: &str = "Agent Message";
13pub const USER_ATTACHMENT_TYPE: &str = "User Attachment";
14pub const AGENT_ATTACHMENT_TYPE: &str = "Agent Attachment";
15pub const TOOL_MESSAGE_TYPE: &str = "Tool Message";
16pub const TOOL_ATTACHMENT_TYPE: &str = "Tool Attachment";
17pub const ATTACHMENT_TYPE: &str = USER_ATTACHMENT_TYPE;
18
19#[derive(Clone, Debug, PartialEq)]
20pub struct ProviderCall {
21 pub tool: String,
22 pub tool_version: String,
23 pub arguments: Value,
24}
25
26#[derive(Clone, Debug, PartialEq)]
27pub struct DispatchedToolCall {
28 pub call: ToolCall,
29 pub call_box_id: BoxId,
30}
31
32#[derive(Clone, Copy, Debug, Eq, PartialEq)]
33pub enum TransitionError {
34 InvalidPhase,
35 BoxIdOverflow,
36 CallIdOverflow,
37 InvalidToolEnvelope,
38 DuplicateToolCall,
39 UnknownToolCall,
40 DuplicateReturn,
41 WrongOriginatingCall,
42 MismatchedTool,
43 MalformedToolConvention,
44}
45
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum RecoveryError {
48 NonContiguousBoxId,
49 InvalidCallSequence,
50 UnknownToolCall,
51 DuplicateReturn,
52 WrongOriginatingCall,
53 MismatchedTool,
54 MalformedToolConvention,
55}
56
57pub struct Chatend {
58 boxes: Vec<ChatBox>,
59 round_active: bool,
60 active_arrivals: Vec<ChatBox>,
61 next_call: Option<u64>,
62}
63
64impl Chatend {
65 pub const fn new() -> Self {
66 Self {
67 boxes: Vec::new(),
68 round_active: false,
69 active_arrivals: Vec::new(),
70 next_call: Some(1),
71 }
72 }
73
74 pub fn boxes(&self) -> &[ChatBox] {
75 &self.boxes
76 }
77
78 pub const fn round_active(&self) -> bool {
79 self.round_active
80 }
81
82 pub fn accept_box(
83 &mut self,
84 box_type: String,
85 contents: String,
86 hidden_type: String,
87 hidden_contents: String,
88 ) -> Result<Option<BoxId>, TransitionError> {
89 self.accept_arrival(ChatBox::new(
90 BoxId::new(0),
91 box_type,
92 contents,
93 hidden_type,
94 hidden_contents,
95 ))
96 }
97
98 pub fn accept_system(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
99 self.accept_box(
100 SYSTEM_MESSAGE_TYPE.into(),
101 contents,
102 String::new(),
103 String::new(),
104 )
105 }
106
107 pub fn accept_user(&mut self, contents: String) -> Result<Option<BoxId>, TransitionError> {
108 self.accept_box(
109 USER_MESSAGE_TYPE.into(),
110 contents,
111 String::new(),
112 String::new(),
113 )
114 }
115
116 pub fn accept_attachment(
117 &mut self,
118 contents: String,
119 hidden_type: String,
120 hidden_contents: String,
121 ) -> Result<Option<BoxId>, TransitionError> {
122 self.accept_box(
123 USER_ATTACHMENT_TYPE.into(),
124 contents,
125 hidden_type,
126 hidden_contents,
127 )
128 }
129
130 pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
131 if self.round_active {
132 return Err(TransitionError::InvalidPhase);
133 }
134 self.round_active = true;
135 Ok(self.boxes.last().map(ChatBox::id))
136 }
137
138 pub fn append_stage(
139 &mut self,
140 agent_contents: String,
141 calls: Vec<ProviderCall>,
142 ) -> Result<Vec<DispatchedToolCall>, TransitionError> {
143 if !self.round_active {
144 return Err(TransitionError::InvalidPhase);
145 }
146
147 let (calls, next_call) = self.assign_calls(calls)?;
148 let includes_agent = !agent_contents.is_empty();
149 self.ensure_capacity(calls.len() + usize::from(includes_agent))?;
150
151 let mut id = self.last_id();
152 let mut additions = Vec::with_capacity(calls.len() + usize::from(includes_agent));
153 if includes_agent {
154 id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
155 additions.push(ChatBox::new(
156 BoxId::new(id),
157 AGENT_MESSAGE_TYPE.into(),
158 agent_contents,
159 String::new(),
160 String::new(),
161 ));
162 }
163
164 let mut dispatched = Vec::with_capacity(calls.len());
165 for call in calls {
166 id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
167 let call_box_id = BoxId::new(id);
168 let box_value = ChatBox::tool_call(call_box_id, call.clone())
169 .map_err(|_| TransitionError::InvalidToolEnvelope)?;
170 additions.push(box_value);
171 dispatched.push(DispatchedToolCall { call, call_box_id });
172 }
173
174 self.boxes.extend(additions);
175 self.next_call = next_call;
176 Ok(dispatched)
177 }
178
179 pub fn accept_async_return(
180 &mut self,
181 result: ToolResult,
182 ) -> Result<Option<BoxId>, TransitionError> {
183 let (call, call_box_id) = self
184 .call_for(result.call_id())?
185 .ok_or(TransitionError::UnknownToolCall)?;
186 if self.has_return(result.call_id())? {
187 return Err(TransitionError::DuplicateReturn);
188 }
189 if result.originating_call_box_id() != call_box_id {
190 return Err(TransitionError::WrongOriginatingCall);
191 }
192 if result.tool() != call.tool() || result.tool_version() != call.tool_version() {
193 return Err(TransitionError::MismatchedTool);
194 }
195
196 let box_value = ChatBox::tool_result(BoxId::new(0), result)
197 .map_err(|_| TransitionError::InvalidToolEnvelope)?;
198 self.accept_arrival(box_value)
199 }
200
201 pub fn flush_active_arrivals(&mut self) -> Result<Vec<ChatBox>, TransitionError> {
202 if !self.round_active {
203 return Err(TransitionError::InvalidPhase);
204 }
205 let arrivals = self.active_arrivals.clone();
206 let appended = self.append_batch(&arrivals)?;
207 self.active_arrivals.clear();
208 Ok(appended)
209 }
210
211 pub fn done(&mut self, final_agent_contents: String) -> Result<Vec<ChatBox>, TransitionError> {
212 if !self.round_active {
213 return Err(TransitionError::InvalidPhase);
214 }
215
216 let mut arrivals = self.active_arrivals.clone();
217 if !final_agent_contents.is_empty() {
218 arrivals.insert(
219 0,
220 ChatBox::new(
221 BoxId::new(0),
222 AGENT_MESSAGE_TYPE.into(),
223 final_agent_contents,
224 String::new(),
225 String::new(),
226 ),
227 );
228 }
229 let appended = self.append_batch(&arrivals)?;
230 self.active_arrivals.clear();
231 self.round_active = false;
232 Ok(appended)
233 }
234
235 pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
236 let mut calls = Vec::<(ToolCall, BoxId)>::new();
237 let mut returned = Vec::<ToolCallId>::new();
238 let mut next_call = Some(1);
239
240 for (index, value) in boxes.iter().enumerate() {
241 let expected_box_id = u64::try_from(index)
242 .ok()
243 .and_then(|value| value.checked_add(1))
244 .ok_or(RecoveryError::NonContiguousBoxId)?;
245 if value.id() != BoxId::new(expected_box_id) {
246 return Err(RecoveryError::NonContiguousBoxId);
247 }
248
249 if let Some(call) = value
250 .tool_call_metadata()
251 .map_err(|_| RecoveryError::MalformedToolConvention)?
252 {
253 let expected_call = next_call.ok_or(RecoveryError::InvalidCallSequence)?;
254 if call.call_id().get() != expected_call {
255 return Err(RecoveryError::InvalidCallSequence);
256 }
257 next_call = expected_call.checked_add(1);
258 calls.push((call, value.id()));
259 }
260
261 if let Some(result) = value
262 .tool_result_metadata()
263 .map_err(|_| RecoveryError::MalformedToolConvention)?
264 {
265 let Some((call, call_box_id)) = calls
266 .iter()
267 .find(|(call, _)| call.call_id() == result.call_id())
268 else {
269 return Err(RecoveryError::UnknownToolCall);
270 };
271 if result.originating_call_box_id() != *call_box_id {
272 return Err(RecoveryError::WrongOriginatingCall);
273 }
274 if result.tool() != call.tool() || result.tool_version() != call.tool_version() {
275 return Err(RecoveryError::MismatchedTool);
276 }
277 if returned.contains(&result.call_id()) {
278 return Err(RecoveryError::DuplicateReturn);
279 }
280 returned.push(result.call_id());
281 }
282 }
283
284 Ok(Self {
285 boxes,
286 round_active: false,
287 active_arrivals: Vec::new(),
288 next_call,
289 })
290 }
291
292 fn assign_calls(
293 &self,
294 calls: Vec<ProviderCall>,
295 ) -> Result<(Vec<ToolCall>, Option<u64>), TransitionError> {
296 let mut next = self.next_call;
297 let mut assigned = Vec::with_capacity(calls.len());
298 for call in calls {
299 let id = next.ok_or(TransitionError::CallIdOverflow)?;
300 let call_id = ToolCallId::new(id).map_err(|_| TransitionError::CallIdOverflow)?;
301 let call = ToolCall::new(call_id, call.tool, call.tool_version, call.arguments)
302 .map_err(|_| TransitionError::InvalidToolEnvelope)?;
303 assigned.push(call);
304 next = id.checked_add(1);
305 }
306 Ok((assigned, next))
307 }
308
309 fn accept_arrival(&mut self, value: ChatBox) -> Result<Option<BoxId>, TransitionError> {
310 if self.round_active {
311 self.active_arrivals.push(value);
312 return Ok(None);
313 }
314 Ok(self.append_batch(&[value])?.pop().map(|value| value.id()))
315 }
316
317 fn append_batch(&mut self, additions: &[ChatBox]) -> Result<Vec<ChatBox>, TransitionError> {
318 self.ensure_capacity(additions.len())?;
319 let mut id = self.last_id();
320 let appended = additions
321 .iter()
322 .map(|value| {
323 id = id.checked_add(1).ok_or(TransitionError::BoxIdOverflow)?;
324 Ok(ChatBox::new(
325 BoxId::new(id),
326 value.box_type().into(),
327 value.contents().into(),
328 value.hidden_type().into(),
329 value.hidden_contents().into(),
330 ))
331 })
332 .collect::<Result<Vec<_>, _>>()?;
333 self.boxes.extend(appended.iter().cloned());
334 Ok(appended)
335 }
336
337 fn last_id(&self) -> u64 {
338 self.boxes.last().map_or(0, |value| value.id().get())
339 }
340
341 fn ensure_capacity(&self, count: usize) -> Result<(), TransitionError> {
342 let count = u64::try_from(count).map_err(|_| TransitionError::BoxIdOverflow)?;
343 self.last_id()
344 .checked_add(count)
345 .ok_or(TransitionError::BoxIdOverflow)
346 .map(|_| ())
347 }
348
349 fn call_for(&self, id: ToolCallId) -> Result<Option<(ToolCall, BoxId)>, TransitionError> {
350 let mut found = None;
351 for value in &self.boxes {
352 if let Some(call) = value
353 .tool_call_metadata()
354 .map_err(|_| TransitionError::MalformedToolConvention)?
355 && call.call_id() == id
356 {
357 if found.is_some() {
358 return Err(TransitionError::DuplicateToolCall);
359 }
360 found = Some((call, value.id()));
361 }
362 }
363 Ok(found)
364 }
365
366 fn has_return(&self, id: ToolCallId) -> Result<bool, TransitionError> {
367 for value in self.boxes.iter().chain(&self.active_arrivals) {
368 let result = value
369 .tool_result_metadata()
370 .map_err(|_| TransitionError::MalformedToolConvention)?;
371 if result.is_some_and(|result| result.call_id() == id) {
372 return Ok(true);
373 }
374 }
375 Ok(false)
376 }
377}
378
379impl Default for Chatend {
380 fn default() -> Self {
381 Self::new()
382 }
383}
384
385#[cfg(test)]
386mod tests;