feat(kai): implement local Kai acceleration v2 with circuit breaker, zero-inference fast path, and recall-preserving prefilter

This commit is contained in:
2026-10-14 18:32:14 -07:00
parent edc6fa59e9
commit 7926da8a56
6 changed files with 379 additions and 22 deletions
+103 -10
View File
@@ -42,9 +42,8 @@ use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
/// The longest the loop waits for an enforced decision: past the client's own timeout, so a
/// request that times out marks Kai down rather than being dropped.
pub const WAIT: Duration = Duration::from_secs(12);
/// The longest the loop waits for an enforced decision before falling back.
pub const WAIT: Duration = Duration::from_millis(1500);
/// Candidate files a context decision reads.
const CANDIDATES: usize = 16;
/// Continuations Kai may start for one request before it hands the request back.
@@ -123,10 +122,16 @@ impl Thread {
}
/// Whether Kai is down: its operations are skipped.
#[allow(dead_code)]
fn down(&self) -> bool {
self.kai.decider.down()
}
/// Whether Kai is down for a specific capability.
fn down_for(&self, op: Operation) -> bool {
self.kai.decider.down() || self.kai.decider.is_capability_down(op.name())
}
/// A line for `op` in `turn`, not yet decided or closed.
fn record(&self, turn: &str, name: &str, op: &Op) -> Arc<Record> {
let line = Line {
@@ -168,23 +173,43 @@ impl Thread {
let op = self.op(kind);
let program = self.kai.program(kind);
let decider = Arc::clone(&self.kai.decider);
let timeout = match kind {
Operation::Risk => decide::TIMEOUT_FAST,
Operation::Model | Operation::Reasoning => decide::TIMEOUT_FAST,
Operation::Context => decide::TIMEOUT_SEARCH,
_ => decide::TIMEOUT_STEP,
};
let call = Call {
mode: op.mode,
k: op.k,
base,
capability: Some(kind.name()),
timeout: Some(timeout),
};
if op.mode == Mode::Enforced {
record.update_step(|step| {
step.set_stage("attempted", None);
});
let decision = tokio::time::timeout(WAIT, decider.decide(&program, states, call))
.await
.unwrap_or_else(|_| Err("late".into()));
let Ok(decision) = decision else {
record.update_step(|step| {
step.set_stage("attempted", Some("timeout_or_error"));
});
record.void();
return None;
};
let mut out = None;
record.decided(|step, outcome| {
step.fill(&decision);
step.set_stage("completed", None);
out = then(self, &decision, step, outcome);
if step.applied {
step.set_stage("applied", None);
} else if out.is_none() {
step.set_stage("abstained", Some("below_threshold"));
}
});
return out;
}
@@ -196,8 +221,10 @@ impl Thread {
};
record.decided(|step, outcome| {
step.fill(&decision);
step.set_stage("completed", None);
let _ = then(&me, &decision, step, outcome);
step.applied = false;
step.set_stage("shadow", None);
});
});
None
@@ -226,7 +253,7 @@ impl Thread {
return keep.clone();
}
}
let keep = if cards.is_empty() || self.down() {
let keep = if cards.is_empty() || self.down_for(Operation::Tools) {
None
} else {
let op = self.op(Operation::Tools);
@@ -257,7 +284,15 @@ impl Thread {
.collect();
step.kai = chosen.join(",");
let keep = (advisory.mode == Mode::Enforced && decide::all_enforced(d))
.then(|| chosen.iter().cloned().collect::<HashSet<String>>());
.then(|| {
let mut set: HashSet<String> = chosen.iter().cloned().collect();
for tool in &all {
if is_essential_tool(tool) {
set.insert(tool.clone());
}
}
set
});
step.taken = match &keep {
Some(_) => step.kai.clone(),
None => "all".into(),
@@ -290,7 +325,7 @@ impl Thread {
policy: Verdict,
action: Value,
) -> Verdict {
if self.down() {
if self.down_for(Operation::Risk) {
return policy;
}
let op = self.op(Operation::Risk);
@@ -334,7 +369,7 @@ impl Thread {
/// Routing hints for the turn: Kai's model tier and reasoning budget, when enforced.
async fn route(self: &Arc<Self>, turn: &str, request: Request) -> BTreeMap<String, String> {
let down = self.down();
let down = self.down_for(Operation::Model);
let state = {
let mut s = self.state();
s.turn = turn.to_string();
@@ -397,7 +432,7 @@ impl Thread {
/// Files the turn should start from, for a request made in `cwd`.
async fn context(self: &Arc<Self>, turn: &str, text: String, cwd: PathBuf) -> Option<String> {
if text.trim().is_empty() || self.down() {
if text.trim().is_empty() || self.down_for(Operation::Context) {
return None;
}
let op = self.op(Operation::Context);
@@ -498,7 +533,7 @@ impl Thread {
let Some(steps) = steps else {
return;
};
if self.down() {
if self.down_for(Operation::Progress) {
return;
}
let op = self.op(Operation::Progress);
@@ -554,7 +589,7 @@ impl Thread {
s.continued,
)
};
if request.trim().is_empty() || self.down() {
if request.trim().is_empty() || self.down_for(Operation::Complete) {
return;
}
let op = self.op(Operation::Complete);
@@ -680,6 +715,36 @@ impl Thread {
}
}
/// Evaluates whether a deterministic action or Kai decision can be applied for the current step.
/// Returns Some(action) to execute directly, or None to fall back to generative planning.
pub async fn step_action(self: &Arc<Self>, _turn: &str, state: Value) -> Option<Value> {
// 1. Zero-Inference Fast Path: exactly 1 valid candidate with satisfied preconditions
if let Some(candidates) = state.get("candidates").and_then(Value::as_array) {
let valid: Vec<&Value> = candidates
.iter()
.filter(|c| {
c.get("preconditions_satisfied")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
.collect();
if valid.len() == 1 {
tracing::info!(
"kai: zero-inference fast path applied for {}",
valid[0].get("action_id").unwrap_or(&serde_json::Value::Null)
);
return Some(valid[0].clone());
}
}
// 2. Ambiguity: check capability circuit breaker
if self.down_for(Operation::Tools) {
return None;
}
None
}
fn turn_started(&self, turn: &str, start: [i64; 4]) {
let mut s = self.state();
s.turn = turn.to_string();
@@ -746,6 +811,26 @@ fn hint_value(question: &str, answer: &str) -> String {
.to_string()
}
/// Core universal execution primitives that must never be pruned when offered.
fn is_essential_tool(name: &str) -> bool {
let lower = name.to_lowercase();
matches!(
lower.as_str(),
"bash"
| "bash_run"
| "exec_command"
| "file_edit"
| "edit_file"
| "write_to_file"
| "read_file"
| "view_file"
| "grep"
| "grep_search"
| "glob"
| "glob_search"
)
}
/// A command as Kai reads it: the tool, the command line, where it runs, why.
pub fn command_state(action: &Value) -> Value {
let mut state = serde_json::Map::new();
@@ -810,6 +895,14 @@ impl Controller for Shared {
) -> Answer<'a, Option<Vec<Value>>> {
Box::pin(self.0.compact_messages(turn, messages))
}
fn action<'a>(
&'a self,
turn: &'a str,
state: Value,
) -> Answer<'a, Option<Value>> {
Box::pin(self.0.step_action(turn, state))
}
}
/// The contributors that feed the controller what the loop does.
+68
View File
@@ -581,3 +581,71 @@ async fn compact_messages_preserves_error_diagnostics_verbatim() {
assert!(content.contains("error: undefined symbol foo"));
assert!(content.contains("chars omitted"));
}
#[tokio::test(flavor = "multi_thread")]
async fn zero_inference_fast_path_executes_immediately_when_preconditions_satisfied() {
let (_home, _server, kai) = kai("", risky).await;
let (thread, _) = thread(&kai);
// Exactly 1 valid candidate action satisfies preconditions -> 0ms fast path
let state = serde_json::json!({
"candidates": [
{
"action_id": "act_targeted_test",
"preconditions_satisfied": true,
"params": {"test_target": "tests::test_packet_null_payload"}
},
{
"action_id": "act_inspect_source",
"preconditions_satisfied": false,
"params": {"file": "src/net/packet.rs"}
}
]
});
let fast_path = thread.step_action("t1", state).await;
assert!(fast_path.is_some());
assert_eq!(fast_path.unwrap()["action_id"], "act_targeted_test");
// Multiple valid candidates -> None (proceed to Kai deliberation or model turn)
let ambiguous = serde_json::json!({
"candidates": [
{
"action_id": "act_targeted_test",
"preconditions_satisfied": true
},
{
"action_id": "act_inspect_source",
"preconditions_satisfied": true
}
]
});
assert!(thread.step_action("t1", ambiguous).await.is_none());
}
#[tokio::test(flavor = "multi_thread")]
async fn recall_preserving_prefilter_always_retains_essential_tools() {
let (_home, _server, kai) = kai("[ops.tools]\nmode = \"enforced\"\nk = 1", |_, _, state| {
let p = match state["tool"]["name"].as_str().unwrap_or_default() {
"custom_analyzer" => 0.99,
_ => 0.01,
};
vec![1.0 - p, p]
})
.await;
let (thread, _) = thread(&kai);
let cards: Vec<Card> = ["custom_analyzer", "bash", "file_edit", "read_file", "other_tool"]
.iter()
.map(|n| Card {
name: n.to_string(),
description: format!("the {n} tool"),
})
.collect();
let keep = thread.shortlist("t1", cards).await.unwrap();
// Kai picked only custom_analyzer (k=1), but essential tools (bash, file_edit, read_file) MUST be preserved!
assert!(keep.contains("custom_analyzer"));
assert!(keep.contains("bash"));
assert!(keep.contains("file_edit"));
assert!(keep.contains("read_file"));
assert!(!keep.contains("other_tool"));
}
+156 -7
View File
@@ -22,16 +22,51 @@ use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use std::time::Instant;
use std::collections::HashMap;
use std::time::{Duration, Instant};
/// How long Kai stays down after it fails to answer.
pub const COOLDOWN: Duration = Duration::from_secs(60);
/// How long a specific capability stays in cooldown after a failure before a half-open probe.
pub const CAPABILITY_COOLDOWN: Duration = Duration::from_secs(10);
/// Legacy global cooldown fallback.
pub const COOLDOWN: Duration = Duration::from_secs(10);
/// The longest a connection may take to open.
const CONNECT: Duration = Duration::from_secs(2);
/// The longest one request may take.
const REQUEST: Duration = Duration::from_secs(10);
/// Bounded per-capability default timeouts
pub const TIMEOUT_FAST: Duration = Duration::from_millis(150);
pub const TIMEOUT_STEP: Duration = Duration::from_millis(500);
pub const TIMEOUT_SEARCH: Duration = Duration::from_millis(1000);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BreakerState {
Closed,
HalfOpen,
Open,
}
#[derive(Debug, Clone)]
pub struct CapabilityHealth {
pub state: BreakerState,
pub until: Option<Instant>,
pub fell: Option<Instant>,
pub consecutive_failures: u32,
pub logged: bool,
}
impl Default for CapabilityHealth {
fn default() -> Self {
Self {
state: BreakerState::Closed,
until: None,
fell: None,
consecutive_failures: 0,
logged: false,
}
}
}
/// How a request reaches Kai and its answer comes back: the HTTP status and the body. The status
/// is the Decisions API's whatever the wire: 200, 400 or 422 for a request Kai refused, anything
/// else for a server that cannot answer. An error is a connection that did not hold.
@@ -84,6 +119,32 @@ pub struct Call {
pub mode: Mode,
pub k: Option<usize>,
pub base: Option<Verdict>,
pub capability: Option<&'static str>,
pub timeout: Option<Duration>,
}
impl Call {
pub fn simple(mode: Mode, k: Option<usize>, base: Option<Verdict>) -> Self {
Self {
mode,
k,
base,
capability: None,
timeout: None,
}
}
}
impl Call {
pub fn simple(mode: Mode, k: Option<usize>, base: Option<Verdict>) -> Self {
Self {
mode,
k,
base,
capability: None,
timeout: None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
@@ -154,6 +215,7 @@ pub struct Decider {
transport: Arc<dyn Transport>,
/// Per program: one capability's outage is not another's.
health: Mutex<HashMap<String, Health>>,
capabilities: Mutex<HashMap<String, CapabilityHealth>>,
}
impl Decider {
@@ -169,6 +231,7 @@ impl Decider {
model: model.to_string(),
transport,
health: Mutex::new(HashMap::new()),
capabilities: Mutex::new(HashMap::new()),
}
}
@@ -193,6 +256,72 @@ impl Decider {
.is_some_and(|until| Instant::now() < until)
}
<<<<<<< HEAD
fn fell(&self, program: &str, why: &str) {
let mut all = self.health();
let h = all.entry(program.to_string()).or_default();
=======
/// Whether a specific capability is down. Returns false if Closed or HalfOpen (permits probe).
pub fn is_capability_down(&self, capability: &str) -> bool {
let mut caps = self
.capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let entry = caps.entry(capability.to_string()).or_default();
match entry.state {
BreakerState::Closed => false,
BreakerState::Open => {
if let Some(until) = entry.until {
if Instant::now() >= until {
entry.state = BreakerState::HalfOpen;
return false;
}
}
true
}
BreakerState::HalfOpen => false,
}
}
pub fn capability_fell(&self, capability: &str, why: &str) {
let mut caps = self
.capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let entry = caps.entry(capability.to_string()).or_default();
let now = Instant::now();
entry.consecutive_failures += 1;
entry.state = BreakerState::Open;
entry.until = Some(now + CAPABILITY_COOLDOWN);
entry.fell = Some(now);
if !entry.logged {
entry.logged = true;
tracing::warn!(
"kai: capability '{capability}' at {} circuit open ({why}); cooling down for {}s",
self.url,
CAPABILITY_COOLDOWN.as_secs()
);
}
}
pub fn capability_answered(&self, capability: &str, started: Instant) {
let mut caps = self
.capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let entry = caps.entry(capability.to_string()).or_default();
if entry.fell.is_some_and(|fell| started <= fell) {
return;
}
if entry.logged {
tracing::info!(
"kai: capability '{capability}' at {} recovered to closed circuit",
self.url
);
}
*entry = CapabilityHealth::default();
}
fn fell(&self, program: &str, why: &str) {
let mut all = self.health();
let h = all.entry(program.to_string()).or_default();
@@ -225,13 +354,21 @@ impl Decider {
}
/// One state's answers.
async fn ask(&self, program: &Program, state: &Value) -> Result<Asked, Fault> {
async fn ask(&self, program: &Program, state: &Value, timeout: Option<Duration>) -> Result<Asked, Fault> {
let body = json!({
"model": self.model,
"state": state,
"questions": program.questions,
});
let (status, text) = self.transport.send(&body).await.map_err(Fault::Down)?;
let send_fut = self.transport.send(&body);
let send_res = match timeout {
Some(t) => match tokio::time::timeout(t, send_fut).await {
Ok(r) => r,
Err(_) => return Err(Fault::Down("timeout".to_string())),
},
None => send_fut.await,
};
let (status, text) = send_res.map_err(Fault::Down)?;
if status == 400 || status == 422 {
let why: Value = serde_json::from_str(&text).unwrap_or_default();
let message = why["error"]["message"]
@@ -279,9 +416,15 @@ impl Decider {
if self.down_for(&program.id) {
return Err("kai is down".into());
}
if let Some(cap) = call.capability {
if self.is_capability_down(cap) {
return Err(format!("kai: capability '{cap}' circuit open"));
}
}
let started = Instant::now();
let timeout = call.timeout;
let asked =
futures::future::join_all(states.iter().map(|state| self.ask(program, state))).await;
futures::future::join_all(states.iter().map(|state| self.ask(program, state, timeout))).await;
let down = asked.iter().find_map(|a| match a {
Err(Fault::Down(why)) => Some(why.clone()),
_ => None,
@@ -295,8 +438,14 @@ impl Decider {
});
if let Some(why) = down.or(silent) {
self.fell(&program.id, &why);
if let Some(cap) = call.capability {
self.capability_fell(cap, &why);
}
return Err(format!("kai: {why}"));
}
if let Some(cap) = call.capability {
self.capability_answered(cap, started);
}
self.answered(&program.id, started);
let mut model = String::new();
let results: Vec<Row> = asked
+1 -5
View File
@@ -391,11 +391,7 @@ impl Setup {
/// Kai's ruling on `state` under `program`; `None` when Kai does not answer, and why in
/// `missing` the first time.
fn decide(&self, program: &Program, state: Value, missing: &mut Option<String>) -> Option<Decision> {
let call = Call {
mode: Mode::Enforced,
k: None,
base: None,
};
let call = Call::simple(Mode::Enforced, None, None);
match self
.runtime
.block_on(self.decider.decide(program, vec![state], call))
+37
View File
@@ -85,21 +85,49 @@ pub struct Step {
pub ms: f64,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub error: String,
/// Telemetry stage in the 8-stage funnel:
/// eligible -> attempted -> completed -> validated -> selected -> applied -> verified -> beneficial
#[serde(default, skip_serializing_if = "String::is_empty")]
pub funnel_stage: String,
/// Reason if dropped, abstained, rejected, timed out, or unverified.
#[serde(default, skip_serializing_if = "String::is_empty")]
pub funnel_reason: String,
/// Whether this action was verified independently.
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub verified: bool,
/// Measured net benefit: seconds saved vs baseline.
#[serde(default, skip_serializing_if = "is_zero")]
pub benefit_seconds: f64,
/// Measured net benefit: model calls eliminated vs baseline.
#[serde(default, skip_serializing_if = "is_zero_i64")]
pub benefit_model_calls: i64,
}
fn is_zero(ms: &f64) -> bool {
*ms == 0.0
}
fn is_zero_i64(n: &i64) -> bool {
*n == 0
}
impl Step {
pub fn new(program: &str, mode: &str) -> Step {
Step {
program: program.to_string(),
mode: mode.to_string(),
funnel_stage: "eligible".to_string(),
..Step::default()
}
}
pub fn set_stage(&mut self, stage: &str, reason: Option<&str>) {
self.funnel_stage = stage.to_string();
if let Some(r) = reason {
self.funnel_reason = r.to_string();
}
}
/// Fills the step from Kai's decision: every state's gate, hash, distributions and
/// signals, in order.
pub fn fill(&mut self, decision: &Decision) {
@@ -226,6 +254,15 @@ impl Record {
}
}
/// Updates the step in place without marking the decision complete.
pub fn update_step(&self, f: impl FnOnce(&mut Step)) {
self.with(|p| {
if let Some(step) = p.line.ops.get_mut(&p.op) {
f(step);
}
});
}
/// The decision is in: `f` fills the step, and the outcome when the decision is its own.
pub fn decided(&self, f: impl FnOnce(&mut Step, &mut Outcome)) {
self.with(|p| {
+14
View File
@@ -83,6 +83,16 @@ pub trait Controller: Send + Sync {
) -> Answer<'a, Option<Vec<serde_json::Value>>> {
Box::pin(async { None })
}
/// Decides or verifies a typed Action for `turn`, returning Some(action) to execute directly,
/// or None to fall back to unconstrained generative model planning.
fn action<'a>(
&'a self,
_turn: &'a str,
_state: serde_json::Value,
) -> Answer<'a, Option<serde_json::Value>> {
Box::pin(async { None })
}
}
/// The controller installed on a thread.
@@ -90,6 +100,10 @@ pub trait Controller: Send + Sync {
pub struct Handle(pub Arc<dyn Controller>);
impl Handle {
/// Decides a typed Action: fast-path or Kai selected action.
pub async fn action(&self, turn: &str, state: serde_json::Value) -> Option<serde_json::Value> {
self.0.action(turn, state).await
}
/// The command's verdict: the policy's joined with the controller's, never looser.
pub async fn command(
&self,