kcode-k1-codex-adapter 0.1.0

mod runtime;
mod shim;

use serde_json::Value;
use std::{
    fmt,
    path::PathBuf,
    sync::{Arc, Mutex, Weak},
};
use tokio::sync::{mpsc, oneshot};

pub use shim::{ASYNC_TOOL_ACKNOWLEDGEMENT, BoxCodec, Shim, ShimItem, ShimOutput};

#[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>,
}

#[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,
    Protocol,
    Server,
    Unavailable,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Error {
    pub kind: ErrorKind,
    pub message: String,
    pub diagnostics: Vec<u8>,
}

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 {}

type Diagnostics = Arc<Mutex<Vec<u8>>>;
type Reply = oneshot::Sender<Result<(), Error>>;
type StartReply = oneshot::Sender<Result<u64, Error>>;
type Events = mpsc::UnboundedSender<Event>;

enum Command {
    Start {
        key: String,
        input: String,
        events: Events,
        reply: StartReply,
    },
    ToolResult {
        key: String,
        turn: u64,
        call_id: String,
        result: ToolResult,
        reply: Reply,
    },
    Abandon {
        key: String,
        turn: u64,
    },
}

struct Client {
    commands: mpsc::UnboundedSender<Command>,
    diagnostics: Diagnostics,
}

#[derive(Clone)]
pub struct Adapter {
    client: Arc<Client>,
}

impl Adapter {
    pub async fn open(config: Config) -> Result<Self, Error> {
        runtime::open(config).await
    }

    pub async fn start_turn(
        &self,
        conversation_key: impl Into<String>,
        input: impl Into<String>,
    ) -> Result<Turn, Error> {
        let key = conversation_key.into();
        let (events, receiver) = mpsc::unbounded_channel();
        let (reply, answer) = oneshot::channel();
        self.client
            .commands
            .send(Command::Start {
                key: key.clone(),
                input: input.into(),
                events,
                reply,
            })
            .map_err(|_| self.unavailable())?;
        let serial = answer.await.map_err(|_| self.unavailable())??;
        Ok(Turn {
            key,
            serial,
            client: Arc::downgrade(&self.client),
            events: receiver,
            diagnostics: self.client.diagnostics.clone(),
            terminal: false,
        })
    }

    pub fn diagnostics(&self) -> Vec<u8> {
        snapshot(&self.client.diagnostics)
    }

    fn unavailable(&self) -> Error {
        fault(
            ErrorKind::Unavailable,
            "Codex app-server is unavailable",
            &self.client.diagnostics,
        )
    }
}

pub struct Turn {
    key: String,
    serial: u64,
    client: Weak<Client>,
    events: mpsc::UnboundedReceiver<Event>,
    diagnostics: Diagnostics,
    terminal: bool,
}

impl fmt::Debug for Turn {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Turn")
            .field("key", &self.key)
            .field("serial", &self.serial)
            .finish_non_exhaustive()
    }
}

impl Turn {
    pub async fn next_event(&mut self) -> Option<Event> {
        let event = self.events.recv().await;
        if event
            .as_ref()
            .is_none_or(|event| matches!(event, Event::Done | Event::Error(_)))
        {
            self.terminal = true;
        }
        event
    }

    pub async fn respond(
        &self,
        call_id: impl Into<String>,
        result: ToolResult,
    ) -> Result<(), Error> {
        let client = self.client.upgrade().ok_or_else(|| {
            fault(
                ErrorKind::Unavailable,
                "Codex app-server is unavailable",
                &self.diagnostics,
            )
        })?;
        let (reply, answer) = oneshot::channel();
        client
            .commands
            .send(Command::ToolResult {
                key: self.key.clone(),
                turn: self.serial,
                call_id: call_id.into(),
                result,
                reply,
            })
            .map_err(|_| {
                fault(
                    ErrorKind::Unavailable,
                    "Codex app-server is unavailable",
                    &self.diagnostics,
                )
            })?;
        answer.await.map_err(|_| {
            fault(
                ErrorKind::Unavailable,
                "Codex app-server is unavailable",
                &self.diagnostics,
            )
        })?
    }
}

impl Drop for Turn {
    fn drop(&mut self) {
        if !self.terminal
            && let Some(client) = self.client.upgrade()
        {
            let _ = client.commands.send(Command::Abandon {
                key: self.key.clone(),
                turn: self.serial,
            });
        }
    }
}

fn snapshot(diagnostics: &Diagnostics) -> Vec<u8> {
    diagnostics
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .clone()
}

fn fault(kind: ErrorKind, message: impl Into<String>, diagnostics: &Diagnostics) -> Error {
    Error {
        kind,
        message: message.into(),
        diagnostics: snapshot(diagnostics),
    }
}