@@ -11,30 +11,20 @@ evidence::EvidencePacket, harness_block::{self, HarnessBlock, DEFAULT_FORMAT_STYLE}, interrupt::Interrupt, memory::MemoryStore, - models::{LlmClient, LlmEvent, Message, Role, ToolCall}, + models::{LlmClient, LlmEvent, Message, Role}, pipeline::{BenchQuery, ContextBudget, MemoryPipeline}, router::{RouteDecision, Router}, tools::{self, ToolContext}, - AgentEvent, AgentMetrics, CompactionRecord, + AgentEvent, CompactionRecord, }; -const MAX_CONSECUTIVE_STREAM_FAILURES: usize = 5; +use crate::tools::format_duration_human; -struct StreamIteration { - response: String, - thinking: String, - tool_calls: Vec, - stream_error: Option, - cancelled_by: Option, -} +mod turn; -enum ToolIterationOutcome { - Continue, - Suspend(Option), - Restart, -} +const MAX_CONSECUTIVE_STREAM_FAILURES: usize = 5; pub struct Agent { config: Config, memory: MemoryStore, @@ -553,615 +543,10 @@ consecutive_waits, ) .await } - - async fn stream_iteration( - &mut self, - llm: &LlmClient, - ctx: &mut Context, - stashed_interrupt: &mut Option, - ) -> StreamIteration { - let (tok_tx, mut tok_rx) = mpsc::channel(256); - let llm2 = llm.clone(); - let msgs = ctx.as_messages_with_refs(&self.memory); - let defs = self.registry.definitions(); - let stream_task = tokio::spawn(async move { llm2.stream(&msgs, &defs, tok_tx).await }); - - let mut response = String::new(); - let mut thinking = String::new(); - let mut tool_calls = vec![]; - let mut cancelled_by = None; - - loop { - tokio::select! { - biased; - interrupt = self.rx.recv() => { - match interrupt { - Some(int) if int.is_preemptive() => { - tracing::info!(source = int.source_tag(), "preemptive interrupt during stream; aborting"); - stream_task.abort(); - cancelled_by = Some(int); - break; - } - Some(int) => { - // non-preemptive (e.g. ExternalEvent) — stash for after turn - if stashed_interrupt.is_none() { - *stashed_interrupt = Some(int); - } - } - None => {} // channel closed = shutdown - } - } - ev = tok_rx.recv() => { - match ev { - None => break, - Some(LlmEvent::ThinkToken(tok)) => { - thinking.push_str(&tok); - let _ = self.output.send(AgentEvent::ThinkToken(tok)); - } - Some(LlmEvent::Token(tok)) => { - response.push_str(&tok); - let _ = self.output.send(AgentEvent::ScratchToken(tok)); - } - Some(LlmEvent::Usage(usage)) => { - ctx.update_tokens(&usage); - } - Some(LlmEvent::ToolCalls(calls)) => { - tool_calls = calls; - } - Some(LlmEvent::ToolCallStarted) => { - let _ = self.output.send(AgentEvent::ToolCallStarted); - } - } - } - } - } - - let stream_error = match stream_task.await { - Ok(Ok(())) => None, - Ok(Err(e)) => { - let msg = format!("llm stream failed: {e}"); - tracing::error!(%msg); - let _ = self.output.send(AgentEvent::Error(msg.clone())); - let _ = self - .output - .send(AgentEvent::Status("llm stream failed".into())); - Some(msg) - } - Err(e) if e.is_cancelled() => { - // intentional abort — already handled by cancelled_by - None - } - Err(e) => { - let msg = format!("llm stream task panicked/cancelled: {e}"); - tracing::error!(%msg); - let _ = self.output.send(AgentEvent::Error(msg.clone())); - let _ = self - .output - .send(AgentEvent::Status("llm stream task failed".into())); - Some(msg) - } - }; - - StreamIteration { - response, - thinking, - tool_calls, - stream_error, - cancelled_by, - } - } - - async fn execute_tool_iteration( - &mut self, - tool_ctx: &ToolContext, - ctx: &mut Context, - turn_count: &mut usize, - runtime_soul: &mut String, - tool_calls: &[ToolCall], - iteration_thinking: &str, - stashed_interrupt: &mut Option, - consecutive_waits: &mut usize, - discord_send_called: &mut bool, - local_send_called: &mut bool, - ) -> Result { - let reasoning_content = if iteration_thinking.is_empty() { - None - } else { - Some(iteration_thinking.to_string()) - }; - ctx.push_assistant_tool_calls(tool_calls.to_vec(), None, reasoning_content); - let _ = self.memory.log_turn_with_tools( - "assistant", - "", - (!iteration_thinking.is_empty()).then_some(iteration_thinking), - Some(tool_calls), - None, - ); - - let mut pending_local_sends = Vec::new(); - let mut should_restart = false; - for call in tool_calls { - let name = call.function.name.clone(); - if name == "restart_harness" { - should_restart = true; - } - if name != "wait_and_continue" { - *consecutive_waits = 0; - } - if name == "discord_send" { - *discord_send_called = true; - } - if name == "local_send" { - *local_send_called = true; - } - let args = call.function.arguments.clone(); - if name != "local_send" { - let _ = self.output.send(AgentEvent::ToolCall { - name: name.clone(), - args: args.clone(), - }); - } - - let wait_result = if name == "wait_and_continue" { - Some(execute_wait_and_continue(&args, &mut self.rx).await) - } else { - None - }; - let mut incoming_after_tool = None; - let mut suspend_after_wait = false; - let mut sent_local_content = None; - let result = if name == "local_send" { - match tools::local_send_content(&args) { - Ok(content) => { - let result = - format!("sent {} chars to local operator", content.chars().count()); - sent_local_content = Some(content); - result - } - Err(err) => { - tracing::warn!(err = %err, "local_send call had invalid arguments"); - format!("error: {err}") - } - } - } else { - match wait_result { - Some(WaitAndContinueResult::Timeout(mut result)) => { - suspend_after_wait = is_wait_timeout_without_event(&result); - if suspend_after_wait { - *consecutive_waits += 1; - if *consecutive_waits >= 5 && *consecutive_waits % 5 == 0 { - let nudge = format!( - "\n\nnudge: you have called wait_and_continue {} consecutive times without taking other actions. to conserve local resources, please increase the wait duration in your next wait_and_continue call (e.g. wait_and_continue(duration=\"10m\")).", - *consecutive_waits - ); - result.push_str(&nudge); - } - } else { - *consecutive_waits = 0; - } - result - } - Some(WaitAndContinueResult::Incoming { result, interrupt }) => { - incoming_after_tool = Some(interrupt); - result - } - Some(WaitAndContinueResult::Reset) => { - ctx.clear(); - *turn_count = 0; - *runtime_soul = self.sys_prompt()?; - ctx.update_soul(runtime_soul, &[]); - save_context_snapshot(&self.memory, ctx); - let _ = self.output.send(AgentEvent::Reset); - "interrupted by reset; context reset".to_string() - } - Some(WaitAndContinueResult::UpdateSoul(content)) => { - self.memory.set_soul_text(&content)?; - *runtime_soul = self.sys_prompt()?; - let pinned = self.memory.pinned_memories().unwrap_or_default(); - ctx.update_soul(runtime_soul, &pinned); - save_context_snapshot(&self.memory, ctx); - let _ = self.output.send(AgentEvent::Status("soul updated".into())); - "interrupted by soul update; soul updated and continuing".to_string() - } - Some(WaitAndContinueResult::Compact) => { - let _ = self.output.send(AgentEvent::Status("compacting...".into())); - if let Err(e) = compact(tool_ctx, ctx, 0, 0, &self.output).await { - tracing::error!(err = %e, "compaction while waiting failed"); - format!("compaction failed while waiting: {e}") - } else { - *runtime_soul = self.sys_prompt()?; - refresh_context_soul(&self.memory, runtime_soul, ctx); - *turn_count = ctx.turn_count(); - "interrupted by compaction; compacted and continuing".to_string() - } - } - None => { - // Check for a preemptive interrupt before running the tool. - // We use try_recv here (not select!) so the tool future is never - // created and dropped — that would risk double-execution on retry. - match self.rx.try_recv() { - Ok(int) if int.is_preemptive() => { - tracing::info!( - source = int.source_tag(), - tool = name, - "preemptive interrupt before tool execution" - ); - let cancel_result = format!( - "tool execution cancelled by incoming {} message", - int.source_tag() - ); - if name != "local_send" { - let _ = self.output.send(AgentEvent::ToolResult { - name: name.clone(), - content: cancel_result.clone(), - }); - } - ctx.push_tool_result(&call.id, &cancel_result); - let _ = self.memory.log_turn_with_tools( - "tool", - &cancel_result, - None, - None, - Some(&call.id), - ); - flush_pending_local_sends( - &self.memory, - ctx, - turn_count, - &mut pending_local_sends, - ); - return Ok(ToolIterationOutcome::Suspend(Some(int))); - } - Ok(int) => { - // non-preemptive — stash - if stashed_interrupt.is_none() { - *stashed_interrupt = Some(int); - } - } - Err(_) => {} - } - - // Now run the tool, allowing preemptive interrupts during its execution. - let mut tool_output = None; - let mut interrupted_by = None; - let mut exec_fut = Box::pin(self.registry.execute(call, tool_ctx)); - loop { - tokio::select! { - res = &mut exec_fut => { - tool_output = Some(res); - break; - } - int_opt = self.rx.recv() => { - match int_opt { - Some(int) if int.is_preemptive() => { - interrupted_by = Some(int); - break; - } - Some(int) => { - if stashed_interrupt.is_none() { - *stashed_interrupt = Some(int); - } - } - None => {} - } - } - } - } - - if let Some(int) = interrupted_by { - tracing::info!( - source = int.source_tag(), - tool = name, - "tool execution interrupted by incoming preemptive event" - ); - let cancel_result = format!( - "tool execution interrupted by incoming {} message", - int.source_tag() - ); - if name != "local_send" { - let _ = self.output.send(AgentEvent::ToolResult { - name: name.clone(), - content: cancel_result.clone(), - }); - } - ctx.push_tool_result(&call.id, &cancel_result); - let _ = self.memory.log_turn_with_tools( - "tool", - &cancel_result, - None, - None, - Some(&call.id), - ); - flush_pending_local_sends( - &self.memory, - ctx, - turn_count, - &mut pending_local_sends, - ); - return Ok(ToolIterationOutcome::Suspend(Some(int))); - } - - tool_output.unwrap_or_default() - } - } - }; - - if name != "local_send" { - let _ = self.output.send(AgentEvent::ToolResult { - name: name.clone(), - content: result.clone(), - }); - } - - ctx.push_tool_result(&call.id, &result); - let _ = self - .memory - .log_turn_with_tools("tool", &result, None, None, Some(&call.id)); - if let Some(content) = sent_local_content { - let _ = self.output.send(AgentEvent::Token(content.clone())); - pending_local_sends.push(content); - } - if let Some(interrupt) = incoming_after_tool { - flush_pending_local_sends(&self.memory, ctx, turn_count, &mut pending_local_sends); - return Ok(ToolIterationOutcome::Suspend(Some(interrupt))); - } - if suspend_after_wait { - flush_pending_local_sends(&self.memory, ctx, turn_count, &mut pending_local_sends); - return Ok(ToolIterationOutcome::Suspend(None)); - } - // Check for interrupts that arrived *during* tool execution. - // Preemptive ones cut the turn short; non-preemptive are stashed. - match self.rx.try_recv() { - Ok(int) if int.is_preemptive() => { - flush_pending_local_sends( - &self.memory, - ctx, - turn_count, - &mut pending_local_sends, - ); - return Ok(ToolIterationOutcome::Suspend(Some(int))); - } - Ok(int) => { - if stashed_interrupt.is_none() { - *stashed_interrupt = Some(int); - } - } - Err(_) => {} - } - } - - flush_pending_local_sends(&self.memory, ctx, turn_count, &mut pending_local_sends); - if should_restart { - Ok(ToolIterationOutcome::Restart) - } else { - Ok(ToolIterationOutcome::Continue) - } - } - - async fn run_turn( - &mut self, - llm: &LlmClient, - tool_ctx: &ToolContext, - ctx: &mut Context, - turn_count: &mut usize, - runtime_soul: &mut String, - recalled_index: Option, - is_local_message: bool, - is_discord_event: bool, - consecutive_waits: &mut usize, - ) -> Result> { - let mut tool_iterations = 0usize; - let mut discord_send_called = false; - let mut discord_nudge_injected = false; - let mut local_send_called = false; - let mut local_send_nudge_injected = false; - // non-preemptive interrupt received during streaming/tool execution; returned at end of turn - let mut stashed_interrupt: Option = None; - const MAX_TOOL_ITERATIONS: usize = 20; - - // emit Started once — all tokens/thinking for the whole turn go into one bubble - let _ = self.output.send(AgentEvent::Started); - - loop { - let StreamIteration { - response: iteration_response, - thinking: iteration_thinking, - tool_calls, - stream_error, - cancelled_by: stream_cancelled_by, - } = self - .stream_iteration(llm, ctx, &mut stashed_interrupt) - .await; - - if let Some(interrupt) = stream_cancelled_by { - // turn was preempted mid-stream — finish cleanly and hand off - let _ = self.output.send(AgentEvent::Done); - cleanup_recalled_context(ctx, recalled_index); - save_context_snapshot(&self.memory, ctx); - return Ok(Some(interrupt)); - } - - if !tool_calls.is_empty() && tool_iterations < MAX_TOOL_ITERATIONS { - tool_iterations += 1; - let outcome = self - .execute_tool_iteration( - tool_ctx, - ctx, - turn_count, - runtime_soul, - &tool_calls, - &iteration_thinking, - &mut stashed_interrupt, - consecutive_waits, - &mut discord_send_called, - &mut local_send_called, - ) - .await?; - match outcome { - ToolIterationOutcome::Continue => {} - ToolIterationOutcome::Suspend(interrupt) => { - self.finish_suspended_turn(ctx, turn_count, recalled_index) - .await; - return Ok(interrupt); - } - ToolIterationOutcome::Restart => { - self.finish_suspended_turn(ctx, turn_count, recalled_index) - .await; - tracing::info!("restart harness requested, performing restart"); - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - perform_restart(); - } - } - - *runtime_soul = self.sys_prompt()?; - refresh_context_soul(&self.memory, runtime_soul, ctx); - continue; - } - - if let Some(err_msg) = &stream_error { - if iteration_response.is_empty() && tool_calls.is_empty() { - let error_content = format!("error: {err_msg}"); - ctx.push_assistant(&error_content, None); - let _ = self.memory.log_turn("assistant", &error_content, None); - cleanup_recalled_context(ctx, recalled_index); - save_context_snapshot(&self.memory, ctx); - let _ = self.output.send(AgentEvent::Done); - break; - } - } - - // plain assistant text is scratchpad, not delivered speech. - let has_final_text = !iteration_response.is_empty(); - if has_final_text || !iteration_thinking.is_empty() { - ctx.push_assistant( - &scratchpad_context(&iteration_response), - (!iteration_thinking.is_empty()).then_some(iteration_thinking.as_str()), - ); - } - - if has_final_text - && is_local_message - && !local_send_called - && !local_send_nudge_injected - { - local_send_nudge_injected = true; - push_visible_nudge( - ctx, - &self.output, - &plain_text_local_nudge_message(&iteration_response), - ); - tracing::warn!("local_send nudge: assistant emitted plain text without local_send"); - continue; - } - - if has_final_text && is_discord_event && !discord_send_called && !discord_nudge_injected - { - discord_nudge_injected = true; - push_visible_nudge( - ctx, - &self.output, - &plain_text_discord_nudge_message(&iteration_response), - ); - tracing::warn!( - "discord_send nudge: assistant emitted plain text without discord_send" - ); - continue; - } - - self.finish_suspended_turn(ctx, turn_count, recalled_index) - .await; - - if ctx.total_tokens > self.config.watermark_tokens { - let _ = self.output.send(AgentEvent::Status("compacting...".into())); - if let Err(e) = compact( - tool_ctx, - ctx, - self.config.compaction_keep, - self.config.watermark_tokens / 4, - &self.output, - ) - .await - { - tracing::error!(err = %e, "automatic compaction failed"); - let _ = self - .output - .send(AgentEvent::Status(format!("compaction failed: {e}"))); - } else { - *runtime_soul = self.sys_prompt()?; - refresh_context_soul(&self.memory, runtime_soul, ctx); - *turn_count = ctx.turn_count(); - } - } - - *consecutive_waits = 0; - break; - } - - Ok(stashed_interrupt) - } - - async fn finish_suspended_turn( - &self, - ctx: &mut Context, - turn_count: &usize, - recalled_index: Option, - ) { - cleanup_recalled_context(ctx, recalled_index); - save_context_snapshot(&self.memory, ctx); - - let _ = self.output.send(AgentEvent::Done); - let metrics = AgentMetrics { - turn_count: *turn_count, - context_tokens: ctx.total_tokens, - watermark: self.config.watermark_tokens, - prompt_tokens: ctx.prompt_tokens, - completion_tokens: ctx.completion_tokens, - cached_prompt_tokens: ctx.cached_prompt_tokens, - prompt_processing_ms: ctx.prompt_processing_ms, - generation_ms: ctx.generation_ms, - bridge_processing_ms: ctx.bridge_processing_ms, - }; - *self.snapshot.write().await = Some(metrics.clone()); - let _ = self.output.send(AgentEvent::Metrics(metrics)); - } } -#[cfg(unix)] -fn perform_restart() -> ! { - let exe = match std::env::current_exe() { - Ok(path) => path, - Err(e) => { - tracing::error!("failed to get current executable path: {e}"); - std::process::exit(1); - } - }; - let args: Vec = std::env::args().skip(1).collect(); - - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or_default(); - - tracing::info!("restarting harness: {:?} with args {:?}", exe, args); - - use std::os::unix::process::CommandExt; - let mut cmd = std::process::Command::new(exe); - cmd.args(&args); - cmd.env("KLBR_RESTART_TIMESTAMP", now.to_string()); - - let err = cmd.exec(); - tracing::error!("exec failed: {err}"); - std::process::exit(1); -} - -#[cfg(not(unix))] -fn perform_restart() -> ! { - tracing::warn!("restart only supported on unix systems"); - std::process::exit(0); -} - fn unix_timestamp() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|duration| duration.as_secs() as i64) @@ -1295,9 +680,9 @@ fn is_completed_wait_tool_result(ctx: &Context) -> bool { let Some(last) = ctx.turns.last() else { return false; }; - if last.role != Role::Tool || !is_wait_timeout_without_event(last.content_str()) { + if last.role != Role::Tool || !tools::wait_timeout_without_event(last.content_str()) { return false; } let Some(tool_call_id) = last.tool_call_id.as_deref() else { return false; @@ -1315,111 +700,8 @@ .any(|call| call.id == tool_call_id && call.function.name == "wait_and_continue") }) } -fn is_wait_timeout_without_event(result: &str) -> bool { - result.contains("no incoming event arrived") -} - -enum WaitAndContinueResult { - Timeout(String), - Incoming { - result: String, - interrupt: Interrupt, - }, - Reset, - Compact, - UpdateSoul(String), -} - -async fn execute_wait_and_continue( - arguments: &str, - rx: &mut mpsc::Receiver, -) -> WaitAndContinueResult { - let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default(); - let timeout = match tools::wait_and_continue_timeout(&args) { - Ok(timeout) => timeout, - Err(err) => return WaitAndContinueResult::Timeout(format!("error: {err}")), - }; - - let start_time = std::time::Instant::now(); - let mut sleep = Box::pin(tokio::time::sleep(timeout)); - tokio::select! { - _ = &mut sleep => { - let elapsed = start_time.elapsed(); - let elapsed_str = format_duration_human(elapsed); - let current_time = chrono::Local::now().format("%Y-%m-%d %H:%M:%S %Z").to_string(); - WaitAndContinueResult::Timeout(format!( - "waited {elapsed_str}; no incoming event arrived; current time: {current_time}" - )) - } - interrupt = rx.recv() => { - let elapsed = start_time.elapsed(); - let elapsed_str = format_duration_human(elapsed); - let current_time = chrono::Local::now().format("%Y-%m-%d %H:%M:%S %Z").to_string(); - match interrupt { - Some(Interrupt::Reset) => WaitAndContinueResult::Reset, - Some(Interrupt::Compact) => WaitAndContinueResult::Compact, - Some(Interrupt::UpdateSoul { content }) => WaitAndContinueResult::UpdateSoul(content), - Some(Interrupt::DebugReflectRequest) => WaitAndContinueResult::Timeout( - format!("debug reflection request received while waiting; current time: {current_time}; elapsed since last turn: {elapsed_str}; try again after this turn") - ), - Some(interrupt) => { - let source = interrupt.source_tag().to_string(); - WaitAndContinueResult::Incoming { - result: format!( - "interrupted early by incoming {source} event after {elapsed_str}; current time: {current_time}" - ), - interrupt, - } - } - None => WaitAndContinueResult::Timeout(format!("interrupt channel closed; current time: {current_time}; elapsed since last turn: {elapsed_str}; continue")), - } - } - } -} - -fn format_duration_human(d: std::time::Duration) -> String { - let secs = d.as_secs(); - if secs == 0 { - return format!("{}ms", d.as_millis()); - } - let h = secs / 3600; - let m = (secs % 3600) / 60; - let s = secs % 60; - let mut parts = vec![]; - if h > 0 { - parts.push(format!("{h}hour")); - } - if m > 0 { - parts.push(format!("{m}min")); - } - if s > 0 || parts.is_empty() { - parts.push(format!("{s}sec")); - } - parts.join(" ") -} - -fn inject_incoming_interrupt( - interrupt: Interrupt, - agent_name: &str, - ctx: &mut Context, - memory: &MemoryStore, - output: &broadcast::Sender, - turn_count: &mut usize, -) { - let prompt_text = interrupt.prompt_text(agent_name); - ctx.push_input(&prompt_text); - if let Ok(entry) = memory.log_turn("user", &prompt_text, None) { - let _ = output.send(AgentEvent::UserTurn(entry)); - } - *turn_count += 1; - let _ = output.send(AgentEvent::Status(format!( - "injected incoming {} event", - interrupt.source_tag() - ))); -} - fn save_context_snapshot(memory: &MemoryStore, ctx: &Context) { if let Err(e) = memory.save_context(&ctx.turns) { tracing::warn!(err = %e, "failed to persist model context snapshot"); }