#![forbid(unsafe_code)]
pub use kcode_k1_chat_codex_state::{
BoxValue, ChatBox, PreparedCall, PreparedMailboxFlush, RestartError, ShimOutput, Start, Status,
ToolCallId,
};
pub use kcode_k1_chat_persistence::EventRecord;
use kcode_k1_chat_codex_state::{AGENT_RESPONSE_TYPE, ConversationState};
use kcode_k1_chat_persistence::{Record, Session};
use kcode_k1_chat_thread_recovery::recover as recover_thread;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub struct TokenBreakdown {
pub input_tokens: i64,
pub cached_input_tokens: i64,
pub cache_write_input_tokens: i64,
pub output_tokens: i64,
pub reasoning_output_tokens: i64,
pub total_tokens: i64,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub struct ModelUsage {
pub provider: String,
pub model: String,
pub context_id: String,
pub provider_turn_id: String,
pub usage: TokenBreakdown,
pub cumulative_usage: Option<TokenBreakdown>,
pub context_limit_tokens: Option<i64>,
}
pub struct DurableTurn {
state: ConversationState,
records: Vec<Record>,
mirrored: usize,
durable: usize,
session: Session,
returned: Vec<ToolCallId>,
}
impl DurableTurn {
pub fn recover(session: Session) -> Result<Self, String> {
let recovered = recover_thread(&session)?;
let returned = returned_ids(recovered.state.boxes())?;
Ok(Self {
state: recovered.state,
records: recovered.records,
mirrored: recovered.mirrored,
durable: recovered.durable,
session,
returned,
})
}
pub fn boxes(&self) -> &[ChatBox] {
self.state.boxes()
}
pub fn events(&self) -> Vec<EventRecord> {
self.records[..self.durable]
.iter()
.filter_map(|record| match record {
Record::Event(event) => Some(event.clone()),
Record::Box(_) => None,
})
.collect()
}
pub fn status(&self) -> Status {
self.state.status()
}
pub fn accept(
&mut self,
box_type: String,
contents: String,
hidden_type: String,
hidden_contents: String,
) -> Result<(), String> {
let result = self
.state
.accept(box_type, contents, hidden_type, hidden_contents);
self.finish(result)
}
pub fn accept_tool_return(
&mut self,
tool_call_id: ToolCallId,
result: Result<String, String>,
) -> Result<(), String> {
if self.returned.contains(&tool_call_id) {
return self.finish(Ok(()));
}
let accepted = self.state.accept_tool_return(tool_call_id, result);
if accepted.is_ok() {
self.returned.push(tool_call_id);
}
self.finish(accepted)
}
pub fn accept_tool_message(
&mut self,
tool_call_id: ToolCallId,
message: String,
) -> Result<(), String> {
let result = self.state.accept_tool_message(tool_call_id, message);
self.finish(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> {
let accepted = self.state.accept_tool_return_v2(
tool_call_id,
result,
metadata_type,
metadata_contents,
);
if accepted.is_ok() {
self.returned.push(tool_call_id);
}
self.finish(accepted)
}
pub fn begin(&mut self) -> Result<Option<Start>, String> {
self.state.begin()
}
pub fn prepare_stage(
&mut self,
job: u64,
text: String,
values: Vec<BoxValue>,
) -> Result<Vec<PreparedCall>, String> {
let result = self.state.prepare_stage(job, text, values);
self.finish(result)
}
pub fn prepare_mailbox_flush(
&mut self,
job: u64,
) -> Result<Option<PreparedMailboxFlush>, String> {
let result = self.state.prepare_mailbox_flush(job);
self.finish(result)
}
pub fn validate_mailbox_flush(&self, prepared: &PreparedMailboxFlush) -> Result<(), String> {
self.state.validate_mailbox_flush(prepared)
}
pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
self.state.commit_mailbox_flush(prepared)
}
pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
if let Err(error) = self.state.complete(job, output) {
return self.finish(Err(error));
}
self.mirror_boxes()?;
let resume = matches!(self.state.status(), Status::Running);
let after_box_id = self.latest_box_id()?;
self.persist_event(EventRecord {
after_box_id,
event_index: self.next_event_index(after_box_id)?,
connected_box_id: 0,
handler: "llm_done".into(),
data: json!({"resume": resume}),
})?;
Ok(resume)
}
pub fn record_model_usage(
&mut self,
connected_box_id: u64,
usage: ModelUsage,
) -> Result<(), String> {
self.mirror_boxes()?;
let after_box_id = self.latest_box_id()?;
if connected_box_id != 0
&& !self.state.boxes().iter().any(|box_| {
box_.id().get() == connected_box_id && box_.box_type() == AGENT_RESPONSE_TYPE
})
{
return Err("model usage must connect to a canonical Agent Response box".to_owned());
}
self.persist_event(EventRecord {
after_box_id,
event_index: self.next_event_index(after_box_id)?,
connected_box_id,
handler: "model_usage".into(),
data: model_usage_data(usage)?,
})
}
pub fn fail(&mut self, job: u64, message: String, restartable_before_launch: bool) {
self.state.fail(job, message, restartable_before_launch);
}
pub fn restart(&mut self) -> Result<(), RestartError> {
self.state.restart()
}
fn finish<T>(&mut self, operation: Result<T, String>) -> Result<T, String> {
let persistence = self.mirror_and_persist();
match (operation, persistence) {
(Ok(value), Ok(())) => Ok(value),
(Err(error), Ok(())) | (Ok(_), Err(error)) => Err(error),
(Err(operation), Err(persistence)) => Err(format!(
"{operation}; additionally failed to persist canonical history: {persistence}"
)),
}
}
fn mirror_and_persist(&mut self) -> Result<(), String> {
self.mirror_boxes()?;
self.persist_pending()
}
fn mirror_boxes(&mut self) -> Result<(), String> {
let boxes = self.state.boxes();
let additions = boxes
.get(self.mirrored..)
.ok_or_else(|| "canonical box frontier moved backwards".to_owned())?;
self.records
.extend(additions.iter().cloned().map(Record::Box));
self.mirrored = boxes.len();
Ok(())
}
fn latest_box_id(&self) -> Result<u64, String> {
self.state
.boxes()
.last()
.map(|box_| box_.id().get())
.ok_or_else(|| "durable event requires a canonical box".to_owned())
}
fn next_event_index(&self, after_box_id: u64) -> Result<u64, String> {
match self.records.last() {
Some(Record::Event(event)) if event.after_box_id == after_box_id => event
.event_index
.checked_add(1)
.ok_or_else(|| "durable event index space was exhausted".to_owned()),
Some(Record::Event(_)) => {
Err("durable event frontier diverged from canonical boxes".to_owned())
}
_ => Ok(1),
}
}
fn persist_event(&mut self, event: EventRecord) -> Result<(), String> {
let suffix = self
.records
.get(self.durable..)
.ok_or_else(|| "durable record frontier moved past canonical records".to_owned())?;
let mut pending = suffix.to_vec();
pending.push(Record::Event(event.clone()));
self.session.persist(pending)?;
self.records.push(Record::Event(event));
self.durable = self.records.len();
Ok(())
}
fn persist_pending(&mut self) -> Result<(), String> {
let suffix = self
.records
.get(self.durable..)
.ok_or_else(|| "durable record frontier moved past canonical records".to_owned())?;
if suffix.is_empty() {
return Ok(());
}
self.session.persist(suffix.to_vec())?;
self.durable = self.records.len();
Ok(())
}
}
fn model_usage_data(usage: ModelUsage) -> Result<Value, String> {
let mut data = serde_json::to_value(usage).map_err(|error| error.to_string())?;
let Value::Object(fields) = &mut data else {
return Err("model usage data did not serialize to an object".to_owned());
};
fields.insert("version".into(), Value::from(1));
Ok(data)
}
fn returned_ids(boxes: &[ChatBox]) -> Result<Vec<ToolCallId>, String> {
let mut returned = Vec::new();
for value in boxes {
if let Some(result) = value
.tool_result_metadata()
.map_err(|error| format!("{error:?}"))?
{
returned.push(result.tool_call_id);
}
}
Ok(returned)
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_chat_codex_state::Call;
use kcode_k1_chat_persistence::K1ChatPersistence;
use kcode_k1_peering::K1Peering;
use kcode_k1_txn_ordering::K1TxnOrdering;
use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
static NEXT: AtomicU64 = AtomicU64::new(0);
struct Fixture {
root: PathBuf,
session: Option<Session>,
}
impl Fixture {
fn new(nonce: u8) -> Self {
let root = std::env::temp_dir().join(format!(
"k1-durable-turn-{}-{}",
std::process::id(),
NEXT.fetch_add(1, Ordering::Relaxed)
));
let _ = fs::remove_dir_all(&root);
let ordering = Arc::new(K1TxnOrdering::open(&root.join("ordering")).unwrap());
let peering =
Arc::new(K1Peering::open(&root.join("peering"), Arc::clone(&ordering)).unwrap());
let persistence =
K1ChatPersistence::open(&root.join("persistence"), ordering, peering).unwrap();
let (session, _) = persistence.session([nonce; 12]).unwrap();
Self {
root,
session: Some(session),
}
}
fn session(&self) -> Session {
self.session.as_ref().unwrap().clone()
}
}
impl Drop for Fixture {
fn drop(&mut self) {
drop(self.session.take());
let _ = fs::remove_dir_all(&self.root);
}
}
fn usage() -> ModelUsage {
let breakdown = TokenBreakdown {
input_tokens: 10,
cached_input_tokens: 2,
cache_write_input_tokens: 3,
output_tokens: 4,
reasoning_output_tokens: 5,
total_tokens: 14,
};
ModelUsage {
provider: "provider".into(),
model: "model".into(),
context_id: "context".into(),
provider_turn_id: "turn".into(),
usage: breakdown.clone(),
cumulative_usage: Some(breakdown),
context_limit_tokens: Some(128),
}
}
fn completed(turn: &mut DurableTurn) -> u64 {
turn.accept(
"User Message".into(),
"hello".into(),
String::new(),
String::new(),
)
.unwrap();
let start = turn.begin().unwrap().unwrap();
turn.complete(start.job, ShimOutput { items: Vec::new() })
.unwrap();
turn.boxes().last().unwrap().id().get()
}
#[test]
fn completion_persists_terminal_box_event_and_recovery_frontiers() {
let fixture = Fixture::new(1);
let session = fixture.session();
let mut turn = DurableTurn::recover(session.clone()).unwrap();
let terminal = completed(&mut turn);
assert_eq!(terminal, 2);
assert_eq!(turn.status(), Status::Quiet);
assert_eq!((turn.mirrored, turn.durable, turn.records.len()), (2, 3, 3));
let log = session.load().unwrap();
let Record::Event(event) = &log.records[2] else {
panic!("expected llm_done event")
};
assert_eq!(
(
event.after_box_id,
event.event_index,
event.connected_box_id
),
(2, 1, 0)
);
assert_eq!(event.handler, "llm_done");
assert_eq!(event.data, json!({"resume": false}));
drop(turn);
let recovered = DurableTurn::recover(session).unwrap();
assert_eq!(recovered.boxes().len(), 2);
assert_eq!(recovered.status(), Status::Quiet);
assert_eq!(
(recovered.mirrored, recovered.durable),
(recovered.boxes().len(), recovered.records.len())
);
}
#[test]
fn model_usage_round_trips_as_ordered_generic_events() {
let fixture = Fixture::new(5);
let session = fixture.session();
let mut turn = DurableTurn::recover(session.clone()).unwrap();
let terminal = completed(&mut turn);
turn.record_model_usage(terminal, usage()).unwrap();
turn.record_model_usage(0, usage()).unwrap();
let events = turn.events();
assert_eq!(events.len(), 3);
assert_eq!(
(
events[1].after_box_id,
events[1].event_index,
events[1].connected_box_id
),
(terminal, 2, terminal)
);
assert_eq!(events[1].handler, "model_usage");
assert_eq!(
events[1].data,
json!({"version": 1, "provider": "provider", "model": "model", "context_id": "context", "provider_turn_id": "turn", "usage": {"input_tokens": 10, "cached_input_tokens": 2, "cache_write_input_tokens": 3, "output_tokens": 4, "reasoning_output_tokens": 5, "total_tokens": 14}, "cumulative_usage": {"input_tokens": 10, "cached_input_tokens": 2, "cache_write_input_tokens": 3, "output_tokens": 4, "reasoning_output_tokens": 5, "total_tokens": 14}, "context_limit_tokens": 128})
);
assert_eq!(events[1].data, events[2].data);
assert_eq!(events[2].event_index, 3);
drop(turn);
let recovered = DurableTurn::recover(session.clone()).unwrap();
assert_eq!(recovered.events(), events);
assert_eq!(
session
.load()
.unwrap()
.records
.iter()
.filter(|record| matches!(record, Record::Event(_)))
.count(),
3
);
}
#[test]
fn model_usage_rejects_invalid_connections_and_resets_index_after_a_new_box() {
let fixture = Fixture::new(6);
let mut turn = DurableTurn::recover(fixture.session()).unwrap();
let terminal = completed(&mut turn);
assert!(
turn.record_model_usage(1, usage())
.unwrap_err()
.contains("Agent Response")
);
assert!(turn.record_model_usage(99, usage()).is_err());
turn.record_model_usage(terminal, usage()).unwrap();
turn.accept(
"User Message".into(),
"later".into(),
String::new(),
String::new(),
)
.unwrap();
turn.record_model_usage(0, usage()).unwrap();
let events = turn.events();
assert_eq!(
events
.iter()
.map(|event| (event.after_box_id, event.event_index))
.collect::<Vec<_>>(),
vec![(2, 1), (2, 2), (3, 1)]
);
}
#[test]
fn active_turn_mailbox_flush_continues_generation_without_resume() {
let fixture = Fixture::new(4);
let session = fixture.session();
let mut turn = DurableTurn::recover(session.clone()).unwrap();
turn.accept(
"User Message".into(),
"first".into(),
String::new(),
String::new(),
)
.unwrap();
let start = turn.begin().unwrap().unwrap();
assert_eq!(
start.values.last(),
Some(&BoxValue::History("[Box 2 | Agent Response]\n".into()))
);
turn.prepare_stage(start.job, "working".into(), Vec::new())
.unwrap();
turn.accept(
"User Message".into(),
"second".into(),
String::new(),
String::new(),
)
.unwrap();
let prepared = turn.prepare_mailbox_flush(start.job).unwrap().unwrap();
assert_eq!(
prepared.values().last(),
Some(&BoxValue::History("[Box 4 | Agent Response]\n".into()))
);
turn.commit_mailbox_flush(prepared).unwrap();
assert!(
!turn
.complete(start.job, ShimOutput { items: Vec::new() })
.unwrap()
);
assert_eq!(turn.status(), Status::Quiet);
assert!(turn.begin().unwrap().is_none());
let log = session.load().unwrap();
let Record::Event(event) = log.records.last().unwrap() else {
panic!("expected final llm_done event")
};
assert_eq!(event.handler, "llm_done");
}
#[test]
fn active_fifo_is_hidden_then_persisted_and_v1_return_is_idempotent() {
let fixture = Fixture::new(2);
let session = fixture.session();
let mut turn = DurableTurn::recover(session.clone()).unwrap();
turn.accept(
"User Message".into(),
"search".into(),
String::new(),
String::new(),
)
.unwrap();
let start = turn.begin().unwrap().unwrap();
let calls = turn
.prepare_stage(
start.job,
"working".into(),
vec![BoxValue::Call(Ok(Call {
name: "WebSearch".into(),
arguments: "{}".into(),
}))],
)
.unwrap();
let tool_call_id = calls[0].tool_call_id;
assert_eq!(session.load().unwrap().records.len(), 3);
turn.accept_tool_message(tool_call_id, "searching".into())
.unwrap();
turn.accept_tool_return_v2(
tool_call_id,
Ok("found".into()),
"k1.web-search-result/v1".into(),
"opaque".into(),
)
.unwrap();
turn.accept_tool_return(tool_call_id, Ok("duplicate".into()))
.unwrap();
assert_eq!(turn.boxes().len(), 3);
assert_eq!(session.load().unwrap().records.len(), 3);
let prepared = turn.prepare_mailbox_flush(start.job).unwrap().unwrap();
turn.validate_mailbox_flush(&prepared).unwrap();
assert_eq!(turn.boxes().len(), 5);
assert_eq!(session.load().unwrap().records.len(), 5);
assert_eq!((turn.mirrored, turn.durable, turn.records.len()), (5, 5, 5));
assert!(turn.boxes()[3].tool_message_metadata().unwrap().is_some());
assert!(turn.boxes()[4].tool_result_v2_metadata().unwrap().is_some());
turn.commit_mailbox_flush(prepared).unwrap();
}
#[test]
fn unresolved_search_messages_are_inert_and_result_v2_begins_once() {
let fixture = Fixture::new(3);
let session = fixture.session();
let mut turn = DurableTurn::recover(session.clone()).unwrap();
turn.accept(
"User Message".into(),
"search".into(),
String::new(),
String::new(),
)
.unwrap();
let start = turn.begin().unwrap().unwrap();
let calls = turn
.prepare_stage(
start.job,
"working".into(),
vec![BoxValue::Call(Ok(Call {
name: "WebSearch".into(),
arguments: "{}".into(),
}))],
)
.unwrap();
let tool_call_id = calls[0].tool_call_id;
let prepared = turn.prepare_mailbox_flush(start.job).unwrap().unwrap();
turn.commit_mailbox_flush(prepared).unwrap();
assert!(
!turn
.complete(start.job, ShimOutput { items: Vec::new() })
.unwrap()
);
assert_eq!(turn.status(), Status::Quiet);
let log = session.load().unwrap();
let Record::Event(event) = log.records.last().unwrap() else {
panic!("expected llm_done event")
};
assert_eq!(event.data, json!({"resume": false}));
turn.accept_tool_message(tool_call_id, "still searching".into())
.unwrap();
assert_eq!(turn.status(), Status::Quiet);
assert!(turn.begin().unwrap().is_none());
turn.accept_tool_return_v2(
tool_call_id,
Ok("found".into()),
"k1.web-search-result/v1".into(),
"opaque".into(),
)
.unwrap();
assert_eq!(turn.status(), Status::Running);
assert!(turn.begin().unwrap().is_some());
assert!(turn.begin().unwrap().is_none());
}
}