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

kcode_k1_codex_shim/
lib.rs
