pub use kcode_k1_codex_conversation_values::{
Config, Diagnostics, DynamicTool, Error, ErrorKind, Event, ToolCall, ToolResult,
validate_config,
};
use serde_json::{Value, json};
use std::collections::{HashMap, HashSet};
pub type ToolToken = (String, u64, String);
pub struct Active<S> {
pub serial: u64,
pub turn: Option<String>,
pub sink: Option<S>,
pub early: Vec<Value>,
pub cancelled: bool,
pub interrupt_sent: bool,
pub failure: Option<Value>,
}
#[derive(Default)]
pub struct Conversation<S> {
pub thread: Option<String>,
pub active: Option<Active<S>>,
pub closing: bool,
}
pub enum Pending<R> {
Thread {
key: String,
serial: u64,
input: String,
},
Turn {
key: String,
serial: u64,
},
Steer {
key: String,
serial: u64,
reply: R,
},
Close {
key: String,
thread: String,
reply: R,
},
Interrupt,
}
#[derive(Clone, Debug, PartialEq)]
pub struct PendingTool {
pub id: Value,
pub rpc_key: String,
}
pub struct State<S, R> {
pub conversations: HashMap<String, Conversation<S>>,
pub by_thread: HashMap<String, String>,
pub pending: HashMap<u64, Pending<R>>,
pub tools: HashMap<ToolToken, PendingTool>,
pub rpc_ids: HashSet<String>,
pub next_id: u64,
pub next_turn: u64,
}
impl<S, R> Default for State<S, R> {
fn default() -> Self {
Self {
conversations: HashMap::new(),
by_thread: HashMap::new(),
pending: HashMap::new(),
tools: HashMap::new(),
rpc_ids: HashSet::new(),
next_id: 1,
next_turn: 1,
}
}
}
impl<S, R> State<S, R> {
pub fn allocate_request_id(&mut self) -> Result<u64, Error> {
allocate(&mut self.next_id, "client request id space exhausted")
}
pub fn allocate_turn_id(&mut self) -> Result<u64, Error> {
allocate(&mut self.next_turn, "turn serial space exhausted")
}
pub fn begin_turn(&mut self, key: impl Into<String>, sink: S) -> Result<u64, Error> {
let key = key.into();
ensure(
!self
.conversations
.get(&key)
.is_some_and(|value| value.active.is_some() || value.closing),
busy("conversation already has an active turn"),
)?;
let serial = self.allocate_turn_id()?;
self.conversations
.entry(key)
.or_insert_with(empty_conversation)
.active = Some(Active {
serial,
turn: None,
sink: Some(sink),
early: Vec::new(),
cancelled: false,
interrupt_sent: false,
failure: None,
});
Ok(serial)
}
pub fn begin_steer(
&mut self,
key: &str,
serial: u64,
reply: R,
) -> Result<(u64, String, String), Error> {
let target = self
.conversations
.get(key)
.filter(|value| !value.closing)
.and_then(|conversation| {
conversation
.active
.as_ref()
.filter(|active| active.serial == serial)
.and_then(|active| conversation.thread.clone().zip(active.turn.clone()))
})
.ok_or_else(|| busy("conversation has no active native turn"))?;
ensure(
!self
.pending
.values()
.any(|value| matches!(value, Pending::Steer { key: owner, .. } if owner == key)),
busy("conversation already has a pending steer"),
)?;
ensure(
!self.pending.contains_key(&self.next_id),
protocol("duplicate client request id"),
)?;
let id = self.allocate_request_id()?;
self.pending.insert(
id,
Pending::Steer {
key: key.to_owned(),
serial,
reply,
},
);
Ok((id, target.0, target.1))
}
pub fn steer_pending(&self, key: &str, serial: u64) -> bool {
self.pending.values().any(
|value| matches!(value, Pending::Steer { key: owner, serial: pending, .. } if owner == key && *pending == serial),
)
}
pub fn take_active(&mut self, key: &str, serial: u64) -> Option<Active<S>> {
let active = &mut self.conversations.get_mut(key)?.active;
(active.as_ref()?.serial == serial)
.then(|| active.take())
.flatten()
}
pub fn set_native_turn(
&mut self,
key: &str,
serial: u64,
turn: impl Into<String>,
) -> Option<Vec<Value>> {
let active = self
.conversations
.get_mut(key)?
.active
.as_mut()
.filter(|value| value.serial == serial)?;
active.turn = Some(turn.into());
Some(std::mem::take(&mut active.early))
}
pub fn interrupt_target(&mut self, key: &str, serial: u64) -> Option<(String, String)> {
let conversation = self.conversations.get_mut(key)?;
let active = conversation.active.as_mut()?;
if active.serial != serial || active.interrupt_sent {
return None;
}
let target = (conversation.thread.clone()?, active.turn.clone()?);
active.interrupt_sent = true;
Some(target)
}
pub fn take_sinks(&mut self) -> Vec<S> {
self.conversations
.values_mut()
.filter_map(|value| value.active.take()?.sink)
.collect()
}
pub fn thread(&self, key: &str) -> Option<&str> {
self.conversations.get(key)?.thread.as_deref()
}
pub fn owner(&self, thread: &str) -> Option<&str> {
self.by_thread.get(thread).map(String::as_str)
}
pub fn set_thread(&mut self, key: &str, thread: impl Into<String>) -> Result<(), Error> {
let thread = thread.into();
ensure(
!self
.by_thread
.get(&thread)
.is_some_and(|owner| owner != key),
protocol("thread/start reused another conversation thread"),
)?;
let conversation = self
.conversations
.entry(key.to_owned())
.or_insert_with(empty_conversation);
if let Some(old) = conversation.thread.replace(thread.clone())
&& old != thread
{
self.by_thread.remove(&old);
}
self.by_thread.insert(thread, key.to_owned());
Ok(())
}
pub fn begin_close(&mut self, key: &str) -> Result<Option<String>, Error> {
let Some(conversation) = self.conversations.get(key) else {
return Ok(None);
};
ensure(
conversation.active.is_none() && !conversation.closing,
busy("conversation is active or already closing"),
)?;
let Some(thread) = conversation.thread.clone() else {
self.conversations.remove(key);
return Ok(None);
};
self.conversations.get_mut(key).unwrap().closing = true;
Ok(Some(thread))
}
pub fn cancel_close(&mut self, key: &str) {
if let Some(conversation) = self.conversations.get_mut(key) {
conversation.closing = false;
}
}
pub fn finish_close(&mut self, key: &str, thread: &str) -> Result<(), Error> {
ensure(
self.conversations.get(key).is_some_and(|value| {
value.thread.as_deref() == Some(thread) && value.active.is_none() && value.closing
}),
protocol("thread/unsubscribe response did not match closing conversation"),
)?;
self.by_thread.remove(thread);
self.conversations.remove(key);
Ok(())
}
pub fn insert_pending(&mut self, id: u64, pending: Pending<R>) -> Result<(), Error> {
ensure(
!self.pending.contains_key(&id),
protocol("duplicate client request id"),
)?;
self.pending.insert(id, pending);
Ok(())
}
pub fn take_pending(&mut self, id: u64) -> Result<Pending<R>, Error> {
self.pending
.remove(&id)
.ok_or_else(|| protocol("unexpected or duplicate app-server response id"))
}
pub fn track_tool(
&mut self,
key: &str,
serial: u64,
call: impl Into<String>,
id: &Value,
) -> Result<ToolToken, Error> {
let (id, rpc_key) = parse_rpc_id(id)?;
let token = (key.to_owned(), serial, call.into());
ensure(
!self.rpc_ids.contains(&rpc_key),
protocol("duplicate app-server request id"),
)?;
ensure(
!self.tools.contains_key(&token),
protocol("duplicate dynamic tool call id"),
)?;
self.rpc_ids.insert(rpc_key.clone());
self.tools
.insert(token.clone(), PendingTool { id, rpc_key });
Ok(token)
}
pub fn take_tool(&mut self, token: &ToolToken) -> Option<PendingTool> {
let pending = self.tools.remove(token)?;
self.rpc_ids.remove(&pending.rpc_key);
Some(pending)
}
pub fn take_turn_tools(&mut self, key: &str, serial: u64) -> Vec<(ToolToken, PendingTool)> {
self.tools
.keys()
.filter(|(owner, turn, _)| owner == key && *turn == serial)
.cloned()
.collect::<Vec<_>>()
.into_iter()
.filter_map(|token| self.take_tool(&token).map(|pending| (token, pending)))
.collect()
}
pub fn resolve_tool(&mut self, id: &Value) -> Result<Option<ToolToken>, Error> {
let (_, rpc_key) = parse_rpc_id(id)?;
let token = self
.tools
.iter()
.find_map(|(token, pending)| (pending.rpc_key == rpc_key).then(|| token.clone()));
if let Some(token) = &token {
self.take_tool(token);
}
Ok(token)
}
}
fn empty_conversation<S>() -> Conversation<S> {
Conversation {
thread: None,
active: None,
closing: false,
}
}
fn allocate(counter: &mut u64, exhausted: &'static str) -> Result<u64, Error> {
let id = *counter;
*counter = id.checked_add(1).ok_or_else(|| protocol(exhausted))?;
Ok(id)
}
fn ensure(valid: bool, error: Error) -> Result<(), Error> {
valid.then_some(()).ok_or(error)
}
fn busy(message: &'static str) -> Error {
Error::new(ErrorKind::Busy, message)
}
fn protocol(message: impl Into<String>) -> Error {
Error::new(ErrorKind::Protocol, message)
}
pub fn thread_start_params(config: &Config) -> Value {
let tools = config
.tools
.iter()
.map(|tool| {
json!({
"name": tool.name,
"description": tool.description,
"inputSchema": tool.input_schema
})
})
.collect::<Vec<_>>();
json!({
"model": config.model,
"cwd": config.working_directory,
"approvalPolicy": "never",
"sandbox": "read-only",
"baseInstructions": config.base_instructions,
"serviceName": "kcode-k1-codex-adapter",
"dynamicTools": tools
})
}
pub fn web_search_thread_start_params(config: &Config) -> Value {
json!({
"model": config.model,
"cwd": config.working_directory,
"approvalPolicy": "never",
"sandbox": "read-only",
"dynamicTools": [],
"ephemeral": true,
"config": {"web_search": "live"}
})
}
pub fn turn_start_params(thread: &str, input: impl Into<String>) -> Value {
json!({"threadId": thread, "input": [{"type": "text", "text": input.into()}]})
}
pub fn turn_steer_params(thread: &str, expected_turn: &str, input: impl Into<String>) -> Value {
json!({
"threadId": thread,
"expectedTurnId": expected_turn,
"input": [{"type": "text", "text": input.into()}]
})
}
pub fn parse_scope(params: Option<&Value>) -> Result<(&str, &str), Error> {
let params = params.ok_or_else(|| protocol("scoped message omitted params"))?;
let thread = params
.get("threadId")
.and_then(Value::as_str)
.ok_or_else(|| protocol("scoped message omitted threadId"))?;
let turn = params
.get("turnId")
.and_then(Value::as_str)
.or_else(|| params.pointer("/turn/id").and_then(Value::as_str))
.ok_or_else(|| protocol("scoped message omitted turn id"))?;
Ok((thread, turn))
}
pub fn parse_rpc_id(id: &Value) -> Result<(Value, String), Error> {
match id {
Value::String(value) => Ok((id.clone(), format!("s:{value}"))),
Value::Number(value) => Ok((id.clone(), format!("n:{value}"))),
_ => Err(protocol("server request id must be a string or number")),
}
}
pub fn is_model_reroute(method: &str) -> bool {
let method = method.to_ascii_lowercase();
method.contains("model") && method.contains("rerout")
}
#[cfg(test)]
mod tests {
use super::*;
fn active() -> (State<(), &'static str>, u64) {
let mut state = State::default();
state.set_thread("key", "thread-a").unwrap();
let serial = state.begin_turn("key", ()).unwrap();
state.set_native_turn("key", serial, "turn-b").unwrap();
(state, serial)
}
#[test]
fn protocol_and_steering_are_exact() {
let config = Config {
executable: "codex".into(),
working_directory: "/work".into(),
model: "model".into(),
reasoning_effort: None,
base_instructions: "ordinary instructions".into(),
tools: vec![DynamicTool {
name: "ordinary_tool".into(),
description: "ordinary description".into(),
input_schema: json!({"type": "object"}),
}],
};
let expected_ordinary = json!({
"model": "model",
"cwd": "/work",
"approvalPolicy": "never",
"sandbox": "read-only",
"baseInstructions": "ordinary instructions",
"serviceName": "kcode-k1-codex-adapter",
"dynamicTools": [{
"name": "ordinary_tool",
"description": "ordinary description",
"inputSchema": {"type": "object"}
}]
});
assert_eq!(thread_start_params(&config), expected_ordinary);
let expected_search = json!({
"model": "model",
"cwd": "/work",
"approvalPolicy": "never",
"sandbox": "read-only",
"dynamicTools": [],
"ephemeral": true,
"config": {"web_search": "live"}
});
assert_eq!(web_search_thread_start_params(&config), expected_search);
assert_eq!(
turn_steer_params("thread-a", "turn-b", "next"),
json!({"threadId": "thread-a", "expectedTurnId": "turn-b", "input": [{"type": "text", "text": "next"}]})
);
let (mut state, serial) = active();
let (id, thread, turn) = state.begin_steer("key", serial, "first").unwrap();
assert_eq!((thread.as_str(), turn.as_str()), ("thread-a", "turn-b"));
let before = (state.next_id, state.pending.len());
for (key, turn) in [("key", serial), ("wrong", serial), ("key", serial + 1)] {
assert_eq!(
state.begin_steer(key, turn, "bad").unwrap_err().kind,
ErrorKind::Busy
);
assert_eq!((state.next_id, state.pending.len()), before);
}
state.set_thread("other", "thread-c").unwrap();
let other = state.begin_turn("other", ()).unwrap();
state.set_native_turn("other", other, "turn-d").unwrap();
assert!(state.begin_steer("other", other, "other").is_ok());
assert!(
matches!(state.take_pending(id).unwrap(), Pending::Steer { key, serial: got, reply: "first" } if key == "key" && got == serial)
);
assert!(state.begin_steer("key", serial, "again").is_ok());
}
#[test]
fn steer_pending_matches_and_clears() {
let (mut state, serial) = active();
let id = state.begin_steer("key", serial, "reply").unwrap().0;
assert!(state.steer_pending("key", serial));
assert!(!state.steer_pending("other", serial));
assert!(!state.steer_pending("key", serial + 1));
state.take_pending(id).unwrap();
assert!(!state.steer_pending("key", serial));
}
#[test]
fn duplicate_pending_is_transactional() {
let mut state: State<(), ()> = State::default();
state.insert_pending(1, Pending::Interrupt).unwrap();
assert!(state.insert_pending(1, Pending::Interrupt).is_err());
assert!(matches!(state.take_pending(1).unwrap(), Pending::Interrupt));
}
}