omk 0.5.0

A Rust runtime for Kimi CLI. Turns prompts into proof-backed engineering runs with gates, worktrees, and replay.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use anyhow::Context;
use std::collections::VecDeque;
use tokio::process::{Child, ChildStdin, ChildStdout};
use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::sync::CancellationToken;

mod client_trait;
mod dispatch;
mod io;
mod process_impl;
mod spawn;

pub use client_trait::{InMemoryWireClient, WireClient};
pub use dispatch::{process_messages, WireMessage, WireResponse};

const LEGACY_NO_HANDSHAKE_PROTOCOL_VERSION: &str = "legacy/no-handshake";

/// Maximum wire-line length in bytes.
///
/// Each Kimi wire message arrives as a single newline-terminated JSON line.
/// Without a hard cap, a peer that never emits a newline can drive
/// `read_line` to allocate until OOM. 16 MiB is generous enough for any
/// realistic LLM turn (most are well under 100 KB) yet still bounds the
/// damage a misbehaving / hostile producer can cause. The same cap is used
/// for inbound MCP requests in `crate::mcp::server`.
pub(crate) const MAX_WIRE_LINE_LENGTH: usize = 16 * 1024 * 1024;

/// A client for communicating with Kimi Code CLI in wire mode via a child process.
#[derive(Debug)]
pub struct ProcessWireClient {
    pub(crate) child: Child,
    pub(crate) stdin: ChildStdin,
    pub(crate) stdout_reader: FramedRead<ChildStdout, LinesCodec>,
    pub(crate) pending_messages: VecDeque<WireMessage>,
    pub(crate) request_id_counter: u64,
    pub(crate) handshake_done: bool,
    pub(crate) stderr_handle: Option<tokio::task::JoinHandle<()>>,
    pub(crate) cancel_token: CancellationToken,
}

impl Drop for ProcessWireClient {
    fn drop(&mut self) {
        self.cancel_token.cancel();
        if let Some(handle) = self.stderr_handle.take() {
            handle.abort();
        }
    }
}

#[cfg(test)]
mod tests;

// Shared helpers used by process_impl.rs and client_trait.rs
pub(crate) fn wire_message_id(msg: &WireMessage) -> Option<&str> {
    match msg {
        WireMessage::SuccessResponse(resp) => Some(resp.id.as_str()),
        WireMessage::ErrorResponse(resp) => Some(resp.id.as_str()),
        WireMessage::Request(req) => Some(req.id.as_str()),
        WireMessage::Event(_) => None,
    }
}

pub(crate) fn decode_response<T: serde::de::DeserializeOwned>(
    msg: WireMessage,
    expected_id: &str,
) -> anyhow::Result<T> {
    match msg {
        WireMessage::SuccessResponse(resp) if resp.id == expected_id => {
            serde_json::from_value(resp.result).context("Failed to decode response result")
        }
        WireMessage::ErrorResponse(resp) if resp.id == expected_id => bail_wire_error(resp),
        other => anyhow::bail!(
            "Buffered wire message did not match expected response id {expected_id}: {other:?}"
        ),
    }
}

pub(crate) fn bail_wire_error<T>(
    resp: crate::wire::protocol::JsonRpcErrorResponse,
) -> anyhow::Result<T> {
    anyhow::bail!(
        "Wire request failed: {} (code: {})",
        resp.error.message,
        resp.error.code
    )
}