kcode-k1-codex-shim 0.2.0

//! Per-conversation K1 bridge for native Codex turns.
//!
//! One [`Shim::infer`] drives one fresh runtime turn through [`Event::Done`],
//! checkpointing and acknowledging each immediately available call wave.

use std::{collections::VecDeque, future::Future, pin::Pin};

pub use kcode_k1_codex_runtime::{
    Adapter, Config, DynamicTool, Error, ErrorKind, Event, ToolCall, ToolResult, Turn,
};

/// Fixed successful response used to release every Codex dynamic-tool request.
pub 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.";

/// Future returned while handing one output stage to K1.
pub type ToolLaunchFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;

/// Accepts one assistant-text and tool-call wave for durable launch.
pub trait ToolCallLauncher<B>: Send {
    /// Persists, accepts, and launches one wave without waiting for completion.
    fn launch_stage<'a>(&'a mut self, text: String, boxes: Vec<B>) -> ToolLaunchFuture<'a>;
}

/// Converts between K1's canonical box type and text visible to Codex.
pub trait BoxCodec {
    /// Canonical box type owned by K1.
    type Box: Clone;

    /// Converts one typed Codex dynamic-tool request into a canonical box.
    ///
    /// A shim invokes this exactly once for each call it receives.
    fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box;

    /// Returns the complete representation of one box for Codex history.
    fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str;
}

/// One ordered item produced by a completed native Codex turn.
#[derive(Clone, Debug, PartialEq)]
pub enum ShimItem<B> {
    /// Assistant text remaining after the final call wave.
    Text(String),
    /// Reserved canonical box output.
    Box(B),
}

/// Atomic terminal output from one native Codex turn.
#[derive(Clone, Debug, PartialEq)]
pub struct ShimOutput<B> {
    /// Empty, or one final assistant-text item.
    pub items: Vec<ShimItem<B>>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Health {
    Ready,
    Unusable,
}

/// A single-owner, per-conversation bridge between K1 boxes and Codex turns.
///
/// Invoke a shim sequentially. Different shims may share cloned [`Adapter`]s.
pub struct Shim<C: BoxCodec> {
    adapter: Adapter,
    conversation_key: String,
    codec: C,
    launcher: Box<dyn ToolCallLauncher<C::Box>>,
    health: Health,
    pending_boxes: VecDeque<C::Box>,
}

impl<C: BoxCodec> Shim<C> {
    /// Creates a ready shim for one conversation.
    pub fn new(
        adapter: Adapter,
        conversation_key: impl Into<String>,
        codec: C,
        launcher: Box<dyn ToolCallLauncher<C::Box>>,
    ) -> Self {
        Self {
            adapter,
            conversation_key: conversation_key.into(),
            codec,
            launcher,
            health: Health::Ready,
            pending_boxes: VecDeque::new(),
        }
    }

    /// Appends one externally produced box in canonical history order.
    pub fn record_box(&mut self, box_: C::Box) {
        self.pending_boxes.push_back(box_);
    }

    /// Appends externally produced boxes without changing their order.
    pub fn record_boxes(&mut self, boxes: impl IntoIterator<Item = C::Box>) {
        self.pending_boxes.extend(boxes);
    }

    /// Returns the number of boxes waiting for the next fresh native turn.
    pub fn pending_box_count(&self) -> usize {
        self.pending_boxes.len()
    }

    /// Closes this shim's runtime conversation while the shim is ready.
    ///
    /// Pending external boxes are retained and can prefix a later fresh thread.
    pub async fn close_conversation(&mut self) -> Result<(), Error> {
        if self.health == Health::Unusable {
            return Err(self.unusable());
        }
        self.adapter
            .close_conversation(self.conversation_key.clone())
            .await
    }

    /// Runs one fresh native turn through terminal completion.
    ///
    /// Pending boxes clear only after start acceptance. Each immediately
    /// buffered call wave is checkpointed before its calls are acknowledged.
    pub async fn infer(&mut self, input: impl Into<String>) -> Result<ShimOutput<C::Box>, Error> {
        if self.health == Health::Unusable {
            return Err(self.unusable());
        }

        let submitted_box_count = self.pending_boxes.len();
        let input = append_section(self.render_pending_boxes(), &input.into());

        self.health = Health::Unusable;
        let mut turn = match self
            .adapter
            .start_turn(self.conversation_key.clone(), input)
            .await
        {
            Ok(turn) => turn,
            Err(error) => {
                self.health = Health::Ready;
                return Err(error);
            }
        };

        for _ in 0..submitted_box_count {
            debug_assert!(self.pending_boxes.pop_front().is_some());
        }

        let mut text = String::new();
        let mut lookahead = None;
        loop {
            let event = match lookahead.take() {
                Some(event) => Some(event),
                None => turn.next_event().await,
            };
            match event {
                Some(Event::TextDelta(delta)) => text.push_str(&delta),
                Some(Event::ToolCall(first)) => {
                    let first_box = self.codec.tool_call_box(&first);
                    let mut calls = vec![first];
                    let mut boxes = vec![first_box];
                    let mut drain_error = None;

                    loop {
                        match turn.try_next_event() {
                            Ok(Some(Event::ToolCall(call))) => {
                                let box_ = self.codec.tool_call_box(&call);
                                calls.push(call);
                                boxes.push(box_);
                            }
                            Ok(Some(event)) => {
                                lookahead = Some(event);
                                break;
                            }
                            Ok(None) => break,
                            Err(error) => {
                                drain_error = Some(error);
                                break;
                            }
                        }
                    }

                    if let Err(message) = self
                        .launcher
                        .launch_stage(std::mem::take(&mut text), boxes)
                        .await
                    {
                        return Err(Error {
                            kind: ErrorKind::LaunchRejected,
                            message,
                            diagnostics: self.adapter.diagnostics(),
                        });
                    }

                    for call in calls {
                        turn.respond(
                            call.call_id,
                            ToolResult {
                                success: true,
                                output: ASYNC_TOOL_ACKNOWLEDGEMENT.to_owned(),
                            },
                        )
                        .await?;
                    }

                    if let Some(error) = drain_error {
                        return Err(error);
                    }
                }
                Some(Event::Done) => {
                    self.health = Health::Ready;
                    let items = if text.is_empty() {
                        Vec::new()
                    } else {
                        vec![ShimItem::Text(text)]
                    };
                    return Ok(ShimOutput { items });
                }
                Some(Event::Error(error)) => return Err(error),
                None => {
                    return Err(
                        self.error("Codex app-server closed before the active turn completed")
                    );
                }
            }
        }
    }

    fn render_pending_boxes(&self) -> String {
        let mut output = String::new();
        for box_ in &self.pending_boxes {
            output = append_section(output, self.codec.box_text(box_));
        }
        output
    }

    fn unusable(&self) -> Error {
        self.error("Codex shim cannot be reused after an active turn failed or was cancelled")
    }

    fn error(&self, message: impl Into<String>) -> Error {
        Error {
            kind: ErrorKind::Unavailable,
            message: message.into(),
            diagnostics: self.adapter.diagnostics(),
        }
    }
}

fn append_section(mut output: String, section: &str) -> String {
    if section.is_empty() {
        return output;
    }
    if !output.is_empty() && !output.ends_with('\n') {
        output.push('\n');
    }
    output.push_str(section);
    output
}