kcode_codex_terra/
lib.rs

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

kcode_codex_terra/
lib.rs
