kcode_k1_codex_shim/
lib.rs1use 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
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.";
15
16pub type ToolLaunchFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;
18
19pub trait ToolCallLauncher<B>: Send {
21 fn launch<'a>(&'a mut self, box_: &'a B) -> ToolLaunchFuture<'a>;
23}
24
25pub trait BoxCodec {
27 type Box: Clone;
29
30 fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box;
34
35 fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str;
37}
38
39#[derive(Clone, Debug, PartialEq)]
41pub enum ShimItem<B> {
42 Text(String),
44 Box(B),
46}
47
48#[derive(Clone, Debug, PartialEq)]
50pub struct ShimOutput<B> {
51 pub items: Vec<ShimItem<B>>,
53}
54
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56enum Health {
57 Ready,
58 Unusable,
59}
60
61pub 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 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 pub fn record_box(&mut self, box_: C::Box) {
95 self.pending_boxes.push_back(box_);
96 }
97
98 pub fn record_boxes(&mut self, boxes: impl IntoIterator<Item = C::Box>) {
100 self.pending_boxes.extend(boxes);
101 }
102
103 pub fn pending_box_count(&self) -> usize {
105 self.pending_boxes.len()
106 }
107
108 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 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 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}