#![forbid(unsafe_code)]
#![doc = include_str!("../Documentation.md")]
use kcode_k1_chat_codex_state::{AGENT_RESPONSE_TYPE, ConversationState};
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_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(|r| {
if let Record::Event(e) = r {
Some(e.clone())
} else {
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 r = self
.state
.accept(box_type, contents, hidden_type, hidden_contents);
self.finish(r)
}
pub fn accept_idle_context_box(
&mut self,
box_type: String,
contents: String,
hidden_type: String,
hidden_contents: String,
) -> Result<(), String> {
let r =
self.state
.accept_idle_context_box(box_type, contents, hidden_type, hidden_contents);
self.finish(r)
}
pub fn accept_idle_context_tool_call(
&mut self,
name: String,
arguments: String,
) -> Result<ToolCallId, String> {
let r = self.state.accept_idle_context_tool_call(name, arguments);
self.finish(r)
}
pub fn accept_idle_context_tool_return(
&mut self,
id: ToolCallId,
result: Result<String, String>,
) -> Result<(), String> {
let r = self.state.accept_idle_context_tool_return(id, result);
if r.is_ok() {
self.returned.push(id);
}
self.finish(r)
}
pub fn accept_tool_return(
&mut self,
id: ToolCallId,
result: Result<String, String>,
) -> Result<(), String> {
if self.returned.contains(&id) {
return self.finish(Ok(()));
}
let r = self.state.accept_tool_return(id, result);
if r.is_ok() {
self.returned.push(id);
}
self.finish(r)
}
pub fn accept_tool_message(&mut self, id: ToolCallId, message: String) -> Result<(), String> {
let r = self.state.accept_tool_message(id, message);
self.finish(r)
}
pub fn accept_tool_return_v2(
&mut self,
id: ToolCallId,
result: Result<String, String>,
metadata_type: String,
metadata_contents: String,
) -> Result<(), String> {
let r = self
.state
.accept_tool_return_v2(id, result, metadata_type, metadata_contents);
if r.is_ok() {
self.returned.push(id);
}
self.finish(r)
}
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 r = self.state.prepare_stage(job, text, values);
self.finish(r)
}
pub fn prepare_mailbox_flush(
&mut self,
job: u64,
) -> Result<Option<PreparedMailboxFlush>, String> {
let r = self.state.prepare_mailbox_flush(job);
self.finish(r)
}
pub fn validate_mailbox_flush(&self, p: &PreparedMailboxFlush) -> Result<(), String> {
self.state.validate_mailbox_flush(p)
}
pub fn commit_mailbox_flush(&mut self, p: PreparedMailboxFlush) -> Result<(), String> {
self.state.commit_mailbox_flush(p)
}
pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
if let Err(e) = self.state.complete(job, output) {
return self.finish(Err(e));
}
self.mirror_boxes()?;
let resume = matches!(self.state.status(), Status::Running);
let after = self.latest_box_id()?;
self.persist_event(EventRecord {
after_box_id: after,
event_index: self.next_event_index(after)?,
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 = self.latest_box_id()?;
if connected_box_id != 0
&& !self
.state
.boxes()
.iter()
.any(|b| b.id().get() == connected_box_id && b.box_type() == AGENT_RESPONSE_TYPE)
{
return Err("model usage must connect to a canonical Agent Response box".into());
}
self.persist_event(EventRecord {
after_box_id: after,
event_index: self.next_event_index(after)?,
connected_box_id,
handler: "model_usage".into(),
data: model_usage_data(usage)?,
})
}
pub fn fail(&mut self, job: u64, message: String, restartable: bool) {
self.state.fail(job, message, restartable)
}
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(v), Ok(())) => Ok(v),
(Err(e), Ok(())) | (Ok(_), Err(e)) => Err(e),
(Err(a), Err(b)) => Err(format!(
"{a}; additionally failed to persist canonical history: {b}"
)),
}
}
fn mirror_and_persist(&mut self) -> Result<(), String> {
self.mirror_boxes()?;
self.persist_pending()
}
fn mirror_boxes(&mut self) -> Result<(), String> {
let additions = self
.state
.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 = self.state.boxes().len();
Ok(())
}
fn latest_box_id(&self) -> Result<u64, String> {
self.state
.boxes()
.last()
.map(|b| b.id().get())
.ok_or_else(|| "durable event requires a canonical box".into())
}
fn next_event_index(&self, after: u64) -> Result<u64, String> {
match self.records.last() {
Some(Record::Event(e)) if e.after_box_id == after => e
.event_index
.checked_add(1)
.ok_or_else(|| "durable event index space was exhausted".into()),
Some(Record::Event(_)) => {
Err("durable event frontier diverged from canonical boxes".into())
}
_ => Ok(1),
}
}
fn persist_event(&mut self, event: EventRecord) -> Result<(), String> {
let mut pending = self
.records
.get(self.durable..)
.ok_or_else(|| "durable record frontier moved past canonical records".to_owned())?
.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(|e| e.to_string())?;
let Value::Object(fields) = &mut data else {
return Err("model usage data did not serialize to an object".into());
};
fields.insert("version".into(), Value::from(1));
Ok(data)
}
fn returned_ids(boxes: &[ChatBox]) -> Result<Vec<ToolCallId>, String> {
let mut out = Vec::new();
for b in boxes {
if let Some(r) = b.tool_result_metadata().map_err(|e| format!("{e:?}"))? {
out.push(r.tool_call_id);
}
if let Some(r) = b.tool_result_v2_metadata().map_err(|e| format!("{e:?}"))? {
out.push(r.tool_call_id);
}
}
Ok(out)
}