use serde_json::{Value, json};
use std::{
collections::{HashMap, HashSet},
fmt,
path::PathBuf,
sync::{Arc, Mutex},
};
pub type ToolToken = (String, u64, String);
#[derive(Clone, Debug)]
pub struct Config {
pub executable: PathBuf,
pub working_directory: String,
pub model: String,
pub reasoning_effort: Option<String>,
pub base_instructions: String,
pub tools: Vec<DynamicTool>,
}
impl Config {
pub fn validate(&self) -> Result<(), Error> {
validate_config(self)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct DynamicTool {
pub name: String,
pub description: String,
pub input_schema: Value,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ToolCall {
pub call_id: String,
pub name: String,
pub arguments: Value,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ToolResult {
pub success: bool,
pub output: String,
}
#[derive(Clone, Debug, PartialEq)]
pub enum Event {
TextDelta(String),
ToolCall(ToolCall),
Done,
Error(Error),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ErrorKind {
Busy,
Interrupted,
InvalidToolResult,
LaunchRejected,
Protocol,
Server,
Unavailable,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Error {
pub kind: ErrorKind,
pub message: String,
pub diagnostics: Vec<u8>,
}
impl Error {
pub fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
Self {
kind,
message: message.into(),
diagnostics: Vec::new(),
}
}
pub fn with_diagnostics(mut self, diagnostics: Vec<u8>) -> Self {
self.diagnostics = diagnostics;
self
}
pub fn server(value: &Value, diagnostics: &Diagnostics) -> Self {
let detail = value
.get("message")
.and_then(Value::as_str)
.or_else(|| value.pointer("/error/message").and_then(Value::as_str));
let message = detail.map_or_else(
|| "Codex app-server error".to_owned(),
|detail| format!("Codex app-server error: {detail}"),
);
diagnostics.error(ErrorKind::Server, message)
}
}
impl fmt::Display for Error {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for Error {}
#[derive(Clone, Debug, Default)]
pub struct Diagnostics(Arc<Mutex<Vec<u8>>>);
impl Diagnostics {
pub fn new(bytes: Vec<u8>) -> Self {
Self(Arc::new(Mutex::new(bytes)))
}
pub fn snapshot(&self) -> Vec<u8> {
self.0
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
pub fn replace(&self, bytes: Vec<u8>) {
*self
.0
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner()) = bytes;
}
pub fn error(&self, kind: ErrorKind, message: impl Into<String>) -> Error {
Error {
kind,
message: message.into(),
diagnostics: self.snapshot(),
}
}
}
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>,
}
pub struct Conversation<S> {
pub thread: Option<String>,
pub active: Option<Active<S>>,
pub closing: bool,
}
impl<S> Default for Conversation<S> {
fn default() -> Self {
Self {
thread: None,
active: None,
closing: false,
}
}
}
pub enum Pending<R> {
Thread {
key: String,
serial: u64,
input: String,
},
Turn {
key: String,
serial: u64,
},
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> {
let id = self.next_id;
self.next_id = id
.checked_add(1)
.ok_or_else(|| Error::new(ErrorKind::Protocol, "client request id space exhausted"))?;
Ok(id)
}
pub fn allocate_turn_id(&mut self) -> Result<u64, Error> {
let id = self.next_turn;
self.next_turn = id
.checked_add(1)
.ok_or_else(|| Error::new(ErrorKind::Protocol, "turn serial space exhausted"))?;
Ok(id)
}
pub fn begin_turn(&mut self, key: impl Into<String>, sink: S) -> Result<u64, Error> {
let key = key.into();
if self
.conversations
.get(&key)
.is_some_and(|conversation| conversation.active.is_some() || conversation.closing)
{
return Err(Error::new(
ErrorKind::Busy,
"conversation already has an active turn",
));
}
let serial = self.allocate_turn_id()?;
self.conversations.entry(key).or_default().active = Some(Active {
serial,
turn: None,
sink: Some(sink),
early: Vec::new(),
cancelled: false,
interrupt_sent: false,
failure: None,
});
Ok(serial)
}
pub fn take_active(&mut self, key: &str, serial: u64) -> Option<Active<S>> {
let conversation = self.conversations.get_mut(key)?;
if conversation
.active
.as_ref()
.is_some_and(|active| active.serial == serial)
{
conversation.active.take()
} else {
None
}
}
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(|active| active.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 thread = conversation.thread.clone()?;
let active = conversation.active.as_mut()?;
if active.serial != serial || active.interrupt_sent {
return None;
}
let turn = active.turn.clone()?;
active.interrupt_sent = true;
Some((thread, turn))
}
pub fn take_sinks(&mut self) -> Vec<S> {
self.conversations
.values_mut()
.filter_map(|conversation| conversation.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();
if self
.by_thread
.get(&thread)
.is_some_and(|owner| owner != key)
{
return Err(Error::new(
ErrorKind::Protocol,
"thread/start reused another conversation thread",
));
}
let old = self
.conversations
.entry(key.to_owned())
.or_default()
.thread
.replace(thread.clone());
if let Some(old) = old.filter(|old| 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);
};
if conversation.active.is_some() || conversation.closing {
return Err(Error::new(
ErrorKind::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)
.expect("conversation exists")
.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> {
let valid = self.conversations.get(key).is_some_and(|conversation| {
conversation.thread.as_deref() == Some(thread)
&& conversation.active.is_none()
&& conversation.closing
});
if !valid {
return Err(Error::new(
ErrorKind::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> {
if self.pending.insert(id, pending).is_some() {
return Err(Error::new(
ErrorKind::Protocol,
"duplicate client request id",
));
}
Ok(())
}
pub fn take_pending(&mut self, id: u64) -> Result<Pending<R>, Error> {
self.pending.remove(&id).ok_or_else(|| {
Error::new(
ErrorKind::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());
if self.rpc_ids.contains(&rpc_key) {
return Err(Error::new(
ErrorKind::Protocol,
"duplicate app-server request id",
));
}
if self.tools.contains_key(&token) {
return Err(Error::new(
ErrorKind::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)> {
let tokens: Vec<_> = self
.tools
.keys()
.filter(|(owner, turn, _)| owner == key && *turn == serial)
.cloned()
.collect();
tokens
.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)
}
}
pub fn validate_config(config: &Config) -> Result<(), Error> {
let mut names = HashSet::new();
for tool in &config.tools {
if tool.name.is_empty() {
return Err(Error::new(
ErrorKind::Protocol,
"dynamic tool names must not be empty",
));
}
if !names.insert(&tool.name) {
return Err(Error::new(
ErrorKind::Protocol,
format!("duplicate dynamic tool name: {}", tool.name),
));
}
}
Ok(())
}
pub fn thread_start_params(config: &Config) -> Value {
let tools: Vec<_> = config
.tools
.iter()
.map(|tool| {
json!({
"name": tool.name,
"description": tool.description,
"inputSchema": tool.input_schema
})
})
.collect();
json!({
"model": config.model,
"cwd": config.working_directory,
"approvalPolicy": "never",
"sandbox": "readOnly",
"baseInstructions": config.base_instructions,
"serviceName": "kcode-k1-codex-adapter",
"dynamicTools": tools
})
}
pub fn turn_start_params(thread: &str, input: impl Into<String>) -> Value {
json!({"threadId": thread, "input": [{"type": "text", "text": input.into()}]})
}
pub fn parse_scope(params: Option<&Value>) -> Result<(&str, &str), Error> {
let params =
params.ok_or_else(|| Error::new(ErrorKind::Protocol, "scoped message omitted params"))?;
let thread = params
.get("threadId")
.and_then(Value::as_str)
.ok_or_else(|| Error::new(ErrorKind::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(|| Error::new(ErrorKind::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(Error::new(
ErrorKind::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")
}