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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use async_trait::async_trait;
use std::sync::Arc;
use tokio::sync::broadcast;

use crate::runtime::classifier::{ClassifierInput, ClassifierOutput};
use crate::runtime::goal::chat_api::{ChildGoalEvent, ChildGoalHandle, CreateChildRequest};

#[async_trait]
pub trait ClassifierBackend: Send + Sync {
    async fn classify(&self, input: ClassifierInput) -> ClassifierOutput;
}

#[async_trait]
pub trait LlmDirectBackend: Send + Sync {
    async fn answer_direct(
        &self,
        prompt: &str,
        context: &[crate::runtime::classifier::ConversationTurn],
    ) -> anyhow::Result<u32>;
}

#[derive(Debug)]
pub struct SmallEditResult {
    pub worker_id: String,
    pub files_touched: u32,
    pub diff_summary: String,
}

#[derive(Debug)]
pub struct MediumPlanResult {
    pub workers: Vec<String>,
    pub steps_completed: u32,
    pub steps_failed: u32,
}

#[async_trait]
pub trait WireWorkerBackend: Send + Sync {
    async fn run_small_edit(&self, task: &str) -> anyhow::Result<SmallEditResult>;
    async fn run_medium_plan(&self, plan: &[String]) -> anyhow::Result<MediumPlanResult>;
}

#[async_trait]
pub trait GoalBridgeBackend: Send + Sync {
    async fn create_child(&self, req: CreateChildRequest) -> anyhow::Result<ChildGoalHandle>;
    async fn subscribe(&self, goal_id: &str)
        -> anyhow::Result<broadcast::Receiver<ChildGoalEvent>>;
}

// Production wrappers

pub struct ProductionClassifierBackend {
    pub inner: Arc<dyn crate::runtime::classifier::LlmClassifierBackend>,
    pub cache: tokio::sync::Mutex<lru::LruCache<u64, ClassifierOutput>>,
}

impl std::fmt::Debug for ProductionClassifierBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProductionClassifierBackend")
            .finish_non_exhaustive()
    }
}

#[async_trait]
impl ClassifierBackend for ProductionClassifierBackend {
    async fn classify(&self, input: ClassifierInput) -> ClassifierOutput {
        let mut cache = self.cache.lock().await;
        crate::runtime::classifier::classify(input, self.inner.as_ref(), &mut cache).await
    }
}

pub struct ProductionLlmDirectBackend<W: crate::llm::LlmClient> {
    pub inner: Arc<tokio::sync::Mutex<W>>,
}

impl<W: crate::llm::LlmClient> std::fmt::Debug for ProductionLlmDirectBackend<W> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProductionLlmDirectBackend")
            .finish_non_exhaustive()
    }
}

#[async_trait]
impl<W: crate::llm::LlmClient + Send + Sync> LlmDirectBackend for ProductionLlmDirectBackend<W> {
    async fn answer_direct(
        &self,
        prompt: &str,
        _context: &[crate::runtime::classifier::ConversationTurn],
    ) -> anyhow::Result<u32> {
        let budget = crate::llm::types::TokenBudget::new(usize::MAX);
        let _resp = self.inner.lock().await.complete(prompt, &budget).await?;
        Ok(0)
    }
}

#[derive(Debug)]
pub struct ProductionGoalBridgeBackend;

#[async_trait]
impl GoalBridgeBackend for ProductionGoalBridgeBackend {
    async fn create_child(&self, req: CreateChildRequest) -> anyhow::Result<ChildGoalHandle> {
        crate::runtime::goal::chat_api::create_child(req).await
    }

    async fn subscribe(&self, id: &str) -> anyhow::Result<broadcast::Receiver<ChildGoalEvent>> {
        crate::runtime::goal::chat_api::subscribe(id)
    }
}