use std::thread::JoinHandle;
use tokio::sync::oneshot;
use std::sync::Arc;
use super::observation::ShellObservation;
use super::{
chat_complete, openai_chat_complete, ChatCtx, CodeSearch, Embedder, NoMcp, PromptDisposition,
SemanticIndex,
};
use crate::{BackendKind, CompactionTriggerPolicy, MemMessage, Role, TokenUsage};
#[derive(Clone)]
pub struct HeadlessCodeSearch {
pub embedder: Arc<dyn Embedder>,
pub index: Arc<dyn SemanticIndex>,
pub top_k: usize,
}
impl std::fmt::Debug for HeadlessCodeSearch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HeadlessCodeSearch")
.field("top_k", &self.top_k)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone)]
pub struct TurnDriverConfig {
pub url: String,
pub model: String,
pub kind: BackendKind,
pub api_key: Option<String>,
pub workspace: String,
pub caveats: crate::caveats::Caveats,
pub max_tool_rounds: usize,
pub narration_nudge_cap: usize,
pub workflow_grace_rounds: usize,
pub tool_output_lines: usize,
pub num_ctx: Option<u32>,
pub connect_timeout_secs: u64,
pub inference_timeout_secs: u64,
pub mid_loop_trim_threshold: usize,
pub compaction_trigger_policy: CompactionTriggerPolicy,
pub mid_loop_trim_tokens: Option<usize>,
pub max_ok_input: Option<u32>,
pub build_check_cmd: Option<String>,
pub safe_context: Option<u32>,
pub input_ceiling_pct: u32,
pub low_budget_pct: usize,
pub code_search: Option<HeadlessCodeSearch>,
}
impl TurnDriverConfig {
pub fn new(
url: impl Into<String>,
model: impl Into<String>,
kind: BackendKind,
workspace: impl Into<String>,
) -> Self {
Self {
url: url.into(),
model: model.into(),
kind,
api_key: None,
workspace: workspace.into(),
caveats: crate::caveats::Caveats::top(),
max_tool_rounds: 40,
narration_nudge_cap: 1,
workflow_grace_rounds: 5,
tool_output_lines: 20,
num_ctx: None,
connect_timeout_secs: 5,
inference_timeout_secs: 120,
mid_loop_trim_threshold: 40,
compaction_trigger_policy: CompactionTriggerPolicy::HeadroomAware,
mid_loop_trim_tokens: None,
max_ok_input: None,
build_check_cmd: None,
safe_context: None,
input_ceiling_pct: 80,
low_budget_pct: 15,
code_search: None,
}
}
#[must_use]
pub fn with_code_search(
mut self,
embedder: Arc<dyn Embedder>,
index: Arc<dyn SemanticIndex>,
top_k: usize,
) -> Self {
self.code_search = Some(HeadlessCodeSearch {
embedder,
index,
top_k,
});
self
}
}
#[derive(Debug, Clone)]
pub struct TurnOutcome {
pub reply: String,
pub was_streamed: bool,
pub usage: Option<TokenUsage>,
pub hallucinations: u32,
pub tool_events: Vec<crate::ToolEvent>,
pub end_reason: Option<crate::TurnEndReason>,
pub error: Option<String>,
}
#[derive(Debug)]
pub enum TurnStatus {
Idle,
Running,
Completed(TurnOutcome),
Failed(String),
}
struct InFlight {
handle: JoinHandle<()>,
rx: oneshot::Receiver<Result<TurnOutcome, String>>,
cancel_tx: Option<oneshot::Sender<()>>,
}
pub struct TurnDriver {
config: TurnDriverConfig,
transcript: Vec<MemMessage>,
in_flight: Option<InFlight>,
}
impl TurnDriver {
pub fn new(config: TurnDriverConfig) -> Self {
Self {
config,
transcript: Vec::new(),
in_flight: None,
}
}
pub fn with_transcript(config: TurnDriverConfig, transcript: Vec<MemMessage>) -> Self {
Self {
config,
transcript,
in_flight: None,
}
}
pub fn transcript(&self) -> &[MemMessage] {
&self.transcript
}
pub fn is_running(&self) -> bool {
self.in_flight.is_some()
}
pub fn submit_observation(&mut self, obs: ShellObservation) {
self.transcript.push(obs.into_mem_message());
}
pub fn submit(&mut self, input: impl Into<String>) -> Result<(), TurnDriverError> {
if self.in_flight.is_some() {
return Err(TurnDriverError::Busy);
}
let task = input.into();
self.transcript.push(MemMessage::user(task.clone()));
self.spawn_turn(task);
Ok(())
}
fn spawn_turn(&mut self, task: String) {
let config = self.config.clone();
let messages = self.transcript.clone();
let (tx, rx) = oneshot::channel();
let (cancel_tx, cancel_rx) = oneshot::channel();
let handle = std::thread::spawn(move || {
let rt = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(rt) => rt,
Err(e) => {
let _ = tx.send(Err(format!("failed to start turn runtime: {e}")));
return;
}
};
let result = rt.block_on(async move {
tokio::select! {
biased;
_ = cancel_rx => Err("turn cancelled".to_string()),
out = run_one_turn(&config, &messages, &task) => out,
}
});
let _ = tx.send(result);
});
self.in_flight = Some(InFlight {
handle,
rx,
cancel_tx: Some(cancel_tx),
});
}
pub fn poll(&mut self) -> TurnStatus {
let Some(in_flight) = self.in_flight.as_mut() else {
return TurnStatus::Idle;
};
match in_flight.rx.try_recv() {
Ok(Ok(outcome)) => {
self.transcript
.push(MemMessage::assistant(outcome.reply.clone()));
self.in_flight = None;
TurnStatus::Completed(outcome)
}
Ok(Err(err)) => {
self.in_flight = None;
TurnStatus::Failed(err)
}
Err(oneshot::error::TryRecvError::Empty) => TurnStatus::Running,
Err(oneshot::error::TryRecvError::Closed) => {
self.in_flight = None;
TurnStatus::Failed("turn task ended without a result".to_string())
}
}
}
pub fn cancel(&mut self) {
if let Some(mut in_flight) = self.in_flight.take() {
if let Some(cancel_tx) = in_flight.cancel_tx.take() {
let _ = cancel_tx.send(());
}
let _ = in_flight.handle.join();
}
}
}
async fn run_one_turn(
config: &TurnDriverConfig,
messages: &[MemMessage],
task: &str,
) -> Result<TurnOutcome, String> {
let mut tool_events: Vec<crate::ToolEvent> = Vec::new();
let mut end_reason: Option<crate::TurnEndReason> = None;
let ctx = ChatCtx {
url: &config.url,
model: &config.model,
kind: config.kind,
api_key: config.api_key.as_deref(),
action_nudges: true,
prompt_disposition: PromptDisposition::Act,
prompt_intake: None,
messages,
task,
workspace: &config.workspace,
color: false,
markdown: false,
tool_offload: false,
spill_store: None,
compaction_store: None,
scratchpad: false,
scratchpad_store: None,
code_search: config.code_search.as_ref().map(|cs| CodeSearch {
embedder: cs.embedder.as_ref(),
index: cs.index.as_ref(),
top_k: cs.top_k,
steer: None,
status: None,
}),
where_is: None,
nav: None,
exposure: crate::agentic::tools::ExposureSettings::default(),
experience_store: None,
step_ledger: None,
caveats: &config.caveats,
persona_tools: None,
max_tool_rounds: config.max_tool_rounds,
narration_nudge_cap: config.narration_nudge_cap,
workflow_grace_rounds: config.workflow_grace_rounds,
tool_output_lines: config.tool_output_lines,
debug: false,
trace: false,
num_ctx: config.num_ctx,
input_ceiling_pct: config.input_ceiling_pct,
low_budget_pct: config.low_budget_pct,
connect_timeout_secs: config.connect_timeout_secs,
inference_timeout_secs: config.inference_timeout_secs,
mid_loop_trim_threshold: config.mid_loop_trim_threshold,
compaction_trigger_policy: config.compaction_trigger_policy,
mid_loop_trim_tokens: config.mid_loop_trim_tokens,
max_ok_input: config.max_ok_input,
build_check_cmd: config.build_check_cmd.clone(),
safe_context: config.safe_context,
recover_cw_400: None,
note_sink: None,
note_nudge: None,
recall_source: None,
memory_source: None,
summarizer: None,
compress_state: None,
tool_events: Some(&mut tool_events),
phantom_reaches: None,
end_reason: Some(&mut end_reason),
permission_gate: None,
on_round_usage: None,
estimate_ratio: None,
estimation: crate::tokens::TokenEstimation::default(),
summary_input_cap_floor_chars: 8_192,
exec_floor: None,
write_ledger: None,
cancel: None,
live_tool_output: None,
git_tool: None,
crew_runner: None,
operating_mode_control: None,
plan_mode_control: None,
};
let mut mcp = NoMcp;
let dispatch = if config.kind == BackendKind::Openai {
openai_chat_complete(ctx, &mut mcp).await
} else {
chat_complete(ctx, &mut mcp).await
};
match dispatch {
Ok((reply, was_streamed, usage, hallucinations)) => Ok(TurnOutcome {
reply,
was_streamed,
usage,
hallucinations,
tool_events,
end_reason,
error: None,
}),
Err(e) => Ok(TurnOutcome {
reply: String::new(),
was_streamed: false,
usage: None,
hallucinations: 0,
tool_events,
end_reason,
error: Some(e.to_string()),
}),
}
}
#[derive(Debug, thiserror::Error)]
pub enum TurnDriverError {
#[error("a turn is already in flight — poll() for it or cancel() before submitting another")]
Busy,
}
pub const VISIBLE_TRANSCRIPT_ROLES: [Role; 2] = [Role::User, Role::Assistant];
#[cfg(test)]
mod tests {
use super::*;
use crate::agentic::SessionSemanticIndex;
use crate::caveats::Caveats;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
struct PlainOllama {
served: Arc<AtomicUsize>,
reply: String,
}
impl Respond for PlainOllama {
fn respond(&self, _req: &Request) -> ResponseTemplate {
self.served.fetch_add(1, Ordering::SeqCst);
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"message": { "content": self.reply }
}))
}
}
fn cfg(url: &str) -> TurnDriverConfig {
let mut c = TurnDriverConfig::new(url, "test-model", BackendKind::Ollama, ".");
c.caveats = Caveats::top();
c
}
async fn pump_to_done(driver: &mut TurnDriver) -> TurnStatus {
for _ in 0..600 {
match driver.poll() {
TurnStatus::Running => tokio::time::sleep(Duration::from_millis(10)).await,
other => return other,
}
}
panic!("turn did not complete within the pump budget");
}
#[tokio::test]
async fn driver_pumps_one_turn_and_yields_the_answer() {
let server = MockServer::start().await;
let served = Arc::new(AtomicUsize::new(0));
Mock::given(method("POST"))
.and(path("/api/chat"))
.respond_with(PlainOllama {
served: served.clone(),
reply: "the driver answered".into(),
})
.mount(&server)
.await;
let mut driver = TurnDriver::new(cfg(&server.uri()));
assert!(matches!(driver.poll(), TurnStatus::Idle));
assert!(!driver.is_running());
driver
.submit("what is 2 + 2?")
.expect("submit starts a turn");
assert!(driver.is_running());
let status = pump_to_done(&mut driver).await;
match status {
TurnStatus::Completed(outcome) => {
assert_eq!(outcome.reply, "the driver answered");
}
other => panic!("expected Completed, got {other:?}"),
}
let t = driver.transcript();
assert_eq!(t.len(), 2);
assert_eq!(t[0].role, Role::User);
assert_eq!(t[0].content, "what is 2 + 2?");
assert_eq!(t[1].role, Role::Assistant);
assert_eq!(t[1].content, "the driver answered");
assert!(served.load(Ordering::SeqCst) >= 1);
assert!(matches!(driver.poll(), TurnStatus::Idle));
assert!(!driver.is_running());
}
struct StubEmbedder;
#[async_trait::async_trait]
impl Embedder for StubEmbedder {
async fn embed(&self, _text: &str) -> anyhow::Result<Vec<f32>> {
Ok(vec![0.0_f32; 4])
}
}
struct CapturingOllama {
bodies: Arc<std::sync::Mutex<Vec<serde_json::Value>>>,
reply: String,
}
impl Respond for CapturingOllama {
fn respond(&self, req: &Request) -> ResponseTemplate {
if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&req.body) {
self.bodies.lock().unwrap().push(v);
}
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"message": { "content": self.reply }
}))
}
}
fn advertised_code_search(bodies: &[serde_json::Value]) -> bool {
bodies.iter().any(|b| {
b["tools"]
.as_array()
.is_some_and(|ts| ts.iter().any(|t| t["function"]["name"] == "code_search"))
})
}
async fn drive_once_capturing(config: TurnDriverConfig) -> Vec<serde_json::Value> {
let server = MockServer::start().await;
let bodies = Arc::new(std::sync::Mutex::new(Vec::new()));
Mock::given(method("POST"))
.and(path("/api/chat"))
.respond_with(CapturingOllama {
bodies: bodies.clone(),
reply: "done".into(),
})
.mount(&server)
.await;
let mut config = config;
config.url = server.uri();
let mut driver = TurnDriver::new(config);
driver.submit("find the retry backoff").expect("submit");
let _ = pump_to_done(&mut driver).await;
let out = bodies.lock().unwrap().clone();
out
}
#[tokio::test]
async fn code_search_advertised_only_when_the_config_carries_retrieval() {
let bare = drive_once_capturing(cfg("http://placeholder")).await;
assert!(
!advertised_code_search(&bare),
"no code_search tool without a configured index"
);
let index: Arc<dyn SemanticIndex> = Arc::new(SessionSemanticIndex::default());
let with = drive_once_capturing(cfg("http://placeholder").with_code_search(
Arc::new(StubEmbedder),
index,
3,
))
.await;
assert!(
advertised_code_search(&with),
"code_search advertised once the driver carries retrieval"
);
}
#[tokio::test]
async fn observation_is_in_context_for_the_next_turn_and_redacted() {
let server = MockServer::start().await;
let seen = Arc::new(std::sync::Mutex::new(String::new()));
struct Capture {
seen: Arc<std::sync::Mutex<String>>,
}
impl Respond for Capture {
fn respond(&self, req: &Request) -> ResponseTemplate {
*self.seen.lock().unwrap() = String::from_utf8_lossy(&req.body).into_owned();
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"message": { "content": "ack" }
}))
}
}
Mock::given(method("POST"))
.and(path("/api/chat"))
.respond_with(Capture { seen: seen.clone() })
.mount(&server)
.await;
let mut driver = TurnDriver::new(cfg(&server.uri()));
driver.submit_observation(ShellObservation::new(
"zsh",
"$ cat creds\nsecret_key=wJalrXUtnFEMIabcdEFGHIJKLMNOPbPxRfiCY",
));
assert!(!driver.is_running());
assert_eq!(driver.transcript().len(), 1);
driver.submit("did you see my shell?").expect("submit");
let _ = pump_to_done(&mut driver).await;
let body = seen.lock().unwrap().clone();
assert!(
body.contains("shell observation"),
"observation missing from request body: {body}"
);
assert!(
!body.contains("wJalrXUtnFEMIabcdEFGHIJKLMNOPbPxRfiCY"),
"secret leaked into the request body: {body}"
);
}
#[tokio::test]
async fn second_submit_while_running_is_busy() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/chat"))
.respond_with(
ResponseTemplate::new(200)
.set_delay(Duration::from_millis(200))
.set_body_json(serde_json::json!({ "message": { "content": "slow" } })),
)
.mount(&server)
.await;
let mut driver = TurnDriver::new(cfg(&server.uri()));
driver.submit("first").expect("first submit");
let err = driver
.submit("second")
.expect_err("second must be rejected");
assert!(matches!(err, TurnDriverError::Busy));
let _ = pump_to_done(&mut driver).await;
}
#[tokio::test]
async fn cancel_aborts_the_in_flight_turn() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/chat"))
.respond_with(
ResponseTemplate::new(200)
.set_delay(Duration::from_secs(30))
.set_body_json(serde_json::json!({ "message": { "content": "never" } })),
)
.mount(&server)
.await;
let mut driver = TurnDriver::new(cfg(&server.uri()));
driver.submit("start a slow turn").expect("submit");
assert!(driver.is_running());
driver.cancel();
assert!(!driver.is_running());
assert!(matches!(driver.poll(), TurnStatus::Idle));
assert_eq!(driver.transcript().len(), 1);
assert_eq!(driver.transcript()[0].content, "start a slow turn");
}
}