kcode_k1_codex_adapter/
shim.rs

1use std::collections::VecDeque;
2
3use crate::{Adapter, Error, Event, ToolCall, ToolResult};
4
5/// Fixed successful response used to release every Codex dynamic-tool request.
6pub const ASYNC_TOOL_ACKNOWLEDGEMENT: &str = "The tool was launched asynchronously. Its result will not be available during this turn. Do not wait for or poll this call; continue the turn without its result. The result will be provided in the next user turn.";
7
8/// Converts between a chat engine's canonical box type and text visible to Codex.
9///
10/// A shim calls `tool_call_box` exactly once for each dynamic tool call it
11/// receives. `box_text` must return the complete representation that should be
12/// placed into Codex context. Implementations must not add ordering of their
13/// own: the shim preserves insertion order.
14pub trait BoxCodec {
15    /// Canonical box type owned by the chat engine.
16    type Box: Clone;
17
18    /// Converts one typed Codex dynamic-tool request into a canonical box.
19    fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box;
20
21    /// Returns the exact textual representation of one box for Codex.
22    fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str;
23}
24
25/// One ordered item produced by a completed native Codex turn.
26#[derive(Clone, Debug, PartialEq)]
27pub enum ShimItem<B> {
28    /// Adjacent streamed assistant-text deltas, coalesced.
29    Text(String),
30    /// A canonical dynamic-tool-call box.
31    Box(B),
32}
33
34/// Atomic terminal output from one native Codex turn.
35#[derive(Clone, Debug, PartialEq)]
36pub struct ShimOutput<B> {
37    /// Assistant text and call boxes in exact provider event order.
38    pub items: Vec<ShimItem<B>>,
39}
40
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42enum Health {
43    Ready,
44    Unusable,
45}
46
47/// A single-owner, per-conversation bridge between K1 boxes and Codex turns.
48///
49/// `Shim` deliberately contains no locking. A caller creates one shim for one
50/// conversation and invokes it sequentially. Different shims may share clones
51/// of the same [`Adapter`], whose actor owns transport concurrency.
52pub struct Shim<C: BoxCodec> {
53    adapter: Adapter,
54    conversation_key: String,
55    codec: C,
56    health: Health,
57    pending_boxes: VecDeque<C::Box>,
58}
59
60impl<C: BoxCodec> Shim<C> {
61    /// Creates a ready shim for one conversation.
62    pub fn new(adapter: Adapter, conversation_key: impl Into<String>, codec: C) -> Self {
63        Self {
64            adapter,
65            conversation_key: conversation_key.into(),
66            codec,
67            health: Health::Ready,
68            pending_boxes: VecDeque::new(),
69        }
70    }
71
72    /// Appends one externally produced box in canonical chat-engine order.
73    ///
74    /// Call boxes returned by [`Shim::infer`] are output only and are not added
75    /// to this pending queue automatically.
76    pub fn record_box(&mut self, box_: C::Box) {
77        self.pending_boxes.push_back(box_);
78    }
79
80    /// Appends externally produced boxes without changing their order.
81    pub fn record_boxes(&mut self, boxes: impl IntoIterator<Item = C::Box>) {
82        self.pending_boxes.extend(boxes);
83    }
84
85    /// Number of boxes waiting to be submitted to a fresh Codex turn.
86    pub fn pending_box_count(&self) -> usize {
87        self.pending_boxes.len()
88    }
89
90    /// Runs exactly one fresh native Codex turn through terminal completion.
91    ///
92    /// Pending K1 boxes prefix ordinary input and are removed only after the
93    /// adapter accepts the turn. Every dynamic call is acknowledged immediately
94    /// with [`ASYNC_TOOL_ACKNOWLEDGEMENT`], then retained as an ordered output
95    /// box. No partial output is returned if the active turn fails.
96    pub async fn infer(&mut self, input: impl Into<String>) -> Result<ShimOutput<C::Box>, Error> {
97        if self.health == Health::Unusable {
98            return Err(self.unusable());
99        }
100
101        let submitted_box_count = self.pending_boxes.len();
102        let input = append_section(self.render_pending_boxes(), &input.into());
103
104        // Cancellation anywhere after this point fails closed. A definite
105        // start rejection restores readiness because no native turn exists.
106        self.health = Health::Unusable;
107        let mut turn = match self
108            .adapter
109            .start_turn(self.conversation_key.clone(), input)
110            .await
111        {
112            Ok(turn) => turn,
113            Err(error) => {
114                self.health = Health::Ready;
115                return Err(error);
116            }
117        };
118
119        for _ in 0..submitted_box_count {
120            let removed = self.pending_boxes.pop_front();
121            debug_assert!(removed.is_some());
122        }
123
124        let mut items = Vec::new();
125        loop {
126            match turn.next_event().await {
127                Some(Event::TextDelta(delta)) => push_text(&mut items, delta),
128                Some(Event::ToolCall(call)) => {
129                    let box_ = self.codec.tool_call_box(&call);
130                    turn.respond(
131                        call.call_id,
132                        ToolResult {
133                            success: true,
134                            output: ASYNC_TOOL_ACKNOWLEDGEMENT.to_owned(),
135                        },
136                    )
137                    .await?;
138                    items.push(ShimItem::Box(box_));
139                }
140                Some(Event::Done) => {
141                    self.health = Health::Ready;
142                    return Ok(ShimOutput { items });
143                }
144                Some(Event::Error(error)) => return Err(error),
145                None => return Err(self.adapter.unavailable()),
146            }
147        }
148    }
149
150    fn render_pending_boxes(&self) -> String {
151        let mut output = String::new();
152        for box_ in &self.pending_boxes {
153            output = append_section(output, self.codec.box_text(box_));
154        }
155        output
156    }
157
158    fn unusable(&self) -> Error {
159        let mut error = self.adapter.unavailable();
160        error.message =
161            "Codex shim cannot be reused after an active turn failed or was cancelled".to_owned();
162        error
163    }
164}
165
166fn push_text<B>(items: &mut Vec<ShimItem<B>>, delta: String) {
167    match items.last_mut() {
168        Some(ShimItem::Text(text)) => text.push_str(&delta),
169        _ => items.push(ShimItem::Text(delta)),
170    }
171}
172
173fn append_section(mut output: String, section: &str) -> String {
174    if section.is_empty() {
175        return output;
176    }
177    if !output.is_empty() && !output.ends_with('\n') {
178        output.push('\n');
179    }
180    output.push_str(section);
181    output
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::{Config, DynamicTool};
188    use serde_json::{Value, json};
189
190    #[derive(Clone, Debug, PartialEq, Eq)]
191    struct TestBox(String);
192
193    #[derive(Default)]
194    struct TestCodec {
195        conversions: Vec<String>,
196    }
197
198    impl BoxCodec for TestCodec {
199        type Box = TestBox;
200
201        fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box {
202            self.conversions.push(call.call_id.clone());
203            TestBox(call.call_id.clone())
204        }
205
206        fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str {
207            &box_.0
208        }
209    }
210
211    #[cfg(unix)]
212    struct TestApp {
213        directory: std::path::PathBuf,
214        executable: std::path::PathBuf,
215    }
216
217    #[cfg(unix)]
218    impl TestApp {
219        fn new(script: &str) -> Self {
220            use std::{
221                os::unix::fs::PermissionsExt,
222                sync::atomic::{AtomicU64, Ordering},
223                time::{SystemTime, UNIX_EPOCH},
224            };
225
226            static NEXT: AtomicU64 = AtomicU64::new(0);
227            let nonce = SystemTime::now()
228                .duration_since(UNIX_EPOCH)
229                .unwrap()
230                .as_nanos();
231            let directory = std::env::temp_dir().join(format!(
232                "kcode-k1-codex-adapter-{}-{nonce}-{}",
233                std::process::id(),
234                NEXT.fetch_add(1, Ordering::Relaxed)
235            ));
236            std::fs::create_dir(&directory).unwrap();
237            let executable = directory.join("codex");
238            std::fs::write(&executable, script).unwrap();
239            let mut permissions = std::fs::metadata(&executable).unwrap().permissions();
240            permissions.set_mode(0o700);
241            std::fs::set_permissions(&executable, permissions).unwrap();
242            Self {
243                directory,
244                executable,
245            }
246        }
247
248        fn path(&self, name: &str) -> std::path::PathBuf {
249            self.directory.join(name)
250        }
251
252        async fn adapter(&self) -> Adapter {
253            Adapter::open(Config {
254                executable: self.executable.clone(),
255                working_directory: self.directory.to_string_lossy().into_owned(),
256                model: "test-model".into(),
257                reasoning_effort: None,
258                base_instructions: String::new(),
259                tools: vec![DynamicTool {
260                    name: "lookup".into(),
261                    description: "Lookup a value".into(),
262                    input_schema: json!({"type":"object"}),
263                }],
264            })
265            .await
266            .unwrap()
267        }
268    }
269
270    #[cfg(unix)]
271    impl Drop for TestApp {
272        fn drop(&mut self) {
273            let _ = std::fs::remove_dir_all(&self.directory);
274        }
275    }
276
277    #[cfg(unix)]
278    async fn wait_for_lines(path: &std::path::Path, minimum: usize) {
279        tokio::time::timeout(std::time::Duration::from_secs(3), async {
280            loop {
281                let count = std::fs::read_to_string(path)
282                    .map(|text| text.lines().count())
283                    .unwrap_or(0);
284                if count >= minimum {
285                    return;
286                }
287                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
288            }
289        })
290        .await
291        .expect("timed out waiting for app-server log");
292    }
293
294    #[cfg(unix)]
295    #[tokio::test]
296    async fn completes_one_turn_with_immediate_ordered_call_responses() {
297        let app = TestApp::new(
298            r#"#!/bin/sh
299set -eu
300IFS= read -r initialize
301echo '{"id":0,"result":{}}'
302IFS= read -r initialized
303IFS= read -r thread_start
304echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
305IFS= read -r turn_start
306printf '%s\n' "$turn_start" > turn-start.log
307echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
308echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"pre"}}'
309echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"face"}}'
310echo '{"id":77,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"A","tool":"lookup","arguments":{"n":1}}}'
311IFS= read -r response
312printf '%s\n' "$response" >> responses.log
313echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"mid"}}'
314echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"dle"}}'
315echo '{"id":78,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"B","tool":"lookup","arguments":{"n":2}}}'
316echo '{"id":79,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"C","tool":"lookup","arguments":{"n":3}}}'
317IFS= read -r response
318printf '%s\n' "$response" >> responses.log
319IFS= read -r response
320printf '%s\n' "$response" >> responses.log
321while [ ! -e release ]; do sleep 0.01; done
322echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"tail"}}'
323echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"end"}}'
324echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
325"#,
326        );
327        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());
328        shim.record_boxes([TestBox("<pending-1>".into()), TestBox("<pending-2>".into())]);
329
330        let task = tokio::spawn(async move {
331            let result = shim.infer("request").await;
332            (shim, result)
333        });
334        wait_for_lines(&app.path("responses.log"), 3).await;
335        assert!(!task.is_finished(), "infer returned before Event::Done");
336        std::fs::write(app.path("release"), "").unwrap();
337        let (shim, output) = tokio::time::timeout(std::time::Duration::from_secs(3), task)
338            .await
339            .unwrap()
340            .unwrap();
341        let output = output.unwrap();
342
343        assert_eq!(
344            output.items,
345            vec![
346                ShimItem::Text("preface".into()),
347                ShimItem::Box(TestBox("A".into())),
348                ShimItem::Text("middle".into()),
349                ShimItem::Box(TestBox("B".into())),
350                ShimItem::Box(TestBox("C".into())),
351                ShimItem::Text("tailend".into()),
352            ]
353        );
354        assert_eq!(shim.codec.conversions, ["A", "B", "C"]);
355        assert_eq!(shim.pending_box_count(), 0);
356
357        let start: Value =
358            serde_json::from_str(&std::fs::read_to_string(app.path("turn-start.log")).unwrap())
359                .unwrap();
360        assert_eq!(
361            start
362                .pointer("/params/input/0/text")
363                .and_then(Value::as_str),
364            Some("<pending-1>\n<pending-2>\nrequest")
365        );
366
367        let responses = std::fs::read_to_string(app.path("responses.log")).unwrap();
368        for (line, id) in responses.lines().zip([77, 78, 79]) {
369            let response: Value = serde_json::from_str(line).unwrap();
370            assert_eq!(response["id"], id);
371            assert_eq!(response["result"]["success"], true);
372            assert_eq!(
373                response
374                    .pointer("/result/contentItems/0/text")
375                    .and_then(Value::as_str),
376                Some(ASYNC_TOOL_ACKNOWLEDGEMENT)
377            );
378        }
379    }
380
381    #[cfg(unix)]
382    #[tokio::test]
383    async fn normal_inference_uses_a_fresh_turn_each_time() {
384        let app = TestApp::new(
385            r#"#!/bin/sh
386set -eu
387read initialize
388echo '{"id":0,"result":{}}'
389read initialized
390read thread_start
391echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
392read first_start
393printf '%s\n' "$first_start" >> starts.log
394echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
395echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"hel"}}'
396echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"lo"}}'
397echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
398read second_start
399printf '%s\n' "$second_start" >> starts.log
400echo '{"id":3,"result":{"turn":{"id":"turn-2"}}}'
401echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-2","delta":"again"}}'
402echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-2","status":"completed"}}}'
403"#,
404        );
405        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());
406
407        assert_eq!(
408            shim.infer("first").await.unwrap().items,
409            [ShimItem::Text("hello".into())]
410        );
411        assert_eq!(
412            shim.infer("second").await.unwrap().items,
413            [ShimItem::Text("again".into())]
414        );
415        let starts = std::fs::read_to_string(app.path("starts.log")).unwrap();
416        assert_eq!(starts.lines().count(), 2);
417    }
418
419    #[cfg(unix)]
420    #[tokio::test]
421    async fn finite_thousand_call_sequence_is_not_truncated() {
422        let app = TestApp::new(
423            r#"#!/bin/sh
424set -eu
425read initialize
426echo '{"id":0,"result":{}}'
427read initialized
428read thread_start
429echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
430read turn_start
431echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
432i=1
433while [ "$i" -le 1000 ]; do
434  echo "{\"id\":$((1000 + i)),\"method\":\"item/tool/call\",\"params\":{\"threadId\":\"thread-1\",\"turnId\":\"turn-1\",\"callId\":\"call-$i\",\"tool\":\"lookup\",\"arguments\":{\"n\":$i}}}"
435  read response
436  i=$((i + 1))
437done
438echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
439"#,
440        );
441        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());
442
443        let output = shim.infer("burst").await.unwrap();
444        assert_eq!(output.items.len(), 1000);
445        for (index, item) in output.items.iter().enumerate() {
446            assert_eq!(item, &ShimItem::Box(TestBox(format!("call-{}", index + 1))));
447        }
448        assert_eq!(shim.codec.conversions.len(), 1000);
449    }
450
451    #[cfg(unix)]
452    #[tokio::test]
453    async fn response_failure_returns_no_output_and_poisons_reuse() {
454        let app = TestApp::new(
455            r#"#!/bin/sh
456set -eu
457read initialize
458echo '{"id":0,"result":{}}'
459read initialized
460read thread_start
461echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
462read turn_start
463echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
464exec 0<&-
465echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"partial"}}'
466echo '{"id":77,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"A","tool":"lookup","arguments":{}}}'
467sleep 1
468"#,
469        );
470        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());
471
472        assert!(shim.infer("request").await.is_err());
473        let reuse = shim.infer("must reject").await.unwrap_err();
474        assert!(
475            reuse.message.contains("cannot be reused"),
476            "unexpected reuse error: {reuse}"
477        );
478    }
479}
kcode_k1_codex_adapter/shim.rs

kcode_k1_codex_adapter/
shim.rs
