kcode-k1-codex-adapter 0.1.0

use std::collections::VecDeque;

use crate::{Adapter, Error, Event, ToolCall, ToolResult};

/// 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 the next user turn.";

/// Converts between a chat engine's canonical box type and text visible to Codex.
///
/// A shim calls `tool_call_box` exactly once for each dynamic tool call it
/// receives. `box_text` must return the complete representation that should be
/// placed into Codex context. Implementations must not add ordering of their
/// own: the shim preserves insertion order.
pub trait BoxCodec {
    /// Canonical box type owned by the chat engine.
    type Box: Clone;

    /// Converts one typed Codex dynamic-tool request into a canonical box.
    fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box;

    /// Returns the exact textual representation of one box for Codex.
    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.
///
/// `Shim` deliberately contains no locking. A caller creates one shim for one
/// conversation and invokes it sequentially. Different shims may share clones
/// of the same [`Adapter`], whose actor owns transport concurrency.
pub struct Shim<C: BoxCodec> {
    adapter: Adapter,
    conversation_key: String,
    codec: C,
    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) -> Self {
        Self {
            adapter,
            conversation_key: conversation_key.into(),
            codec,
            health: Health::Ready,
            pending_boxes: VecDeque::new(),
        }
    }

    /// Appends one externally produced box in canonical chat-engine order.
    ///
    /// Call boxes returned by [`Shim::infer`] are output only and are not added
    /// to this pending queue 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);
    }

    /// Number of boxes waiting to be submitted to a fresh Codex turn.
    pub fn pending_box_count(&self) -> usize {
        self.pending_boxes.len()
    }

    /// Runs exactly one fresh native Codex turn through terminal completion.
    ///
    /// Pending K1 boxes prefix ordinary input and are removed only after the
    /// adapter accepts the turn. Every dynamic call is acknowledged immediately
    /// with [`ASYNC_TOOL_ACKNOWLEDGEMENT`], then retained as an ordered output
    /// box. No partial output is returned if the active turn fails.
    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());

        // Cancellation anywhere after this point fails closed. A definite
        // start rejection restores readiness because no native turn exists.
        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);
                    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.adapter.unavailable()),
            }
        }
    }

    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 {
        let mut error = self.adapter.unavailable();
        error.message =
            "Codex shim cannot be reused after an active turn failed or was cancelled".to_owned();
        error
    }
}

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
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Config, DynamicTool};
    use serde_json::{Value, json};

    #[derive(Clone, Debug, PartialEq, Eq)]
    struct TestBox(String);

    #[derive(Default)]
    struct TestCodec {
        conversions: Vec<String>,
    }

    impl BoxCodec for TestCodec {
        type Box = TestBox;

        fn tool_call_box(&mut self, call: &ToolCall) -> Self::Box {
            self.conversions.push(call.call_id.clone());
            TestBox(call.call_id.clone())
        }

        fn box_text<'a>(&self, box_: &'a Self::Box) -> &'a str {
            &box_.0
        }
    }

    #[cfg(unix)]
    struct TestApp {
        directory: std::path::PathBuf,
        executable: std::path::PathBuf,
    }

    #[cfg(unix)]
    impl TestApp {
        fn new(script: &str) -> Self {
            use std::{
                os::unix::fs::PermissionsExt,
                sync::atomic::{AtomicU64, Ordering},
                time::{SystemTime, UNIX_EPOCH},
            };

            static NEXT: AtomicU64 = AtomicU64::new(0);
            let nonce = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos();
            let directory = std::env::temp_dir().join(format!(
                "kcode-k1-codex-adapter-{}-{nonce}-{}",
                std::process::id(),
                NEXT.fetch_add(1, Ordering::Relaxed)
            ));
            std::fs::create_dir(&directory).unwrap();
            let executable = directory.join("codex");
            std::fs::write(&executable, script).unwrap();
            let mut permissions = std::fs::metadata(&executable).unwrap().permissions();
            permissions.set_mode(0o700);
            std::fs::set_permissions(&executable, permissions).unwrap();
            Self {
                directory,
                executable,
            }
        }

        fn path(&self, name: &str) -> std::path::PathBuf {
            self.directory.join(name)
        }

        async fn adapter(&self) -> Adapter {
            Adapter::open(Config {
                executable: self.executable.clone(),
                working_directory: self.directory.to_string_lossy().into_owned(),
                model: "test-model".into(),
                reasoning_effort: None,
                base_instructions: String::new(),
                tools: vec![DynamicTool {
                    name: "lookup".into(),
                    description: "Lookup a value".into(),
                    input_schema: json!({"type":"object"}),
                }],
            })
            .await
            .unwrap()
        }
    }

    #[cfg(unix)]
    impl Drop for TestApp {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.directory);
        }
    }

    #[cfg(unix)]
    async fn wait_for_lines(path: &std::path::Path, minimum: usize) {
        tokio::time::timeout(std::time::Duration::from_secs(3), async {
            loop {
                let count = std::fs::read_to_string(path)
                    .map(|text| text.lines().count())
                    .unwrap_or(0);
                if count >= minimum {
                    return;
                }
                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
            }
        })
        .await
        .expect("timed out waiting for app-server log");
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn completes_one_turn_with_immediate_ordered_call_responses() {
        let app = TestApp::new(
            r#"#!/bin/sh
set -eu
IFS= read -r initialize
echo '{"id":0,"result":{}}'
IFS= read -r initialized
IFS= read -r thread_start
echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
IFS= read -r turn_start
printf '%s\n' "$turn_start" > turn-start.log
echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"pre"}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"face"}}'
echo '{"id":77,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"A","tool":"lookup","arguments":{"n":1}}}'
IFS= read -r response
printf '%s\n' "$response" >> responses.log
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"mid"}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"dle"}}'
echo '{"id":78,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"B","tool":"lookup","arguments":{"n":2}}}'
echo '{"id":79,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"C","tool":"lookup","arguments":{"n":3}}}'
IFS= read -r response
printf '%s\n' "$response" >> responses.log
IFS= read -r response
printf '%s\n' "$response" >> responses.log
while [ ! -e release ]; do sleep 0.01; done
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"tail"}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"end"}}'
echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
"#,
        );
        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());
        shim.record_boxes([TestBox("<pending-1>".into()), TestBox("<pending-2>".into())]);

        let task = tokio::spawn(async move {
            let result = shim.infer("request").await;
            (shim, result)
        });
        wait_for_lines(&app.path("responses.log"), 3).await;
        assert!(!task.is_finished(), "infer returned before Event::Done");
        std::fs::write(app.path("release"), "").unwrap();
        let (shim, output) = tokio::time::timeout(std::time::Duration::from_secs(3), task)
            .await
            .unwrap()
            .unwrap();
        let output = output.unwrap();

        assert_eq!(
            output.items,
            vec![
                ShimItem::Text("preface".into()),
                ShimItem::Box(TestBox("A".into())),
                ShimItem::Text("middle".into()),
                ShimItem::Box(TestBox("B".into())),
                ShimItem::Box(TestBox("C".into())),
                ShimItem::Text("tailend".into()),
            ]
        );
        assert_eq!(shim.codec.conversions, ["A", "B", "C"]);
        assert_eq!(shim.pending_box_count(), 0);

        let start: Value =
            serde_json::from_str(&std::fs::read_to_string(app.path("turn-start.log")).unwrap())
                .unwrap();
        assert_eq!(
            start
                .pointer("/params/input/0/text")
                .and_then(Value::as_str),
            Some("<pending-1>\n<pending-2>\nrequest")
        );

        let responses = std::fs::read_to_string(app.path("responses.log")).unwrap();
        for (line, id) in responses.lines().zip([77, 78, 79]) {
            let response: Value = serde_json::from_str(line).unwrap();
            assert_eq!(response["id"], id);
            assert_eq!(response["result"]["success"], true);
            assert_eq!(
                response
                    .pointer("/result/contentItems/0/text")
                    .and_then(Value::as_str),
                Some(ASYNC_TOOL_ACKNOWLEDGEMENT)
            );
        }
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn normal_inference_uses_a_fresh_turn_each_time() {
        let app = TestApp::new(
            r#"#!/bin/sh
set -eu
read initialize
echo '{"id":0,"result":{}}'
read initialized
read thread_start
echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
read first_start
printf '%s\n' "$first_start" >> starts.log
echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"hel"}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"lo"}}'
echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
read second_start
printf '%s\n' "$second_start" >> starts.log
echo '{"id":3,"result":{"turn":{"id":"turn-2"}}}'
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-2","delta":"again"}}'
echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-2","status":"completed"}}}'
"#,
        );
        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());

        assert_eq!(
            shim.infer("first").await.unwrap().items,
            [ShimItem::Text("hello".into())]
        );
        assert_eq!(
            shim.infer("second").await.unwrap().items,
            [ShimItem::Text("again".into())]
        );
        let starts = std::fs::read_to_string(app.path("starts.log")).unwrap();
        assert_eq!(starts.lines().count(), 2);
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn finite_thousand_call_sequence_is_not_truncated() {
        let app = TestApp::new(
            r#"#!/bin/sh
set -eu
read initialize
echo '{"id":0,"result":{}}'
read initialized
read thread_start
echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
read turn_start
echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
i=1
while [ "$i" -le 1000 ]; do
  echo "{\"id\":$((1000 + i)),\"method\":\"item/tool/call\",\"params\":{\"threadId\":\"thread-1\",\"turnId\":\"turn-1\",\"callId\":\"call-$i\",\"tool\":\"lookup\",\"arguments\":{\"n\":$i}}}"
  read response
  i=$((i + 1))
done
echo '{"method":"turn/completed","params":{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed"}}}'
"#,
        );
        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());

        let output = shim.infer("burst").await.unwrap();
        assert_eq!(output.items.len(), 1000);
        for (index, item) in output.items.iter().enumerate() {
            assert_eq!(item, &ShimItem::Box(TestBox(format!("call-{}", index + 1))));
        }
        assert_eq!(shim.codec.conversions.len(), 1000);
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn response_failure_returns_no_output_and_poisons_reuse() {
        let app = TestApp::new(
            r#"#!/bin/sh
set -eu
read initialize
echo '{"id":0,"result":{}}'
read initialized
read thread_start
echo '{"id":1,"result":{"thread":{"id":"thread-1"}}}'
read turn_start
echo '{"id":2,"result":{"turn":{"id":"turn-1"}}}'
exec 0<&-
echo '{"method":"item/agentMessage/delta","params":{"threadId":"thread-1","turnId":"turn-1","delta":"partial"}}'
echo '{"id":77,"method":"item/tool/call","params":{"threadId":"thread-1","turnId":"turn-1","callId":"A","tool":"lookup","arguments":{}}}'
sleep 1
"#,
        );
        let mut shim = Shim::new(app.adapter().await, "conversation-1", TestCodec::default());

        assert!(shim.infer("request").await.is_err());
        let reuse = shim.infer("must reject").await.unwrap_err();
        assert!(
            reuse.message.contains("cannot be reused"),
            "unexpected reuse error: {reuse}"
        );
    }
}