1use serde_json::{Value, json};
2use std::{
3 collections::{HashMap, HashSet},
4 fmt,
5 path::PathBuf,
6 sync::{Arc, Mutex},
7};
8
9pub type ToolToken = (String, u64, String);
10
11#[derive(Clone, Debug)]
12pub struct Config {
13 pub executable: PathBuf,
14 pub working_directory: String,
15 pub model: String,
16 pub reasoning_effort: Option<String>,
17 pub base_instructions: String,
18 pub tools: Vec<DynamicTool>,
19}
20
21impl Config {
22 pub fn validate(&self) -> Result<(), Error> {
23 validate_config(self)
24 }
25}
26
27#[derive(Clone, Debug, PartialEq)]
28pub struct DynamicTool {
29 pub name: String,
30 pub description: String,
31 pub input_schema: Value,
32}
33
34#[derive(Clone, Debug, PartialEq)]
35pub struct ToolCall {
36 pub call_id: String,
37 pub name: String,
38 pub arguments: Value,
39}
40
41#[derive(Clone, Debug, PartialEq, Eq)]
42pub struct ToolResult {
43 pub success: bool,
44 pub output: String,
45}
46
47#[derive(Clone, Debug, PartialEq)]
48pub enum Event {
49 TextDelta(String),
50 ToolCall(ToolCall),
51 Done,
52 Error(Error),
53}
54
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub enum ErrorKind {
57 Busy,
58 Interrupted,
59 InvalidToolResult,
60 LaunchRejected,
61 Protocol,
62 Server,
63 Unavailable,
64}
65
66#[derive(Clone, Debug, PartialEq, Eq)]
67pub struct Error {
68 pub kind: ErrorKind,
69 pub message: String,
70 pub diagnostics: Vec<u8>,
71}
72
73impl Error {
74 pub fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
75 Self {
76 kind,
77 message: message.into(),
78 diagnostics: Vec::new(),
79 }
80 }
81
82 pub fn with_diagnostics(mut self, diagnostics: Vec<u8>) -> Self {
83 self.diagnostics = diagnostics;
84 self
85 }
86
87 pub fn server(value: &Value, diagnostics: &Diagnostics) -> Self {
88 let detail = value
89 .get("message")
90 .and_then(Value::as_str)
91 .or_else(|| value.pointer("/error/message").and_then(Value::as_str));
92 let message = detail.map_or_else(
93 || "Codex app-server error".to_owned(),
94 |detail| format!("Codex app-server error: {detail}"),
95 );
96 diagnostics.error(ErrorKind::Server, message)
97 }
98}
99
100impl fmt::Display for Error {
101 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
102 formatter.write_str(&self.message)
103 }
104}
105
106impl std::error::Error for Error {}
107
108#[derive(Clone, Debug, Default)]
109pub struct Diagnostics(Arc<Mutex<Vec<u8>>>);
110
111impl Diagnostics {
112 pub fn new(bytes: Vec<u8>) -> Self {
113 Self(Arc::new(Mutex::new(bytes)))
114 }
115
116 pub fn snapshot(&self) -> Vec<u8> {
117 self.0
118 .lock()
119 .unwrap_or_else(|poisoned| poisoned.into_inner())
120 .clone()
121 }
122
123 pub fn replace(&self, bytes: Vec<u8>) {
124 *self
125 .0
126 .lock()
127 .unwrap_or_else(|poisoned| poisoned.into_inner()) = bytes;
128 }
129
130 pub fn error(&self, kind: ErrorKind, message: impl Into<String>) -> Error {
131 Error {
132 kind,
133 message: message.into(),
134 diagnostics: self.snapshot(),
135 }
136 }
137}
138
139pub struct Active<S> {
140 pub serial: u64,
141 pub turn: Option<String>,
142 pub sink: Option<S>,
143 pub early: Vec<Value>,
144 pub cancelled: bool,
145 pub interrupt_sent: bool,
146 pub failure: Option<Value>,
147}
148
149pub struct Conversation<S> {
150 pub thread: Option<String>,
151 pub active: Option<Active<S>>,
152 pub closing: bool,
153}
154
155impl<S> Default for Conversation<S> {
156 fn default() -> Self {
157 Self {
158 thread: None,
159 active: None,
160 closing: false,
161 }
162 }
163}
164
165pub enum Pending<R> {
166 Thread {
167 key: String,
168 serial: u64,
169 input: String,
170 },
171 Turn {
172 key: String,
173 serial: u64,
174 },
175 Close {
176 key: String,
177 thread: String,
178 reply: R,
179 },
180 Interrupt,
181}
182
183#[derive(Clone, Debug, PartialEq)]
184pub struct PendingTool {
185 pub id: Value,
186 pub rpc_key: String,
187}
188
189pub struct State<S, R> {
190 pub conversations: HashMap<String, Conversation<S>>,
191 pub by_thread: HashMap<String, String>,
192 pub pending: HashMap<u64, Pending<R>>,
193 pub tools: HashMap<ToolToken, PendingTool>,
194 pub rpc_ids: HashSet<String>,
195 pub next_id: u64,
196 pub next_turn: u64,
197}
198
199impl<S, R> Default for State<S, R> {
200 fn default() -> Self {
201 Self {
202 conversations: HashMap::new(),
203 by_thread: HashMap::new(),
204 pending: HashMap::new(),
205 tools: HashMap::new(),
206 rpc_ids: HashSet::new(),
207 next_id: 1,
208 next_turn: 1,
209 }
210 }
211}
212
213impl<S, R> State<S, R> {
214 pub fn allocate_request_id(&mut self) -> Result<u64, Error> {
215 let id = self.next_id;
216 self.next_id = id
217 .checked_add(1)
218 .ok_or_else(|| Error::new(ErrorKind::Protocol, "client request id space exhausted"))?;
219 Ok(id)
220 }
221
222 pub fn allocate_turn_id(&mut self) -> Result<u64, Error> {
223 let id = self.next_turn;
224 self.next_turn = id
225 .checked_add(1)
226 .ok_or_else(|| Error::new(ErrorKind::Protocol, "turn serial space exhausted"))?;
227 Ok(id)
228 }
229
230 pub fn begin_turn(&mut self, key: impl Into<String>, sink: S) -> Result<u64, Error> {
231 let key = key.into();
232 if self
233 .conversations
234 .get(&key)
235 .is_some_and(|conversation| conversation.active.is_some() || conversation.closing)
236 {
237 return Err(Error::new(
238 ErrorKind::Busy,
239 "conversation already has an active turn",
240 ));
241 }
242 let serial = self.allocate_turn_id()?;
243 self.conversations.entry(key).or_default().active = Some(Active {
244 serial,
245 turn: None,
246 sink: Some(sink),
247 early: Vec::new(),
248 cancelled: false,
249 interrupt_sent: false,
250 failure: None,
251 });
252 Ok(serial)
253 }
254
255 pub fn take_active(&mut self, key: &str, serial: u64) -> Option<Active<S>> {
256 let conversation = self.conversations.get_mut(key)?;
257 if conversation
258 .active
259 .as_ref()
260 .is_some_and(|active| active.serial == serial)
261 {
262 conversation.active.take()
263 } else {
264 None
265 }
266 }
267
268 pub fn set_native_turn(
269 &mut self,
270 key: &str,
271 serial: u64,
272 turn: impl Into<String>,
273 ) -> Option<Vec<Value>> {
274 let active = self
275 .conversations
276 .get_mut(key)?
277 .active
278 .as_mut()
279 .filter(|active| active.serial == serial)?;
280 active.turn = Some(turn.into());
281 Some(std::mem::take(&mut active.early))
282 }
283
284 pub fn interrupt_target(&mut self, key: &str, serial: u64) -> Option<(String, String)> {
285 let conversation = self.conversations.get_mut(key)?;
286 let thread = conversation.thread.clone()?;
287 let active = conversation.active.as_mut()?;
288 if active.serial != serial || active.interrupt_sent {
289 return None;
290 }
291 let turn = active.turn.clone()?;
292 active.interrupt_sent = true;
293 Some((thread, turn))
294 }
295
296 pub fn take_sinks(&mut self) -> Vec<S> {
297 self.conversations
298 .values_mut()
299 .filter_map(|conversation| conversation.active.take()?.sink)
300 .collect()
301 }
302
303 pub fn thread(&self, key: &str) -> Option<&str> {
304 self.conversations.get(key)?.thread.as_deref()
305 }
306
307 pub fn owner(&self, thread: &str) -> Option<&str> {
308 self.by_thread.get(thread).map(String::as_str)
309 }
310
311 pub fn set_thread(&mut self, key: &str, thread: impl Into<String>) -> Result<(), Error> {
312 let thread = thread.into();
313 if self
314 .by_thread
315 .get(&thread)
316 .is_some_and(|owner| owner != key)
317 {
318 return Err(Error::new(
319 ErrorKind::Protocol,
320 "thread/start reused another conversation thread",
321 ));
322 }
323 let old = self
324 .conversations
325 .entry(key.to_owned())
326 .or_default()
327 .thread
328 .replace(thread.clone());
329 if let Some(old) = old.filter(|old| old != &thread) {
330 self.by_thread.remove(&old);
331 }
332 self.by_thread.insert(thread, key.to_owned());
333 Ok(())
334 }
335
336 pub fn begin_close(&mut self, key: &str) -> Result<Option<String>, Error> {
337 let Some(conversation) = self.conversations.get(key) else {
338 return Ok(None);
339 };
340 if conversation.active.is_some() || conversation.closing {
341 return Err(Error::new(
342 ErrorKind::Busy,
343 "conversation is active or already closing",
344 ));
345 }
346 let Some(thread) = conversation.thread.clone() else {
347 self.conversations.remove(key);
348 return Ok(None);
349 };
350 self.conversations
351 .get_mut(key)
352 .expect("conversation exists")
353 .closing = true;
354 Ok(Some(thread))
355 }
356
357 pub fn cancel_close(&mut self, key: &str) {
358 if let Some(conversation) = self.conversations.get_mut(key) {
359 conversation.closing = false;
360 }
361 }
362
363 pub fn finish_close(&mut self, key: &str, thread: &str) -> Result<(), Error> {
364 let valid = self.conversations.get(key).is_some_and(|conversation| {
365 conversation.thread.as_deref() == Some(thread)
366 && conversation.active.is_none()
367 && conversation.closing
368 });
369 if !valid {
370 return Err(Error::new(
371 ErrorKind::Protocol,
372 "thread/unsubscribe response did not match closing conversation",
373 ));
374 }
375 self.by_thread.remove(thread);
376 self.conversations.remove(key);
377 Ok(())
378 }
379
380 pub fn insert_pending(&mut self, id: u64, pending: Pending<R>) -> Result<(), Error> {
381 if self.pending.insert(id, pending).is_some() {
382 return Err(Error::new(
383 ErrorKind::Protocol,
384 "duplicate client request id",
385 ));
386 }
387 Ok(())
388 }
389
390 pub fn take_pending(&mut self, id: u64) -> Result<Pending<R>, Error> {
391 self.pending.remove(&id).ok_or_else(|| {
392 Error::new(
393 ErrorKind::Protocol,
394 "unexpected or duplicate app-server response id",
395 )
396 })
397 }
398
399 pub fn track_tool(
400 &mut self,
401 key: &str,
402 serial: u64,
403 call: impl Into<String>,
404 id: &Value,
405 ) -> Result<ToolToken, Error> {
406 let (id, rpc_key) = parse_rpc_id(id)?;
407 let token = (key.to_owned(), serial, call.into());
408 if self.rpc_ids.contains(&rpc_key) {
409 return Err(Error::new(
410 ErrorKind::Protocol,
411 "duplicate app-server request id",
412 ));
413 }
414 if self.tools.contains_key(&token) {
415 return Err(Error::new(
416 ErrorKind::Protocol,
417 "duplicate dynamic tool call id",
418 ));
419 }
420 self.rpc_ids.insert(rpc_key.clone());
421 self.tools
422 .insert(token.clone(), PendingTool { id, rpc_key });
423 Ok(token)
424 }
425
426 pub fn take_tool(&mut self, token: &ToolToken) -> Option<PendingTool> {
427 let pending = self.tools.remove(token)?;
428 self.rpc_ids.remove(&pending.rpc_key);
429 Some(pending)
430 }
431
432 pub fn take_turn_tools(&mut self, key: &str, serial: u64) -> Vec<(ToolToken, PendingTool)> {
433 let tokens: Vec<_> = self
434 .tools
435 .keys()
436 .filter(|(owner, turn, _)| owner == key && *turn == serial)
437 .cloned()
438 .collect();
439 tokens
440 .into_iter()
441 .filter_map(|token| self.take_tool(&token).map(|pending| (token, pending)))
442 .collect()
443 }
444
445 pub fn resolve_tool(&mut self, id: &Value) -> Result<Option<ToolToken>, Error> {
446 let (_, rpc_key) = parse_rpc_id(id)?;
447 let token = self
448 .tools
449 .iter()
450 .find_map(|(token, pending)| (pending.rpc_key == rpc_key).then(|| token.clone()));
451 if let Some(token) = &token {
452 self.take_tool(token);
453 }
454 Ok(token)
455 }
456}
457
458pub fn validate_config(config: &Config) -> Result<(), Error> {
459 let mut names = HashSet::new();
460 for tool in &config.tools {
461 if tool.name.is_empty() {
462 return Err(Error::new(
463 ErrorKind::Protocol,
464 "dynamic tool names must not be empty",
465 ));
466 }
467 if !names.insert(&tool.name) {
468 return Err(Error::new(
469 ErrorKind::Protocol,
470 format!("duplicate dynamic tool name: {}", tool.name),
471 ));
472 }
473 }
474 Ok(())
475}
476
477pub fn thread_start_params(config: &Config) -> Value {
478 let tools: Vec<_> = config
479 .tools
480 .iter()
481 .map(|tool| {
482 json!({
483 "name": tool.name,
484 "description": tool.description,
485 "inputSchema": tool.input_schema
486 })
487 })
488 .collect();
489 json!({
490 "model": config.model,
491 "cwd": config.working_directory,
492 "approvalPolicy": "never",
493 "sandbox": "readOnly",
494 "baseInstructions": config.base_instructions,
495 "serviceName": "kcode-k1-codex-adapter",
496 "dynamicTools": tools
497 })
498}
499
500pub fn turn_start_params(thread: &str, input: impl Into<String>) -> Value {
501 json!({"threadId": thread, "input": [{"type": "text", "text": input.into()}]})
502}
503
504pub fn parse_scope(params: Option<&Value>) -> Result<(&str, &str), Error> {
505 let params =
506 params.ok_or_else(|| Error::new(ErrorKind::Protocol, "scoped message omitted params"))?;
507 let thread = params
508 .get("threadId")
509 .and_then(Value::as_str)
510 .ok_or_else(|| Error::new(ErrorKind::Protocol, "scoped message omitted threadId"))?;
511 let turn = params
512 .get("turnId")
513 .and_then(Value::as_str)
514 .or_else(|| params.pointer("/turn/id").and_then(Value::as_str))
515 .ok_or_else(|| Error::new(ErrorKind::Protocol, "scoped message omitted turn id"))?;
516 Ok((thread, turn))
517}
518
519pub fn parse_rpc_id(id: &Value) -> Result<(Value, String), Error> {
520 match id {
521 Value::String(value) => Ok((id.clone(), format!("s:{value}"))),
522 Value::Number(value) => Ok((id.clone(), format!("n:{value}"))),
523 _ => Err(Error::new(
524 ErrorKind::Protocol,
525 "server request id must be a string or number",
526 )),
527 }
528}
529
530pub fn is_model_reroute(method: &str) -> bool {
531 let method = method.to_ascii_lowercase();
532 method.contains("model") && method.contains("rerout")
533}