//! Provider boundary stream parser and workbench tool bridge.
//!
//! Lucid Code Posture:
//! - Model stream is strictly typed scratchpad (thoughts and prose deltas).
//! - The model has exactly one tool: Python code execution in workbench.
//! - Diagnostic prints and expression results are cell outputs.
//! - Explicit speech ONLY occurs via `local.send()` passing `delivery.review`.
//! - Terminal wait: `await runtime.wait(...)` is a hard host yield and state commit,
//! halting inference cleanly until external wake.
use std::{sync::Arc, time::Duration};
use async_stream::stream;
use futures::Stream;
use klbr_runtime::{
protocol::{CellStatus, Outcome},
session::SessionRuntime,
store::{now_ms as store_now_ms, ReceiptStatus},
};
use serde_json::{json, Value};
use tokio::sync::mpsc;
use crate::{
models::{LlmClient, LlmEvent, Message, Role, ToolCall, ToolDef},
runtime_turn::state::{
CellOutputRecord, ScratchpadEvent, SpeechReceipt, SpeechStatus, TurnEvent,
TurnStateMachine, YieldRecord,
},
};
/// Wall-clock milliseconds, or 0 when the host clock cannot answer.
fn now_ms() -> i64 {
store_now_ms().unwrap_or(0)
}
/// The leading line of a tool result: when the observation was made, and how long the cell
/// took when one actually ran.
///
/// Both halves are fixed facts about the past, so the line stays true when the transcript is
/// replayed later. A rendered "N minutes ago" would not: it decays into a false claim, which
/// is exactly the failure mode that makes a reader unable to tell a current observation from
/// a stale one. Currency is stated by the host clock (see `runtime.inspect`), not by history.
pub fn observation_header(observed_ms: i64, took_ms: Option<u64>) -> String {
let when = crate::harness_block::format_instant(observed_ms);
match took_ms {
Some(took_ms) => format!("[{when} · ran {}]", ran_for(took_ms)),
None => format!("[{when}]"),
}
}
/// A tool result whose cell never ran: an unexecuted call still happened at a time.
fn not_run_header() -> String {
observation_header(now_ms(), None)
}
pub fn stamped(header: &str, body: &str) -> String {
format!("{header}\n{body}")
}
pub fn ran_for(took_ms: u64) -> String {
match took_ms {
0..=999 => format!("{took_ms}ms"),
1_000..=9_999 => format!("{:.2}s", took_ms as f64 / 1000.0),
10_000..=59_999 => format!("{:.1}s", took_ms as f64 / 1000.0),
_ => {
let seconds = took_ms / 1000;
format!("{}m{:02}s", seconds / 60, seconds % 60)
}
}
}
/// The single tool exposed to the model: Python code execution in the supervised workbench.
pub fn python_tool_def() -> ToolDef {
ToolDef::function(
"python",
"Execute Python code in the persistent workbench kernel. Standard output, errors, and expression results are returned. Operator speech MUST use `await local.send(...)`. Waiting for external input or timers MUST use `await runtime.wait(...)`.",
json!({
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "The Python code snippet to execute."
}
},
"required": ["code"],
"additionalProperties": false
}),
)
}
/// An owned provider streaming task handle that retains a Tokio JoinHandle and automatically
/// aborts the background task on drop or explicit cancellation.
pub struct ProviderTaskHandle {
handle: tokio::task::JoinHandle<Result<(), anyhow::Error>>,
}
impl ProviderTaskHandle {
pub fn new(handle: tokio::task::JoinHandle<Result<(), anyhow::Error>>) -> Self {
Self { handle }
}
pub fn abort(&self) {
self.handle.abort();
}
pub fn is_finished(&self) -> bool {
self.handle.is_finished()
}
}
impl Drop for ProviderTaskHandle {
fn drop(&mut self) {
self.handle.abort();
}
}
/// Execution options for running a turn step.
/// Content identity of the exact messages a provider request was built from.
fn request_digest(messages: &[Message]) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
for message in messages {
hasher.update(serde_json::to_vec(message).unwrap_or_default());
hasher.update([0u8]);
}
format!("{:x}", hasher.finalize())
}
#[derive(Debug, Clone)]
pub struct TurnStepOptions {
pub cell_budget: Duration,
pub step_timeout: Option<Duration>,
pub purpose: klbr_runtime::protocol::Purpose,
/// The provider route this step's request actually goes to. Recorded with the
/// step's usage so a receipt names the route it measured, not the default one.
pub route_model: String,
pub route_url: String,
}
impl Default for TurnStepOptions {
fn default() -> Self {
Self {
cell_budget: Duration::from_secs(30),
step_timeout: None,
purpose: klbr_runtime::protocol::Purpose::Interactive,
route_model: String::new(),
route_url: String::new(),
}
}
}
/// A parsed tool call candidate from the model stream.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedToolCall {
pub call_id: String,
pub name: String,
pub arguments: String,
pub code: Result<String, String>,
}
/// Parse a raw model `ToolCall` into a validated workbench call.
pub fn parse_tool_call(call: &ToolCall) -> ParsedToolCall {
let call_id = call.id.clone();
let name = call.function.name.clone();
let arguments = call.function.arguments.clone();
if name != "python" {
return ParsedToolCall {
call_id,
name: name.clone(),
arguments,
code: Err(format!(
"Unsupported tool '{name}'. The only available tool is 'python'."
)),
};
}
let parsed_json: Result<Value, serde_json::Error> = serde_json::from_str(&arguments);
let code = match parsed_json {
Ok(Value::Object(map)) => match map.get("code") {
Some(Value::String(code_str)) => Ok(code_str.clone()),
Some(_) => Err("Malformed 'code' argument: expected a string.".into()),
None => Err("Missing required argument 'code'.".into()),
},
Ok(_) => Err("Malformed tool arguments: expected a JSON object.".into()),
Err(err) => Err(format!("Malformed JSON in tool arguments: {err}")),
};
ParsedToolCall {
call_id,
name,
arguments,
code,
}
}
/// Result of executing a workbench cell for a tool call.
#[derive(Debug, Clone)]
pub struct WorkbenchExecutionResult {
pub call_id: String,
pub execution_id: Option<String>,
pub cell_outputs: Vec<CellOutputRecord>,
pub speech_receipts: Vec<SpeechReceipt>,
pub yield_record: Option<YieldRecord>,
pub tool_result_message: Message,
}
/// Supervised Workbench Execution Bridge:
/// Executes a Python cell in the session runtime, extracts cell outputs,
/// collects any speech receipts emitted during execution, and detects terminal yields.
pub async fn execute_workbench_cell(
session: &Arc<SessionRuntime>,
call_id: &str,
code: String,
budget: Duration,
) -> Result<WorkbenchExecutionResult, String> {
execute_workbench_cell_with_purpose(
session,
call_id,
code,
budget,
klbr_runtime::protocol::Purpose::Interactive,
)
.await
}
pub async fn execute_workbench_cell_with_purpose(
session: &Arc<SessionRuntime>,
call_id: &str,
code: String,
budget: Duration,
purpose: klbr_runtime::protocol::Purpose,
) -> Result<WorkbenchExecutionResult, String> {
let started = std::time::Instant::now();
let report = session
.execute_with_purpose(code, budget, purpose)
.await
.map_err(|e| format!("Workbench execution failure: {e:#}"))?;
let header = observation_header(now_ms(), Some(started.elapsed().as_millis() as u64));
let mut cell_outputs = Vec::new();
let mut speech_receipts = Vec::new();
let mut yield_record = None;
let tool_result_message;
// Collect speech receipts recorded during this execution
if let Ok(effects) = session.store.execution_effects(&report.id).await {
for effect in effects {
speech_receipts.push(SpeechReceipt {
effect_id: effect.effect_id.to_string(),
status: match effect.status {
ReceiptStatus::Confirmed => SpeechStatus::Confirmed,
ReceiptStatus::Held => SpeechStatus::Held,
},
event_id: effect.event_id,
reason: effect.reason,
text: effect.text,
});
}
}
match &report.outcome {
Outcome::Cell {
status,
output,
truncated_bytes,
error,
} => {
for item in output {
cell_outputs.push(item.clone().into());
}
if let Some(err_text) = error {
cell_outputs.push(CellOutputRecord::stderr(err_text.clone()));
}
if *status == CellStatus::Yielded {
let y = session
.store
.execution_yield(&report.id)
.await
.unwrap_or(None);
let (seconds, reason) = y.unwrap_or((None, "wait".into()));
let y_rec = YieldRecord { seconds, reason };
yield_record = Some(y_rec.clone());
tool_result_message = Message::tool_result(
call_id,
stamped(&header, &format!("Yielded: {}", y_rec.reason)),
);
} else {
let mut body = String::new();
for item in output {
body.push_str(&item.text);
}
if let Some(err_text) = error {
if !body.is_empty() && !body.ends_with('\n') {
body.push('\n');
}
body.push_str("Error: ");
body.push_str(err_text);
}
if *truncated_bytes > 0 {
if !body.is_empty() && !body.ends_with('\n') {
body.push('\n');
}
body.push_str(&format!("[truncated {truncated_bytes} bytes]"));
}
if body.is_empty() {
body = match status {
CellStatus::Ok => "(executed with no output)".into(),
_ => "(no output)".into(),
};
}
tool_result_message = Message::tool_result(call_id, stamped(&header, &body));
}
}
Outcome::Failed { error } => {
cell_outputs.push(CellOutputRecord::stderr(error.clone()));
tool_result_message = Message::tool_result(
call_id,
stamped(&header, &format!("Execution failed: {error}")),
);
}
Outcome::Unknown { reason } => {
// The execution's outcome is genuinely unobserved: the operation may or may
// not have run, and its effects are withheld rather than confirmed or
// dropped. This is NOT a failure and must never be reported as one - a model
// that reads "failed" will reasonably retry, which is exactly the blind
// replay an unknown outcome forbids.
cell_outputs.push(CellOutputRecord::stderr(format!(
"Execution outcome unknown: {reason}"
)));
tool_result_message = Message::tool_result(
call_id,
stamped(
&header,
&format!(
"Execution outcome UNKNOWN: {reason}. The operation may or may not have \
run. Do not repeat it. Any messages it prepared are held, not sent; \
inspect the execution's effects to establish what actually happened \
before acting."
),
),
);
}
Outcome::Remote { .. } | Outcome::RemoteError { .. } => {
// An auxiliary remote namespace answers a host_request; it is never the
// completion of a workbench cell. Reaching this arm means a remote worker's
// outcome was routed into the cell-integration path, which is a routing
// defect rather than a turn result - report it rather than inventing output.
return Err("Unexpected remote-namespace outcome from workbench execution".into());
}
Outcome::Hook { .. } => {
return Err("Unexpected hook outcome from workbench execution".into());
}
}
Ok(WorkbenchExecutionResult {
call_id: call_id.to_string(),
execution_id: Some(report.id.to_string()),
cell_outputs,
speech_receipts,
yield_record,
tool_result_message,
})
}
/// Fully supervised turn step runner with provider ownership, cancellation, timeout,
/// non-quadratic streaming, and durable transcript persistence.
pub fn run_turn_step_supervised(
mut stream_rx: mpsc::Receiver<LlmEvent>,
provider_task: Option<ProviderTaskHandle>,
session: Arc<SessionRuntime>,
mut state: TurnStateMachine,
mut conversation: Vec<Message>,
options: TurnStepOptions,
mut cancel_rx: Option<tokio::sync::oneshot::Receiver<()>>,
) -> impl Stream<Item = (TurnEvent, TurnStateMachine, Vec<Message>)> {
stream! {
let _provider_guard = provider_task;
if let Err(e) = state.begin_stream() {
let err_msg = e.to_string();
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
return;
}
yield (TurnEvent::BeginStream, state.clone(), conversation.clone());
let mut step_thoughts = String::new();
let mut step_prose = String::new();
let mut proposed_calls = Vec::new();
let mut stream_completed = false;
// The provider's own final usage for THIS request. Absent stays absent: a
// provider that reported nothing is not a provider that reported zero.
let mut step_usage: Option<crate::models::Usage> = None;
let mut provider_error: Option<String> = None;
let stream_deadline = options.step_timeout.map(|d| tokio::time::Instant::now() + d);
loop {
let next_event = tokio::select! {
biased;
// 1. Operator cancellation / preemption
_ = async {
if let Some(rx) = cancel_rx.as_mut() {
rx.await.ok()
} else {
futures::future::pending().await
}
} => {
if let Some(ref pt) = _provider_guard {
pt.abort();
}
let _ = state.cancel();
yield (TurnEvent::Cancelled, state, conversation);
return;
}
// 2. Step timeout
_ = async {
if let Some(deadline) = stream_deadline {
tokio::time::sleep_until(deadline).await
} else {
futures::future::pending().await
}
} => {
if let Some(ref pt) = _provider_guard {
pt.abort();
}
let _ = state.timeout();
yield (TurnEvent::Timeout, state, conversation);
return;
}
// 3. Next stream item
ev = stream_rx.recv() => ev,
};
match next_event {
Some(LlmEvent::ThinkToken(t)) => {
step_thoughts.push_str(&t);
let ev = ScratchpadEvent::ThoughtDelta(t);
let _ = state.record_scratchpad(ev.clone());
yield (TurnEvent::Scratchpad(ev), state.clone(), Vec::new());
}
Some(LlmEvent::Token(t)) => {
step_prose.push_str(&t);
let ev = ScratchpadEvent::ProseDelta(t);
let _ = state.record_scratchpad(ev.clone());
yield (TurnEvent::Scratchpad(ev), state.clone(), Vec::new());
}
Some(LlmEvent::ToolCallStarted) => {
let ev = ScratchpadEvent::ToolCallStarted;
let _ = state.record_scratchpad(ev.clone());
yield (TurnEvent::Scratchpad(ev), state.clone(), Vec::new());
}
Some(LlmEvent::ToolCalls(calls)) => {
for call in &calls {
let ev = ScratchpadEvent::ToolCallProposed {
id: call.id.clone(),
name: call.function.name.clone(),
arguments: call.function.arguments.clone(),
};
let _ = state.record_scratchpad(ev.clone());
yield (TurnEvent::Scratchpad(ev), state.clone(), Vec::new());
}
proposed_calls = calls;
}
Some(LlmEvent::Usage(usage)) => {
step_usage = Some(usage);
}
Some(LlmEvent::Completed) => {
stream_completed = true;
}
Some(LlmEvent::Error(err)) => {
provider_error = Some(err);
break;
}
None => {
// Channel EOF
break;
}
}
}
if let Some(err) = provider_error {
if let Some(ref pt) = _provider_guard {
pt.abort();
}
let _ = state.fail_turn(err.clone());
yield (TurnEvent::Failed(err), state, conversation);
return;
}
// Terminal outcome: channel EOF without Completed signal is an unexpected network/stream drop
if !stream_completed {
if let Some(ref pt) = _provider_guard {
pt.abort();
}
let err = "provider stream closed unexpectedly before completion".to_string();
let _ = state.fail_turn(err.clone());
yield (TurnEvent::Failed(err), state, conversation);
return;
}
let reasoning_content = if step_thoughts.is_empty() {
None
} else {
Some(step_thoughts)
};
let prose_content = if step_prose.is_empty() {
None
} else {
Some(step_prose)
};
let scratchpad_text = prose_content.as_deref().unwrap_or("").to_string();
let tool_calls_json: Vec<Value> = proposed_calls
.iter()
.map(|c| {
json!({
"id": c.id,
"name": c.function.name,
"arguments": c.function.arguments,
})
})
.collect();
// Persist validated model response and tool-call group to Store before executing cells
let response_payload = json!({
"type": "model_response",
"role": "assistant",
"text": scratchpad_text,
"content": prose_content,
"reasoning": reasoning_content,
"tool_calls": tool_calls_json,
});
let response_event_id = match session
.store
.append(&session.id, "scratchpad", response_payload)
.await
{
Ok(id) => id,
Err(err) => {
let err_msg = format!("failed to commit model response to store: {err}");
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
return;
}
};
// P17: one durable usage record per actual provider request, written next to
// the response it belongs to. Tool-loop continuations, maintenance and peer
// turns each get their own, so a per-outer-turn receipt cannot satisfy this,
// and the request content it was measured against is named by digest rather
// than left to be inferred from position.
{
let request_sha256 = request_digest(&conversation);
let usage_payload = json!({
"type": "request_usage",
"purpose": format!("{:?}", options.purpose),
"step": state.step_index(),
"request_sha256": request_sha256,
"message_count": conversation.len(),
"response_event_id": response_event_id,
"route": {
"url": options.route_url.clone(),
"model": options.route_model.clone(),
},
"usage_reported": step_usage
.as_ref()
.map(|usage| usage.tokens_reported)
.unwrap_or(false),
"usage": step_usage,
});
let _ = session
.store
.append(&session.id, "model.usage", usage_payload)
.await;
}
if proposed_calls.is_empty() {
conversation.push(Message {
role: Role::Assistant,
content: prose_content,
reasoning_content,
tool_calls: None,
tool_call_id: None,
});
let _ = state.complete_turn();
// Turn completed boundary: yield full conversation snapshot
yield (TurnEvent::TurnCompleted, state, conversation);
return;
}
conversation.push(Message::with_tool_calls(
proposed_calls.clone(),
prose_content,
reasoning_content,
));
let mut peer_cancellation_reason: Option<String> = None;
let mut final_yield_record: Option<YieldRecord> = None;
for call in proposed_calls {
let parsed = parse_tool_call(&call);
let _ = state.begin_tool_execution(&parsed.call_id);
yield (
TurnEvent::BeginToolExecution {
call_id: parsed.call_id.clone(),
},
state.clone(),
conversation.clone(),
);
if let Some(ref reason) = peer_cancellation_reason {
// Peer tool call is cleanly closed (cancelled/skipped) without execution
let skip_text = stamped(¬_run_header(), &format!("Skipped: {reason}"));
let output = CellOutputRecord::stderr(skip_text.clone());
let _ = state.record_cell_output(output.clone());
yield (TurnEvent::CellOutput(output), state.clone(), Vec::new());
let _ = state.complete_tool_execution(&parsed.call_id);
yield (
TurnEvent::ToolExecutionCompleted {
call_id: parsed.call_id.clone(),
},
state.clone(),
conversation.clone(),
);
conversation.push(Message::tool_result(&parsed.call_id, skip_text.clone()));
if let Err(err) = session
.store
.append(
&session.id,
"tool.result",
json!({
"type": "tool_result",
"call_id": parsed.call_id,
"response_event_id": response_event_id,
"status": "skipped",
"reason": reason,
"content": skip_text,
}),
)
.await
{
let err_msg = format!("failed to commit skipped tool result to store: {err}");
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
return;
}
continue;
}
let exec_res = match parsed.code {
Ok(code) => execute_workbench_cell_with_purpose(&session, &parsed.call_id, code, options.cell_budget, options.purpose).await,
Err(err_msg) => {
peer_cancellation_reason = Some(format!("tool '{}' failed: {err_msg}", parsed.call_id));
Ok(WorkbenchExecutionResult {
call_id: parsed.call_id.clone(),
execution_id: None,
cell_outputs: vec![CellOutputRecord::stderr(err_msg.clone())],
speech_receipts: Vec::new(),
yield_record: None,
tool_result_message: Message::tool_result(
&parsed.call_id,
stamped(¬_run_header(), &format!("Error: {err_msg}")),
),
})
}
};
match exec_res {
Ok(res) => {
for output in res.cell_outputs {
let _ = state.record_cell_output(output.clone());
yield (TurnEvent::CellOutput(output), state.clone(), Vec::new());
}
for receipt in res.speech_receipts {
let _ = state.record_speech(receipt.clone());
yield (TurnEvent::Speech(receipt), state.clone(), Vec::new());
}
let _ = state.complete_tool_execution(&parsed.call_id);
yield (
TurnEvent::ToolExecutionCompleted {
call_id: parsed.call_id.clone(),
},
state.clone(),
Vec::new(),
);
conversation.push(res.tool_result_message.clone());
let tool_content = res.tool_result_message.content.clone().unwrap_or_default();
// Link tool call and model response in Store
if let Err(err) = session
.store
.append(
&session.id,
"tool.result",
json!({
"type": "tool_result",
"call_id": parsed.call_id,
"response_event_id": response_event_id,
"execution_id": res.execution_id,
"status": if res.yield_record.is_some() { "yielded" } else { "completed" },
"content": tool_content,
}),
)
.await
{
let err_msg = format!("failed to commit tool result to store: {err}");
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
return;
}
if let Some(y) = res.yield_record {
final_yield_record = Some(y);
peer_cancellation_reason = Some("previous tool call yielded".into());
}
}
Err(err) => {
let err_record = CellOutputRecord::stderr(err.clone());
let _ = state.record_cell_output(err_record.clone());
yield (TurnEvent::CellOutput(err_record), state.clone(), Vec::new());
let _ = state.complete_tool_execution(&parsed.call_id);
yield (
TurnEvent::ToolExecutionCompleted {
call_id: parsed.call_id.clone(),
},
state.clone(),
Vec::new(),
);
let err_text = stamped(¬_run_header(), &format!("Error: {err}"));
conversation.push(Message::tool_result(&parsed.call_id, err_text.clone()));
if let Err(append_err) = session
.store
.append(
&session.id,
"tool.result",
json!({
"type": "tool_result",
"call_id": parsed.call_id,
"response_event_id": response_event_id,
"status": "error",
"error": err,
"content": err_text,
}),
)
.await
{
let err_msg = format!("failed to commit error tool result to store: {append_err}");
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
return;
}
peer_cancellation_reason = Some(format!("previous tool call errored: {err}"));
}
}
}
if let Some(y) = final_yield_record {
let _ = state.transition_yield(y.clone());
yield (TurnEvent::Yielded(y), state, conversation);
return;
}
let step = state.current_step();
match state.step_completed() {
Ok(true) => {
yield (TurnEvent::StepCompleted { step }, state, conversation);
}
Ok(false) => {
// The slice is over. Whether that is completion or exhaustion
// decides whether the host may continue the task, so the two are
// emitted as different events.
yield (TurnEvent::StepCompleted { step }, state.clone(), conversation.clone());
if state.is_exhausted() {
yield (
TurnEvent::exhausted(step, state.max_steps()),
state,
conversation,
);
} else {
yield (TurnEvent::TurnCompleted, state, conversation);
}
}
Err(crate::runtime_turn::state::StateError::MaxStepsExceeded {
current_step,
max_steps,
}) => {
// A spent allowance is a resumable outcome, not a failure.
yield (
TurnEvent::exhausted(current_step, max_steps),
state,
conversation,
);
}
Err(e) => {
let err_msg = e.to_string();
let _ = state.fail_turn(err_msg.clone());
yield (TurnEvent::Failed(err_msg), state, conversation);
}
}
}
}
/// Process a single turn step from a stream of `LlmEvent`s.
pub fn run_turn_step_from_stream(
stream_rx: mpsc::Receiver<LlmEvent>,
session: Arc<SessionRuntime>,
state: TurnStateMachine,
conversation: Vec<Message>,
cell_budget: Duration,
) -> impl Stream<Item = (TurnEvent, TurnStateMachine, Vec<Message>)> {
run_turn_step_supervised(
stream_rx,
None,
session,
state,
conversation,
TurnStepOptions {
cell_budget,
step_timeout: None,
purpose: klbr_runtime::protocol::Purpose::Interactive,
route_model: String::new(),
route_url: String::new(),
},
None,
)
}
/// Run a single step using the real `LlmClient` to stream model completions with task ownership.
pub fn run_turn_step(
llm: &LlmClient,
session: &Arc<SessionRuntime>,
state: TurnStateMachine,
conversation: Vec<Message>,
cell_budget: Duration,
) -> impl Stream<Item = (TurnEvent, TurnStateMachine, Vec<Message>)> {
let (tx, rx) = mpsc::channel(100);
let llm_clone = llm.clone();
let conv_clone = conversation.clone();
let tools = vec![python_tool_def()];
let join_handle = tokio::spawn(async move { llm_clone.stream(&conv_clone, &tools, tx).await });
let provider_task = ProviderTaskHandle::new(join_handle);
run_turn_step_supervised(
rx,
Some(provider_task),
session.clone(),
state,
conversation,
TurnStepOptions {
cell_budget,
step_timeout: None,
purpose: klbr_runtime::protocol::Purpose::Interactive,
route_model: String::new(),
route_url: String::new(),
},
None,
)
}
/// Multi-step Turn Runner.
///
/// Runs model-workbench interaction loops until:
/// - The model completes its response without tool calls.
/// - The workbench yields cleanly via `await runtime.wait(...)`.
/// - The maximum configured steps are reached.
/// - An error occurs.
pub fn run_turn_loop(
llm: &LlmClient,
session: &Arc<SessionRuntime>,
initial_messages: Vec<Message>,
max_steps: usize,
cell_budget: Duration,
) -> impl Stream<Item = TurnEvent> {
run_turn_loop_with_purpose(
llm,
session,
initial_messages,
max_steps,
cell_budget,
klbr_runtime::protocol::Purpose::Interactive,
)
}
/// True for the events that end a turn and name why it ended.
fn is_terminal_event(event: &TurnEvent) -> bool {
matches!(
event,
TurnEvent::TurnCompleted
| TurnEvent::Yielded(_)
| TurnEvent::Exhausted { .. }
| TurnEvent::Failed(_)
| TurnEvent::Interrupted
| TurnEvent::Cancelled
| TurnEvent::Timeout
)
}
pub fn run_turn_loop_with_purpose(
llm: &LlmClient,
session: &Arc<SessionRuntime>,
initial_messages: Vec<Message>,
max_steps: usize,
cell_budget: Duration,
purpose: klbr_runtime::protocol::Purpose,
) -> impl Stream<Item = TurnEvent> {
let llm = llm.clone();
let session = session.clone();
stream! {
let mut state = TurnStateMachine::new(max_steps);
let mut conversation = initial_messages;
while !state.is_terminal() {
let (tx, rx) = mpsc::channel(100);
let llm_clone = llm.clone();
let conv_clone = conversation.clone();
let tools = vec![python_tool_def()];
let join_handle = tokio::spawn(async move {
llm_clone.stream(&conv_clone, &tools, tx).await
});
let provider_task = ProviderTaskHandle::new(join_handle);
let step_stream = run_turn_step_supervised(
rx,
Some(provider_task),
session.clone(),
state.clone(),
conversation.clone(),
TurnStepOptions {
cell_budget,
step_timeout: None,
purpose,
route_model: llm.config.llm.model.clone(),
route_url: llm.config.llm.url.clone(),
},
None,
);
tokio::pin!(step_stream);
use futures::StreamExt;
while let Some((event, next_state, next_conv)) = step_stream.next().await {
let state_is_terminal = next_state.is_terminal();
let event_is_terminal = is_terminal_event(&event);
if !matches!(event, TurnEvent::Scratchpad(_)) {
state = next_state;
}
if !next_conv.is_empty() {
conversation = next_conv;
}
yield event;
// The state reaches terminal one event *before* the event that says
// why: a spent allowance is reported by the event after the final
// StepCompleted. Ending on the state alone would swallow that
// reason, and the host would never learn the task is unfinished.
if state_is_terminal && event_is_terminal {
return;
}
}
}
}
}