kcode-codex-terra 0.1.1

use std::{
    collections::{BTreeMap, VecDeque},
    fmt,
    path::PathBuf,
    time::Duration,
};

use kcode_codex_terra_protocol as protocol;
use kcode_codex_terra_usage::UsageAccumulator;
use kcode_jsonrpc_stdio::{Error as TransportError, IncomingMessage, PeerId, RequestId, StdioRpc};
use kcode_k1_accounting::{Accounting, AccountingEvent, UsageValue};
use serde_json::Value;

const OPERATION: &str = "run tool";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ErrorKind {
    InvalidInput,
    Unavailable,
    Timeout,
    Protocol,
}

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

impl Error {
    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: message.into(),
        }
    }

    pub const fn kind(&self) -> ErrorKind {
        self.kind
    }

    pub fn message(&self) -> &str {
        &self.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 {}

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Clone, Debug, PartialEq)]
pub struct ToolRun {
    pub input: String,
    pub tool_name: String,
    pub tool_description: String,
    pub input_schema: Value,
}

#[derive(Clone, Debug, PartialEq)]
pub struct ToolRunResult {
    pub arguments: Value,
    pub thread_id: String,
    pub turn_id: String,
    pub usage: BTreeMap<String, UsageValue>,
}

#[derive(Clone)]
pub struct CodexTerra {
    accounting: Accounting,
    executable: PathBuf,
    working_directory: PathBuf,
    timeout: Duration,
}

impl CodexTerra {
    pub fn new(
        accounting: Accounting,
        executable: impl Into<PathBuf>,
        working_directory: impl Into<PathBuf>,
        timeout: Duration,
    ) -> Result<Self> {
        let executable = executable.into();
        if executable.as_os_str().is_empty() || timeout.is_zero() {
            return Err(invalid("executable must be nonempty and timeout nonzero"));
        }
        Ok(Self {
            accounting,
            executable,
            working_directory: working_directory.into(),
            timeout,
        })
    }

    pub async fn run(&self, run: ToolRun) -> Result<ToolRunResult> {
        validate_run(&run)?;
        match tokio::time::timeout(self.timeout, execute(self, run)).await {
            Ok(result) => result,
            Err(_) => Err(Error::new(ErrorKind::Timeout, "Codex tool run timed out")),
        }
    }
}

struct AttemptAccounting {
    accounting: Accounting,
    usage: UsageAccumulator,
}

impl AttemptAccounting {
    fn new(accounting: Accounting) -> Self {
        Self {
            accounting,
            usage: UsageAccumulator::new(),
        }
    }

    fn apply(&mut self, value: &Value) -> Result<()> {
        self.usage
            .apply(value)
            .map_err(|error| protocol_error(error.to_string()))
    }

    fn reconciled_rounds(&self) -> usize {
        self.usage.reconciled_rounds()
    }

    fn snapshot(&self) -> BTreeMap<String, UsageValue> {
        self.usage.snapshot()
    }
}

impl Drop for AttemptAccounting {
    fn drop(&mut self) {
        self.accounting.record(&AccountingEvent {
            source: protocol::MODEL.to_owned(),
            operation: OPERATION.to_owned(),
            usage: self.usage.snapshot(),
        });
    }
}

fn validate_run(run: &ToolRun) -> Result<()> {
    let name = run.tool_name.as_bytes();
    if name.is_empty()
        || name.len() > 64
        || !name
            .iter()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
    {
        return Err(invalid("dynamic tool name is invalid"));
    }
    if run.tool_description.is_empty() || !run.input_schema.is_object() {
        return Err(invalid("dynamic tool metadata is invalid"));
    }
    Ok(())
}

