use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;
use crate::agent::{FinishReason, PlanningMode};
use crate::compact::Compactor;
use crate::event::AgentEvent;
use crate::hooks::HookRegistry;
use crate::llm::{LlmProvider, TokenUsage, ToolSpec};
use crate::message::Message;
use crate::permissions::PermissionMode;
use crate::storage::{NoopSessionStore, SessionStore, StorageBackend};
use crate::tool_set_provider::ToolSetProvider;
use crate::tools::PermissionHook;
use crate::tools::ToolRegistry;
pub struct TurnContext {
pub messages: Vec<Message>,
pub step_events_tx: Option<tokio::sync::mpsc::UnboundedSender<AgentEvent>>,
pub plan_confirmed: bool,
pub plan_buffer: Option<Vec<crate::llm::ToolCall>>,
pub tool_specs: Vec<ToolSpec>,
pub streaming: bool,
pub permission_hook: Option<Arc<dyn PermissionHook>>,
pub planning_mode: PlanningMode,
pub exploring_plan_mode: Arc<AtomicBool>,
pub permission_mode: PermissionMode,
pub mailbox: Option<crate::tools::send_message::WorkerMailbox>,
}
#[derive(Debug)]
pub struct TurnOutcome {
pub new_messages: Vec<Message>,
pub final_text: Option<String>,
pub finish_reason: FinishReason,
pub usage: TokenUsage,
pub llm_latency_ms: u64,
pub steps: usize,
pub side_effects: Vec<SideEffect>,
pub plan_buffer: Option<Vec<crate::llm::ToolCall>>,
pub tool_audits: std::collections::HashMap<String, crate::tools::AuditMeta>,
pub plan_confirmed: bool,
}
#[derive(Debug, Clone)]
pub enum SideEffect {
BackgroundJob {
id: String,
pid: u32,
command: String,
},
ScheduleWakeup { delay: Duration, prompt: String },
}
#[derive(Clone)]
pub struct AgentKernel {
pub(crate) llm: Arc<dyn LlmProvider>,
pub(crate) tools: ToolRegistry,
pub(crate) max_steps: usize,
pub(crate) max_transcript_chars: Option<usize>,
pub(crate) compactor: Option<Compactor>,
pub(crate) hooks: HookRegistry,
pub(crate) shutdown_token: Option<tokio_util::sync::CancellationToken>,
pub(crate) storage: Arc<dyn StorageBackend>,
pub(crate) session_store: Arc<dyn SessionStore>,
}
impl std::fmt::Debug for AgentKernel {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let tools_count = self.tools.names().len();
let hooks_count = self.hooks.len();
f.debug_struct("AgentKernel")
.field("llm", &"<LlmProvider>")
.field("tools_count", &tools_count)
.field("max_steps", &self.max_steps)
.field("max_transcript_chars", &self.max_transcript_chars)
.field("compactor", &self.compactor)
.field("hooks_count", &hooks_count)
.field("storage", &"<StorageBackend>")
.field("session_store", &"<SessionStore>")
.finish()
}
}
impl AgentKernel {
pub fn builder() -> AgentKernelBuilder {
AgentKernelBuilder::default()
}
pub fn llm(&self) -> &Arc<dyn LlmProvider> {
&self.llm
}
pub fn tools(&self) -> &ToolRegistry {
&self.tools
}
pub fn tools_mut(&mut self) -> &mut ToolRegistry {
&mut self.tools
}
pub fn shutdown_token(&self) -> Option<&tokio_util::sync::CancellationToken> {
self.shutdown_token.as_ref()
}
pub fn storage(&self) -> &Arc<dyn StorageBackend> {
&self.storage
}
pub fn session_store(&self) -> &Arc<dyn SessionStore> {
&self.session_store
}
pub fn hooks(&self) -> &HookRegistry {
&self.hooks
}
pub fn with_tools(&self, tools: ToolRegistry) -> Self {
Self {
llm: self.llm.clone(),
tools,
max_steps: self.max_steps,
max_transcript_chars: self.max_transcript_chars,
compactor: self.compactor.clone(),
hooks: self.hooks.clone(),
shutdown_token: self.shutdown_token.clone(),
storage: self.storage.clone(),
session_store: self.session_store.clone(),
}
}
pub async fn run(&self, ctx: TurnContext) -> crate::error::Result<TurnOutcome> {
use crate::run_core::RunCore;
let input_len = ctx.messages.len();
let core = RunCore {
messages: ctx.messages,
llm: self.llm.clone(),
tools: Arc::new(self.tools.clone()),
max_steps: self.max_steps,
max_transcript_chars: self.max_transcript_chars,
events: ctx.step_events_tx,
streaming: ctx.streaming,
compactor: self.compactor.clone(),
permission_hook: ctx.permission_hook,
hooks: &self.hooks,
planning_mode: ctx.planning_mode,
total_llm_latency_ms: 0,
plan_buffer: ctx.plan_buffer,
plan_confirmed: ctx.plan_confirmed,
exploring_plan_mode: ctx.exploring_plan_mode,
permission_mode: ctx.permission_mode,
shutdown_token: self.shutdown_token.clone(),
mailbox: ctx.mailbox,
};
let inner = core.run_inner().await?;
let mut new_messages = if inner.messages.len() > input_len {
inner.messages[input_len..].to_vec()
} else {
Vec::new()
};
if !inner.messages.is_empty()
&& inner.messages[0].role == crate::message::Role::System
&& inner.messages[0].content.contains("[compacted:")
{
new_messages.insert(0, inner.messages[0].clone());
}
Ok(TurnOutcome {
new_messages,
final_text: inner.final_message,
finish_reason: inner.finish_reason,
usage: inner.total_usage,
llm_latency_ms: inner.total_llm_latency_ms,
steps: inner.steps,
side_effects: Vec::new(),
plan_buffer: inner.plan_buffer,
plan_confirmed: inner.plan_confirmed,
tool_audits: inner.tool_audits,
})
}
}
#[derive(Default)]
pub struct AgentKernelBuilder {
llm: Option<Arc<dyn LlmProvider>>,
tools: Option<ToolRegistry>,
max_steps: Option<usize>,
max_transcript_chars: Option<usize>,
compactor: Option<Compactor>,
hooks: Option<HookRegistry>,
shutdown_token: Option<tokio_util::sync::CancellationToken>,
storage: Option<Arc<dyn StorageBackend>>,
session_store: Option<Arc<dyn SessionStore>>,
tool_set_provider: Option<Arc<dyn ToolSetProvider>>,
}
impl std::fmt::Debug for AgentKernelBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let tools_desc = self.tools.as_ref().map(|t| t.names().len());
let hooks_desc = self.hooks.as_ref().map(|h| h.len());
f.debug_struct("AgentKernelBuilder")
.field("llm", &self.llm.as_ref().map(|_| "<LlmProvider>"))
.field("tools", &tools_desc)
.field("max_steps", &self.max_steps)
.field("max_transcript_chars", &self.max_transcript_chars)
.field("compactor", &self.compactor)
.field("hooks", &hooks_desc)
.field(
"storage",
&self.storage.as_ref().map(|_| "<StorageBackend>"),
)
.field(
"session_store",
&self.session_store.as_ref().map(|_| "<SessionStore>"),
)
.field(
"tool_set_provider",
&self.tool_set_provider.as_ref().map(|_| "<ToolSetProvider>"),
)
.finish()
}
}
impl AgentKernelBuilder {
pub fn llm(mut self, llm: Arc<dyn LlmProvider>) -> Self {
self.llm = Some(llm);
self
}
pub fn tools(mut self, tools: ToolRegistry) -> Self {
self.tools = Some(tools);
self
}
pub fn max_steps(mut self, n: usize) -> Self {
self.max_steps = Some(n);
self
}
pub fn max_transcript_chars(mut self, n: usize) -> Self {
self.max_transcript_chars = Some(n);
self
}
pub fn compactor(mut self, compactor: Compactor) -> Self {
self.compactor = Some(compactor);
self
}
pub fn hooks(mut self, hooks: HookRegistry) -> Self {
self.hooks = Some(hooks);
self
}
pub fn shutdown_token(mut self, token: tokio_util::sync::CancellationToken) -> Self {
self.shutdown_token = Some(token);
self
}
pub fn with_storage(mut self, backend: Arc<dyn StorageBackend>) -> Self {
self.storage = Some(backend);
self
}
pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
self.session_store = Some(store);
self
}
pub fn with_tool_set_provider(mut self, provider: Arc<dyn ToolSetProvider>) -> Self {
self.tool_set_provider = Some(provider);
self
}
pub fn build(self) -> crate::error::Result<AgentKernel> {
let llm = self.llm.ok_or_else(|| crate::error::Error::Config {
message: "llm provider is required".into(),
})?;
let tools = if let Some(registry) = self.tools {
registry
} else if let Some(ref provider) = self.tool_set_provider {
provider.build_registry()
} else {
ToolRegistry::local()
};
let max_steps = self.max_steps.unwrap_or(32);
let hooks = self.hooks.unwrap_or_default();
let storage: Arc<dyn StorageBackend> = self.storage.unwrap_or_else(|| {
Arc::new(crate::storage::local::LocalStorageBackend::new(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
))
});
let session_store: Arc<dyn SessionStore> = self
.session_store
.unwrap_or_else(|| Arc::new(NoopSessionStore));
Ok(AgentKernel {
llm,
tools,
max_steps,
max_transcript_chars: self.max_transcript_chars,
compactor: self.compactor,
hooks,
shutdown_token: self.shutdown_token,
storage,
session_store,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llm::MockProvider;
#[test]
fn kernel_builder_requires_llm() {
let result = AgentKernel::builder().build();
assert!(result.is_err());
match result {
Err(e) => assert!(e.to_string().contains("llm provider is required")),
Ok(_) => panic!("expected Err"),
}
}
#[test]
fn kernel_builder_happy_path() {
let mock = MockProvider::default();
let tools = ToolRegistry::local();
let kernel = AgentKernel::builder()
.llm(Arc::new(mock))
.tools(tools)
.max_steps(16)
.build()
.expect("build should succeed");
assert_eq!(kernel.max_steps, 16);
}
#[test]
fn kernel_builder_default_max_steps() {
let mock = MockProvider::default();
let tools = ToolRegistry::local();
let kernel = AgentKernel::builder()
.llm(Arc::new(mock))
.tools(tools)
.build()
.expect("build should succeed");
assert_eq!(kernel.max_steps, 32);
}
#[test]
fn kernel_clone_is_independent() {
let mock = MockProvider::default();
let tools1 = ToolRegistry::local();
let kernel = AgentKernel::builder()
.llm(Arc::new(mock))
.tools(tools1)
.build()
.expect("build should succeed");
let mut cloned = kernel.clone();
let new_tools = ToolRegistry::local();
cloned.tools = new_tools;
assert!(!Arc::ptr_eq(
kernel.tools().transport(),
cloned.tools().transport()
));
}
#[test]
fn kernel_with_tools_preserves_llm() {
let mock = MockProvider::default();
let mock_arc = Arc::new(mock);
let tools1 = ToolRegistry::local();
let kernel = AgentKernel::builder()
.llm(mock_arc.clone())
.tools(tools1)
.build()
.expect("build should succeed");
let tools2 = ToolRegistry::local();
let new_kernel = kernel.with_tools(tools2);
assert!(Arc::ptr_eq(&kernel.llm, &new_kernel.llm));
assert_eq!(new_kernel.max_steps, kernel.max_steps);
}
#[test]
fn turn_outcome_default_values() {
let outcome = TurnOutcome {
new_messages: vec![],
final_text: None,
finish_reason: FinishReason::NoMoreToolCalls,
usage: TokenUsage::default(),
llm_latency_ms: 0,
steps: 0,
side_effects: vec![],
plan_buffer: None,
plan_confirmed: false,
tool_audits: std::collections::HashMap::new(),
};
assert!(outcome.new_messages.is_empty());
assert!(outcome.final_text.is_none());
assert_eq!(outcome.finish_reason, FinishReason::NoMoreToolCalls);
assert_eq!(outcome.usage, TokenUsage::default());
assert_eq!(outcome.llm_latency_ms, 0);
assert_eq!(outcome.steps, 0);
assert!(outcome.side_effects.is_empty());
}
#[test]
fn side_effect_variants() {
let bg = SideEffect::BackgroundJob {
id: "job-1".into(),
pid: 12345,
command: "echo hello".into(),
};
match &bg {
SideEffect::BackgroundJob { id, pid, command } => {
assert_eq!(id, "job-1");
assert_eq!(*pid, 12345);
assert_eq!(command, "echo hello");
}
_ => panic!("expected BackgroundJob"),
}
let wake = SideEffect::ScheduleWakeup {
delay: Duration::from_secs(60),
prompt: "check status".into(),
};
match &wake {
SideEffect::ScheduleWakeup { delay, prompt } => {
assert_eq!(delay.as_secs(), 60);
assert_eq!(prompt, "check status");
}
_ => panic!("expected ScheduleWakeup"),
}
}
}