kcode-k1-codex-shim 0.1.0

//! Per-conversation K1 bridge for native Codex turns.
//!
//! One [`Shim::infer`] drives one fresh runtime turn through [`Event::Done`],
//! launching and acknowledging every dynamic call without exposing partial or
//! cross-round suspended output.

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 canonical tool-call box to K1.
pub type ToolLaunchFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;

/// Accepts canonical tool-call boxes for asynchronous execution.
pub trait ToolCallLauncher<B>: Send {
    /// Accepts, starts, or queues one tool call without waiting for completion.
    fn launch<'a>(&'a mut self, box_: &'a 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> {
    /// Adjacent streamed assistant-text deltas, coalesced.
    Text(String),
    /// A canonical dynamic-tool-call box.
    Box(B),
}

/// Atomic terminal output from one native Codex turn.
#[derive(Clone, Debug, PartialEq)]
pub struct ShimOutput<B> {
    /// Assistant text and call boxes in exact provider event order.
    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.
    ///
    /// Call boxes returned by [`Shim::infer`] are not queued automatically.
    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. Calls are converted
    /// once, launched, acknowledged, and retained in provider event order.
    /// Active-turn failure or cancellation returns no partial output and makes
    /// this shim unusable.
    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 = input.into();
        let input = append_section(self.render_pending_boxes(), &input);

        // Cancellation after this assignment fails closed. A returned start
        // error certifies that no active native turn was accepted.
        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 {
            let removed = self.pending_boxes.pop_front();
            debug_assert!(removed.is_some());
        }

        let mut items = Vec::new();
        loop {
            match turn.next_event().await {
                Some(Event::TextDelta(delta)) => push_text(&mut items, delta),
                Some(Event::ToolCall(call)) => {
                    let box_ = self.codec.tool_call_box(&call);
                    if let Err(message) = self.launcher.launch(&box_).await {
                        return Err(Error {
                            kind: ErrorKind::LaunchRejected,
                            message,
                            diagnostics: self.adapter.diagnostics(),
                        });
                    }
                    turn.respond(
                        call.call_id,
                        ToolResult {
                            success: true,
                            output: ASYNC_TOOL_ACKNOWLEDGEMENT.to_owned(),
                        },
                    )
                    .await?;
                    items.push(ShimItem::Box(box_));
                }
                Some(Event::Done) => {
                    self.health = Health::Ready;
                    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 push_text<B>(items: &mut Vec<ShimItem<B>>, delta: String) {
    match items.last_mut() {
        Some(ShimItem::Text(text)) => text.push_str(&delta),
        _ => items.push(ShimItem::Text(delta)),
    }
}

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
}