async fn execute(client: &CodexTerra, run: ToolRun) -> Result<ToolRunResult> {
    let validator = jsonschema::validator_for(&run.input_schema)
        .map_err(|_| invalid("dynamic tool schema could not be compiled"))?;
    let mut rpc = StdioRpc::spawn(protocol::app_server_command(
        &client.executable,
        &client.working_directory,
    ))
    .map_err(|_| unavailable("Codex app-server could not be started"))?;
    let mut inbox = Inbox::default();
    let result = async {
        let request = rpc
            .send_request(
                "initialize",
                protocol::initialize_params(env!("CARGO_PKG_VERSION")),
            )
            .await
            .map_err(map_transport)?;
        inbox.response(&mut rpc, request, "initialize").await?;
        rpc.send_notification("initialized", serde_json::json!({}))
            .await
            .map_err(map_transport)?;

        let request = rpc
            .send_request(
                "thread/start",
                protocol::thread_start_params(
                    &client.working_directory,
                    run.tool_name.clone(),
                    run.tool_description.clone(),
                    run.input_schema.clone(),
                ),
            )
            .await
            .map_err(map_transport)?;
        let response = inbox.response(&mut rpc, request, "thread/start").await?;
        let thread_id = required(&response, "/thread/id", "thread ID")?.to_owned();

        let request = rpc
            .send_request(
                "turn/start",
                protocol::turn_start_params(&thread_id, run.input),
            )
            .await
            .map_err(map_transport)?;
        let mut accounting = AttemptAccounting::new(client.accounting.clone());
        let response = inbox.response(&mut rpc, request, "turn/start").await?;
        let turn_id = required(&response, "/turn/id", "turn ID")?.to_owned();
        let mut arguments = None;
        let mut rounds_at_call = None;

        loop {
            match inbox.next(&mut rpc).await? {
                IncomingMessage::Request { id, method, params } => {
                    if method != "item/tool/call" || arguments.is_some() {
                        return Err(protocol_error("unexpected or repeated Codex request"));
                    }
                    require_scope(&params, &thread_id, &turn_id)?;
                    let call_id = required(&params, "/callId", "tool call ID")?;
                    if call_id.is_empty()
                        || required(&params, "/tool", "tool name")? != run.tool_name
                    {
                        return Err(protocol_error(
                            "Codex supplied an invalid dynamic tool call",
                        ));
                    }
                    let value = params
                        .get("arguments")
                        .cloned()
                        .ok_or_else(|| protocol_error("Codex omitted dynamic tool arguments"))?;
                    if !validator.is_valid(&value) {
                        return Err(protocol_error("Codex tool arguments failed schema"));
                    }
                    let rounds = accounting.reconciled_rounds();
                    if rounds == 0 {
                        return Err(protocol_error(
                            "Codex called the tool before reporting usage",
                        ));
                    }
                    rpc.respond(id, protocol::tool_success_result())
                        .await
                        .map_err(map_transport)?;
                    arguments = Some(value);
                    rounds_at_call = Some(rounds);
                }
                IncomingMessage::Notification { method, params } => {
                    match protocol::classify_notification(&method, &params, &thread_id, &turn_id)
                        .map_err(|error| protocol_error(error.to_string()))?
                    {
                        protocol::NotificationKind::Usage => {
                            accounting.apply(&params["tokenUsage"])?;
                        }
                        protocol::NotificationKind::TurnCompleted => {
                            let arguments = arguments.ok_or_else(|| {
                                protocol_error("Codex completed without calling the tool")
                            })?;
                            if rounds_at_call
                                .is_none_or(|rounds| accounting.reconciled_rounds() <= rounds)
                            {
                                return Err(protocol_error(
                                    "Codex completed without usage after the tool response",
                                ));
                            }
                            let usage = accounting.snapshot();
                            if usage.is_empty() {
                                return Err(protocol_error(
                                    "Codex completed without terminal usage",
                                ));
                            }
                            return Ok(ToolRunResult {
                                arguments,
                                thread_id,
                                turn_id,
                                usage,
                            });
                        }
                        protocol::NotificationKind::Continue => {}
                    }
                }
                IncomingMessage::Response { .. } => {
                    return Err(protocol_error("Codex emitted an unexpected response"));
                }
            }
        }
    }
    .await;
    let shutdown = rpc.shutdown().await.map_err(map_transport);
    match result {
        Err(error) => Err(error),
        Ok(value) => shutdown.map(|_| value),
    }
}

#[derive(Default)]
struct Inbox {
    queued: VecDeque<IncomingMessage>,
}

impl Inbox {
    async fn response(
        &mut self,
        rpc: &mut StdioRpc,
        expected: RequestId,
        label: &str,
    ) -> Result<Value> {
        loop {
            match rpc.next().await.map_err(map_transport)? {
                IncomingMessage::Response { id, result } => {
                    if id != PeerId::Number(expected.0.into()) {
                        return Err(protocol_error(format!("wrong response ID for {label}")));
                    }
                    return result.map_err(|_| protocol_error(format!("Codex {label} failed")));
                }
                message => self.queued.push_back(message),
            }
        }
    }

    async fn next(&mut self, rpc: &mut StdioRpc) -> Result<IncomingMessage> {
        match self.queued.pop_front() {
            Some(message) => Ok(message),
            None => rpc.next().await.map_err(map_transport),
        }
    }
}

fn require_id(value: &Value, pointer: &str, expected: &str) -> Result<()> {
    (value.pointer(pointer).and_then(Value::as_str) == Some(expected))
        .then_some(())
        .ok_or_else(|| protocol_error("Codex used a mismatched identifier"))
}

fn require_scope(value: &Value, thread: &str, turn: &str) -> Result<()> {
    require_id(value, "/threadId", thread)?;
    require_id(value, "/turnId", turn)
}

fn required<'a>(value: &'a Value, pointer: &str, label: &str) -> Result<&'a str> {
    value
        .pointer(pointer)
        .and_then(Value::as_str)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| protocol_error(format!("Codex omitted or emptied {label}")))
}

fn map_transport(error: TransportError) -> Error {
    match error {
        TransportError::Json(_)
        | TransportError::InvalidMessage(_)
        | TransportError::InboundLineTooLong => protocol_error("Codex emitted invalid JSON-RPC"),
        _ => unavailable("Codex app-server transport became unavailable"),
    }
}

fn invalid(message: impl Into<String>) -> Error {
    Error::new(ErrorKind::InvalidInput, message)
}

fn unavailable(message: impl Into<String>) -> Error {
    Error::new(ErrorKind::Unavailable, message)
}

fn protocol_error(message: impl Into<String>) -> Error {
    Error::new(ErrorKind::Protocol, message)
}