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::{PreflightGenerated, ProviderCall, StateError};
18use recovery::recovered_sequence;
19
20pub const MALFORMED_NATIVE_CALL_METADATA_TYPE: &str = "k1.malformed-native-call.v1";
21
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum PreflightMode {
24 Blocking,
25 NonBlocking,
26}
27
28#[derive(Clone, Debug, Eq, PartialEq)]
29pub enum PreflightItem {
30 SystemMessage {
31 contents: String,
32 },
33 KtoolCall {
34 name: String,
35 arguments: String,
36 mode: PreflightMode,
37 },
38}
39
40#[derive(Clone, Debug, Eq, PartialEq)]
41pub struct PreparedPreflightCall {
42 pub tool_call_id: ToolCallId,
43 pub call_box_id: BoxId,
44 pub name: String,
45 pub arguments: String,
46 pub mode: PreflightMode,
47}
48
49#[derive(Clone, Debug)]
50pub struct Start {
51 pub job: u64,
52 pub values: Vec<BoxValue>,
53 pub attempt: Arc<AtomicU8>,
54}
55
56#[derive(Clone, Debug, Eq, PartialEq)]
57pub struct PreparedCall {
58 pub tool_call_id: ToolCallId,
59 pub name: String,
60 pub arguments: String,
61 disposition: PreparedCallDisposition,
62}
63
64impl PreparedCall {
65 pub fn disposition(&self) -> &PreparedCallDisposition {
66 &self.disposition
67 }
68}
69
70#[derive(Clone, Debug, Eq, PartialEq)]
71pub enum PreparedCallDisposition {
72 External,
73 ImmediateError(Box<ImmediateToolError>),
74}
75
76#[derive(Clone, Debug, Eq, PartialEq)]
77pub struct ImmediateToolError {
78 pub message: String,
79 pub native_tool: String,
80 pub attempted_ktool: Option<String>,
81 pub validation_code: String,
82 pub path: String,
83 pub expected: String,
84 pub received: String,
85 pub native_arguments: String,
86 pub metadata_type: String,
87 pub metadata_contents: String,
88}
89
90#[derive(Clone, Debug)]
91pub struct PreparedMailboxFlush(Arc<Prepared>);
92
93#[derive(Debug)]
94struct Prepared {
95 token: u64,
96 values: Vec<BoxValue>,
97 job: u64,
98 external_count: usize,
99}
100
101impl PreparedMailboxFlush {
102 pub fn values(&self) -> &[BoxValue] {
103 &self.0.values
104 }
105}
106
107#[derive(Clone, Debug, Eq, PartialEq)]
108pub enum Status {
109 Running,
110 Quiet,
111 Stalled { message: String, restartable: bool },
112}
113
114#[derive(Clone, Copy, Debug, Eq, PartialEq)]
115pub enum RestartError {
116 NotStalled,
117 ProviderActionAccepted,
118}
119
120#[derive(Clone, Copy, Debug, Eq, PartialEq)]
121enum Phase {
122 ProviderActive,
123 ChatendBoundary,
124 PendingGeneration,
125}
126
127struct Round {
128 job: u64,
129 accepted_provider_action: bool,
130 phase: Phase,
131 mailbox_flush_needed: bool,
132 restartable: bool,
133}
134
135enum Mode {
136 Idle,
137 Running(Round),
138 Stalled(Status, bool),
139}
140
141pub struct ConversationState {
142 state: ActorState,
143 session: [u8; 12],
144 sequence: u64,
145 unsubmitted: Vec<BoxValue>,
146 queued_trigger: bool,
147 token: u64,
148 prepared: Option<Arc<Prepared>>,
149 mode: Mode,
150}
151
152impl ConversationState {
153 pub fn new(session: [u8; 12]) -> Self {
154 Self {
155 state: ActorState::new(false),
156 session,
157 sequence: 0,
158 unsubmitted: Vec::new(),
159 queued_trigger: false,
160 token: 0,
161 prepared: None,
162 mode: Mode::Idle,
163 }
164 }
165
166 pub fn recover(session: [u8; 12], boxes: Vec<ChatBox>, force: bool) -> Result<Self, String> {
167 let sequence = recovered_sequence(session, &boxes)?;
168 let state = ActorState::recover(boxes, force).map_err(debug)?;
169 Ok(Self {
170 unsubmitted: state.boxes().iter().map(project).collect(),
171 state,
172 sequence,
173 ..Self::new(session)
174 })
175 }
176
177 pub fn boxes(&self) -> &[ChatBox] {
178 self.state.boxes()
179 }
180
181 pub fn status(&self) -> Status {
182 match &self.mode {
183 Mode::Running(_) => Status::Running,
184 Mode::Idle if self.state.quiet() => Status::Quiet,
185 Mode::Idle => Status::Running,
186 Mode::Stalled(status, _) => status.clone(),
187 }
188 }
189
190 pub fn prepare_preflight(
191 &mut self,
192 items: Vec<PreflightItem>,
193 ) -> Result<Vec<PreparedPreflightCall>, String> {
194 if items.is_empty()
195 || !matches!(self.mode, Mode::Idle)
196 || !self.state.boxes().is_empty()
197 || self.prepared.is_some()
198 {
199 return Err("preflight requires one fresh nonempty session".to_owned());
200 }
201 let mut sequence = self.sequence;
202 let mut generated = Vec::with_capacity(items.len());
203 let mut calls = Vec::new();
204 for item in items {
205 match item {
206 PreflightItem::SystemMessage { contents } => {
207 generated.push(PreflightGenerated::SystemMessage { contents });
208 }
209 PreflightItem::KtoolCall {
210 name,
211 arguments,
212 mode,
213 } => {
214 sequence = next_sequence(sequence)?;
215 let tool_call_id = ToolCallId::new(self.session, sequence);
216 generated.push(PreflightGenerated::ToolCall(ProviderCall {
217 tool_call_id,
218 name: name.clone(),
219 arguments: arguments.clone(),
220 }));
221 calls.push((tool_call_id, name, arguments, mode));
222 }
223 }
224 }
225 let before = self.state.boxes().len();
226 let dispatched = self.state.append_preflight(generated).map_err(debug)?;
227 if dispatched.len() != calls.len() {
228 return Err("preflight call correlation diverged".to_owned());
229 }
230 self.unsubmitted
231 .extend(self.state.boxes()[before..].iter().map(project));
232 self.sequence = sequence;
233 Ok(dispatched
234 .into_iter()
235 .zip(calls)
236 .map(|(dispatched, (tool_call_id, name, arguments, mode))| {
237 debug_assert_eq!(dispatched.tool_call_id, tool_call_id);
238 PreparedPreflightCall {
239 tool_call_id,
240 call_box_id: dispatched.call_box_id,
241 name,
242 arguments,
243 mode,
244 }
245 })
246 .collect())
247 }
248
249 pub fn accept(
250 &mut self,
251 box_type: String,
252 contents: String,
253 hidden_type: String,
254 hidden_contents: String,
255 ) -> Result<(), String> {
256 self.accept_arrival(true, |state| {
257 state.accept_box(box_type, contents, hidden_type, hidden_contents)
258 })
259 }
260
261 pub fn accept_tool_message(
262 &mut self,
263 tool_call_id: ToolCallId,
264 message: String,
265 ) -> Result<(), String> {
266 self.accept_arrival(false, |state| {
267 state.accept_tool_message(tool_call_id, message)
268 })
269 }
270
271 pub fn accept_tool_return(
272 &mut self,
273 tool_call_id: ToolCallId,
274 result: Result<String, String>,
275 ) -> Result<(), String> {
276 self.accept_arrival(true, |state| {
277 state.accept_async_return(tool_call_id, result)
278 })
279 }
280
281 pub fn accept_tool_return_v2(
282 &mut self,
283 tool_call_id: ToolCallId,
284 result: Result<String, String>,
285 metadata_type: String,
286 metadata_contents: String,
287 ) -> Result<(), String> {
288 self.accept_arrival(true, |state| {
289 state.accept_async_return_v2(tool_call_id, result, metadata_type, metadata_contents)
290 })
291 }
292
293 pub fn begin(&mut self) -> Result<Option<Start>, String> {
294 if !matches!(self.mode, Mode::Idle) {
295 return Ok(None);
296 }
297 let Some(start) = self.state.begin_inference().map_err(debug)? else {
298 return Ok(None);
299 };
300 let promised_id = self.promised_id()?;
301 let mut values = std::mem::take(&mut self.unsubmitted);
302 values.push(open_agent_response(promised_id));
303 let mailbox_flush_needed = std::mem::take(&mut self.queued_trigger);
304 self.mode = Mode::Running(Round {
305 job: start.job,
306 accepted_provider_action: false,
307 phase: Phase::ProviderActive,
308 mailbox_flush_needed,
309 restartable: true,
310 });
311 Ok(Some(Start {
312 job: start.job,
313 values,
314 attempt: start.attempt,
315 }))
316 }
317
318 pub fn prepare_stage(
319 &mut self,
320 job: u64,
321 text: String,
322 values: Vec<BoxValue>,
323 ) -> Result<Vec<PreparedCall>, String> {
324 if !matches!(
325 &self.mode,
326 Mode::Running(round) if round.job == job && round.phase == Phase::ProviderActive
327 ) {
328 return Err("stale Codex inference stage".to_owned());
329 }
330 if self.prepared.is_some() {
331 return Err("a Codex mailbox flush remains uncommitted".to_owned());
332 }
333 let accepted_provider_action = !values.is_empty();
334 let mut sequence = self.sequence;
335 let mut generated = Vec::with_capacity(values.len());
336 let mut prepared = Vec::new();
337 for value in values {
338 match value {
339 BoxValue::AgentMessage(Ok(contents)) => {
340 generated.push(ProviderGenerated::AgentMessage { contents });
341 }
342 BoxValue::Call(Ok(call)) => {
343 sequence = next_sequence(sequence)?;
344 let tool_call_id = ToolCallId::new(self.session, sequence);
345 generated.push(ProviderGenerated::ToolCall(ProviderCall {
346 tool_call_id,
347 name: call.name.clone(),
348 arguments: call.arguments.clone(),
349 }));
350 prepared.push(PreparedCall {
351 tool_call_id,
352 name: call.name,
353 arguments: call.arguments,
354 disposition: PreparedCallDisposition::External,
355 });
356 }
357 BoxValue::MalformedNativeAction(action) => {
358 sequence = next_sequence(sequence)?;
359 let tool_call_id = ToolCallId::new(self.session, sequence);
360 let name = action
361 .attempted_ktool()
362 .unwrap_or_else(|| action.native_tool())
363 .to_owned();
364 let arguments = action.native_arguments_json();
365 let error = ImmediateToolError {
366 message: action.diagnostic(),
367 native_tool: action.native_tool().to_owned(),
368 attempted_ktool: action.attempted_ktool().map(str::to_owned),
369 validation_code: action.validation_code().to_owned(),
370 path: action.path().to_owned(),
371 expected: action.expected().to_owned(),
372 received: action.received().to_owned(),
373 native_arguments: arguments.clone(),
374 metadata_type: MALFORMED_NATIVE_CALL_METADATA_TYPE.to_owned(),
375 metadata_contents: action.diagnostic_json(),
376 };
377 generated.push(ProviderGenerated::ToolCall(ProviderCall {
378 tool_call_id,
379 name: name.clone(),
380 arguments: arguments.clone(),
381 }));
382 prepared.push(PreparedCall {
383 tool_call_id,
384 name,
385 arguments,
386 disposition: PreparedCallDisposition::ImmediateError(Box::new(error)),
387 });
388 }
389 _ => return Err("stage contains a malformed provider action".to_owned()),
390 }
391 }
392 let before = self.state.boxes().len();
393 self.state
394 .append_stage(job, text, generated)
395 .map_err(debug)?;
396 self.unsubmitted.extend(
397 self.state.boxes()[before..]
398 .iter()
399 .filter(|box_| box_.box_type() == TOOL_CALL_TYPE)
400 .map(project),
401 );
402 self.sequence = sequence;
403 if let Mode::Running(round) = &mut self.mode {
404 round.accepted_provider_action |= accepted_provider_action;
405 round.phase = Phase::ChatendBoundary;
406 round.mailbox_flush_needed = true;
407 round.restartable = false;
408 }
409 Ok(prepared)
410 }
411
412 pub fn mailbox_flush(&mut self, job: u64) -> Result<Vec<ChatBox>, String> {
413 if !matches!(
414 &self.mode,
415 Mode::Running(round) if round.job == job && round.phase == Phase::ChatendBoundary
416 ) {
417 return Err("stale Codex active-arrival mailbox flush".to_owned());
418 }
419 let boxes = self.state.flush_active_arrivals(job).map_err(debug)?;
420 self.unsubmitted.extend(boxes.iter().map(project));
421 if let Mode::Running(round) = &mut self.mode {
422 round.phase = Phase::PendingGeneration;
423 }
424 Ok(boxes)
425 }
426
427 pub fn prepare_mailbox_flush(
428 &mut self,
429 job: u64,
430 ) -> Result<Option<PreparedMailboxFlush>, String> {
431 let (mailbox_flush_needed, phase) = match &self.mode {
432 Mode::Running(round) if round.job == job => (round.mailbox_flush_needed, round.phase),
433 _ => return Err("stale Codex inference mailbox flush".to_owned()),
434 };
435 if let Some(prepared) = &self.prepared {
436 return Ok(Some(PreparedMailboxFlush(Arc::clone(prepared))));
437 }
438 if !mailbox_flush_needed {
439 return Ok(None);
440 }
441 if phase == Phase::ProviderActive {
442 return Ok(None);
443 }
444 if phase == Phase::ChatendBoundary {
445 self.mailbox_flush(job)?;
446 }
447 let token = self
448 .token
449 .checked_add(1)
450 .ok_or_else(|| "Codex mailbox-flush token space was exhausted".to_owned())?;
451 let external_count = self.unsubmitted.len();
452 let mut values = self.unsubmitted.clone();
453 values.push(open_agent_response(self.promised_id()?));
454 let prepared = Arc::new(Prepared {
455 token,
456 values,
457 job,
458 external_count,
459 });
460 self.token = token;
461 self.prepared = Some(Arc::clone(&prepared));
462 if let Mode::Running(round) = &mut self.mode {
463 round.mailbox_flush_needed = false;
464 }
465 Ok(Some(PreparedMailboxFlush(prepared)))
466 }
467
468 pub fn validate_mailbox_flush(&self, prepared: &PreparedMailboxFlush) -> Result<(), String> {
469 let prepared = &prepared.0;
470 let prefix_matches = self.unsubmitted.get(..prepared.external_count)
471 == Some(&prepared.values[..prepared.external_count]);
472 let valid = matches!(
473 &self.mode,
474 Mode::Running(round)
475 if round.job == prepared.job && round.phase == Phase::PendingGeneration
476 ) && self.token == prepared.token
477 && prefix_matches
478 && self
479 .prepared
480 .as_ref()
481 .is_some_and(|current| Arc::ptr_eq(current, prepared));
482 if valid {
483 Ok(())
484 } else {
485 Err("stale or invalid Codex mailbox flush".to_owned())
486 }
487 }
488
489 pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
490 self.validate_mailbox_flush(&prepared)?;
491 self.unsubmitted.drain(..prepared.0.external_count);
492 self.prepared = None;
493 if let Mode::Running(round) = &mut self.mode {
494 round.phase = Phase::ProviderActive;
495 }
496 Ok(())
497 }
498
499 pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<(), String> {
500 let mut round = self.take_round(job)?;
501 if round.phase != Phase::ProviderActive {
502 let message = "Codex inference completed outside a provider generation".to_owned();
503 self.preserve(round, message.clone(), false);
504 return Err(message);
505 }
506 if self.prepared.is_some() {
507 let message = "Codex inference completed with an uncommitted mailbox flush".to_owned();
508 self.preserve(round, message.clone(), false);
509 return Err(message);
510 }
511 let mut text = String::new();
512 for item in output.items {
513 match item {
514 ShimItem::Text(value) => text.push_str(&value),
515 ShimItem::Box(_) => {
516 round.accepted_provider_action = true;
517 let message = "terminal Codex output contains a box".to_owned();
518 self.preserve(round, message.clone(), false);
519 return Err(message);
520 }
521 }
522 }
523 self.complete_round(round, text)
524 }
525
526 pub fn complete_recoverable_failure(
527 &mut self,
528 job: u64,
529 message: String,
530 ) -> Result<(), String> {
531 let round = self.take_round(job)?;
532 self.prepared = None;
533 self.complete_round(round, message)
534 }
535
536 pub fn reset_provider_context(&mut self) -> Result<(), String> {
537 if !matches!(self.mode, Mode::Idle) || self.prepared.is_some() {
538 return Err("provider context reset requires an idle conversation".to_owned());
539 }
540 self.unsubmitted = self.state.boxes().iter().map(project).collect();
541 Ok(())
542 }
543
544 pub fn halt_critical(&mut self, message: String) {
545 if matches!(self.mode, Mode::Stalled(_, _)) {
546 return;
547 }
548 match std::mem::replace(&mut self.mode, Mode::Idle) {
549 Mode::Running(round) => self.preserve(round, message, false),
550 Mode::Idle => {
551 self.prepared = None;
552 self.queued_trigger = false;
553 self.mode = Mode::Stalled(
554 Status::Stalled {
555 message,
556 restartable: false,
557 },
558 false,
559 );
560 }
561 Mode::Stalled(status, accepted_provider_action) => {
562 self.mode = Mode::Stalled(status, accepted_provider_action);
563 }
564 }
565 }
566
567 pub fn fail(&mut self, job: u64, message: String, restartable_before_launch: bool) {
568 if let Ok(round) = self.take_round(job) {
569 self.preserve(round, message, restartable_before_launch);
570 }
571 }
572
573 pub fn restart(&mut self) -> Result<(), RestartError> {
574 match &self.mode {
575 Mode::Stalled(_, true) => return Err(RestartError::ProviderActionAccepted),
576 Mode::Stalled(
577 Status::Stalled {
578 restartable: true, ..
579 },
580 false,
581 ) => {}
582 _ => return Err(RestartError::NotStalled),
583 }
584 self.state.restart().map_err(|_| RestartError::NotStalled)?;
585 self.unsubmitted = self.state.boxes().iter().map(project).collect();
586 self.queued_trigger = false;
587 self.prepared = None;
588 self.mode = Mode::Idle;
589 Ok(())
590 }
591
592 fn accept_arrival<F>(&mut self, triggering: bool, accept: F) -> Result<(), String>
593 where
594 F: FnOnce(&mut ActorState) -> Result<(), StateError>,
595 {
596 let before = self.state.boxes().len();
597 accept(&mut self.state).map_err(debug)?;
598 let appended = &self.state.boxes()[before..];
599 self.unsubmitted.extend(appended.iter().map(project));
600 match &mut self.mode {
601 Mode::Running(round) => round.mailbox_flush_needed |= triggering,
602 Mode::Idle if appended.is_empty() => self.queued_trigger |= triggering,
603 Mode::Idle => self.queued_trigger = false,
604 Mode::Stalled(_, _) => {}
605 }
606 Ok(())
607 }
608
609 fn complete_round(&mut self, round: Round, text: String) -> Result<(), String> {
610 let before = self.state.boxes().len();
611 if let Err(error) = self
612 .state
613 .complete_inference(round.job, text)
614 .map_err(debug)
615 {
616 self.preserve(round, error.clone(), false);
617 return Err(error);
618 }
619 self.unsubmitted
620 .extend(self.state.boxes()[before..].iter().skip(1).map(project));
621 self.queued_trigger = false;
622 self.mode = Mode::Idle;
623 Ok(())
624 }
625
626 fn promised_id(&self) -> Result<BoxId, String> {
627 let previous = self.state.boxes().last().map_or(0, |box_| box_.id().get());
628 let value = previous
629 .checked_add(1)
630 .ok_or_else(|| "BoxId space was exhausted".to_owned())?;
631 Ok(BoxId::new(value))
632 }
633
634 fn take_round(&mut self, job: u64) -> Result<Round, String> {
635 match std::mem::replace(&mut self.mode, Mode::Idle) {
636 Mode::Running(round) if round.job == job => Ok(round),
637 other => {
638 self.mode = other;
639 Err("stale Codex inference completion".to_owned())
640 }
641 }
642 }
643
644 fn preserve(&mut self, round: Round, message: String, restartable_before_launch: bool) {
645 self.prepared = None;
646 let restartable = restartable_before_launch
647 && round.restartable
648 && !round.accepted_provider_action
649 && self
650 .state
651 .stall_inference(round.job, message.clone())
652 .is_ok();
653 if !restartable {
654 let _ = self.state.halt(message.clone());
655 }
656 self.mode = Mode::Stalled(
657 Status::Stalled {
658 message,
659 restartable,
660 },
661 round.accepted_provider_action,
662 );
663 }
664}
665
666fn next_sequence(sequence: u64) -> Result<u64, String> {
667 sequence
668 .checked_add(1)
669 .ok_or_else(|| "ToolCallId space was exhausted".to_owned())
670}
671
672fn debug(error: impl std::fmt::Debug) -> String {
673 format!("{error:?}")
674}
675
676#[cfg(test)]
677mod tests {
678 use super::*;
679 use kcode_k1_chat_codex_codec::Codec;
680 use kcode_k1_codex_adapter::{BoxCodec, ToolCall};
681
682 fn active() -> (ConversationState, u64) {
683 let mut state = ConversationState::new([7; 12]);
684 state
685 .accept(
686 USER_MESSAGE_TYPE.into(),
687 "start".into(),
688 String::new(),
689 String::new(),
690 )
691 .unwrap();
692 let job = state.begin().unwrap().unwrap().job;
693 (state, job)
694 }
695
696 fn native(name: &str, arguments: &str) -> BoxValue {
697 let mut codec = Codec;
698 codec.tool_call_box(&ToolCall {
699 call_id: "native".into(),
700 name: name.into(),
701 arguments: arguments.parse().unwrap(),
702 })
703 }
704
705 fn active_after_mailbox_flush() -> (ConversationState, u64, ToolCallId) {
706 let (mut state, job) = active();
707 let calls = state
708 .prepare_stage(
709 job,
710 String::new(),
711 vec![BoxValue::Call(Ok(Call {
712 name: "tool".into(),
713 arguments: "{}".into(),
714 }))],
715 )
716 .unwrap();
717 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
718 state.commit_mailbox_flush(prepared).unwrap();
719 (state, job, calls[0].tool_call_id)
720 }
721
722 fn empty_output() -> ShimOutput<BoxValue> {
723 ShimOutput { items: Vec::new() }
724 }
725
726 #[test]
727 fn fresh_preflight_is_ordered_unscheduled_and_part_of_first_input() {
728 let mut state = ConversationState::new([4; 12]);
729 let calls = state
730 .prepare_preflight(vec![
731 PreflightItem::SystemMessage {
732 contents: "context".into(),
733 },
734 PreflightItem::KtoolCall {
735 name: "CurrentTime".into(),
736 arguments: "{}".into(),
737 mode: PreflightMode::Blocking,
738 },
739 ])
740 .unwrap();
741 assert_eq!(calls.len(), 1);
742 assert_eq!(calls[0].call_box_id.get(), 2);
743 assert_eq!(calls[0].tool_call_id.sequence(), 1);
744 assert!(state.begin().unwrap().is_none());
745 state
746 .accept(
747 USER_MESSAGE_TYPE.into(),
748 "kickoff".into(),
749 String::new(),
750 String::new(),
751 )
752 .unwrap();
753 let start = state.begin().unwrap().unwrap();
754 assert_eq!(start.values.len(), 4);
755 assert!(
756 state
757 .prepare_preflight(vec![PreflightItem::SystemMessage {
758 contents: "again".into()
759 }])
760 .is_err()
761 );
762 }
763
764 #[test]
765 fn malformed_only_is_persisted_and_prepared_as_deterministic_immediate_error() {
766 let (mut state, job) = active();
767 let malformed = native(
768 "call_ktool",
769 r#"{"z":0,"name":"recoverable","arguments":{"b":2,"a":1}}"#,
770 );
771 let (diagnostic, diagnostic_json) = match &malformed {
772 BoxValue::MalformedNativeAction(action) => {
773 (action.diagnostic(), action.diagnostic_json())
774 }
775 _ => panic!("test value must be malformed"),
776 };
777 let calls = state
778 .prepare_stage(job, String::new(), vec![malformed])
779 .unwrap();
780 assert_eq!(state.status(), Status::Running);
781 assert_eq!(calls.len(), 1);
782 assert_eq!(calls[0].name, "recoverable");
783 assert_eq!(
784 calls[0].arguments,
785 r#"{"arguments":{"a":1,"b":2},"name":"recoverable","z":0}"#
786 );
787 let PreparedCallDisposition::ImmediateError(error) = calls[0].disposition() else {
788 panic!("malformed call must be local");
789 };
790 assert_eq!(error.message, diagnostic);
791 assert_eq!(error.native_tool, "call_ktool");
792 assert_eq!(error.attempted_ktool.as_deref(), Some("recoverable"));
793 assert_eq!(error.validation_code, "invalid_wrapper_fields");
794 assert_eq!(error.path, "$");
795 assert_eq!(error.expected, "object with exactly name and arguments");
796 assert_eq!(error.received, calls[0].arguments);
797 assert_eq!(error.native_arguments, calls[0].arguments);
798 assert_eq!(error.metadata_type, MALFORMED_NATIVE_CALL_METADATA_TYPE);
799 assert_eq!(error.metadata_contents, diagnostic_json);
800 let persisted = state
801 .boxes()
802 .last()
803 .unwrap()
804 .tool_call_metadata()
805 .unwrap()
806 .unwrap();
807 assert_eq!(persisted.name, calls[0].name);
808 assert_eq!(persisted.arguments, calls[0].arguments);
809 }
810
811 #[test]
812 fn mixed_stage_preserves_order_contiguous_ids_name_selection_and_external_calls() {
813 let (mut state, job) = active();
814 let before = state.boxes().len();
815 let values = vec![
816 BoxValue::AgentMessage(Ok("first".into())),
817 native(
818 "call_ktool",
819 r#"{"name":"attempted","arguments":{},"extra":1}"#,
820 ),
821 BoxValue::Call(Ok(Call {
822 name: "valid".into(),
823 arguments: "{\"ok\":true}".into(),
824 })),
825 native("future_native", r#"{"z":0,"a":true}"#),
826 ];
827 let calls = state.prepare_stage(job, String::new(), values).unwrap();
828 assert_eq!(
829 calls
830 .iter()
831 .map(|call| call.tool_call_id.sequence())
832 .collect::<Vec<_>>(),
833 vec![1, 2, 3]
834 );
835 assert_eq!(
836 calls
837 .iter()
838 .map(|call| call.name.as_str())
839 .collect::<Vec<_>>(),
840 vec!["attempted", "valid", "future_native"]
841 );
842 assert!(matches!(
843 calls[0].disposition(),
844 PreparedCallDisposition::ImmediateError(_)
845 ));
846 assert_eq!(calls[1].disposition(), &PreparedCallDisposition::External);
847 assert!(matches!(
848 calls[2].disposition(),
849 PreparedCallDisposition::ImmediateError(_)
850 ));
851 assert_eq!(calls[2].arguments, r#"{"a":true,"z":0}"#);
852 assert_eq!(
853 state.boxes()[before..]
854 .iter()
855 .map(ChatBox::box_type)
856 .collect::<Vec<_>>(),
857 vec![
858 AGENT_RESPONSE_TYPE,
859 AGENT_MESSAGE_TYPE,
860 TOOL_CALL_TYPE,
861 TOOL_CALL_TYPE,
862 TOOL_CALL_TYPE,
863 ]
864 );
865 }
866
867 #[test]
868 fn normal_call_precedes_correlated_result_and_open_response_in_stable_flush() {
869 let (mut state, job) = active();
870 let calls = state
871 .prepare_stage(
872 job,
873 String::new(),
874 vec![BoxValue::Call(Ok(Call {
875 name: "tool".into(),
876 arguments: "{\"input\":1}".into(),
877 }))],
878 )
879 .unwrap();
880 state
881 .accept_tool_return(calls[0].tool_call_id, Ok("done".into()))
882 .unwrap();
883
884 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
885 let expected = vec![
886 project(
887 state
888 .boxes()
889 .iter()
890 .find(|box_| box_.box_type() == TOOL_CALL_TYPE)
891 .unwrap(),
892 ),
893 project(
894 state
895 .boxes()
896 .iter()
897 .find(|box_| box_.box_type() == TOOL_RESULT_TYPE)
898 .unwrap(),
899 ),
900 open_agent_response(state.promised_id().unwrap()),
901 ];
902 assert_eq!(prepared.values(), expected);
903
904 let repeated = state.prepare_mailbox_flush(job).unwrap().unwrap();
905 assert_eq!(repeated.values(), prepared.values());
906 state.validate_mailbox_flush(&prepared).unwrap();
907 state.commit_mailbox_flush(repeated).unwrap();
908 assert!(state.unsubmitted.is_empty());
909 }
910
911 #[test]
912 fn malformed_native_call_is_queued_as_its_canonical_tool_call() {
913 let (mut state, job) = active();
914 let malformed = native(
915 "call_ktool",
916 r#"{"name":"recoverable","arguments":{},"extra":1}"#,
917 );
918 state
919 .prepare_stage(job, String::new(), vec![malformed])
920 .unwrap();
921
922 let expected_call = project(
923 state
924 .boxes()
925 .iter()
926 .find(|box_| box_.box_type() == TOOL_CALL_TYPE)
927 .unwrap(),
928 );
929 let expected_response = open_agent_response(state.promised_id().unwrap());
930 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
931 assert_eq!(prepared.values(), &[expected_call, expected_response]);
932 }
933
934 #[test]
935 fn provider_agent_message_and_stage_response_are_not_requeued() {
936 let (mut state, job) = active();
937 state
938 .prepare_stage(
939 job,
940 "stage response".into(),
941 vec![BoxValue::AgentMessage(Ok("provider message".into()))],
942 )
943 .unwrap();
944
945 let expected_response = open_agent_response(state.promised_id().unwrap());
946 let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
947 assert_eq!(prepared.values(), &[expected_response]);
948 }
949
950 #[test]
951 fn committed_mailbox_flush_does_not_schedule_a_fresh_turn() {
952 let (mut state, job, _) = active_after_mailbox_flush();
953 state.complete(job, empty_output()).unwrap();
954 assert_eq!(state.status(), Status::Quiet);
955 assert!(state.begin().unwrap().is_none());
956 }
957
958 #[test]
959 fn tool_messages_after_a_mailbox_flush_remain_inert() {
960 let (mut state, job, tool_call_id) = active_after_mailbox_flush();
961 state
962 .accept_tool_message(tool_call_id, "still running".into())
963 .unwrap();
964 state.complete(job, empty_output()).unwrap();
965 assert_eq!(state.status(), Status::Quiet);
966 assert!(state.begin().unwrap().is_none());
967 }
968
969 #[test]
970 fn one_tool_result_schedules_exactly_one_fresh_turn() {
971 let (mut state, job, tool_call_id) = active_after_mailbox_flush();
972 state
973 .accept_tool_return(tool_call_id, Ok("done".into()))
974 .unwrap();
975 state.complete(job, empty_output()).unwrap();
976 let followup = state.begin().unwrap().unwrap();
977 assert!(state.begin().unwrap().is_none());
978 state.complete(followup.job, empty_output()).unwrap();
979 assert_eq!(state.status(), Status::Quiet);
980 assert!(state.begin().unwrap().is_none());
981 }
982
983 #[test]
984 fn recoverable_failure_is_visible_terminal_and_allows_later_turn() {
985 let (mut state, job) = active();
986 state
987 .complete_recoverable_failure(job, "The response failed. Please try again.".into())
988 .unwrap();
989 assert_eq!(state.status(), Status::Quiet);
990 assert_eq!(
991 state.boxes().last().unwrap().box_type(),
992 AGENT_RESPONSE_TYPE
993 );
994 state.reset_provider_context().unwrap();
995 state
996 .accept(
997 USER_MESSAGE_TYPE.into(),
998 "retry".into(),
999 String::new(),
1000 String::new(),
1001 )
1002 .unwrap();
1003 let start = state.begin().unwrap().unwrap();
1004 assert!(start.values.iter().any(|value| {
1005 matches!(value, BoxValue::History(text) if text.contains("The response failed. Please try again."))
1006 }));
1007 }
1008
1009 #[test]
1010 fn failed_mailbox_transport_can_terminalize_and_replay() {
1011 let (mut state, job) = active();
1012 state
1013 .prepare_stage(job, "partial".into(), Vec::new())
1014 .unwrap();
1015 assert!(state.prepare_mailbox_flush(job).unwrap().is_some());
1016 state
1017 .complete_recoverable_failure(job, "safe failure".into())
1018 .unwrap();
1019 assert_eq!(state.status(), Status::Quiet);
1020 state.reset_provider_context().unwrap();
1021 state
1022 .accept(
1023 USER_MESSAGE_TYPE.into(),
1024 "later".into(),
1025 String::new(),
1026 String::new(),
1027 )
1028 .unwrap();
1029 let start = state.begin().unwrap().unwrap();
1030 assert!(start.values.iter().any(
1031 |value| matches!(value, BoxValue::History(text) if text.contains("safe failure"))
1032 ));
1033 }
1034
1035 #[test]
1036 fn idle_critical_halt_is_nonrestartable_and_first_write_wins() {
1037 let mut state = ConversationState::new([9; 12]);
1038 state.halt_critical("integrity failure".into());
1039 state.halt_critical("replacement".into());
1040 assert_eq!(
1041 state.status(),
1042 Status::Stalled {
1043 message: "integrity failure".into(),
1044 restartable: false,
1045 }
1046 );
1047 assert_eq!(state.restart(), Err(RestartError::NotStalled));
1048 }
1049}