1#![forbid(unsafe_code)]
2
3mod recovery;
4
5pub use kcode_k1_chat_codex_codec::{BoxValue, Call};
6pub use kcode_k1_chat_state::{
7 AGENT_ATTACHMENT_TYPE, AGENT_MESSAGE_TYPE, AGENT_RESPONSE_TYPE, ActorState, BoxId, ChatBox,
8 ProviderGenerated, SYSTEM_MESSAGE_TYPE, TOOL_ATTACHMENT_TYPE, TOOL_CALL_TYPE,
9 TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE, ToolCallId, USER_ATTACHMENT_TYPE, USER_MESSAGE_TYPE,
10};
11pub use kcode_k1_codex_adapter::{ShimItem, ShimOutput};
12
13use std::sync::Arc;
14use std::sync::atomic::AtomicU8;
15
16use kcode_k1_chat_codex_codec::{open_agent_response, project};
17use kcode_k1_chat_state::{ProviderCall, StateError};
18use recovery::recovered_sequence;
19
20pub const MALFORMED_NATIVE_CALL_METADATA_TYPE: &str = "k1.malformed-native-call.v1";
21
22#[derive(Clone, Debug)]
23pub struct Start {
24 pub job: u64,
25 pub values: Vec<BoxValue>,
26 pub attempt: Arc<AtomicU8>,
27}
28
29#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct PreparedCall {
31 pub tool_call_id: ToolCallId,
32 pub name: String,
33 pub arguments: String,
34 disposition: PreparedCallDisposition,
35}
36
37impl PreparedCall {
38 pub fn disposition(&self) -> &PreparedCallDisposition {
39 &self.disposition
40 }
41}
42
43#[derive(Clone, Debug, Eq, PartialEq)]
44pub enum PreparedCallDisposition {
45 External,
46 ImmediateError(Box<ImmediateToolError>),
47}
48
49#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct ImmediateToolError {
51 pub message: String,
52 pub native_tool: String,
53 pub attempted_ktool: Option<String>,
54 pub validation_code: String,
55 pub path: String,
56 pub expected: String,
57 pub received: String,
58 pub native_arguments: String,
59 pub metadata_type: String,
60 pub metadata_contents: String,
61}
62
63#[derive(Clone, Debug)]
64pub struct PreparedMailboxFlush(Arc<Prepared>);
65
66#[derive(Debug)]
67struct Prepared {
68 token: u64,
69 values: Vec<BoxValue>,
70 job: u64,
71 external_count: usize,
72}
73
74impl PreparedMailboxFlush {
75 pub fn values(&self) -> &[BoxValue] {
76 &self.0.values
77 }
78}
79
80#[derive(Clone, Debug, Eq, PartialEq)]
81pub enum Status {
82 Running,
83 Quiet,
84 Stalled { message: String, restartable: bool },
85}
86
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum RestartError {
89 NotStalled,
90 ProviderActionAccepted,
91}
92
93#[derive(Clone, Copy, Debug, Eq, PartialEq)]
94enum Phase {
95 ProviderActive,
96 ChatendBoundary,
97 PendingGeneration,
98}
99
100struct Round {
101 job: u64,
102 accepted_provider_action: bool,
103 phase: Phase,
104 mailbox_flush_needed: bool,
105 restartable: bool,
106}
107
108enum Mode {
109 Idle,
110 Running(Round),
111 Stalled(Status, bool),
112}
113
114pub struct ConversationState {
115 state: ActorState,
116 session: [u8; 12],
117 sequence: u64,
118 unsubmitted: Vec<BoxValue>,
119 queued_trigger: bool,
120 token: u64,
121 prepared: Option<Arc<Prepared>>,
122 mode: Mode,
123}
124
125impl ConversationState {
126 pub fn new(session: [u8; 12]) -> Self {
127 Self {
128 state: ActorState::new(false),
129 session,
130 sequence: 0,
131 unsubmitted: Vec::new(),
132 queued_trigger: false,
133 token: 0,
134 prepared: None,
135 mode: Mode::Idle,
136 }
137 }
138
139 pub fn recover(session: [u8; 12], boxes: Vec<ChatBox>, force: bool) -> Result<Self, String> {
140 let sequence = recovered_sequence(session, &boxes)?;
141 let state = ActorState::recover(boxes, force).map_err(debug)?;
142 Ok(Self {
143 unsubmitted: state.boxes().iter().map(project).collect(),
144 state,
145 sequence,
146 ..Self::new(session)
147 })
148 }
149
150 pub fn boxes(&self) -> &[ChatBox] {
151 self.state.boxes()
152 }
153
154 pub fn status(&self) -> Status {
155 match &self.mode {
156 Mode::Running(_) => Status::Running,
157 Mode::Idle if self.state.quiet() => Status::Quiet,
158 Mode::Idle => Status::Running,
159 Mode::Stalled(status, _) => status.clone(),
160 }
161 }
162
163 pub fn accept(
164 &mut self,
165 box_type: String,
166 contents: String,
167 hidden_type: String,
168 hidden_contents: String,
169 ) -> Result<(), String> {
170 self.accept_arrival(true, |state| {
171 state.accept_box(box_type, contents, hidden_type, hidden_contents)
172 })
173 }
174
175 pub fn accept_tool_message(
176 &mut self,
177 tool_call_id: ToolCallId,
178 message: String,
179 ) -> Result<(), String> {
180 self.accept_arrival(false, |state| {
181 state.accept_tool_message(tool_call_id, message)
182 })
183 }
184
185 pub fn accept_tool_return(
186 &mut self,
187 tool_call_id: ToolCallId,
188 result: Result<String, String>,
189 ) -> Result<(), String> {
190 self.accept_arrival(true, |state| {
191 state.accept_async_return(tool_call_id, result)
192 })
193 }
194
195 pub fn accept_tool_return_v2(
196 &mut self,
197 tool_call_id: ToolCallId,
198 result: Result<String, String>,
199 metadata_type: String,
200 metadata_contents: String,
201 ) -> Result<(), String> {
202 self.accept_arrival(true, |state| {
203 state.accept_async_return_v2(tool_call_id, result, metadata_type, metadata_contents)
204 })
205 }
206
207 pub fn begin(&mut self) -> Result<Option<Start>, String> {
208 if !matches!(self.mode, Mode::Idle) {
209 return Ok(None);
210 }
211 let Some(start) = self.state.begin_inference().map_err(debug)? else {
212 return Ok(None);
213 };
214 let promised_id = self.promised_id()?;
215 let mut values = std::mem::take(&mut self.unsubmitted);
216 values.push(open_agent_response(promised_id));
217 let mailbox_flush_needed = std::mem::take(&mut self.queued_trigger);
218 self.mode = Mode::Running(Round {
219 job: start.job,
220 accepted_provider_action: false,
221 phase: Phase::ProviderActive,
222 mailbox_flush_needed,
223 restartable: true,
224 });
225 Ok(Some(Start {
226 job: start.job,
227 values,
228 attempt: start.attempt,
229 }))
230 }
231
232 pub fn prepare_stage(
233 &mut self,
234 job: u64,
235 text: String,
236 values: Vec<BoxValue>,
237 ) -> Result<Vec<PreparedCall>, String> {
238 if !matches!(
239 &self.mode,
240 Mode::Running(round) if round.job == job && round.phase == Phase::ProviderActive
241 ) {
242 return Err("stale Codex inference stage".to_owned());
243 }
244 if self.prepared.is_some() {
245 return Err("a Codex mailbox flush remains uncommitted".to_owned());
246 }
247 let accepted_provider_action = !values.is_empty();
248 let mut sequence = self.sequence;
249 let mut generated = Vec::with_capacity(values.len());
250 let mut prepared = Vec::new();
251 for value in values {
252 match value {
253 BoxValue::AgentMessage(Ok(contents)) => {
254 generated.push(ProviderGenerated::AgentMessage { contents });
255 }
256 BoxValue::Call(Ok(call)) => {
257 sequence = next_sequence(sequence)?;
258 let tool_call_id = ToolCallId::new(self.session, sequence);
259 generated.push(ProviderGenerated::ToolCall(ProviderCall {
260 tool_call_id,
261 name: call.name.clone(),
262 arguments: call.arguments.clone(),
263 }));
264 prepared.push(PreparedCall {
265 tool_call_id,
266 name: call.name,
267 arguments: call.arguments,
268 disposition: PreparedCallDisposition::External,
269 });
270 }
271 BoxValue::MalformedNativeAction(action) => {
272 sequence = next_sequence(sequence)?;
273 let tool_call_id = ToolCallId::new(self.session, sequence);
274 let name = action
275 .attempted_ktool()
276 .unwrap_or_else(|| action.native_tool())
277 .to_owned();
278 let arguments = action.native_arguments_json();
279 let error = ImmediateToolError {
280 message: action.diagnostic(),
281 native_tool: action.native_tool().to_owned(),
282 attempted_ktool: action.attempted_ktool().map(str::to_owned),
283 validation_code: action.validation_code().to_owned(),
284 path: action.path().to_owned(),
285 expected: action.expected().to_owned(),
286 received: action.received().to_owned(),
287 native_arguments: arguments.clone(),
288 metadata_type: MALFORMED_NATIVE_CALL_METADATA_TYPE.to_owned(),
289 metadata_contents: action.diagnostic_json(),
290 };
291 generated.push(ProviderGenerated::ToolCall(ProviderCall {
292 tool_call_id,
293 name: name.clone(),
294 arguments: arguments.clone(),
295 }));
296 prepared.push(PreparedCall {
297 tool_call_id,
298 name,
299 arguments,
300 disposition: PreparedCallDisposition::ImmediateError(Box::new(error)),
301 });
302 }
303 _ => return Err("stage contains a malformed provider action".to_owned()),
304 }
305 }
306 let before = self.state.boxes().len();
307 self.state
308 .append_stage(job, text, generated)
309 .map_err(debug)?;
310 self.unsubmitted.extend(
311 self.state.boxes()[before..]
312 .iter()
313 .filter(|box_| box_.box_type() == TOOL_CALL_TYPE)
314 .map(project),
315 );
316 self.sequence = sequence;
317 if let Mode::Running(round) = &mut self.mode {
318 round.accepted_provider_action |= accepted_provider_action;
319 round.phase = Phase::ChatendBoundary;
320 round.mailbox_flush_needed = true;
321 round.restartable = false;
322 }
323 Ok(prepared)
324 }
325
326 pub fn mailbox_flush(&mut self, job: u64) -> Result<Vec<ChatBox>, String> {
327 if !matches!(
328 &self.mode,
329 Mode::Running(round) if round.job == job && round.phase == Phase::ChatendBoundary
330 ) {
331 return Err("stale Codex active-arrival mailbox flush".to_owned());
332 }
333 let boxes = self.state.flush_active_arrivals(job).map_err(debug)?;
334 self.unsubmitted.extend(boxes.iter().map(project));
335 if let Mode::Running(round) = &mut self.mode {
336 round.phase = Phase::PendingGeneration;
337 }
338 Ok(boxes)
339 }
340
341 pub fn prepare_mailbox_flush(
342 &mut self,
343 job: u64,
344 ) -> Result<Option<PreparedMailboxFlush>, String> {
345 let (mailbox_flush_needed, phase) = match &self.mode {
346 Mode::Running(round) if round.job == job => (round.mailbox_flush_needed, round.phase),
347 _ => return Err("stale Codex inference mailbox flush".to_owned()),
348 };
349 if let Some(prepared) = &self.prepared {
350 return Ok(Some(PreparedMailboxFlush(Arc::clone(prepared))));
351 }
352 if !mailbox_flush_needed {
353 return Ok(None);
354 }
355 if phase == Phase::ProviderActive {
356 return Ok(None);
357 }
358 if phase == Phase::ChatendBoundary {
359 self.mailbox_flush(job)?;
360 }
361 let token = self
362 .token
363 .checked_add(1)
364 .ok_or_else(|| "Codex mailbox-flush token space was exhausted".to_owned())?;
365 let external_count = self.unsubmitted.len();
366 let mut values = self.unsubmitted.clone();
367 values.push(open_agent_response(self.promised_id()?));
368 let prepared = Arc::new(Prepared {
369 token,
370 values,
371 job,
372 external_count,
373 });
374 self.token = token;
375 self.prepared = Some(Arc::clone(&prepared));
376 if let Mode::Running(round) = &mut self.mode {
377 round.mailbox_flush_needed = false;
378 }
379 Ok(Some(PreparedMailboxFlush(prepared)))
380 }
381
382 pub fn validate_mailbox_flush(&self, prepared: &PreparedMailboxFlush) -> Result<(), String> {
383 let prepared = &prepared.0;
384 let prefix_matches = self.unsubmitted.get(..prepared.external_count)
385 == Some(&prepared.values[..prepared.external_count]);
386 let valid = matches!(
387 &self.mode,
388 Mode::Running(round)
389 if round.job == prepared.job && round.phase == Phase::PendingGeneration
390 ) && self.token == prepared.token
391 && prefix_matches
392 && self
393 .prepared
394 .as_ref()
395 .is_some_and(|current| Arc::ptr_eq(current, prepared));
396 if valid {
397 Ok(())
398 } else {
399 Err("stale or invalid Codex mailbox flush".to_owned())
400 }
401 }
402
403 pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
404 self.validate_mailbox_flush(&prepared)?;
405 self.unsubmitted.drain(..prepared.0.external_count);
406 self.prepared = None;
407 if let Mode::Running(round) = &mut self.mode {
408 round.phase = Phase::ProviderActive;
409 }
410 Ok(())
411 }
412
413 pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<(), String> {
414 let mut round = self.take_round(job)?;
415 if round.phase != Phase::ProviderActive {
416 let message = "Codex inference completed outside a provider generation".to_owned();
417 self.preserve(round, message.clone(), false);
418 return Err(message);
419 }
420 if self.prepared.is_some() {
421 let message = "Codex inference completed with an uncommitted mailbox flush".to_owned();
422 self.preserve(round, message.clone(), false);
423 return Err(message);
424 }
425 let mut text = String::new();
426 for item in output.items {
427 match item {
428 ShimItem::Text(value) => text.push_str(&value),
429 ShimItem::Box(_) => {
430 round.accepted_provider_action = true;
431 let message = "terminal Codex output contains a box".to_owned();
432 self.preserve(round, message.clone(), false);
433 return Err(message);
434 }
435 }
436 }
437 let before = self.state.boxes().len();
438 if let Err(error) = self.state.complete_inference(job, text).map_err(debug) {
439 self.preserve(round, error.clone(), false);
440 return Err(error);
441 }
442 self.unsubmitted
443 .extend(self.state.boxes()[before..].iter().skip(1).map(project));
444 self.queued_trigger = false;
445 self.mode = Mode::Idle;
446 Ok(())
447 }
448
449 pub fn fail(&mut self, job: u64, message: String, restartable_before_launch: bool) {
450 if let Ok(round) = self.take_round(job) {
451 self.preserve(round, message, restartable_before_launch);
452 }
453 }
454
455 pub fn restart(&mut self) -> Result<(), RestartError> {
456 match &self.mode {
457 Mode::Stalled(_, true) => return Err(RestartError::ProviderActionAccepted),
458 Mode::Stalled(
459 Status::Stalled {
460 restartable: true, ..
461 },
462 false,
463 ) => {}
464 _ => return Err(RestartError::NotStalled),
465 }
466 self.state.restart().map_err(|_| RestartError::NotStalled)?;
467 self.unsubmitted = self.state.boxes().iter().map(project).collect();
468 self.queued_trigger = false;
469 self.prepared = None;
470 self.mode = Mode::Idle;
471 Ok(())
472 }
473
474 fn accept_arrival<F>(&mut self, triggering: bool, accept: F) -> Result<(), String>
475 where
476 F: FnOnce(&mut ActorState) -> Result<(), StateError>,
477 {
478 let before = self.state.boxes().len();
479 accept(&mut self.state).map_err(debug)?;
480 let appended = &self.state.boxes()[before..];
481 self.unsubmitted.extend(appended.iter().map(project));
482 match &mut self.mode {
483 Mode::Running(round) => round.mailbox_flush_needed |= triggering,
484 Mode::Idle if appended.is_empty() => self.queued_trigger |= triggering,
485 Mode::Idle => self.queued_trigger = false,
486 Mode::Stalled(_, _) => {}
487 }
488 Ok(())
489 }
490
491 fn promised_id(&self) -> Result<BoxId, String> {
492 let previous = self.state.boxes().last().map_or(0, |box_| box_.id().get());
493 let value = previous
494 .checked_add(1)
495 .ok_or_else(|| "BoxId space was exhausted".to_owned())?;
496 Ok(BoxId::new(value))
497 }
498
499 fn take_round(&mut self, job: u64) -> Result<Round, String> {
500 match std::mem::replace(&mut self.mode, Mode::Idle) {
501 Mode::Running(round) if round.job == job => Ok(round),
502 other => {
503 self.mode = other;
504 Err("stale Codex inference completion".to_owned())
505 }
506 }
507 }
508
509 fn preserve(&mut self, round: Round, message: String, restartable_before_launch: bool) {
510 self.prepared = None;
511 let restartable = restartable_before_launch
512 && round.restartable
513 && !round.accepted_provider_action
514 && self
515 .state
516 .stall_inference(round.job, message.clone())
517 .is_ok();
518 if !restartable {
519 let _ = self.state.halt(message.clone());
520 }
521 self.mode = Mode::Stalled(
522 Status::Stalled {
523 message,
524 restartable,
525 },
526 round.accepted_provider_action,
527 );
528 }
529}
530
531fn next_sequence(sequence: u64) -> Result<u64, String> {
532 sequence
533 .checked_add(1)
534 .ok_or_else(|| "ToolCallId space was exhausted".to_owned())
535}
536
537fn debug(error: impl std::fmt::Debug) -> String {
538 format!("{error:?}")
539}
540
541#[cfg(test)]
542mod tests {
543 use super::*;
544 use kcode_k1_chat_codex_codec::Codec;
545 use kcode_k1_codex_adapter::{BoxCodec, ToolCall};
546
547 fn active() -> (ConversationState, u64) {
548 let mut state = ConversationState::new([7; 12]);
549 state
550 .accept(
551 USER_MESSAGE_TYPE.into(),
552 "start".into(),
553 String::new(),
554 String::new(),
555 )
556 .unwrap();
557 let job = state.begin().unwrap().unwrap().job;
558 (state, job)
559 }
560
561 fn native(name: &str, arguments: &str) -> BoxValue {
562 let mut codec = Codec;
563 codec.tool_call_box(&ToolCall {
564 call_id: "native".into(),
565 name: name.into(),
566 arguments: arguments.parse().unwrap(),
567 })
568 }
569
570 fn active_after_mailbox_flush() -> (ConversationState, u64, ToolCallId) {
571 let (mut state, job) = active();
572 let calls = state
573 .prepare_stage(
574 job,
575 String::new(),
576 vec![BoxValue::Call(Ok(Call {
577 name: "tool".into(),
578 arguments: "{}".into(),
579 }))],
580 )
581 .unwrap();
582 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
583 state.commit_mailbox_flush(prepared).unwrap();
584 (state, job, calls[0].tool_call_id)
585 }
586
587 fn empty_output() -> ShimOutput<BoxValue> {
588 ShimOutput { items: Vec::new() }
589 }
590
591 #[test]
592 fn malformed_only_is_persisted_and_prepared_as_deterministic_immediate_error() {
593 let (mut state, job) = active();
594 let malformed = native(
595 "call_ktool",
596 r#"{"z":0,"name":"recoverable","arguments":{"b":2,"a":1}}"#,
597 );
598 let (diagnostic, diagnostic_json) = match &malformed {
599 BoxValue::MalformedNativeAction(action) => {
600 (action.diagnostic(), action.diagnostic_json())
601 }
602 _ => panic!("test value must be malformed"),
603 };
604 let calls = state
605 .prepare_stage(job, String::new(), vec![malformed])
606 .unwrap();
607 assert_eq!(state.status(), Status::Running);
608 assert_eq!(calls.len(), 1);
609 assert_eq!(calls[0].name, "recoverable");
610 assert_eq!(
611 calls[0].arguments,
612 r#"{"arguments":{"a":1,"b":2},"name":"recoverable","z":0}"#
613 );
614 let PreparedCallDisposition::ImmediateError(error) = calls[0].disposition() else {
615 panic!("malformed call must be local");
616 };
617 assert_eq!(error.message, diagnostic);
618 assert_eq!(error.native_tool, "call_ktool");
619 assert_eq!(error.attempted_ktool.as_deref(), Some("recoverable"));
620 assert_eq!(error.validation_code, "invalid_wrapper_fields");
621 assert_eq!(error.path, "$");
622 assert_eq!(error.expected, "object with exactly name and arguments");
623 assert_eq!(error.received, calls[0].arguments);
624 assert_eq!(error.native_arguments, calls[0].arguments);
625 assert_eq!(error.metadata_type, MALFORMED_NATIVE_CALL_METADATA_TYPE);
626 assert_eq!(error.metadata_contents, diagnostic_json);
627 let persisted = state
628 .boxes()
629 .last()
630 .unwrap()
631 .tool_call_metadata()
632 .unwrap()
633 .unwrap();
634 assert_eq!(persisted.name, calls[0].name);
635 assert_eq!(persisted.arguments, calls[0].arguments);
636 }
637
638 #[test]
639 fn mixed_stage_preserves_order_contiguous_ids_name_selection_and_external_calls() {
640 let (mut state, job) = active();
641 let before = state.boxes().len();
642 let values = vec![
643 BoxValue::AgentMessage(Ok("first".into())),
644 native(
645 "call_ktool",
646 r#"{"name":"attempted","arguments":{},"extra":1}"#,
647 ),
648 BoxValue::Call(Ok(Call {
649 name: "valid".into(),
650 arguments: "{\"ok\":true}".into(),
651 })),
652 native("future_native", r#"{"z":0,"a":true}"#),
653 ];
654 let calls = state.prepare_stage(job, String::new(), values).unwrap();
655 assert_eq!(
656 calls
657 .iter()
658 .map(|call| call.tool_call_id.sequence())
659 .collect::<Vec<_>>(),
660 vec![1, 2, 3]
661 );
662 assert_eq!(
663 calls
664 .iter()
665 .map(|call| call.name.as_str())
666 .collect::<Vec<_>>(),
667 vec!["attempted", "valid", "future_native"]
668 );
669 assert!(matches!(
670 calls[0].disposition(),
671 PreparedCallDisposition::ImmediateError(_)
672 ));
673 assert_eq!(calls[1].disposition(), &PreparedCallDisposition::External);
674 assert!(matches!(
675 calls[2].disposition(),
676 PreparedCallDisposition::ImmediateError(_)
677 ));
678 assert_eq!(calls[2].arguments, r#"{"a":true,"z":0}"#);
679 assert_eq!(
680 state.boxes()[before..]
681 .iter()
682 .map(ChatBox::box_type)
683 .collect::<Vec<_>>(),
684 vec![
685 AGENT_RESPONSE_TYPE,
686 AGENT_MESSAGE_TYPE,
687 TOOL_CALL_TYPE,
688 TOOL_CALL_TYPE,
689 TOOL_CALL_TYPE,
690 ]
691 );
692 }
693
694 #[test]
695 fn normal_call_precedes_correlated_result_and_open_response_in_stable_flush() {
696 let (mut state, job) = active();
697 let calls = state
698 .prepare_stage(
699 job,
700 String::new(),
701 vec![BoxValue::Call(Ok(Call {
702 name: "tool".into(),
703 arguments: "{\"input\":1}".into(),
704 }))],
705 )
706 .unwrap();
707 state
708 .accept_tool_return(calls[0].tool_call_id, Ok("done".into()))
709 .unwrap();
710
711 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
712 let expected = vec![
713 project(
714 state
715 .boxes()
716 .iter()
717 .find(|box_| box_.box_type() == TOOL_CALL_TYPE)
718 .unwrap(),
719 ),
720 project(
721 state
722 .boxes()
723 .iter()
724 .find(|box_| box_.box_type() == TOOL_RESULT_TYPE)
725 .unwrap(),
726 ),
727 open_agent_response(state.promised_id().unwrap()),
728 ];
729 assert_eq!(prepared.values(), expected);
730
731 let repeated = state.prepare_mailbox_flush(job).unwrap().unwrap();
732 assert_eq!(repeated.values(), prepared.values());
733 state.validate_mailbox_flush(&prepared).unwrap();
734 state.commit_mailbox_flush(repeated).unwrap();
735 assert!(state.unsubmitted.is_empty());
736 }
737
738 #[test]
739 fn malformed_native_call_is_queued_as_its_canonical_tool_call() {
740 let (mut state, job) = active();
741 let malformed = native(
742 "call_ktool",
743 r#"{"name":"recoverable","arguments":{},"extra":1}"#,
744 );
745 state
746 .prepare_stage(job, String::new(), vec![malformed])
747 .unwrap();
748
749 let expected_call = project(
750 state
751 .boxes()
752 .iter()
753 .find(|box_| box_.box_type() == TOOL_CALL_TYPE)
754 .unwrap(),
755 );
756 let expected_response = open_agent_response(state.promised_id().unwrap());
757 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
758 assert_eq!(prepared.values(), &[expected_call, expected_response]);
759 }
760
761 #[test]
762 fn provider_agent_message_and_stage_response_are_not_requeued() {
763 let (mut state, job) = active();
764 state
765 .prepare_stage(
766 job,
767 "stage response".into(),
768 vec![BoxValue::AgentMessage(Ok("provider message".into()))],
769 )
770 .unwrap();
771
772 let expected_response = open_agent_response(state.promised_id().unwrap());
773 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
774 assert_eq!(prepared.values(), &[expected_response]);
775 }
776
777 #[test]
778 fn committed_mailbox_flush_does_not_schedule_a_fresh_turn() {
779 let (mut state, job, _) = active_after_mailbox_flush();
780 state.complete(job, empty_output()).unwrap();
781 assert_eq!(state.status(), Status::Quiet);
782 assert!(state.begin().unwrap().is_none());
783 }
784
785 #[test]
786 fn tool_messages_after_a_mailbox_flush_remain_inert() {
787 let (mut state, job, tool_call_id) = active_after_mailbox_flush();
788 state
789 .accept_tool_message(tool_call_id, "still running".into())
790 .unwrap();
791 state.complete(job, empty_output()).unwrap();
792 assert_eq!(state.status(), Status::Quiet);
793 assert!(state.begin().unwrap().is_none());
794 }
795
796 #[test]
797 fn one_tool_result_schedules_exactly_one_fresh_turn() {
798 let (mut state, job, tool_call_id) = active_after_mailbox_flush();
799 state
800 .accept_tool_return(tool_call_id, Ok("done".into()))
801 .unwrap();
802 state.complete(job, empty_output()).unwrap();
803 let followup = state.begin().unwrap().unwrap();
804 assert!(state.begin().unwrap().is_none());
805 state.complete(followup.job, empty_output()).unwrap();
806 assert_eq!(state.status(), Status::Quiet);
807 assert!(state.begin().unwrap().is_none());
808 }
809}