kcode-codex-terra 0.1.0

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

use kcode_jsonrpc_stdio::{Error as TransportError, IncomingMessage, PeerId, RequestId, StdioRpc};
use kcode_k1_accounting::{Accounting, UsageValue};
use serde_json::{Value, json};
use tokio::process::Command;

mod usage;

use usage::AttemptAccounting;

const MODEL: &str = "gpt-5.6-terra";
const OPERATION: &str = "run tool";
const DEVELOPER_INSTRUCTION: &str =
    "Call the supplied dynamic function exactly once. Do not produce assistant prose.";
const CODEX_CONFIG: &str = "web_search=\"disabled\"|mcp_servers={}|features.shell_tool=false|features.apps=false|features.browser_use=false|features.computer_use=false|features.goals=false|features.hooks=false|features.image_generation=false|features.multi_agent=false|features.plugins=false|features.tool_suggest=false|features.remote_plugin=false|model_auto_compact_token_limit=9223372036854775807";
const DISABLED_EVENTS: &str = "commandExecution|fileChange|mcpToolCall|webSearch|imageView|imageGeneration|collabAgentToolCall|subAgentActivity";

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

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(app_server_command(client))
        .map_err(|_| unavailable("Codex app-server could not be started"))?;
    let mut inbox = Inbox::default();
    let result = async {
        let request = rpc
            .send_request(
                "initialize",
                json!({
                    "clientInfo":{"name":"kcode-codex-terra","version":env!("CARGO_PKG_VERSION")},
                    "capabilities":{"experimentalApi":true}
                }),
            )
            .await
            .map_err(map_transport)?;
        inbox.response(&mut rpc, request, "initialize").await?;
        rpc.send_notification("initialized", json!({}))
            .await
            .map_err(map_transport)?;

        let request = rpc
            .send_request(
                "thread/start",
                json!({
                    "model":MODEL,
                    "cwd":client.working_directory,
                    "approvalPolicy":"never",
                    "sandbox":"readOnly",
                    "baseInstructions":"",
                    "developerInstructions":DEVELOPER_INSTRUCTION,
                    "dynamicTools":[{
                        "type":"function",
                        "name":run.tool_name,
                        "description":run.tool_description,
                        "inputSchema":run.input_schema
                    }],
                    "ephemeral":true,
                    "environments":[]
                }),
            )
            .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",
                json!({
                    "threadId":thread_id,
                    "input":[{"type":"text","text":run.input}],
                    "approvalPolicy":"never"
                }),
            )
            .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("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("Codex supplied an invalid dynamic tool call"));
                    }
                    let value = params
                        .get("arguments")
                        .cloned()
                        .ok_or_else(|| protocol("Codex omitted dynamic tool arguments"))?;
                    if !validator.is_valid(&value) {
                        return Err(protocol("Codex tool arguments failed schema"));
                    }
                    let rounds = accounting.reconciled_rounds();
                    if rounds == 0 {
                        return Err(protocol("Codex called the tool before reporting usage"));
                    }
                    rpc.respond(
                        id,
                        json!({"success":true,"contentItems":[{"type":"inputText","text":"ok"}]}),
                    )
                    .await
                    .map_err(map_transport)?;
                    arguments = Some(value);
                    rounds_at_call = Some(rounds);
                }
                IncomingMessage::Notification { method, params } => match method.as_str() {
                    "thread/started" => require_id(&params, "/thread/id", &thread_id)?,
                    "thread/tokenUsage/updated" => {
                        require_scope(&params, &thread_id, &turn_id)?;
                        accounting.apply(
                            params
                                .get("tokenUsage")
                                .ok_or_else(|| protocol("Codex omitted token usage"))?,
                        )?;
                    }
                    "turn/started" => require_turn_object(&params, &thread_id, &turn_id)?,
                    "item/started" | "item/completed" => {
                        require_scope(&params, &thread_id, &turn_id)?;
                        validate_item(&params)?;
                    }
                    "turn/completed" => {
                        require_turn_object(&params, &thread_id, &turn_id)?;
                        if params.pointer("/turn/status").and_then(Value::as_str)
                            != Some("completed")
                        {
                            return Err(protocol("Codex turn failed"));
                        }
                        let arguments = arguments
                            .ok_or_else(|| protocol("Codex completed without calling the tool"))?;
                        if rounds_at_call
                            .is_none_or(|rounds| accounting.reconciled_rounds() <= rounds)
                        {
                            return Err(protocol(
                                "Codex completed without usage after the tool response",
                            ));
                        }
                        let usage = accounting.snapshot();
                        if usage.is_empty() {
                            return Err(protocol("Codex completed without terminal usage"));
                        }
                        return Ok(ToolRunResult {
                            arguments,
                            thread_id,
                            turn_id,
                            usage,
                        });
                    }
                    _ => validate_notification(&method, &params, &thread_id, &turn_id)?,
                },
                IncomingMessage::Response { .. } => {
                    return Err(protocol("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(format!("wrong response ID for {label}")));
                    }
                    return result.map_err(|_| protocol(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 app_server_command(client: &CodexTerra) -> Command {
    let mut command = Command::new(&client.executable);
    for value in CODEX_CONFIG.split('|') {
        command.arg("-c").arg(value);
    }
    command
        .args(["app-server", "--stdio"])
        .current_dir(&client.working_directory)
        .env_remove("OPENAI_API_KEY")
        .env_remove("CODEX_API_KEY");
    command
}

fn validate_notification(method: &str, params: &Value, thread: &str, turn: &str) -> Result<()> {
    if method.contains("rerout") || DISABLED_EVENTS.split('|').any(|kind| method.contains(kind)) {
        return Err(protocol("Codex attempted a disabled capability"));
    }
    if matches!(
        method,
        "model/safetyBuffering/updated" | "model/verification"
    ) {
        return require_scope(params, thread, turn);
    }
    if method.starts_with("thread/") {
        return require_id(params, "/threadId", thread);
    }
    if method.starts_with("turn/") || method.starts_with("item/") {
        require_scope(params, thread, turn)?;
        if params.get("item").is_some() {
            validate_item(params)?;
        }
        return Ok(());
    }
    Err(protocol("Codex emitted an unexpected event"))
}

fn validate_item(params: &Value) -> Result<()> {
    match params.pointer("/item/type").and_then(Value::as_str) {
        Some("userMessage" | "agentMessage" | "reasoning" | "dynamicToolCall") => Ok(()),
        Some(_) => Err(protocol("Codex attempted a disabled built-in tool")),
        None => Err(protocol("Codex item event omitted its item type")),
    }
}

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("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 require_turn_object(value: &Value, thread: &str, turn: &str) -> Result<()> {
    require_id(value, "/threadId", thread)?;
    require_id(value, "/turn/id", 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(format!("Codex omitted or emptied {label}")))
}

fn map_transport(error: TransportError) -> Error {
    match error {
        TransportError::Json(_)
        | TransportError::InvalidMessage(_)
        | TransportError::InboundLineTooLong => protocol("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(message: impl Into<String>) -> Error {
    Error::new(ErrorKind::Protocol, message)
}