#![forbid(unsafe_code)]
mod recovery;
pub use kcode_k1_chat_codex_codec::{BoxValue, Call};
pub use kcode_k1_chat_state::{
AGENT_ATTACHMENT_TYPE, AGENT_MESSAGE_TYPE, AGENT_RESPONSE_TYPE, ActorState, BoxId, ChatBox,
ProviderGenerated, SYSTEM_MESSAGE_TYPE, TOOL_ATTACHMENT_TYPE, TOOL_CALL_TYPE,
TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE, ToolCallId, USER_ATTACHMENT_TYPE, USER_MESSAGE_TYPE,
};
pub use kcode_k1_codex_adapter::{ShimItem, ShimOutput};
use std::sync::Arc;
use std::sync::atomic::AtomicU8;
use kcode_k1_chat_codex_codec::{open_agent_response, project};
use kcode_k1_chat_state::{ProviderCall, StateError};
use recovery::recovered_sequence;
pub const MALFORMED_NATIVE_CALL_METADATA_TYPE: &str = "k1.malformed-native-call.v1";
#[derive(Clone, Debug)]
pub struct Start {
pub job: u64,
pub values: Vec<BoxValue>,
pub attempt: Arc<AtomicU8>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PreparedCall {
pub tool_call_id: ToolCallId,
pub name: String,
pub arguments: String,
disposition: PreparedCallDisposition,
}
impl PreparedCall {
pub fn disposition(&self) -> &PreparedCallDisposition {
&self.disposition
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PreparedCallDisposition {
External,
ImmediateError(Box<ImmediateToolError>),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ImmediateToolError {
pub message: String,
pub native_tool: String,
pub attempted_ktool: Option<String>,
pub validation_code: String,
pub path: String,
pub expected: String,
pub received: String,
pub native_arguments: String,
pub metadata_type: String,
pub metadata_contents: String,
}
#[derive(Clone, Debug)]
pub struct PreparedMailboxFlush(Arc<Prepared>);
#[derive(Debug)]
struct Prepared {
token: u64,
values: Vec<BoxValue>,
job: u64,
external_count: usize,
}
impl PreparedMailboxFlush {
pub fn values(&self) -> &[BoxValue] {
&self.0.values
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Status {
Running,
Quiet,
Stalled { message: String, restartable: bool },
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RestartError {
NotStalled,
ProviderActionAccepted,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Phase {
ProviderActive,
ChatendBoundary,
PendingGeneration,
}
struct Round {
job: u64,
accepted_provider_action: bool,
phase: Phase,
mailbox_flush_needed: bool,
restartable: bool,
}
enum Mode {
Idle,
Running(Round),
Stalled(Status, bool),
}
pub struct ConversationState {
state: ActorState,
session: [u8; 12],
sequence: u64,
unsubmitted: Vec<BoxValue>,
queued_trigger: bool,
token: u64,
prepared: Option<Arc<Prepared>>,
mode: Mode,
}
impl ConversationState {
pub fn new(session: [u8; 12]) -> Self {
Self {
state: ActorState::new(false),
session,
sequence: 0,
unsubmitted: Vec::new(),
queued_trigger: false,
token: 0,
prepared: None,
mode: Mode::Idle,
}
}
pub fn recover(session: [u8; 12], boxes: Vec<ChatBox>, force: bool) -> Result<Self, String> {
let sequence = recovered_sequence(session, &boxes)?;
let state = ActorState::recover(boxes, force).map_err(debug)?;
Ok(Self {
unsubmitted: state.boxes().iter().map(project).collect(),
state,
sequence,
..Self::new(session)
})
}
pub fn boxes(&self) -> &[ChatBox] {
self.state.boxes()
}
pub fn status(&self) -> Status {
match &self.mode {
Mode::Running(_) => Status::Running,
Mode::Idle if self.state.quiet() => Status::Quiet,
Mode::Idle => Status::Running,
Mode::Stalled(status, _) => status.clone(),
}
}
pub fn accept(
&mut self,
box_type: String,
contents: String,
hidden_type: String,
hidden_contents: String,
) -> Result<(), String> {
self.accept_arrival(true, |state| {
state.accept_box(box_type, contents, hidden_type, hidden_contents)
})
}
pub fn accept_tool_message(
&mut self,
tool_call_id: ToolCallId,
message: String,
) -> Result<(), String> {
self.accept_arrival(false, |state| {
state.accept_tool_message(tool_call_id, message)
})
}
pub fn accept_tool_return(
&mut self,
tool_call_id: ToolCallId,
result: Result<String, String>,
) -> Result<(), String> {
self.accept_arrival(true, |state| {
state.accept_async_return(tool_call_id, result)
})
}
pub fn accept_tool_return_v2(
&mut self,
tool_call_id: ToolCallId,
result: Result<String, String>,
metadata_type: String,
metadata_contents: String,
) -> Result<(), String> {
self.accept_arrival(true, |state| {
state.accept_async_return_v2(tool_call_id, result, metadata_type, metadata_contents)
})
}
pub fn begin(&mut self) -> Result<Option<Start>, String> {
if !matches!(self.mode, Mode::Idle) {
return Ok(None);
}
let Some(start) = self.state.begin_inference().map_err(debug)? else {
return Ok(None);
};
let promised_id = self.promised_id()?;
let mut values = std::mem::take(&mut self.unsubmitted);
values.push(open_agent_response(promised_id));
let mailbox_flush_needed = std::mem::take(&mut self.queued_trigger);
self.mode = Mode::Running(Round {
job: start.job,
accepted_provider_action: false,
phase: Phase::ProviderActive,
mailbox_flush_needed,
restartable: true,
});
Ok(Some(Start {
job: start.job,
values,
attempt: start.attempt,
}))
}
pub fn prepare_stage(
&mut self,
job: u64,
text: String,
values: Vec<BoxValue>,
) -> Result<Vec<PreparedCall>, String> {
if !matches!(
&self.mode,
Mode::Running(round) if round.job == job && round.phase == Phase::ProviderActive
) {
return Err("stale Codex inference stage".to_owned());
}
if self.prepared.is_some() {
return Err("a Codex mailbox flush remains uncommitted".to_owned());
}
let accepted_provider_action = !values.is_empty();
let mut sequence = self.sequence;
let mut generated = Vec::with_capacity(values.len());
let mut prepared = Vec::new();
for value in values {
match value {
BoxValue::AgentMessage(Ok(contents)) => {
generated.push(ProviderGenerated::AgentMessage { contents });
}
BoxValue::Call(Ok(call)) => {
sequence = next_sequence(sequence)?;
let tool_call_id = ToolCallId::new(self.session, sequence);
generated.push(ProviderGenerated::ToolCall(ProviderCall {
tool_call_id,
name: call.name.clone(),
arguments: call.arguments.clone(),
}));
prepared.push(PreparedCall {
tool_call_id,
name: call.name,
arguments: call.arguments,
disposition: PreparedCallDisposition::External,
});
}
BoxValue::MalformedNativeAction(action) => {
sequence = next_sequence(sequence)?;
let tool_call_id = ToolCallId::new(self.session, sequence);
let name = action
.attempted_ktool()
.unwrap_or_else(|| action.native_tool())
.to_owned();
let arguments = action.native_arguments_json();
let error = ImmediateToolError {
message: action.diagnostic(),
native_tool: action.native_tool().to_owned(),
attempted_ktool: action.attempted_ktool().map(str::to_owned),
validation_code: action.validation_code().to_owned(),
path: action.path().to_owned(),
expected: action.expected().to_owned(),
received: action.received().to_owned(),
native_arguments: arguments.clone(),
metadata_type: MALFORMED_NATIVE_CALL_METADATA_TYPE.to_owned(),
metadata_contents: action.diagnostic_json(),
};
generated.push(ProviderGenerated::ToolCall(ProviderCall {
tool_call_id,
name: name.clone(),
arguments: arguments.clone(),
}));
prepared.push(PreparedCall {
tool_call_id,
name,
arguments,
disposition: PreparedCallDisposition::ImmediateError(Box::new(error)),
});
}
_ => return Err("stage contains a malformed provider action".to_owned()),
}
}
let before = self.state.boxes().len();
self.state
.append_stage(job, text, generated)
.map_err(debug)?;
self.unsubmitted.extend(
self.state.boxes()[before..]
.iter()
.filter(|box_| box_.box_type() == TOOL_CALL_TYPE)
.map(project),
);
self.sequence = sequence;
if let Mode::Running(round) = &mut self.mode {
round.accepted_provider_action |= accepted_provider_action;
round.phase = Phase::ChatendBoundary;
round.mailbox_flush_needed = true;
round.restartable = false;
}
Ok(prepared)
}
pub fn mailbox_flush(&mut self, job: u64) -> Result<Vec<ChatBox>, String> {
if !matches!(
&self.mode,
Mode::Running(round) if round.job == job && round.phase == Phase::ChatendBoundary
) {
return Err("stale Codex active-arrival mailbox flush".to_owned());
}
let boxes = self.state.flush_active_arrivals(job).map_err(debug)?;
self.unsubmitted.extend(boxes.iter().map(project));
if let Mode::Running(round) = &mut self.mode {
round.phase = Phase::PendingGeneration;
}
Ok(boxes)
}
pub fn prepare_mailbox_flush(
&mut self,
job: u64,
) -> Result<Option<PreparedMailboxFlush>, String> {
let (mailbox_flush_needed, phase) = match &self.mode {
Mode::Running(round) if round.job == job => (round.mailbox_flush_needed, round.phase),
_ => return Err("stale Codex inference mailbox flush".to_owned()),
};
if let Some(prepared) = &self.prepared {
return Ok(Some(PreparedMailboxFlush(Arc::clone(prepared))));
}
if !mailbox_flush_needed {
return Ok(None);
}
if phase == Phase::ProviderActive {
return Ok(None);
}
if phase == Phase::ChatendBoundary {
self.mailbox_flush(job)?;
}
let token = self
.token
.checked_add(1)
.ok_or_else(|| "Codex mailbox-flush token space was exhausted".to_owned())?;
let external_count = self.unsubmitted.len();
let mut values = self.unsubmitted.clone();
values.push(open_agent_response(self.promised_id()?));
let prepared = Arc::new(Prepared {
token,
values,
job,
external_count,
});
self.token = token;
self.prepared = Some(Arc::clone(&prepared));
if let Mode::Running(round) = &mut self.mode {
round.mailbox_flush_needed = false;
}
Ok(Some(PreparedMailboxFlush(prepared)))
}
pub fn validate_mailbox_flush(&self, prepared: &PreparedMailboxFlush) -> Result<(), String> {
let prepared = &prepared.0;
let prefix_matches = self.unsubmitted.get(..prepared.external_count)
== Some(&prepared.values[..prepared.external_count]);
let valid = matches!(
&self.mode,
Mode::Running(round)
if round.job == prepared.job && round.phase == Phase::PendingGeneration
) && self.token == prepared.token
&& prefix_matches
&& self
.prepared
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, prepared));
if valid {
Ok(())
} else {
Err("stale or invalid Codex mailbox flush".to_owned())
}
}
pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
self.validate_mailbox_flush(&prepared)?;
self.unsubmitted.drain(..prepared.0.external_count);
self.prepared = None;
if let Mode::Running(round) = &mut self.mode {
round.phase = Phase::ProviderActive;
}
Ok(())
}
pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<(), String> {
let mut round = self.take_round(job)?;
if round.phase != Phase::ProviderActive {
let message = "Codex inference completed outside a provider generation".to_owned();
self.preserve(round, message.clone(), false);
return Err(message);
}
if self.prepared.is_some() {
let message = "Codex inference completed with an uncommitted mailbox flush".to_owned();
self.preserve(round, message.clone(), false);
return Err(message);
}
let mut text = String::new();
for item in output.items {
match item {
ShimItem::Text(value) => text.push_str(&value),
ShimItem::Box(_) => {
round.accepted_provider_action = true;
let message = "terminal Codex output contains a box".to_owned();
self.preserve(round, message.clone(), false);
return Err(message);
}
}
}
let before = self.state.boxes().len();
if let Err(error) = self.state.complete_inference(job, text).map_err(debug) {
self.preserve(round, error.clone(), false);
return Err(error);
}
self.unsubmitted
.extend(self.state.boxes()[before..].iter().skip(1).map(project));
self.queued_trigger = false;
self.mode = Mode::Idle;
Ok(())
}
pub fn fail(&mut self, job: u64, message: String, restartable_before_launch: bool) {
if let Ok(round) = self.take_round(job) {
self.preserve(round, message, restartable_before_launch);
}
}
pub fn restart(&mut self) -> Result<(), RestartError> {
match &self.mode {
Mode::Stalled(_, true) => return Err(RestartError::ProviderActionAccepted),
Mode::Stalled(
Status::Stalled {
restartable: true, ..
},
false,
) => {}
_ => return Err(RestartError::NotStalled),
}
self.state.restart().map_err(|_| RestartError::NotStalled)?;
self.unsubmitted = self.state.boxes().iter().map(project).collect();
self.queued_trigger = false;
self.prepared = None;
self.mode = Mode::Idle;
Ok(())
}
fn accept_arrival<F>(&mut self, triggering: bool, accept: F) -> Result<(), String>
where
F: FnOnce(&mut ActorState) -> Result<(), StateError>,
{
let before = self.state.boxes().len();
accept(&mut self.state).map_err(debug)?;
let appended = &self.state.boxes()[before..];
self.unsubmitted.extend(appended.iter().map(project));
match &mut self.mode {
Mode::Running(round) => round.mailbox_flush_needed |= triggering,
Mode::Idle if appended.is_empty() => self.queued_trigger |= triggering,
Mode::Idle => self.queued_trigger = false,
Mode::Stalled(_, _) => {}
}
Ok(())
}
fn promised_id(&self) -> Result<BoxId, String> {
let previous = self.state.boxes().last().map_or(0, |box_| box_.id().get());
let value = previous
.checked_add(1)
.ok_or_else(|| "BoxId space was exhausted".to_owned())?;
Ok(BoxId::new(value))
}
fn take_round(&mut self, job: u64) -> Result<Round, String> {
match std::mem::replace(&mut self.mode, Mode::Idle) {
Mode::Running(round) if round.job == job => Ok(round),
other => {
self.mode = other;
Err("stale Codex inference completion".to_owned())
}
}
}
fn preserve(&mut self, round: Round, message: String, restartable_before_launch: bool) {
self.prepared = None;
let restartable = restartable_before_launch
&& round.restartable
&& !round.accepted_provider_action
&& self
.state
.stall_inference(round.job, message.clone())
.is_ok();
if !restartable {
let _ = self.state.halt(message.clone());
}
self.mode = Mode::Stalled(
Status::Stalled {
message,
restartable,
},
round.accepted_provider_action,
);
}
}
fn next_sequence(sequence: u64) -> Result<u64, String> {
sequence
.checked_add(1)
.ok_or_else(|| "ToolCallId space was exhausted".to_owned())
}
fn debug(error: impl std::fmt::Debug) -> String {
format!("{error:?}")
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_chat_codex_codec::Codec;
use kcode_k1_codex_adapter::{BoxCodec, ToolCall};
fn active() -> (ConversationState, u64) {
let mut state = ConversationState::new([7; 12]);
state
.accept(
USER_MESSAGE_TYPE.into(),
"start".into(),
String::new(),
String::new(),
)
.unwrap();
let job = state.begin().unwrap().unwrap().job;
(state, job)
}
fn native(name: &str, arguments: &str) -> BoxValue {
let mut codec = Codec;
codec.tool_call_box(&ToolCall {
call_id: "native".into(),
name: name.into(),
arguments: arguments.parse().unwrap(),
})
}
fn active_after_mailbox_flush() -> (ConversationState, u64, ToolCallId) {
let (mut state, job) = active();
let calls = state
.prepare_stage(
job,
String::new(),
vec![BoxValue::Call(Ok(Call {
name: "tool".into(),
arguments: "{}".into(),
}))],
)
.unwrap();
let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
state.commit_mailbox_flush(prepared).unwrap();
(state, job, calls[0].tool_call_id)
}
fn empty_output() -> ShimOutput<BoxValue> {
ShimOutput { items: Vec::new() }
}
#[test]
fn malformed_only_is_persisted_and_prepared_as_deterministic_immediate_error() {
let (mut state, job) = active();
let malformed = native(
"call_ktool",
r#"{"z":0,"name":"recoverable","arguments":{"b":2,"a":1}}"#,
);
let (diagnostic, diagnostic_json) = match &malformed {
BoxValue::MalformedNativeAction(action) => {
(action.diagnostic(), action.diagnostic_json())
}
_ => panic!("test value must be malformed"),
};
let calls = state
.prepare_stage(job, String::new(), vec![malformed])
.unwrap();
assert_eq!(state.status(), Status::Running);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "recoverable");
assert_eq!(
calls[0].arguments,
r#"{"arguments":{"a":1,"b":2},"name":"recoverable","z":0}"#
);
let PreparedCallDisposition::ImmediateError(error) = calls[0].disposition() else {
panic!("malformed call must be local");
};
assert_eq!(error.message, diagnostic);
assert_eq!(error.native_tool, "call_ktool");
assert_eq!(error.attempted_ktool.as_deref(), Some("recoverable"));
assert_eq!(error.validation_code, "invalid_wrapper_fields");
assert_eq!(error.path, "$");
assert_eq!(error.expected, "object with exactly name and arguments");
assert_eq!(error.received, calls[0].arguments);
assert_eq!(error.native_arguments, calls[0].arguments);
assert_eq!(error.metadata_type, MALFORMED_NATIVE_CALL_METADATA_TYPE);
assert_eq!(error.metadata_contents, diagnostic_json);
let persisted = state
.boxes()
.last()
.unwrap()
.tool_call_metadata()
.unwrap()
.unwrap();
assert_eq!(persisted.name, calls[0].name);
assert_eq!(persisted.arguments, calls[0].arguments);
}
#[test]
fn mixed_stage_preserves_order_contiguous_ids_name_selection_and_external_calls() {
let (mut state, job) = active();
let before = state.boxes().len();
let values = vec![
BoxValue::AgentMessage(Ok("first".into())),
native(
"call_ktool",
r#"{"name":"attempted","arguments":{},"extra":1}"#,
),
BoxValue::Call(Ok(Call {
name: "valid".into(),
arguments: "{\"ok\":true}".into(),
})),
native("future_native", r#"{"z":0,"a":true}"#),
];
let calls = state.prepare_stage(job, String::new(), values).unwrap();
assert_eq!(
calls
.iter()
.map(|call| call.tool_call_id.sequence())
.collect::<Vec<_>>(),
vec![1, 2, 3]
);
assert_eq!(
calls
.iter()
.map(|call| call.name.as_str())
.collect::<Vec<_>>(),
vec!["attempted", "valid", "future_native"]
);
assert!(matches!(
calls[0].disposition(),
PreparedCallDisposition::ImmediateError(_)
));
assert_eq!(calls[1].disposition(), &PreparedCallDisposition::External);
assert!(matches!(
calls[2].disposition(),
PreparedCallDisposition::ImmediateError(_)
));
assert_eq!(calls[2].arguments, r#"{"a":true,"z":0}"#);
assert_eq!(
state.boxes()[before..]
.iter()
.map(ChatBox::box_type)
.collect::<Vec<_>>(),
vec![
AGENT_RESPONSE_TYPE,
AGENT_MESSAGE_TYPE,
TOOL_CALL_TYPE,
TOOL_CALL_TYPE,
TOOL_CALL_TYPE,
]
);
}
#[test]
fn normal_call_precedes_correlated_result_and_open_response_in_stable_flush() {
let (mut state, job) = active();
let calls = state
.prepare_stage(
job,
String::new(),
vec![BoxValue::Call(Ok(Call {
name: "tool".into(),
arguments: "{\"input\":1}".into(),
}))],
)
.unwrap();
state
.accept_tool_return(calls[0].tool_call_id, Ok("done".into()))
.unwrap();
let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
let expected = vec![
project(
state
.boxes()
.iter()
.find(|box_| box_.box_type() == TOOL_CALL_TYPE)
.unwrap(),
),
project(
state
.boxes()
.iter()
.find(|box_| box_.box_type() == TOOL_RESULT_TYPE)
.unwrap(),
),
open_agent_response(state.promised_id().unwrap()),
];
assert_eq!(prepared.values(), expected);
let repeated = state.prepare_mailbox_flush(job).unwrap().unwrap();
assert_eq!(repeated.values(), prepared.values());
state.validate_mailbox_flush(&prepared).unwrap();
state.commit_mailbox_flush(repeated).unwrap();
assert!(state.unsubmitted.is_empty());
}
#[test]
fn malformed_native_call_is_queued_as_its_canonical_tool_call() {
let (mut state, job) = active();
let malformed = native(
"call_ktool",
r#"{"name":"recoverable","arguments":{},"extra":1}"#,
);
state
.prepare_stage(job, String::new(), vec![malformed])
.unwrap();
let expected_call = project(
state
.boxes()
.iter()
.find(|box_| box_.box_type() == TOOL_CALL_TYPE)
.unwrap(),
);
let expected_response = open_agent_response(state.promised_id().unwrap());
let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
assert_eq!(prepared.values(), &[expected_call, expected_response]);
}
#[test]
fn provider_agent_message_and_stage_response_are_not_requeued() {
let (mut state, job) = active();
state
.prepare_stage(
job,
"stage response".into(),
vec![BoxValue::AgentMessage(Ok("provider message".into()))],
)
.unwrap();
let expected_response = open_agent_response(state.promised_id().unwrap());
let prepared = state.prepare_mailbox_flush(job).unwrap().unwrap();
assert_eq!(prepared.values(), &[expected_response]);
}
#[test]
fn committed_mailbox_flush_does_not_schedule_a_fresh_turn() {
let (mut state, job, _) = active_after_mailbox_flush();
state.complete(job, empty_output()).unwrap();
assert_eq!(state.status(), Status::Quiet);
assert!(state.begin().unwrap().is_none());
}
#[test]
fn tool_messages_after_a_mailbox_flush_remain_inert() {
let (mut state, job, tool_call_id) = active_after_mailbox_flush();
state
.accept_tool_message(tool_call_id, "still running".into())
.unwrap();
state.complete(job, empty_output()).unwrap();
assert_eq!(state.status(), Status::Quiet);
assert!(state.begin().unwrap().is_none());
}
#[test]
fn one_tool_result_schedules_exactly_one_fresh_turn() {
let (mut state, job, tool_call_id) = active_after_mailbox_flush();
state
.accept_tool_return(tool_call_id, Ok("done".into()))
.unwrap();
state.complete(job, empty_output()).unwrap();
let followup = state.begin().unwrap().unwrap();
assert!(state.begin().unwrap().is_none());
state.complete(followup.job, empty_output()).unwrap();
assert_eq!(state.status(), Status::Quiet);
assert!(state.begin().unwrap().is_none());
}
}