kcode_codex_terra/
lib.rs

1use std::{
2    collections::{BTreeMap, VecDeque},
3    fmt,
4    path::PathBuf,
5    time::Duration,
6};
7
8use kcode_jsonrpc_stdio::{Error as TransportError, IncomingMessage, PeerId, RequestId, StdioRpc};
9use kcode_k1_accounting::{Accounting, UsageValue};
10use serde_json::{Value, json};
11use tokio::process::Command;
12
13mod usage;
14
15use usage::AttemptAccounting;
16
17const MODEL: &str = "gpt-5.6-terra";
18const OPERATION: &str = "run tool";
19const DEVELOPER_INSTRUCTION: &str =
20    "Call the supplied dynamic function exactly once. Do not produce assistant prose.";
21const 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";
22const DISABLED_EVENTS: &str = "commandExecution|fileChange|mcpToolCall|webSearch|imageView|imageGeneration|collabAgentToolCall|subAgentActivity";
23
24#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub enum ErrorKind {
26    InvalidInput,
27    Unavailable,
28    Timeout,
29    Protocol,
30}
31
32#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct Error {
34    kind: ErrorKind,
35    message: String,
36}
37
38impl Error {
39    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
40        Self {
41            kind,
42            message: message.into(),
43        }
44    }
45
46    pub const fn kind(&self) -> ErrorKind {
47        self.kind
48    }
49
50    pub fn message(&self) -> &str {
51        &self.message
52    }
53}
54
55impl fmt::Display for Error {
56    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
57        formatter.write_str(&self.message)
58    }
59}
60
61impl std::error::Error for Error {}
62
63pub type Result<T> = std::result::Result<T, Error>;
64
65#[derive(Clone, Debug, PartialEq)]
66pub struct ToolRun {
67    pub input: String,
68    pub tool_name: String,
69    pub tool_description: String,
70    pub input_schema: Value,
71}
72
73#[derive(Clone, Debug, PartialEq)]
74pub struct ToolRunResult {
75    pub arguments: Value,
76    pub thread_id: String,
77    pub turn_id: String,
78    pub usage: BTreeMap<String, UsageValue>,
79}
80
81#[derive(Clone)]
82pub struct CodexTerra {
83    accounting: Accounting,
84    executable: PathBuf,
85    working_directory: PathBuf,
86    timeout: Duration,
87}
88
89impl CodexTerra {
90    pub fn new(
91        accounting: Accounting,
92        executable: impl Into<PathBuf>,
93        working_directory: impl Into<PathBuf>,
94        timeout: Duration,
95    ) -> Result<Self> {
96        let executable = executable.into();
97        if executable.as_os_str().is_empty() || timeout.is_zero() {
98            return Err(invalid("executable must be nonempty and timeout nonzero"));
99        }
100        Ok(Self {
101            accounting,
102            executable,
103            working_directory: working_directory.into(),
104            timeout,
105        })
106    }
107
108    pub async fn run(&self, run: ToolRun) -> Result<ToolRunResult> {
109        validate_run(&run)?;
110        match tokio::time::timeout(self.timeout, execute(self, run)).await {
111            Ok(result) => result,
112            Err(_) => Err(Error::new(ErrorKind::Timeout, "Codex tool run timed out")),
113        }
114    }
115}
116
117fn validate_run(run: &ToolRun) -> Result<()> {
118    let name = run.tool_name.as_bytes();
119    if name.is_empty()
120        || name.len() > 64
121        || !name
122            .iter()
123            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
124    {
125        return Err(invalid("dynamic tool name is invalid"));
126    }
127    if run.tool_description.is_empty() || !run.input_schema.is_object() {
128        return Err(invalid("dynamic tool metadata is invalid"));
129    }
130    Ok(())
131}
132
133async fn execute(client: &CodexTerra, run: ToolRun) -> Result<ToolRunResult> {
134    let validator = jsonschema::validator_for(&run.input_schema)
135        .map_err(|_| invalid("dynamic tool schema could not be compiled"))?;
136    let mut rpc = StdioRpc::spawn(app_server_command(client))
137        .map_err(|_| unavailable("Codex app-server could not be started"))?;
138    let mut inbox = Inbox::default();
139    let result = async {
140        let request = rpc
141            .send_request(
142                "initialize",
143                json!({
144                    "clientInfo":{"name":"kcode-codex-terra","version":env!("CARGO_PKG_VERSION")},
145                    "capabilities":{"experimentalApi":true}
146                }),
147            )
148            .await
149            .map_err(map_transport)?;
150        inbox.response(&mut rpc, request, "initialize").await?;
151        rpc.send_notification("initialized", json!({}))
152            .await
153            .map_err(map_transport)?;
154
155        let request = rpc
156            .send_request(
157                "thread/start",
158                json!({
159                    "model":MODEL,
160                    "cwd":client.working_directory,
161                    "approvalPolicy":"never",
162                    "sandbox":"readOnly",
163                    "baseInstructions":"",
164                    "developerInstructions":DEVELOPER_INSTRUCTION,
165                    "dynamicTools":[{
166                        "type":"function",
167                        "name":run.tool_name,
168                        "description":run.tool_description,
169                        "inputSchema":run.input_schema
170                    }],
171                    "ephemeral":true,
172                    "environments":[]
173                }),
174            )
175            .await
176            .map_err(map_transport)?;
177        let response = inbox.response(&mut rpc, request, "thread/start").await?;
178        let thread_id = required(&response, "/thread/id", "thread ID")?.to_owned();
179
180        let request = rpc
181            .send_request(
182                "turn/start",
183                json!({
184                    "threadId":thread_id,
185                    "input":[{"type":"text","text":run.input}],
186                    "approvalPolicy":"never"
187                }),
188            )
189            .await
190            .map_err(map_transport)?;
191        let mut accounting = AttemptAccounting::new(client.accounting.clone());
192        let response = inbox.response(&mut rpc, request, "turn/start").await?;
193        let turn_id = required(&response, "/turn/id", "turn ID")?.to_owned();
194        let mut arguments = None;
195        let mut rounds_at_call = None;
196
197        loop {
198            match inbox.next(&mut rpc).await? {
199                IncomingMessage::Request { id, method, params } => {
200                    if method != "item/tool/call" || arguments.is_some() {
201                        return Err(protocol("unexpected or repeated Codex request"));
202                    }
203                    require_scope(&params, &thread_id, &turn_id)?;
204                    let call_id = required(&params, "/callId", "tool call ID")?;
205                    if call_id.is_empty()
206                        || required(&params, "/tool", "tool name")? != run.tool_name
207                    {
208                        return Err(protocol("Codex supplied an invalid dynamic tool call"));
209                    }
210                    let value = params
211                        .get("arguments")
212                        .cloned()
213                        .ok_or_else(|| protocol("Codex omitted dynamic tool arguments"))?;
214                    if !validator.is_valid(&value) {
215                        return Err(protocol("Codex tool arguments failed schema"));
216                    }
217                    let rounds = accounting.reconciled_rounds();
218                    if rounds == 0 {
219                        return Err(protocol("Codex called the tool before reporting usage"));
220                    }
221                    rpc.respond(
222                        id,
223                        json!({"success":true,"contentItems":[{"type":"inputText","text":"ok"}]}),
224                    )
225                    .await
226                    .map_err(map_transport)?;
227                    arguments = Some(value);
228                    rounds_at_call = Some(rounds);
229                }
230                IncomingMessage::Notification { method, params } => match method.as_str() {
231                    "thread/started" => require_id(&params, "/thread/id", &thread_id)?,
232                    "thread/tokenUsage/updated" => {
233                        require_scope(&params, &thread_id, &turn_id)?;
234                        accounting.apply(
235                            params
236                                .get("tokenUsage")
237                                .ok_or_else(|| protocol("Codex omitted token usage"))?,
238                        )?;
239                    }
240                    "turn/started" => require_turn_object(&params, &thread_id, &turn_id)?,
241                    "item/started" | "item/completed" => {
242                        require_scope(&params, &thread_id, &turn_id)?;
243                        validate_item(&params)?;
244                    }
245                    "turn/completed" => {
246                        require_turn_object(&params, &thread_id, &turn_id)?;
247                        if params.pointer("/turn/status").and_then(Value::as_str)
248                            != Some("completed")
249                        {
250                            return Err(protocol("Codex turn failed"));
251                        }
252                        let arguments = arguments
253                            .ok_or_else(|| protocol("Codex completed without calling the tool"))?;
254                        if rounds_at_call
255                            .is_none_or(|rounds| accounting.reconciled_rounds() <= rounds)
256                        {
257                            return Err(protocol(
258                                "Codex completed without usage after the tool response",
259                            ));
260                        }
261                        let usage = accounting.snapshot();
262                        if usage.is_empty() {
263                            return Err(protocol("Codex completed without terminal usage"));
264                        }
265                        return Ok(ToolRunResult {
266                            arguments,
267                            thread_id,
268                            turn_id,
269                            usage,
270                        });
271                    }
272                    _ => validate_notification(&method, &params, &thread_id, &turn_id)?,
273                },
274                IncomingMessage::Response { .. } => {
275                    return Err(protocol("Codex emitted an unexpected response"));
276                }
277            }
278        }
279    }
280    .await;
281    let shutdown = rpc.shutdown().await.map_err(map_transport);
282    match result {
283        Err(error) => Err(error),
284        Ok(value) => shutdown.map(|_| value),
285    }
286}
287
288#[derive(Default)]
289struct Inbox {
290    queued: VecDeque<IncomingMessage>,
291}
292
293impl Inbox {
294    async fn response(
295        &mut self,
296        rpc: &mut StdioRpc,
297        expected: RequestId,
298        label: &str,
299    ) -> Result<Value> {
300        loop {
301            match rpc.next().await.map_err(map_transport)? {
302                IncomingMessage::Response { id, result } => {
303                    if id != PeerId::Number(expected.0.into()) {
304                        return Err(protocol(format!("wrong response ID for {label}")));
305                    }
306                    return result.map_err(|_| protocol(format!("Codex {label} failed")));
307                }
308                message => self.queued.push_back(message),
309            }
310        }
311    }
312
313    async fn next(&mut self, rpc: &mut StdioRpc) -> Result<IncomingMessage> {
314        match self.queued.pop_front() {
315            Some(message) => Ok(message),
316            None => rpc.next().await.map_err(map_transport),
317        }
318    }
319}
320
321fn app_server_command(client: &CodexTerra) -> Command {
322    let mut command = Command::new(&client.executable);
323    for value in CODEX_CONFIG.split('|') {
324        command.arg("-c").arg(value);
325    }
326    command
327        .args(["app-server", "--stdio"])
328        .current_dir(&client.working_directory)
329        .env_remove("OPENAI_API_KEY")
330        .env_remove("CODEX_API_KEY");
331    command
332}
333
334fn validate_notification(method: &str, params: &Value, thread: &str, turn: &str) -> Result<()> {
335    if method.contains("rerout") || DISABLED_EVENTS.split('|').any(|kind| method.contains(kind)) {
336        return Err(protocol("Codex attempted a disabled capability"));
337    }
338    if matches!(
339        method,
340        "model/safetyBuffering/updated" | "model/verification"
341    ) {
342        return require_scope(params, thread, turn);
343    }
344    if method.starts_with("thread/") {
345        return require_id(params, "/threadId", thread);
346    }
347    if method.starts_with("turn/") || method.starts_with("item/") {
348        require_scope(params, thread, turn)?;
349        if params.get("item").is_some() {
350            validate_item(params)?;
351        }
352        return Ok(());
353    }
354    Err(protocol("Codex emitted an unexpected event"))
355}
356
357fn validate_item(params: &Value) -> Result<()> {
358    match params.pointer("/item/type").and_then(Value::as_str) {
359        Some("userMessage" | "agentMessage" | "reasoning" | "dynamicToolCall") => Ok(()),
360        Some(_) => Err(protocol("Codex attempted a disabled built-in tool")),
361        None => Err(protocol("Codex item event omitted its item type")),
362    }
363}
364
365fn require_id(value: &Value, pointer: &str, expected: &str) -> Result<()> {
366    (value.pointer(pointer).and_then(Value::as_str) == Some(expected))
367        .then_some(())
368        .ok_or_else(|| protocol("Codex used a mismatched identifier"))
369}
370
371fn require_scope(value: &Value, thread: &str, turn: &str) -> Result<()> {
372    require_id(value, "/threadId", thread)?;
373    require_id(value, "/turnId", turn)
374}
375
376fn require_turn_object(value: &Value, thread: &str, turn: &str) -> Result<()> {
377    require_id(value, "/threadId", thread)?;
378    require_id(value, "/turn/id", turn)
379}
380
381fn required<'a>(value: &'a Value, pointer: &str, label: &str) -> Result<&'a str> {
382    value
383        .pointer(pointer)
384        .and_then(Value::as_str)
385        .filter(|value| !value.is_empty())
386        .ok_or_else(|| protocol(format!("Codex omitted or emptied {label}")))
387}
388
389fn map_transport(error: TransportError) -> Error {
390    match error {
391        TransportError::Json(_)
392        | TransportError::InvalidMessage(_)
393        | TransportError::InboundLineTooLong => protocol("Codex emitted invalid JSON-RPC"),
394        _ => unavailable("Codex app-server transport became unavailable"),
395    }
396}
397
398fn invalid(message: impl Into<String>) -> Error {
399    Error::new(ErrorKind::InvalidInput, message)
400}
401
402fn unavailable(message: impl Into<String>) -> Error {
403    Error::new(ErrorKind::Unavailable, message)
404}
405
406fn protocol(message: impl Into<String>) -> Error {
407    Error::new(ErrorKind::Protocol, message)
408}
kcode_codex_terra/lib.rs

kcode_codex_terra/
lib.rs
