use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::sync::{Mutex, RwLock};
use tokio_util::sync::CancellationToken;
use edgecrab_state::SessionDb;
use edgecrab_tools::ProcessTable;
use edgecrab_tools::TodoStore;
use edgecrab_tools::config_ref::{AppConfigRef, LspServerConfigRef};
use edgecrab_tools::registry::{GatewaySender, ToolInventoryEntry, ToolRegistry};
use edgecrab_types::{AgentError, ApiMode, Cost, Message, Platform, Role, Usage};
use edgequake_llm::LLMProvider;
use crate::config::AppConfig;
pub struct Agent {
pub(crate) config: RwLock<AgentConfig>,
pub(crate) provider: RwLock<Arc<dyn LLMProvider>>,
#[allow(dead_code)] pub(crate) state_db: Option<Arc<SessionDb>>,
pub(crate) tool_registry: Option<Arc<ToolRegistry>>,
pub(crate) gateway_sender: RwLock<Option<Arc<dyn GatewaySender>>>,
pub(crate) process_table: Arc<ProcessTable>,
pub(crate) conversation_lock: Mutex<()>,
pub(crate) session: RwLock<SessionState>,
pub(crate) budget: Arc<IterationBudget>,
pub(crate) cancel: std::sync::Mutex<CancellationToken>,
pub(crate) gc_cancel: CancellationToken,
pub(crate) todo_store: Arc<edgecrab_tools::TodoStore>,
}
#[derive(Debug, Clone, Default)]
pub struct IsolatedAgentOptions {
pub session_id: Option<String>,
pub platform: Option<Platform>,
pub quiet_mode: Option<bool>,
pub origin_chat: Option<(String, String)>,
}
#[derive(Debug, Clone)]
pub struct AgentConfig {
pub model: String,
pub max_iterations: u32,
pub enabled_toolsets: Vec<String>,
pub disabled_toolsets: Vec<String>,
pub enabled_tools: Vec<String>,
pub disabled_tools: Vec<String>,
pub streaming: bool,
pub temperature: Option<f32>,
pub platform: Platform,
pub api_mode: ApiMode,
pub session_id: Option<String>,
pub quiet_mode: bool,
pub save_trajectories: bool,
pub skip_context_files: bool,
pub skip_memory: bool,
pub reasoning_effort: Option<String>,
pub personality_addon: Option<String>,
pub model_config: crate::config::ModelConfig,
pub skills_config: crate::config::SkillsConfig,
pub delegation_enabled: bool,
pub delegation_model: Option<String>,
pub delegation_provider: Option<String>,
pub delegation_max_subagents: u32,
pub delegation_max_iterations: u32,
pub origin_chat: Option<(String, String)>,
pub browser: crate::config::BrowserConfig,
pub checkpoints_enabled: bool,
pub checkpoints_max_snapshots: u32,
pub terminal_backend: edgecrab_tools::tools::backends::BackendKind,
pub terminal_docker: edgecrab_tools::tools::backends::DockerBackendConfig,
pub terminal_ssh: edgecrab_tools::tools::backends::SshBackendConfig,
pub terminal_modal: edgecrab_tools::tools::backends::ModalBackendConfig,
pub terminal_daytona: edgecrab_tools::tools::backends::DaytonaBackendConfig,
pub terminal_singularity: edgecrab_tools::tools::backends::SingularityBackendConfig,
pub compression: crate::config::CompressionConfig,
pub auxiliary: crate::config::AuxiliaryConfig,
pub moa: crate::config::MoaConfig,
pub tts: crate::config::TtsConfig,
pub stt: crate::config::SttConfig,
pub image_generation: crate::config::ImageGenerationConfig,
pub terminal_env_passthrough: Vec<String>,
pub file_allowed_roots: Vec<std::path::PathBuf>,
pub path_restrictions: Vec<std::path::PathBuf>,
pub lsp: crate::config::LspConfig,
}
impl Default for AgentConfig {
fn default() -> Self {
Self {
model: "ollama/gemma4:latest".into(),
max_iterations: 90,
enabled_toolsets: Vec::new(),
disabled_toolsets: Vec::new(),
enabled_tools: Vec::new(),
disabled_tools: Vec::new(),
streaming: true,
temperature: None,
platform: Platform::Cli,
api_mode: ApiMode::ChatCompletions,
session_id: None,
quiet_mode: false,
save_trajectories: false,
skip_context_files: false,
skip_memory: false,
reasoning_effort: None,
personality_addon: None,
model_config: crate::config::ModelConfig::default(),
skills_config: crate::config::SkillsConfig::default(),
delegation_enabled: true,
delegation_model: None,
delegation_provider: None,
delegation_max_subagents: 3,
delegation_max_iterations: 50,
origin_chat: None,
browser: crate::config::BrowserConfig::default(),
checkpoints_enabled: true,
checkpoints_max_snapshots: 50,
terminal_backend: edgecrab_tools::tools::backends::BackendKind::Local,
terminal_docker: edgecrab_tools::tools::backends::DockerBackendConfig::default(),
terminal_ssh: edgecrab_tools::tools::backends::SshBackendConfig::default(),
terminal_modal: edgecrab_tools::tools::backends::ModalBackendConfig::default(),
terminal_daytona: edgecrab_tools::tools::backends::DaytonaBackendConfig::default(),
terminal_singularity:
edgecrab_tools::tools::backends::SingularityBackendConfig::default(),
compression: crate::config::CompressionConfig::default(),
auxiliary: crate::config::AuxiliaryConfig::default(),
moa: crate::config::MoaConfig::default(),
tts: crate::config::TtsConfig::default(),
stt: crate::config::SttConfig::default(),
image_generation: crate::config::ImageGenerationConfig::default(),
terminal_env_passthrough: Vec::new(),
file_allowed_roots: Vec::new(),
path_restrictions: Vec::new(),
lsp: crate::config::LspConfig::default(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ResolvedToolPolicy {
pub expanded_enabled: Vec<String>,
pub expanded_disabled: Vec<String>,
pub parent_active_toolsets: Vec<String>,
}
pub(crate) fn resolve_tool_policy(config: &AgentConfig) -> ResolvedToolPolicy {
let expanded_enabled = edgecrab_tools::toolsets::expand_toolset_names(&config.enabled_toolsets);
let expanded_disabled =
edgecrab_tools::toolsets::expand_toolset_names(&config.disabled_toolsets);
let parent_active_toolsets = if config.enabled_toolsets.is_empty()
|| edgecrab_tools::toolsets::contains_all_sentinel(&config.enabled_toolsets)
|| expanded_enabled.is_empty()
{
Vec::new()
} else {
expanded_enabled
.iter()
.filter(|toolset| {
!expanded_disabled
.iter()
.any(|disabled| disabled == *toolset)
})
.cloned()
.collect()
};
ResolvedToolPolicy {
expanded_enabled,
expanded_disabled,
parent_active_toolsets,
}
}
impl AgentConfig {
fn lsp_server_refs(&self) -> std::collections::HashMap<String, LspServerConfigRef> {
self.lsp
.servers
.iter()
.map(|(name, server)| {
(
name.clone(),
LspServerConfigRef {
command: server.command.clone(),
args: server.args.clone(),
file_extensions: server.file_extensions.clone(),
language_id: server.language_id.clone(),
root_markers: server.root_markers.clone(),
env: server.env.clone(),
initialization_options: server.initialization_options.clone(),
},
)
})
.collect()
}
fn disabled_skills_for_platform(&self) -> Vec<String> {
let mut disabled = self.skills_config.disabled.clone();
let platform_key = self.platform.to_string();
if let Some(platform_disabled) = self.skills_config.platform_disabled.get(&platform_key) {
disabled.extend(platform_disabled.iter().cloned());
}
disabled
}
pub(crate) fn to_app_config_ref(
&self,
gateway_running: bool,
tool_policy: &ResolvedToolPolicy,
) -> AppConfigRef {
AppConfigRef {
edgecrab_home: crate::config::edgecrab_home(),
file_allowed_roots: self.file_allowed_roots.clone(),
path_restrictions: self.path_restrictions.clone(),
lsp_enabled: self.lsp.enabled,
lsp_file_size_limit_bytes: self.lsp.file_size_limit_bytes,
lsp_servers: self.lsp_server_refs(),
delegation_enabled: self.delegation_enabled,
delegation_model: self.delegation_model.clone(),
delegation_provider: self.delegation_provider.clone(),
delegation_max_subagents: self.delegation_max_subagents,
delegation_max_iterations: self.delegation_max_iterations,
parent_active_toolsets: tool_policy.parent_active_toolsets.clone(),
disabled_toolsets: tool_policy.expanded_disabled.clone(),
enabled_tools: self.enabled_tools.clone(),
disabled_tools: self.disabled_tools.clone(),
external_skill_dirs: self.skills_config.external_dirs.clone(),
disabled_skills: self.disabled_skills_for_platform(),
preloaded_skills: self.skills_config.preloaded.clone(),
browser_record_sessions: self.browser.record_sessions,
browser_command_timeout: self.browser.command_timeout,
browser_recording_max_age_hours: self.browser.recording_max_age_hours,
checkpoints_enabled: self.checkpoints_enabled,
checkpoints_max_snapshots: self.checkpoints_max_snapshots,
terminal_backend: self.terminal_backend.clone(),
terminal_docker: self.terminal_docker.clone(),
terminal_ssh: self.terminal_ssh.clone(),
terminal_modal: self.terminal_modal.clone(),
terminal_daytona: self.terminal_daytona.clone(),
terminal_singularity: self.terminal_singularity.clone(),
terminal_env_passthrough: self.terminal_env_passthrough.clone(),
auxiliary_provider: self.auxiliary.provider.clone(),
auxiliary_model: self.auxiliary.model.clone(),
auxiliary_base_url: self.auxiliary.base_url.clone(),
auxiliary_api_key_env: self.auxiliary.api_key_env.clone(),
moa_enabled: self.moa.enabled,
moa_reference_models: self.moa.reference_models.clone(),
moa_aggregator_model: Some(self.moa.aggregator_model.clone()),
tts_provider: Some(self.tts.provider.clone()),
tts_voice: Some(self.tts.voice.clone()),
tts_rate: self.tts.rate.clone(),
tts_model: self.tts.model.clone(),
tts_elevenlabs_voice_id: self.tts.elevenlabs_voice_id.clone(),
tts_elevenlabs_model_id: self.tts.elevenlabs_model_id.clone(),
tts_elevenlabs_api_key_env: Some(self.tts.elevenlabs_api_key_env.clone()),
stt_provider: Some(self.stt.provider.clone()),
stt_whisper_model: Some(self.stt.whisper_model.clone()),
image_provider: Some(self.image_generation.provider.clone()),
image_model: Some(self.image_generation.model.clone()),
gateway_running,
..Default::default()
}
}
}
#[derive(Clone, Default)]
pub struct SessionState {
pub session_id: Option<String>,
pub messages: Vec<Message>,
pub cached_system_prompt: Option<String>,
pub user_turn_count: u32,
pub api_call_count: u32,
pub session_input_tokens: u64,
pub session_output_tokens: u64,
pub session_cache_read_tokens: u64,
pub session_cache_write_tokens: u64,
pub session_reasoning_tokens: u64,
pub last_prompt_tokens: u64,
pub native_tool_streaming_disabled: bool,
pub session_tool_call_count: u32,
}
pub struct IterationBudget {
remaining: AtomicU32,
max: u32,
}
impl IterationBudget {
pub fn new(max: u32) -> Self {
Self {
remaining: AtomicU32::new(max),
max,
}
}
pub fn try_consume(&self) -> bool {
loop {
let current = self.remaining.load(Ordering::Relaxed);
if current == 0 {
return false;
}
if self
.remaining
.compare_exchange_weak(current, current - 1, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
return true;
}
}
}
pub fn remaining(&self) -> u32 {
self.remaining.load(Ordering::Relaxed)
}
pub fn max(&self) -> u32 {
self.max
}
pub fn used(&self) -> u32 {
self.max.saturating_sub(self.remaining())
}
pub fn reset(&self) {
self.remaining.store(self.max, Ordering::Relaxed);
}
}
#[derive(Debug, Clone)]
pub struct ConversationResult {
pub final_response: String,
pub messages: Vec<Message>,
pub session_id: String,
pub api_calls: u32,
pub interrupted: bool,
pub budget_exhausted: bool,
pub model: String,
pub usage: Usage,
pub cost: Cost,
pub tool_errors: Vec<edgecrab_types::ToolErrorRecord>,
}
impl Agent {
fn build_runtime_clone(
config: AgentConfig,
provider: Arc<dyn LLMProvider>,
state_db: Option<Arc<SessionDb>>,
tool_registry: Option<Arc<ToolRegistry>>,
) -> Self {
let budget = Arc::new(IterationBudget::new(config.max_iterations));
let gc_cancel = CancellationToken::new();
let process_table = Arc::new(ProcessTable::new());
process_table.spawn_gc_task(gc_cancel.clone());
Self {
config: RwLock::new(config),
provider: RwLock::new(provider),
state_db,
tool_registry,
gateway_sender: RwLock::new(None),
process_table,
conversation_lock: Mutex::new(()),
session: RwLock::new(SessionState::default()),
budget,
cancel: std::sync::Mutex::new(CancellationToken::new()),
gc_cancel,
todo_store: Arc::new(TodoStore::new()),
}
}
pub async fn chat(&self, message: &str) -> Result<String, AgentError> {
let result = self.run_conversation(message, None, None).await?;
Ok(result.final_response)
}
pub async fn chat_in_cwd(&self, message: &str, cwd: &Path) -> Result<String, AgentError> {
let result = self
.run_conversation_in_cwd(message, None, None, cwd)
.await?;
Ok(result.final_response)
}
pub async fn set_gateway_sender(&self, sender: Arc<dyn GatewaySender>) {
*self.gateway_sender.write().await = Some(sender);
}
pub async fn chat_with_origin(
&self,
message: &str,
platform: &str,
chat_id: &str,
) -> Result<String, AgentError> {
{
let mut cfg = self.config.write().await;
cfg.origin_chat = Some((platform.to_string(), chat_id.to_string()));
cfg.platform = platform_from_str(platform).unwrap_or(cfg.platform);
}
let result = self.run_conversation(message, None, None).await?;
{
let mut cfg = self.config.write().await;
cfg.origin_chat = None;
}
Ok(result.final_response)
}
pub async fn chat_streaming_with_origin(
&self,
message: &str,
platform: &str,
chat_id: &str,
chunk_tx: tokio::sync::mpsc::UnboundedSender<StreamEvent>,
) -> Result<(), AgentError> {
{
let mut cfg = self.config.write().await;
cfg.origin_chat = Some((platform.to_string(), chat_id.to_string()));
cfg.platform = platform_from_str(platform).unwrap_or(cfg.platform);
}
let result = self.chat_streaming(message, chunk_tx).await;
{
let mut cfg = self.config.write().await;
cfg.origin_chat = None;
}
result
}
pub async fn chat_streaming(
&self,
message: &str,
chunk_tx: tokio::sync::mpsc::UnboundedSender<StreamEvent>,
) -> Result<(), AgentError> {
let streaming_enabled = {
let config = self.config.read().await;
config.streaming
};
let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel();
let saw_live_content = Arc::new(AtomicBool::new(false));
let saw_live_content_forwarder = saw_live_content.clone();
let chunk_tx_forwarder = chunk_tx.clone();
let forwarder = tokio::spawn(async move {
while let Some(event) = event_rx.recv().await {
if matches!(event, StreamEvent::Token(_) | StreamEvent::Reasoning(_)) {
saw_live_content_forwarder.store(true, AtomicOrdering::Relaxed);
}
let _ = chunk_tx_forwarder.send(event);
}
});
match self
.execute_loop(message, None, None, Some(&event_tx), None)
.await
{
Ok(result) => {
drop(event_tx);
let _ = forwarder.await;
if !saw_live_content.load(AtomicOrdering::Relaxed) {
if let Some(reasoning) = result
.messages
.iter()
.rev()
.find(|msg| msg.role == Role::Assistant)
.and_then(|msg| msg.reasoning.clone())
.filter(|reasoning| !reasoning.trim().is_empty())
{
let _ = chunk_tx.send(StreamEvent::Reasoning(reasoning));
}
if streaming_enabled {
for chunk in result.final_response.as_bytes().chunks(50) {
let text = String::from_utf8_lossy(chunk).into_owned();
let _ = chunk_tx.send(StreamEvent::Token(text));
}
} else if !result.final_response.is_empty() {
let _ = chunk_tx.send(StreamEvent::Token(result.final_response.clone()));
}
}
let _ = chunk_tx.send(StreamEvent::Done);
Ok(())
}
Err(e) => {
drop(event_tx);
let _ = forwarder.await;
let _ = chunk_tx.send(StreamEvent::Error(e.to_string()));
Err(e)
}
}
}
pub async fn run_conversation(
&self,
user_message: &str,
system_message: Option<&str>,
conversation_history: Option<Vec<Message>>,
) -> Result<ConversationResult, AgentError> {
self.execute_loop(
user_message,
system_message,
conversation_history,
None,
None,
)
.await
}
pub async fn run_conversation_in_cwd(
&self,
user_message: &str,
system_message: Option<&str>,
conversation_history: Option<Vec<Message>>,
cwd: &Path,
) -> Result<ConversationResult, AgentError> {
self.execute_loop(
user_message,
system_message,
conversation_history,
None,
Some(cwd),
)
.await
}
pub async fn fork_isolated(&self, options: IsolatedAgentOptions) -> Result<Self, AgentError> {
let mut config = self.config.read().await.clone();
let provider = self.provider.read().await.clone();
let gateway_sender = self.gateway_sender.read().await.clone();
if let Some(session_id) = options.session_id {
config.session_id = Some(session_id);
} else {
config.session_id = None;
}
if let Some(platform) = options.platform {
config.platform = platform;
}
if let Some(quiet_mode) = options.quiet_mode {
config.quiet_mode = quiet_mode;
}
config.origin_chat = options.origin_chat;
let child = Self::build_runtime_clone(
config,
provider,
self.state_db.clone(),
self.tool_registry.clone(),
);
if let Some(sender) = gateway_sender {
child.set_gateway_sender(sender).await;
}
Ok(child)
}
pub fn interrupt(&self) {
self.cancel
.lock()
.expect("cancel mutex not poisoned")
.cancel();
}
pub fn is_cancelled(&self) -> bool {
self.cancel
.lock()
.expect("cancel mutex not poisoned")
.is_cancelled()
}
pub async fn new_session(&self) {
let passthrough = {
let cfg = self.config.read().await;
cfg.terminal_env_passthrough.clone()
};
edgecrab_tools::tools::backends::local::clear_env_passthrough();
if !passthrough.is_empty() {
edgecrab_tools::tools::backends::local::register_env_passthrough(&passthrough);
}
let mut session = self.session.write().await;
*session = SessionState::default();
self.budget.reset();
}
pub async fn swap_model(&self, model: String, provider: Arc<dyn LLMProvider>) {
{
let mut cfg = self.config.write().await;
cfg.model = model;
}
{
let mut prov = self.provider.write().await;
*prov = provider;
}
}
pub async fn model(&self) -> String {
self.config.read().await.model.clone()
}
pub async fn session_snapshot(&self) -> SessionSnapshot {
let session = self.session.read().await;
let config = self.config.read().await;
SessionSnapshot {
session_id: session.session_id.clone(),
model: config.model.clone(),
message_count: session.messages.len(),
user_turn_count: session.user_turn_count,
api_call_count: session.api_call_count,
input_tokens: session.session_input_tokens,
output_tokens: session.session_output_tokens,
cache_read_tokens: session.session_cache_read_tokens,
cache_write_tokens: session.session_cache_write_tokens,
reasoning_tokens: session.session_reasoning_tokens,
last_prompt_tokens: session.last_prompt_tokens,
budget_remaining: self.budget.remaining(),
budget_max: self.budget.max(),
}
}
pub fn try_session_snapshot(&self) -> Option<SessionSnapshot> {
let session = self.session.try_read().ok()?;
let config = self.config.try_read().ok()?;
Some(SessionSnapshot {
session_id: session.session_id.clone(),
model: config.model.clone(),
message_count: session.messages.len(),
user_turn_count: session.user_turn_count,
api_call_count: session.api_call_count,
input_tokens: session.session_input_tokens,
output_tokens: session.session_output_tokens,
cache_read_tokens: session.session_cache_read_tokens,
cache_write_tokens: session.session_cache_write_tokens,
reasoning_tokens: session.session_reasoning_tokens,
last_prompt_tokens: session.last_prompt_tokens,
budget_remaining: self.budget.remaining(),
budget_max: self.budget.max(),
})
}
pub async fn system_prompt(&self) -> Option<String> {
self.session.read().await.cached_system_prompt.clone()
}
pub(crate) async fn publish_session_state(&self, session: &SessionState) {
*self.session.write().await = session.clone();
}
pub async fn append_to_system_prompt(&self, note: &str) {
let mut session = self.session.write().await;
if let Some(ref mut prompt) = session.cached_system_prompt {
prompt.push_str("\n\n");
prompt.push_str(note);
}
}
pub async fn invalidate_system_prompt(&self) {
let mut session = self.session.write().await;
session.cached_system_prompt = None;
}
pub async fn set_personality_addon(&self, addon: Option<String>) {
{
let mut config = self.config.write().await;
config.personality_addon = addon;
}
self.invalidate_system_prompt().await;
}
pub async fn inject_assistant_context(&self, text: &str) {
let mut session = self.session.write().await;
session.messages.push(Message::assistant(text));
}
pub async fn messages(&self) -> Vec<Message> {
self.session.read().await.messages.clone()
}
pub fn try_messages(&self) -> Option<Vec<Message>> {
Some(self.session.try_read().ok()?.messages.clone())
}
pub async fn undo_last_turn(&self) -> usize {
let mut session = self.session.write().await;
let mut removed = 0;
while let Some(m) = session.messages.last() {
if m.role == Role::User {
session.messages.pop();
removed += 1;
break;
}
session.messages.pop();
removed += 1;
}
removed
}
pub async fn force_compress(&self) {
let provider = self.provider.read().await.clone();
let config = self.config.read().await.clone();
let mut session = self.session.write().await;
let params = crate::compression::CompressionParams::from_model_config(
&config.model,
&config.compression,
);
session.messages =
crate::compression::compress_with_llm(&session.messages, ¶ms, &provider).await;
}
pub async fn set_session_title(&self, title: String) {
let session = self.session.read().await;
if let (Some(db), Some(sid)) = (&self.state_db, &session.session_id) {
let _ = db.update_session_title(sid, &title);
}
}
pub async fn restore_session(&self, id: &str) -> Result<usize, AgentError> {
let db = self
.state_db
.as_ref()
.ok_or_else(|| AgentError::Config("No state database configured".into()))?;
let record = db
.get_session(id)?
.ok_or_else(|| AgentError::Config(format!("Session not found: {id}")))?;
let messages = db.get_messages(id)?;
{
let mut session = self.session.write().await;
session.session_id = Some(record.id.clone());
session.cached_system_prompt = record.system_prompt.clone();
session.user_turn_count = messages
.iter()
.filter(|m| matches!(m.role, edgecrab_types::Role::User))
.count() as u32;
session.api_call_count = 0;
session.session_input_tokens = record.input_tokens.max(0) as u64;
session.session_output_tokens = record.output_tokens.max(0) as u64;
session.session_cache_read_tokens = record.cache_read_tokens.max(0) as u64;
session.session_cache_write_tokens = record.cache_write_tokens.max(0) as u64;
session.session_reasoning_tokens = record.reasoning_tokens.max(0) as u64;
session.last_prompt_tokens = 0;
session.messages = messages;
}
if let Some(model) = record.model {
let mut config = self.config.write().await;
config.model = model;
}
self.budget.reset();
Ok(self.session.read().await.messages.len())
}
pub fn list_sessions(
&self,
limit: usize,
) -> Result<Vec<edgecrab_state::SessionSummary>, AgentError> {
match &self.state_db {
Some(db) => db.list_sessions(limit),
None => Ok(Vec::new()),
}
}
pub fn delete_session(&self, id: &str) -> Result<(), AgentError> {
match &self.state_db {
Some(db) => db.delete_session(id),
None => Err(AgentError::Config("No state database configured".into())),
}
}
pub fn rename_session(&self, id: &str, title: &str) -> Result<(), AgentError> {
match &self.state_db {
Some(db) => db.update_session_title(id, title),
None => Err(AgentError::Config("No state database configured".into())),
}
}
pub fn prune_sessions(
&self,
older_than_days: u32,
source: Option<&str>,
) -> Result<usize, AgentError> {
match &self.state_db {
Some(db) => db.prune_sessions(older_than_days, source),
None => Err(AgentError::Config("No state database configured".into())),
}
}
pub fn has_state_db(&self) -> bool {
self.state_db.is_some()
}
pub async fn state_db(&self) -> Option<Arc<SessionDb>> {
self.state_db.clone()
}
pub fn state_db_handle(&self) -> Option<Arc<SessionDb>> {
self.state_db.clone()
}
pub async fn provider_handle(&self) -> Arc<dyn LLMProvider> {
self.provider.read().await.clone()
}
pub async fn auxiliary_config(&self) -> crate::config::AuxiliaryConfig {
self.config.read().await.auxiliary.clone()
}
pub async fn smart_routing_config(&self) -> crate::config::SmartRoutingYaml {
self.config.read().await.model_config.smart_routing.clone()
}
pub async fn moa_config(&self) -> crate::config::MoaConfig {
self.config.read().await.moa.clone()
}
pub async fn media_config(
&self,
) -> (
crate::config::TtsConfig,
crate::config::SttConfig,
crate::config::ImageGenerationConfig,
) {
let config = self.config.read().await;
(
config.tts.clone(),
config.stt.clone(),
config.image_generation.clone(),
)
}
pub async fn tool_names(&self) -> Vec<String> {
match &self.tool_registry {
Some(reg) => reg
.tool_names()
.into_iter()
.map(|s| s.to_string())
.collect(),
None => Vec::new(),
}
}
pub async fn toolset_summary(&self) -> Vec<(String, usize)> {
match &self.tool_registry {
Some(reg) => reg.toolset_summary(),
None => Vec::new(),
}
}
pub async fn tool_inventory(&self) -> Vec<ToolInventoryEntry> {
let Some(registry) = &self.tool_registry else {
return Vec::new();
};
let config = self.config.read().await.clone();
let tool_policy = resolve_tool_policy(&config);
let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
let ctx = edgecrab_tools::registry::ToolContext {
task_id: "tool-inventory".into(),
cwd,
session_id: config
.session_id
.clone()
.unwrap_or_else(|| "tool-inventory".into()),
user_task: None,
cancel: CancellationToken::new(),
config: config
.to_app_config_ref(self.gateway_sender.read().await.is_some(), &tool_policy),
state_db: self.state_db.clone(),
platform: config.platform,
process_table: Some(self.process_table.clone()),
provider: Some(self.provider.read().await.clone()),
tool_registry: self.tool_registry.clone(),
delegate_depth: 0,
sub_agent_runner: None,
delegation_event_tx: None,
clarify_tx: None,
approval_tx: None,
on_skills_changed: None,
gateway_sender: self.gateway_sender.read().await.clone(),
origin_chat: config.origin_chat.clone(),
session_key: config.session_id.clone(),
todo_store: Some(self.todo_store.clone()),
current_tool_call_id: None,
current_tool_name: None,
tool_progress_tx: None,
};
registry.tool_inventory(&ctx)
}
pub async fn set_reasoning_effort(&self, level: Option<String>) {
let mut config = self.config.write().await;
config.reasoning_effort = level;
}
pub async fn set_streaming(&self, enabled: bool) {
let mut config = self.config.write().await;
config.streaming = enabled;
config.model_config.streaming = enabled;
}
pub async fn set_auxiliary_config(&self, auxiliary: crate::config::AuxiliaryConfig) {
let mut config = self.config.write().await;
config.auxiliary = auxiliary;
}
pub async fn set_smart_routing_config(&self, smart_routing: crate::config::SmartRoutingYaml) {
let mut config = self.config.write().await;
config.model_config.smart_routing = smart_routing;
}
pub async fn set_image_generation_config(
&self,
image_generation: crate::config::ImageGenerationConfig,
) {
let mut config = self.config.write().await;
config.image_generation = image_generation;
}
pub async fn set_moa_config(&self, moa: crate::config::MoaConfig) {
let mut config = self.config.write().await;
config.moa = moa.sanitized();
}
pub async fn set_toolset_filters(&self, enabled: Vec<String>, disabled: Vec<String>) {
let mut config = self.config.write().await;
config.enabled_toolsets = enabled;
config.disabled_toolsets = disabled;
}
pub async fn set_tool_filters(
&self,
enabled_toolsets: Vec<String>,
disabled_toolsets: Vec<String>,
enabled_tools: Vec<String>,
disabled_tools: Vec<String>,
) {
let mut config = self.config.write().await;
config.enabled_toolsets = enabled_toolsets;
config.disabled_toolsets = disabled_toolsets;
config.enabled_tools = enabled_tools;
config.disabled_tools = disabled_tools;
}
pub async fn toolset_filters(&self) -> (Vec<String>, Vec<String>) {
let config = self.config.read().await;
(
config.enabled_toolsets.clone(),
config.disabled_toolsets.clone(),
)
}
pub async fn tool_filters(&self) -> (Vec<String>, Vec<String>, Vec<String>, Vec<String>) {
let config = self.config.read().await;
(
config.enabled_toolsets.clone(),
config.disabled_toolsets.clone(),
config.enabled_tools.clone(),
config.disabled_tools.clone(),
)
}
}
#[derive(Debug, Clone)]
pub struct SessionSnapshot {
pub session_id: Option<String>,
pub model: String,
pub message_count: usize,
pub user_turn_count: u32,
pub api_call_count: u32,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_tokens: u64,
pub cache_write_tokens: u64,
pub reasoning_tokens: u64,
pub last_prompt_tokens: u64,
pub budget_remaining: u32,
pub budget_max: u32,
}
impl SessionSnapshot {
pub fn prompt_tokens(&self) -> u64 {
self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
}
pub fn total_tokens(&self) -> u64 {
self.prompt_tokens() + self.output_tokens + self.reasoning_tokens
}
pub fn context_pressure_tokens(&self) -> u64 {
if self.last_prompt_tokens > 0 {
self.last_prompt_tokens
} else {
self.prompt_tokens()
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApprovalChoice {
Once,
Session,
Always,
Deny,
}
pub enum StreamEvent {
Token(String),
Reasoning(String),
ToolExec {
tool_call_id: String,
name: String,
args_json: String,
},
ToolProgress {
tool_call_id: String,
name: String,
message: String,
},
ToolDone {
tool_call_id: String,
name: String,
args_json: String,
result_preview: Option<String>,
duration_ms: u64,
is_error: bool,
},
SubAgentStart {
task_index: usize,
task_count: usize,
goal: String,
},
SubAgentReasoning {
task_index: usize,
task_count: usize,
text: String,
},
SubAgentToolExec {
task_index: usize,
task_count: usize,
name: String,
args_json: String,
},
SubAgentFinish {
task_index: usize,
task_count: usize,
status: String,
duration_ms: u64,
summary: String,
api_calls: u32,
model: Option<String>,
},
Done,
Error(String),
Clarify {
question: String,
choices: Option<Vec<String>>,
response_tx: tokio::sync::oneshot::Sender<String>,
},
Approval {
command: String,
full_command: String,
reasons: Vec<String>,
response_tx: tokio::sync::oneshot::Sender<ApprovalChoice>,
},
SecretRequest {
var_name: String,
prompt: String,
is_sudo: bool,
response_tx: tokio::sync::oneshot::Sender<String>,
},
HookEvent {
event: String,
context_json: String,
},
ContextPressure {
estimated_tokens: usize,
threshold_tokens: usize,
},
}
impl std::fmt::Debug for StreamEvent {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Token(t) => write!(f, "Token({t:?})"),
Self::Reasoning(t) => write!(f, "Reasoning({t:?})"),
Self::ToolExec { name, .. } => write!(f, "ToolExec({name:?})"),
Self::ToolProgress { name, message, .. } => {
write!(f, "ToolProgress({name:?}, {message:?})")
}
Self::ToolDone {
name,
duration_ms,
is_error,
..
} => {
write!(f, "ToolDone({name:?}, {duration_ms}ms, err={is_error})")
}
Self::SubAgentStart {
task_index,
task_count,
goal,
} => write!(
f,
"SubAgentStart({}/{}, {:?})",
task_index + 1,
task_count,
goal
),
Self::SubAgentReasoning {
task_index,
task_count,
text,
} => write!(
f,
"SubAgentReasoning({}/{}, {:?})",
task_index + 1,
task_count,
text
),
Self::SubAgentToolExec {
task_index,
task_count,
name,
..
} => write!(
f,
"SubAgentToolExec({}/{}, {:?})",
task_index + 1,
task_count,
name
),
Self::SubAgentFinish {
task_index,
task_count,
status,
duration_ms,
..
} => write!(
f,
"SubAgentFinish({}/{}, {:?}, {}ms)",
task_index + 1,
task_count,
status,
duration_ms
),
Self::Done => write!(f, "Done"),
Self::Error(e) => write!(f, "Error({e:?})"),
Self::Clarify {
question, choices, ..
} => {
if choices.is_some() {
write!(f, "Clarify({question:?}, multiple-choice)")
} else {
write!(f, "Clarify({question:?})")
}
}
Self::Approval { command, .. } => write!(f, "Approval({command:?})"),
Self::SecretRequest {
var_name, is_sudo, ..
} => write!(f, "SecretRequest({var_name:?}, sudo={is_sudo})"),
Self::HookEvent { event, .. } => write!(f, "HookEvent({event:?})"),
Self::ContextPressure {
estimated_tokens,
threshold_tokens,
} => write!(
f,
"ContextPressure(est={estimated_tokens}, threshold={threshold_tokens})"
),
}
}
}
fn platform_from_str(s: &str) -> Option<Platform> {
match s {
"cli" => Some(Platform::Cli),
"telegram" => Some(Platform::Telegram),
"discord" => Some(Platform::Discord),
"slack" => Some(Platform::Slack),
"whatsapp" => Some(Platform::Whatsapp),
"feishu" => Some(Platform::Feishu),
"wecom" => Some(Platform::Wecom),
"signal" => Some(Platform::Signal),
"email" => Some(Platform::Email),
"matrix" => Some(Platform::Matrix),
"mattermost" => Some(Platform::Mattermost),
"dingtalk" => Some(Platform::DingTalk),
"sms" => Some(Platform::Sms),
"webhook" => Some(Platform::Webhook),
"api" | "api_server" => Some(Platform::Api),
"homeassistant" => Some(Platform::HomeAssistant),
"cron" => Some(Platform::Cron),
_ => None,
}
}
pub struct AgentBuilder {
config: AgentConfig,
provider: Option<Arc<dyn LLMProvider>>,
state_db: Option<Arc<SessionDb>>,
tool_registry: Option<Arc<ToolRegistry>>,
}
impl AgentBuilder {
pub fn new(model: &str) -> Self {
Self {
config: AgentConfig {
model: model.to_string(),
..Default::default()
},
provider: None,
state_db: None,
tool_registry: None,
}
}
pub fn from_config(config: &AppConfig) -> Self {
let personality_addon =
crate::config::resolve_personality(config, &config.display.personality);
Self {
config: AgentConfig {
model: config.model.default_model.clone(),
enabled_toolsets: config.tools.enabled_toolsets.clone().unwrap_or_default(),
disabled_toolsets: config.tools.disabled_toolsets.clone().unwrap_or_default(),
enabled_tools: config.tools.enabled_tools.clone().unwrap_or_default(),
disabled_tools: config.tools.disabled_tools.clone().unwrap_or_default(),
max_iterations: config.model.max_iterations,
streaming: config.model.streaming,
save_trajectories: config.save_trajectories,
skip_context_files: config.skip_context_files,
skip_memory: config.skip_memory,
temperature: config.model.temperature,
model_config: config.model.clone(),
skills_config: config.skills.clone(),
delegation_enabled: config.delegation.enabled,
delegation_model: config.delegation.model.clone(),
delegation_provider: config.delegation.provider.clone(),
delegation_max_subagents: config.delegation.max_subagents,
delegation_max_iterations: config.delegation.max_iterations,
personality_addon,
browser: config.browser.clone(),
checkpoints_enabled: config.checkpoints.enabled,
checkpoints_max_snapshots: config.checkpoints.max_snapshots,
terminal_backend: config.terminal.backend.clone(),
terminal_docker: config.terminal.docker.clone(),
terminal_ssh: config.terminal.ssh.clone(),
terminal_modal: config.terminal.modal.clone(),
terminal_daytona: config.terminal.daytona.clone(),
terminal_singularity: config.terminal.singularity.clone(),
compression: config.compression.clone(),
auxiliary: config.auxiliary.clone(),
moa: config.moa.clone(),
tts: config.tts.clone(),
stt: config.stt.clone(),
image_generation: config.image_generation.clone(),
terminal_env_passthrough: config.terminal.env_passthrough.clone(),
file_allowed_roots: config.tools.file.allowed_roots.clone(),
path_restrictions: config.security.path_restrictions.clone(),
lsp: config.lsp.clone(),
..Default::default()
},
provider: None,
state_db: None,
tool_registry: None,
}
}
pub fn provider(mut self, p: Arc<dyn LLMProvider>) -> Self {
self.provider = Some(p);
self
}
pub fn state_db(mut self, db: Arc<SessionDb>) -> Self {
self.state_db = Some(db);
self
}
pub fn tools(mut self, registry: Arc<ToolRegistry>) -> Self {
self.tool_registry = Some(registry);
self
}
pub fn max_iterations(mut self, n: u32) -> Self {
self.config.max_iterations = n;
self
}
pub fn streaming(mut self, enabled: bool) -> Self {
self.config.streaming = enabled;
self
}
pub fn platform(mut self, p: Platform) -> Self {
self.config.platform = p;
self
}
pub fn session_id(mut self, id: String) -> Self {
self.config.session_id = Some(id);
self
}
pub fn origin_chat(mut self, platform: String, chat_id: String) -> Self {
self.config.origin_chat = Some((platform, chat_id));
self
}
pub fn temperature(mut self, t: f32) -> Self {
self.config.temperature = Some(t);
self
}
pub fn quiet_mode(mut self, enabled: bool) -> Self {
self.config.quiet_mode = enabled;
self
}
pub fn build(self) -> Result<Agent, AgentError> {
let provider = self
.provider
.ok_or_else(|| AgentError::Config("provider is required".into()))?;
Ok(Agent::build_runtime_clone(
self.config,
provider,
self.state_db,
self.tool_registry,
))
}
}
impl Drop for Agent {
fn drop(&mut self) {
self.gc_cancel.cancel();
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use edgecrab_tools::registry::GatewaySender;
use edgequake_llm::traits::{
ChatMessage, CompletionOptions, StreamChunk, ToolChoice, ToolDefinition,
};
use futures::StreamExt;
use tokio::sync::Notify;
struct ReasoningStreamProvider;
struct SlowChatProvider {
started: Arc<Notify>,
release: Arc<Notify>,
}
struct MockGatewaySender;
#[test]
fn platform_from_str_accepts_gateway_api_server_alias() {
assert_eq!(platform_from_str("api"), Some(Platform::Api));
assert_eq!(platform_from_str("api_server"), Some(Platform::Api));
}
#[async_trait]
impl LLMProvider for ReasoningStreamProvider {
fn name(&self) -> &str {
"reasoning-stream"
}
fn model(&self) -> &str {
"reasoning-stream/mock"
}
fn max_context_length(&self) -> usize {
128_000
}
async fn complete(
&self,
_prompt: &str,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
Ok(edgequake_llm::LLMResponse::new(
"fallback complete",
self.model(),
))
}
async fn complete_with_options(
&self,
prompt: &str,
_options: &CompletionOptions,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
self.complete(prompt).await
}
async fn chat(
&self,
messages: &[ChatMessage],
options: Option<&CompletionOptions>,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
self.chat_with_tools(messages, &[], None, options).await
}
async fn chat_with_tools(
&self,
_messages: &[ChatMessage],
_tools: &[ToolDefinition],
_tool_choice: Option<ToolChoice>,
_options: Option<&CompletionOptions>,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
let mut response = edgequake_llm::LLMResponse::new("nonstreamed answer", self.model());
response.thinking_content = Some("hidden reasoning".to_string());
Ok(response)
}
async fn stream(
&self,
_prompt: &str,
) -> edgequake_llm::Result<futures::stream::BoxStream<'static, edgequake_llm::Result<String>>>
{
Ok(futures::stream::iter(vec![Ok("plain stream".to_string())]).boxed())
}
async fn chat_with_tools_stream(
&self,
_messages: &[ChatMessage],
_tools: &[ToolDefinition],
_tool_choice: Option<ToolChoice>,
_options: Option<&CompletionOptions>,
) -> edgequake_llm::Result<
futures::stream::BoxStream<'static, edgequake_llm::Result<StreamChunk>>,
> {
let chunks = vec![
Ok(StreamChunk::ThinkingContent {
text: "live reasoning".to_string(),
tokens_used: Some(3),
budget_total: None,
}),
Ok(StreamChunk::Content("streamed answer".to_string())),
Ok(StreamChunk::Finished {
reason: "stop".to_string(),
ttft_ms: None,
usage: None,
}),
];
Ok(futures::stream::iter(chunks).boxed())
}
fn supports_streaming(&self) -> bool {
true
}
fn supports_tool_streaming(&self) -> bool {
true
}
fn supports_function_calling(&self) -> bool {
true
}
}
#[async_trait]
impl LLMProvider for SlowChatProvider {
fn name(&self) -> &str {
"slow-chat"
}
fn model(&self) -> &str {
"slow-chat/mock"
}
fn max_context_length(&self) -> usize {
128_000
}
async fn complete(
&self,
prompt: &str,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
Ok(edgequake_llm::LLMResponse::new(prompt, self.model()))
}
async fn complete_with_options(
&self,
prompt: &str,
_options: &CompletionOptions,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
self.complete(prompt).await
}
async fn chat(
&self,
messages: &[ChatMessage],
options: Option<&CompletionOptions>,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
self.chat_with_tools(messages, &[], None, options).await
}
async fn chat_with_tools(
&self,
_messages: &[ChatMessage],
_tools: &[ToolDefinition],
_tool_choice: Option<ToolChoice>,
_options: Option<&CompletionOptions>,
) -> edgequake_llm::Result<edgequake_llm::LLMResponse> {
self.started.notify_waiters();
self.release.notified().await;
Ok(edgequake_llm::LLMResponse::new("slow answer", self.model()))
}
}
#[async_trait]
impl GatewaySender for MockGatewaySender {
async fn send_message(
&self,
_platform: &str,
_recipient: &str,
_message: &str,
) -> Result<(), String> {
Ok(())
}
async fn list_targets(&self) -> Result<Vec<String>, String> {
Ok(vec!["telegram".into()])
}
}
#[test]
fn iteration_budget_counts_down() {
let budget = IterationBudget::new(3);
assert!(budget.try_consume());
assert!(budget.try_consume());
assert!(budget.try_consume());
assert!(!budget.try_consume());
assert_eq!(budget.used(), 3);
}
#[test]
fn iteration_budget_reset() {
let budget = IterationBudget::new(2);
budget.try_consume();
budget.try_consume();
assert!(!budget.try_consume());
budget.reset();
assert!(budget.try_consume());
assert_eq!(budget.remaining(), 1);
}
#[test]
fn builder_requires_provider() {
let result = AgentBuilder::new("test/model").build();
assert!(result.is_err());
}
#[tokio::test]
async fn builder_with_mock_provider() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.max_iterations(10)
.build()
.expect("build agent");
let cfg = agent.config.read().await;
assert_eq!(cfg.model, "mock");
assert_eq!(cfg.max_iterations, 10);
assert_eq!(agent.budget.remaining(), 10);
}
#[tokio::test]
async fn chat_with_mock_provider() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
let result = agent.chat("hello").await;
assert!(result.is_ok());
let response = result.expect("response");
assert!(!response.is_empty());
}
#[tokio::test]
async fn try_session_snapshot_remains_readable_during_slow_turn() {
let started = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let provider: Arc<dyn LLMProvider> = Arc::new(SlowChatProvider {
started: started.clone(),
release: release.clone(),
});
let agent = Arc::new(
AgentBuilder::new("slow-chat/mock")
.provider(provider)
.build()
.expect("build agent"),
);
let agent_task = agent.clone();
let chat_task = tokio::spawn(async move { agent_task.chat("hello").await });
tokio::time::timeout(std::time::Duration::from_secs(2), started.notified())
.await
.expect("provider started");
let snapshot = agent
.try_session_snapshot()
.expect("session snapshot should stay readable mid-turn");
assert_eq!(snapshot.message_count, 1);
assert_eq!(snapshot.user_turn_count, 1);
release.notify_waiters();
let result = chat_task.await.expect("join").expect("chat");
assert_eq!(result, "slow answer");
}
#[test]
fn from_config_wires_agent_flags() {
let config = AppConfig {
save_trajectories: true,
skip_context_files: true,
skip_memory: true,
..Default::default()
};
let builder = AgentBuilder::from_config(&config);
assert!(builder.config.save_trajectories);
assert!(builder.config.skip_context_files);
assert!(builder.config.skip_memory);
}
#[test]
fn resolve_tool_policy_expands_aliases_once() {
let config = AgentConfig {
enabled_toolsets: vec!["core".into()],
disabled_toolsets: vec!["browser".into()],
..Default::default()
};
let policy = resolve_tool_policy(&config);
assert!(
policy
.expanded_enabled
.iter()
.any(|toolset| toolset == "file")
);
assert!(
policy
.expanded_disabled
.iter()
.any(|toolset| toolset == "browser")
);
assert!(
!policy
.parent_active_toolsets
.iter()
.any(|toolset| toolset == "browser")
);
}
#[test]
fn app_config_ref_projection_keeps_runtime_policy_consistent() {
let config = AgentConfig {
platform: Platform::Telegram,
enabled_toolsets: vec!["core".into()],
disabled_toolsets: vec!["browser".into()],
..Default::default()
};
let policy = resolve_tool_policy(&config);
let projected = config.to_app_config_ref(true, &policy);
assert!(projected.gateway_running);
assert_eq!(
projected.parent_active_toolsets,
policy.parent_active_toolsets
);
assert_eq!(projected.disabled_toolsets, policy.expanded_disabled);
}
#[tokio::test]
async fn new_session_resets_state() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.max_iterations(5)
.build()
.expect("build agent");
agent.budget.try_consume();
agent.budget.try_consume();
assert_eq!(agent.budget.remaining(), 3);
let _ = agent.chat("hi").await;
agent.new_session().await;
assert_eq!(agent.budget.remaining(), 5);
let session = agent.session.read().await;
assert!(session.messages.is_empty());
}
#[tokio::test]
async fn interrupt_triggers_cancellation() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
assert!(!agent.is_cancelled());
agent.interrupt();
assert!(agent.is_cancelled());
}
#[tokio::test]
async fn cancel_resets_between_turns() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
agent.interrupt();
assert!(agent.is_cancelled());
let result = agent.chat("hello after cancel").await;
assert!(result.is_ok(), "expected success after cancel: {result:?}");
assert!(!agent.is_cancelled());
}
#[tokio::test]
async fn conversation_result_structure() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.session_id("test-session".into())
.build()
.expect("build agent");
let result = agent
.run_conversation("hello", Some("You are helpful."), None)
.await
.expect("conversation");
assert_eq!(result.session_id, "test-session");
assert_eq!(result.api_calls, 1);
assert!(!result.interrupted);
assert_eq!(result.model, "mock");
assert_eq!(result.messages.len(), 2);
}
#[tokio::test]
async fn session_state_gets_session_id_after_chat() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
{
let session = agent.session.read().await;
assert!(session.session_id.is_none());
}
let _ = agent.chat("hello").await;
{
let session = agent.session.read().await;
assert!(session.session_id.is_some());
}
}
#[tokio::test]
async fn fork_isolated_creates_fresh_session_state() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let parent = AgentBuilder::new("mock")
.provider(provider)
.max_iterations(7)
.build()
.expect("build agent");
let _ = parent.chat("parent turn").await.expect("parent chat");
let parent_before = parent.session_snapshot().await;
let child = parent
.fork_isolated(IsolatedAgentOptions {
session_id: Some("bg-test".into()),
quiet_mode: Some(true),
..Default::default()
})
.await
.expect("fork isolated");
let child_before = child.session_snapshot().await;
let child_cfg = child.config.read().await;
assert_eq!(child_before.message_count, 0);
assert_eq!(child_before.api_call_count, 0);
assert_eq!(child_before.model, parent_before.model);
assert_eq!(child_cfg.session_id.as_deref(), Some("bg-test"));
assert_eq!(child.budget.remaining(), 7);
drop(child_cfg);
let _ = child.chat("background turn").await.expect("child chat");
let parent_after = parent.session_snapshot().await;
let child_after = child.session_snapshot().await;
assert_eq!(parent_after.message_count, parent_before.message_count);
assert!(child_after.message_count > 0);
assert_eq!(child_after.session_id.as_deref(), Some("bg-test"));
assert_ne!(parent_after.session_id, child_after.session_id);
}
#[tokio::test]
async fn fork_isolated_preserves_gateway_sender() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let parent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
parent.set_gateway_sender(Arc::new(MockGatewaySender)).await;
let child = parent
.fork_isolated(IsolatedAgentOptions::default())
.await
.expect("fork isolated");
assert!(child.gateway_sender.read().await.is_some());
}
#[tokio::test]
async fn model_config_propagates_to_agent() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
let cfg = agent.config.read().await;
assert!(!cfg.model_config.smart_routing.enabled);
assert!(cfg.model_config.smart_routing.cheap_model.is_empty());
}
#[tokio::test]
async fn chat_streaming_emits_reasoning_when_enabled() {
let provider: Arc<dyn LLMProvider> = Arc::new(ReasoningStreamProvider);
let agent = AgentBuilder::new("reasoning-stream/mock")
.provider(provider)
.streaming(true)
.build()
.expect("build agent");
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
agent
.chat_streaming("hello", tx)
.await
.expect("streaming chat");
let mut saw_reasoning = false;
let mut saw_token = false;
let mut saw_done = false;
while let Some(event) = rx.recv().await {
match event {
StreamEvent::Reasoning(text) => saw_reasoning |= !text.is_empty(),
StreamEvent::Token(text) => saw_token |= !text.is_empty(),
StreamEvent::Done => {
saw_done = true;
break;
}
_ => {}
}
}
assert!(
saw_reasoning,
"expected a live reasoning event when streaming is enabled"
);
assert!(saw_token, "expected streamed answer tokens");
assert!(saw_done, "expected the stream to terminate cleanly");
}
#[tokio::test]
async fn chat_streaming_sends_single_final_answer_when_streaming_disabled() {
let provider: Arc<dyn LLMProvider> = Arc::new(ReasoningStreamProvider);
let agent = AgentBuilder::new("reasoning-stream/mock")
.provider(provider)
.streaming(false)
.build()
.expect("build agent");
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
agent
.chat_streaming("hello", tx)
.await
.expect("streaming chat");
let mut saw_reasoning = false;
let mut token_events = 0usize;
let mut collected_tokens = String::new();
while let Some(event) = rx.recv().await {
match event {
StreamEvent::Reasoning(text) => saw_reasoning |= !text.is_empty(),
StreamEvent::Token(text) => {
token_events += 1;
collected_tokens.push_str(&text);
}
StreamEvent::Done => break,
_ => {}
}
}
assert!(
saw_reasoning,
"final reasoning content should still be available for think mode"
);
assert_eq!(
token_events, 1,
"streaming-off mode should emit one complete answer instead of pseudo-streaming chunks"
);
assert_eq!(collected_tokens, "nonstreamed answer");
}
#[tokio::test]
async fn session_personality_overlay_replaces_previous_overlay() {
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.build()
.expect("build agent");
{
let mut session = agent.session.write().await;
session.cached_system_prompt = Some("Base prompt".to_string());
}
agent
.set_personality_addon(Some("First overlay".to_string()))
.await;
assert!(agent.system_prompt().await.is_none());
{
let mut session = agent.session.write().await;
session.cached_system_prompt = Some("Base prompt".to_string());
}
agent
.set_personality_addon(Some("Second overlay".to_string()))
.await;
assert!(agent.system_prompt().await.is_none());
agent.set_personality_addon(None).await;
assert!(agent.system_prompt().await.is_none());
}
#[test]
fn session_snapshot_prompt_tokens_include_cache_buckets() {
let snap = SessionSnapshot {
session_id: None,
model: "mock/model".into(),
message_count: 0,
user_turn_count: 0,
api_call_count: 0,
input_tokens: 3,
output_tokens: 10,
cache_read_tokens: 15_000,
cache_write_tokens: 2_000,
reasoning_tokens: 7,
last_prompt_tokens: 1_234,
budget_remaining: 0,
budget_max: 0,
};
assert_eq!(snap.prompt_tokens(), 17_003);
assert_eq!(snap.total_tokens(), 17_020);
assert_eq!(snap.context_pressure_tokens(), 1_234);
}
#[tokio::test]
async fn restore_session_reuses_persisted_system_prompt() {
let tmp = tempfile::tempdir().expect("tempdir");
let db = Arc::new(SessionDb::open(&tmp.path().join("sessions.db")).expect("db"));
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let session = edgecrab_state::SessionRecord {
id: "restore-me".into(),
source: "cli".into(),
user_id: None,
model: Some("mock/model".into()),
system_prompt: Some("Persisted system prompt".into()),
parent_session_id: None,
started_at: 1.0,
ended_at: Some(2.0),
end_reason: None,
message_count: 2,
tool_call_count: 0,
input_tokens: 10,
output_tokens: 4,
cache_read_tokens: 0,
cache_write_tokens: 0,
reasoning_tokens: 0,
estimated_cost_usd: None,
title: Some("restore".into()),
};
db.save_session(&session).expect("save session");
db.save_message("restore-me", &Message::user("hello"), 1.0)
.expect("save user");
db.save_message("restore-me", &Message::assistant("hi"), 2.0)
.expect("save assistant");
let agent = AgentBuilder::new("mock/model")
.provider(provider)
.state_db(db)
.build()
.expect("build agent");
let restored = agent.restore_session("restore-me").await.expect("restore");
assert_eq!(restored, 2);
assert_eq!(
agent.system_prompt().await.as_deref(),
Some("Persisted system prompt")
);
}
}