use std::{
collections::VecDeque,
future::Future,
sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use kcode_k1_chat_core::{
Call, ChatError, ChatEvent, ChatView, CompactFuture, CompactRequest, Inference, Llm, LlmError,
LlmFuture, LlmThread, PendingAction, Runtime, ToolMode, ToolOutput, ToolRequest, ToolStart,
Updates, WorkerRequest, WorkerStart,
};
use tokio::sync::{Notify, mpsc};
pub trait Candidate {
type Chat: Send + Sync + 'static;
fn open(
runtime: Arc<dyn Runtime>,
initial_primary: String,
llm: Arc<dyn Llm>,
) -> (Self::Chat, mpsc::UnboundedReceiver<ChatEvent>);
fn append(
chat: &Self::Chat,
text: String,
) -> impl Future<Output = Result<(), ChatError>> + Send;
fn restart(chat: &Self::Chat) -> impl Future<Output = Result<(), ChatError>> + Send;
fn view(chat: &Self::Chat) -> impl Future<Output = Result<ChatView, ChatError>> + Send;
fn finalize(
chat: Self::Chat,
) -> impl Future<Output = Result<ChatView, ChatError>> + Send + 'static;
}
struct Reply {
gate: Option<Arc<Notify>>,
result: Result<Inference, LlmError>,
}
struct ScriptedInner {
starts: AtomicUsize,
replies: Mutex<VecDeque<Reply>>,
deltas: Mutex<Vec<(usize, String)>>,
}
#[derive(Clone)]
struct Scripted(Arc<ScriptedInner>);
impl Scripted {
fn new(replies: Vec<Reply>) -> Arc<Self> {
Arc::new(Self(Arc::new(ScriptedInner {
starts: AtomicUsize::new(0),
replies: Mutex::new(replies.into()),
deltas: Mutex::new(Vec::new()),
})))
}
fn deltas(&self) -> Vec<(usize, String)> {
self.0.deltas.lock().unwrap().clone()
}
fn starts(&self) -> usize {
self.0.starts.load(Ordering::SeqCst)
}
}
struct ScriptedThread {
owner: Scripted,
id: usize,
}
impl Llm for Scripted {
fn start(&self) -> Box<dyn LlmThread> {
let id = self.0.starts.fetch_add(1, Ordering::SeqCst) + 1;
Box::new(ScriptedThread {
owner: self.clone(),
id,
})
}
}
impl LlmThread for ScriptedThread {
fn infer<'a>(&'a mut self, delta: &'a str) -> LlmFuture<'a> {
self.owner
.0
.deltas
.lock()
.unwrap()
.push((self.id, delta.to_owned()));
let reply = self
.owner
.0
.replies
.lock()
.unwrap()
.pop_front()
.expect("scripted LLM reply exhausted");
Box::pin(async move {
if let Some(gate) = reply.gate {
gate.notified().await;
}
reply.result
})
}
}
struct ToolPlan {
mode: ToolMode,
queued: String,
result: String,
gate: Option<Arc<Notify>>,
activity: Option<String>,
}
#[derive(Default)]
struct FixtureRuntime {
tools: Mutex<VecDeque<(String, ToolPlan)>>,
}
impl FixtureRuntime {
fn install(&self, name: &str, plan: ToolPlan) {
self.tools
.lock()
.unwrap()
.push_back((name.to_owned(), plan));
}
}
impl Runtime for FixtureRuntime {
fn start_tool(&self, request: ToolRequest, updates: Updates) -> Result<ToolStart, String> {
let mut tools = self.tools.lock().unwrap();
let index = tools
.iter()
.position(|(name, _)| name == &request.name)
.ok_or_else(|| format!("missing tool plan: {}", request.name))?;
let (_, plan) = tools.remove(index).expect("tool plan index disappeared");
Ok(ToolStart {
mode: plan.mode,
queued: plan.queued,
future: Box::pin(async move {
if let Some(activity) = plan.activity {
let _ = updates.activity(activity).send();
}
if let Some(gate) = plan.gate {
gate.notified().await;
}
ToolOutput {
text: plan.result,
cost_cents: Default::default(),
}
}),
})
}
fn start_worker(&self, _request: &WorkerRequest) -> Result<WorkerStart, String> {
Err("workers are unavailable in this verifier".to_owned())
}
fn compact(&self, _request: CompactRequest, _primary: String) -> CompactFuture {
Box::pin(async { Err("compaction is unavailable in this verifier".to_owned()) })
}
}
fn success(text: &str, calls: Vec<Call>) -> Reply {
Reply {
gate: None,
result: Ok(Inference {
text: text.to_owned(),
calls,
continue_inference: false,
}),
}
}
fn gated(gate: Arc<Notify>, text: &str, calls: Vec<Call>) -> Reply {
Reply {
gate: Some(gate),
result: success(text, calls).result,
}
}
fn transient(text: &str) -> Reply {
Reply {
gate: None,
result: Err(LlmError::Transient(text.to_owned())),
}
}
fn permanent(text: &str) -> Reply {
Reply {
gate: None,
result: Err(LlmError::Permanent(text.to_owned())),
}
}
fn tool(name: &str) -> Call {
Call::Tool(ToolRequest {
name: name.to_owned(),
input: String::new(),
})
}
fn plan(
mode: ToolMode,
queued: &str,
result: &str,
gate: Option<Arc<Notify>>,
activity: Option<&str>,
) -> ToolPlan {
ToolPlan {
mode,
queued: queued.to_owned(),
result: result.to_owned(),
gate,
activity: activity.map(str::to_owned),
}
}
async fn settle() {
for _ in 0..30 {
tokio::task::yield_now().await;
}
}
pub fn verify_initial_primary_delta_output_order_and_no_self_trigger<C: Candidate>() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("failed to build verifier runtime");
runtime.block_on(async {
let gate = Arc::new(Notify::new());
let llm = Scripted::new(vec![
gated(gate.clone(), "o", Vec::new()),
success("", Vec::new()),
success("", Vec::new()),
]);
let (chat, mut events) = C::open(
Arc::new(FixtureRuntime::default()),
"i".to_owned(),
llm.clone(),
);
settle().await;
assert!(llm.deltas().is_empty());
assert_eq!(C::view(&chat).await.unwrap().primary, "i");
assert_eq!(C::append(&chat, String::new()).await, Err(ChatError::Empty));
C::append(&chat, "u".to_owned()).await.unwrap();
settle().await;
assert_eq!(llm.deltas(), vec![(1, "iu".to_owned())]);
C::append(&chat, "a".to_owned()).await.unwrap();
assert_eq!(C::view(&chat).await.unwrap().pending, "a");
gate.notify_one();
settle().await;
assert_eq!(llm.deltas()[1], (1, "a".to_owned()));
assert_eq!(C::view(&chat).await.unwrap().primary, "iuoa");
assert_eq!(events.try_recv(), Ok(ChatEvent::Text("o".to_owned())));
settle().await;
assert_eq!(llm.deltas().len(), 2);
C::append(&chat, "v".to_owned()).await.unwrap();
settle().await;
assert_eq!(llm.deltas()[2], (1, "v".to_owned()));
assert_eq!(C::view(&chat).await.unwrap().primary, "iuoav");
assert_eq!(C::finalize(chat).await.unwrap().primary, "iuoav");
assert_eq!(events.recv().await, None);
});
}
pub fn verify_retry_schedule_stall_pending_and_fresh_restart<C: Candidate>() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("failed to build verifier runtime");
runtime.block_on(async {
let mut replies = (1..=5)
.map(|number| transient(&number.to_string()))
.collect::<Vec<_>>();
replies.push(success("", Vec::new()));
let llm = Scripted::new(replies);
let (chat, mut events) = C::open(
Arc::new(FixtureRuntime::default()),
"i".to_owned(),
llm.clone(),
);
C::append(&chat, "u".to_owned()).await.unwrap();
settle().await;
assert_eq!(
C::view(&chat).await.unwrap().actions,
vec![PendingAction::Inference { attempt: 1 }]
);
for (wait, count) in [(10, 2), (20, 3), (40, 4), (80, 5)] {
tokio::time::advance(Duration::from_secs(wait - 1)).await;
settle().await;
assert_eq!(llm.deltas().len(), count - 1);
assert_eq!(
C::view(&chat).await.unwrap().actions,
vec![PendingAction::Inference {
attempt: (count - 1) as u8
}]
);
tokio::time::advance(Duration::from_secs(1)).await;
settle().await;
assert_eq!(llm.deltas().len(), count);
if count < 5 {
assert_eq!(
C::view(&chat).await.unwrap().actions,
vec![PendingAction::Inference {
attempt: count as u8
}]
);
}
}
assert_eq!(llm.deltas(), vec![(1, "iu".to_owned()); 5]);
assert_eq!(events.try_recv(), Ok(ChatEvent::Stalled("5".to_owned())));
C::append(&chat, "later".to_owned()).await.unwrap();
assert_eq!(C::view(&chat).await.unwrap().pending, "later");
C::restart(&chat).await.unwrap();
settle().await;
assert_eq!(llm.starts(), 2);
assert_eq!(llm.deltas().last(), Some(&(2, "iulater".to_owned())));
assert_eq!(C::finalize(chat).await.unwrap().primary, "iulater");
assert_eq!(events.recv().await, None);
});
}
pub fn verify_blocked_threads_are_independent<C: Candidate>() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("failed to build verifier runtime");
runtime.block_on(async {
let gate = Arc::new(Notify::new());
let blocked = Scripted::new(vec![gated(gate.clone(), "", Vec::new())]);
let free = Scripted::new(vec![success("x", Vec::new())]);
let fixtures: Arc<dyn Runtime> = Arc::new(FixtureRuntime::default());
let (first, mut first_events) = C::open(fixtures.clone(), String::new(), blocked);
let (second, mut second_events) = C::open(fixtures, String::new(), free);
C::append(&first, "a".to_owned()).await.unwrap();
C::append(&second, "b".to_owned()).await.unwrap();
settle().await;
assert_eq!(C::view(&second).await.unwrap().primary, "bx");
assert_eq!(
C::view(&first).await.unwrap().actions,
vec![PendingAction::Inference { attempt: 1 }]
);
gate.notify_one();
settle().await;
assert_eq!(C::finalize(first).await.unwrap().primary, "a");
assert_eq!(C::finalize(second).await.unwrap().primary, "bx");
assert_eq!(first_events.recv().await, None);
assert_eq!(
second_events.recv().await,
Some(ChatEvent::Text("x".to_owned()))
);
assert_eq!(second_events.recv().await, None);
});
}
pub fn verify_finalize_waits_closes_and_stalled_finalize_returns<C: Candidate>() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("failed to build verifier runtime");
runtime.block_on(async {
let gate = Arc::new(Notify::new());
let fixtures = Arc::new(FixtureRuntime::default());
fixtures.install(
"q",
plan(ToolMode::Queued, "q", "r", Some(gate.clone()), None),
);
let llm = Scripted::new(vec![
success("", vec![tool("q")]),
success("", Vec::new()),
success("", Vec::new()),
]);
let (chat, mut events) = C::open(fixtures, String::new(), llm);
C::append(&chat, "u".to_owned()).await.unwrap();
settle().await;
assert_eq!(C::view(&chat).await.unwrap().primary, "uq");
let task = tokio::spawn(C::finalize(chat));
settle().await;
assert!(!task.is_finished());
gate.notify_one();
settle().await;
assert_eq!(task.await.unwrap().unwrap().primary, "uqr");
assert_eq!(events.recv().await, None);
let llm = Scripted::new(vec![permanent("stop")]);
let (stalled, mut stalled_events) =
C::open(Arc::new(FixtureRuntime::default()), String::new(), llm);
C::append(&stalled, "u".to_owned()).await.unwrap();
settle().await;
let view = tokio::time::timeout(Duration::from_secs(1), C::finalize(stalled))
.await
.expect("stalled finalization timed out")
.expect("stalled finalization failed");
assert_eq!(view.primary, "u");
assert_eq!(
stalled_events.try_recv(),
Ok(ChatEvent::Stalled("stop".to_owned()))
);
assert_eq!(stalled_events.recv().await, None);
});
}
pub fn verify_dropped_text_and_activity_receivers_stall_cleanly<C: Candidate>() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("failed to build verifier runtime");
runtime.block_on(async {
let llm = Scripted::new(vec![success("text", Vec::new()), success("", Vec::new())]);
let (chat, events) = C::open(
Arc::new(FixtureRuntime::default()),
"i".to_owned(),
llm.clone(),
);
drop(events);
C::append(&chat, "u".to_owned()).await.unwrap();
settle().await;
C::restart(&chat).await.unwrap();
settle().await;
assert_eq!(llm.starts(), 2);
assert_eq!(llm.deltas()[1], (2, "iutext".to_owned()));
assert_eq!(C::finalize(chat).await.unwrap().primary, "iutext");
let fixtures = Arc::new(FixtureRuntime::default());
fixtures.install(
"status",
plan(ToolMode::Fast, "unused", "R", None, Some("status")),
);
let llm = Scripted::new(vec![
success("", vec![tool("status")]),
success("", Vec::new()),
]);
let (chat, events) = C::open(fixtures, String::new(), llm.clone());
drop(events);
C::append(&chat, "u".to_owned()).await.unwrap();
settle().await;
let view = C::view(&chat).await.unwrap();
assert!(view.actions.is_empty());
assert_eq!(view.pending, "R");
C::restart(&chat).await.unwrap();
settle().await;
assert_eq!(llm.starts(), 2);
assert_eq!(llm.deltas()[1], (2, "uR".to_owned()));
assert_eq!(C::finalize(chat).await.unwrap().primary, "uR");
});
}