Skip to main content

kcode_k1_codex_shim/
lib.rs

1//! Per-conversation K1 bridge for native Codex turns.
2//!
3//! One [`Shim::infer`] drives one fresh runtime turn through [`Event::Done`],
4//! launching and acknowledging every dynamic call without exposing partial or
5//! cross-round suspended output.
6
7use std::{collections::VecDeque, future::Future, pin::Pin};
8
9pub use kcode_k1_codex_runtime::{
10    Adapter, Config, DynamicTool, Error, ErrorKind, Event, ToolCall, ToolResult, Turn,
11};
12
13/// Fixed successful response used to release every Codex dynamic-tool request.
14pub 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 a subsequent turn when it becomes available.";
15
16/// Future returned while handing one canonical tool-call box to K1.
17pub type ToolLaunchFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;
18
19/// Accepts canonical tool-call boxes for asynchronous execution.
20pub trait ToolCallLauncher<B>: Send {
21    /// Accepts, starts, or queues one tool call without waiting for completion.
22    fn launch<'a>(&'a mut self, box_: &'a B) -> ToolLaunchFuture<'a>;
23}
24
25/// Converts between K1's canonical box type and text visible to Codex.
26pub trait BoxCodec {
27    /// Canonical box type owned by K1.
28    type Box: Clone;
29
30    /// Converts one typed Codex dynamic-tool request into a canonical box.
31    ///
32    /// A shim invokes this exactly once for each call it receives.
33    fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box;
34
35    /// Returns the complete representation of one box for Codex history.
36    fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str;
37}
38
39/// One ordered item produced by a completed native Codex turn.
40#[derive(Clone, Debug, PartialEq)]
41pub enum ShimItem<B> {
42    /// Adjacent streamed assistant-text deltas, coalesced.
43    Text(String),
44    /// A canonical dynamic-tool-call box.
45    Box(B),
46}
47
48/// Atomic terminal output from one native Codex turn.
49#[derive(Clone, Debug, PartialEq)]
50pub struct ShimOutput<B> {
51    /// Assistant text and call boxes in exact provider event order.
52    pub items: Vec<ShimItem<B>>,
53}
54
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56enum Health {
57    Ready,
58    Unusable,
59}
60
61/// A single-owner, per-conversation bridge between K1 boxes and Codex turns.
62///
63/// Invoke a shim sequentially. Different shims may share cloned [`Adapter`]s.
64pub struct Shim<C: BoxCodec> {
65    adapter: Adapter,
66    conversation_key: String,
67    codec: C,
68    launcher: Box<dyn ToolCallLauncher<C::Box>>,
69    health: Health,
70    pending_boxes: VecDeque<C::Box>,
71}
72
73impl<C: BoxCodec> Shim<C> {
74    /// Creates a ready shim for one conversation.
75    pub fn new(
76        adapter: Adapter,
77        conversation_key: impl Into<String>,
78        codec: C,
79        launcher: Box<dyn ToolCallLauncher<C::Box>>,
80    ) -> Self {
81        Self {
82            adapter,
83            conversation_key: conversation_key.into(),
84            codec,
85            launcher,
86            health: Health::Ready,
87            pending_boxes: VecDeque::new(),
88        }
89    }
90
91    /// Appends one externally produced box in canonical history order.
92    ///
93    /// Call boxes returned by [`Shim::infer`] are not queued automatically.
94    pub fn record_box(&mut self, box_: C::Box) {
95        self.pending_boxes.push_back(box_);
96    }
97
98    /// Appends externally produced boxes without changing their order.
99    pub fn record_boxes(&mut self, boxes: impl IntoIterator<Item = C::Box>) {
100        self.pending_boxes.extend(boxes);
101    }
102
103    /// Returns the number of boxes waiting for the next fresh native turn.
104    pub fn pending_box_count(&self) -> usize {
105        self.pending_boxes.len()
106    }
107
108    /// Closes this shim's runtime conversation while the shim is ready.
109    ///
110    /// Pending external boxes are retained and can prefix a later fresh thread.
111    pub async fn close_conversation(&mut self) -> Result<(), Error> {
112        if self.health == Health::Unusable {
113            return Err(self.unusable());
114        }
115        self.adapter
116            .close_conversation(self.conversation_key.clone())
117            .await
118    }
119
120    /// Runs one fresh native turn through terminal completion.
121    ///
122    /// Pending boxes clear only after start acceptance. Calls are converted
123    /// once, launched, acknowledged, and retained in provider event order.
124    /// Active-turn failure or cancellation returns no partial output and makes
125    /// this shim unusable.
126    pub async fn infer(&mut self, input: impl Into<String>) -> Result<ShimOutput<C::Box>, Error> {
127        if self.health == Health::Unusable {
128            return Err(self.unusable());
129        }
130
131        let submitted_box_count = self.pending_boxes.len();
132        let input = input.into();
133        let input = append_section(self.render_pending_boxes(), &input);
134
135        // Cancellation after this assignment fails closed. A returned start
136        // error certifies that no active native turn was accepted.
137        self.health = Health::Unusable;
138        let mut turn = match self
139            .adapter
140            .start_turn(self.conversation_key.clone(), input)
141            .await
142        {
143            Ok(turn) => turn,
144            Err(error) => {
145                self.health = Health::Ready;
146                return Err(error);
147            }
148        };
149
150        for _ in 0..submitted_box_count {
151            let removed = self.pending_boxes.pop_front();
152            debug_assert!(removed.is_some());
153        }
154
155        let mut items = Vec::new();
156        loop {
157            match turn.next_event().await {
158                Some(Event::TextDelta(delta)) => push_text(&mut items, delta),
159                Some(Event::ToolCall(call)) => {
160                    let box_ = self.codec.tool_call_box(&call);
161                    if let Err(message) = self.launcher.launch(&box_).await {
162                        return Err(Error {
163                            kind: ErrorKind::LaunchRejected,
164                            message,
165                            diagnostics: self.adapter.diagnostics(),
166                        });
167                    }
168                    turn.respond(
169                        call.call_id,
170                        ToolResult {
171                            success: true,
172                            output: ASYNC_TOOL_ACKNOWLEDGEMENT.to_owned(),
173                        },
174                    )
175                    .await?;
176                    items.push(ShimItem::Box(box_));
177                }
178                Some(Event::Done) => {
179                    self.health = Health::Ready;
180                    return Ok(ShimOutput { items });
181                }
182                Some(Event::Error(error)) => return Err(error),
183                None => {
184                    return Err(
185                        self.error("Codex app-server closed before the active turn completed")
186                    );
187                }
188            }
189        }
190    }
191
192    fn render_pending_boxes(&self) -> String {
193        let mut output = String::new();
194        for box_ in &self.pending_boxes {
195            output = append_section(output, self.codec.box_text(box_));
196        }
197        output
198    }
199
200    fn unusable(&self) -> Error {
201        self.error("Codex shim cannot be reused after an active turn failed or was cancelled")
202    }
203
204    fn error(&self, message: impl Into<String>) -> Error {
205        Error {
206            kind: ErrorKind::Unavailable,
207            message: message.into(),
208            diagnostics: self.adapter.diagnostics(),
209        }
210    }
211}
212
213fn push_text<B>(items: &mut Vec<ShimItem<B>>, delta: String) {
214    match items.last_mut() {
215        Some(ShimItem::Text(text)) => text.push_str(&delta),
216        _ => items.push(ShimItem::Text(delta)),
217    }
218}
219
220fn append_section(mut output: String, section: &str) -> String {
221    if section.is_empty() {
222        return output;
223    }
224    if !output.is_empty() && !output.ends_with('\n') {
225        output.push('\n');
226    }
227    output.push_str(section);
228    output
229}