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 lingshu_plugins::{discover_plugins, hermes_supports_hook, invoke_hermes_hook};
use lingshu_state::SessionDb;
use lingshu_tools::ProcessTable;
use lingshu_tools::TodoStore;
use lingshu_tools::config_ref::{AppConfigRef, LspServerConfigRef};
use lingshu_tools::registry::{GatewaySender, ToolInventoryEntry, ToolRegistry};
use lingshu_types::{
AgentError, ApiMode, Cost, Message, OriginChat, Platform, Role, RunOutcome, 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: RwLock<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<lingshu_tools::TodoStore>,
pub(crate) goal_store: Arc<dyn crate::goals::GoalStore>,
pub(crate) context_engine: Option<Arc<dyn crate::context_engine::ContextEngine>>,
pub(crate) steer_tx: crate::steering::SteeringSender,
pub(crate) steer_rx: std::sync::Mutex<Option<crate::steering::SteeringReceiver>>,
pub(crate) steer_event_tx:
std::sync::Mutex<Option<tokio::sync::mpsc::UnboundedSender<crate::StreamEvent>>>,
pub(crate) steer_pending: std::sync::atomic::AtomicUsize,
}
#[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<OriginChat>,
pub kanban_task_id: Option<String>,
pub profile: Option<String>,
pub install_root: Option<std::path::PathBuf>,
}
#[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 custom_system_prompt: Option<String>,
pub model_config: crate::config::ModelConfig,
pub skills_config: crate::config::SkillsConfig,
pub memory_config: crate::config::MemoryConfig,
pub approvals_config: crate::config::ApprovalsConfig,
pub kanban_config: crate::config::KanbanConfig,
pub kanban_task_id: Option<String>,
pub plugins_config: crate::config::PluginsConfig,
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<OriginChat>,
pub browser: crate::config::BrowserConfig,
pub checkpoints_enabled: bool,
pub checkpoints_max_snapshots: u32,
pub checkpoints_max_total_size_mb: u32,
pub checkpoints_max_file_size_mb: u32,
pub computer_use_enabled: bool,
pub computer_use_keep_last_n_screenshots: u32,
pub computer_use_confirm_destructive: bool,
pub computer_use_cua_cmd: String,
pub terminal_backend: lingshu_tools::tools::backends::BackendKind,
pub terminal_docker: lingshu_tools::tools::backends::DockerBackendConfig,
pub terminal_ssh: lingshu_tools::tools::backends::SshBackendConfig,
pub terminal_modal: lingshu_tools::tools::backends::ModalBackendConfig,
pub terminal_daytona: lingshu_tools::tools::backends::DaytonaBackendConfig,
pub terminal_singularity: lingshu_tools::tools::backends::SingularityBackendConfig,
pub compression: crate::config::CompressionConfig,
pub auxiliary: crate::config::AuxiliaryConfig,
pub shadow_judge: crate::config::ShadowJudgeConfig,
pub goals: crate::config::GoalsConfig,
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,
pub result_spill: bool,
pub result_spill_threshold: usize,
pub result_spill_preview_lines: usize,
pub result_turn_budget_chars: usize,
pub max_write_payload_kib: Option<u32>,
pub local_write_create_dirs: bool,
pub local_max_tool_turn_tokens: usize,
pub file_mutation_verifier: bool,
pub cache: crate::config::CacheConfig,
pub web_search: crate::config::WebSearchConfig,
pub web: crate::config::WebToolsConfig,
}
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,
custom_system_prompt: None,
model_config: crate::config::ModelConfig::default(),
skills_config: crate::config::SkillsConfig::default(),
memory_config: crate::config::MemoryConfig::default(),
approvals_config: crate::config::ApprovalsConfig::default(),
kanban_config: crate::config::KanbanConfig::default(),
kanban_task_id: None,
plugins_config: crate::config::PluginsConfig::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: 20,
checkpoints_max_total_size_mb: 200,
checkpoints_max_file_size_mb: 10,
computer_use_enabled: false,
computer_use_keep_last_n_screenshots: 1,
computer_use_confirm_destructive: true,
computer_use_cua_cmd: "cua-driver".into(),
terminal_backend: lingshu_tools::tools::backends::BackendKind::Local,
terminal_docker: lingshu_tools::tools::backends::DockerBackendConfig::default(),
terminal_ssh: lingshu_tools::tools::backends::SshBackendConfig::default(),
terminal_modal: lingshu_tools::tools::backends::ModalBackendConfig::default(),
terminal_daytona: lingshu_tools::tools::backends::DaytonaBackendConfig::default(),
terminal_singularity:
lingshu_tools::tools::backends::SingularityBackendConfig::default(),
compression: crate::config::CompressionConfig::default(),
auxiliary: crate::config::AuxiliaryConfig::default(),
shadow_judge: crate::config::ShadowJudgeConfig::default(),
goals: crate::config::GoalsConfig::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(),
result_spill: true,
result_spill_threshold: 16_384,
result_spill_preview_lines: 80,
result_turn_budget_chars: 200_000,
max_write_payload_kib: None,
local_write_create_dirs: true,
local_max_tool_turn_tokens:
lingshu_tools::mutation_turn_policy::LOCAL_TOOL_TURN_ABS_MAX_TOKENS,
file_mutation_verifier: true,
cache: crate::config::CacheConfig::default(),
web_search: crate::config::WebSearchConfig::default(),
web: crate::config::WebToolsConfig::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 = lingshu_tools::toolsets::expand_toolset_names(&config.enabled_toolsets);
let expanded_disabled =
lingshu_tools::toolsets::expand_toolset_names(&config.disabled_toolsets);
let parent_active_toolsets = if config.enabled_toolsets.is_empty()
|| lingshu_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
}
fn disabled_plugins_for_platform(&self) -> Vec<String> {
let mut disabled = self.plugins_config.disabled.clone();
let platform_key = self.platform.to_string();
if let Some(platform_disabled) = self.plugins_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 {
lingshu_home: crate::config::lingshu_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_post_write_timeout_ms: self.lsp.timeout_ms,
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(),
disabled_plugins: self.disabled_plugins_for_platform(),
plugin_install_dir: self.plugins_config.install_dir.clone(),
preloaded_skills: self.skills_config.preloaded.clone(),
skills_write_approval: self.skills_config.write_approval,
memory_write_approval: self.memory_config.write_approval,
approval_mode: self.approvals_config.mode,
approvals_smart_model: self.approvals_config.smart_model.clone(),
kanban_enabled: self.kanban_config.enabled,
kanban_claim_ttl_secs: self.kanban_config.claim_ttl_secs,
kanban_default_max_runtime_secs: self.kanban_config.default_max_runtime_secs,
skills_hub_url: self.skills_config.hub_url.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,
checkpoints_max_total_size_mb: self.checkpoints_max_total_size_mb,
checkpoints_max_file_size_mb: self.checkpoints_max_file_size_mb,
computer_use_enabled: self.computer_use_enabled,
computer_use_keep_last_n_screenshots: self.computer_use_keep_last_n_screenshots,
computer_use_confirm_destructive: self.computer_use_confirm_destructive,
computer_use_cua_cmd: self.computer_use_cua_cmd.clone(),
active_model: self.model.clone(),
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()),
result_spill: self.result_spill,
result_spill_threshold: self.result_spill_threshold,
result_spill_preview_lines: self.result_spill_preview_lines,
result_turn_budget_chars: self.result_turn_budget_chars,
max_write_payload_kib: self.max_write_payload_kib.map_or(
lingshu_tools::edit_contract::DEFAULT_MAX_MUTATION_PAYLOAD_KIB,
|kib| kib as usize,
),
local_write_create_dirs: crate::local_provider_policy::effective_local_write_create_dirs(
self.local_write_create_dirs,
&self.model,
),
local_max_tool_turn_tokens: self.local_max_tool_turn_tokens,
web_search: lingshu_tools::config_ref::WebSearchConfigRef {
primary: self.web_search.primary.clone(),
fallbacks: self.web_search.fallbacks.clone(),
timeout_secs: self.web_search.timeout_secs,
backends: self
.web_search
.backends
.iter()
.map(|(k, v)| {
(
k.clone(),
lingshu_tools::config_ref::WebSearchBackendConfigRef {
api_key: v.api_key.clone(),
endpoint: v.endpoint.clone(),
rps: v.rps,
timeout_secs: v.timeout_secs,
},
)
})
.collect(),
},
web: lingshu_tools::config_ref::WebToolsConfigRef {
search_backend: self.web.search_backend.clone(),
extract_backend: self.web.extract_backend.clone(),
backend: self.web.backend.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 cached_stable_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 tool_result_image_downgrades: std::collections::HashSet<(String, String)>,
pub last_run_outcome: Option<RunOutcome>,
pub session_tool_call_count: u32,
pub invalid_tool_call_retries: u32,
pub first_compression_done: bool,
}
impl SessionState {
pub(crate) fn apply_model_transfer_outcome(
&mut self,
outcome: &crate::model_transfer::ModelTransferOutcome,
) {
self.cached_system_prompt = None;
self.cached_stable_prompt = None;
if outcome.compressed {
self.first_compression_done = true;
}
self.messages.push(Message::user(
&crate::model_transfer::format_model_transfer_user_message(
&outcome.from_model,
&outcome.to_model,
&outcome.brief,
outcome.compressed,
),
));
}
}
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 run_outcome: RunOutcome,
pub model: String,
pub usage: Usage,
pub cost: Cost,
pub tool_errors: Vec<lingshu_types::ToolErrorRecord>,
}
impl Agent {
fn build_runtime_clone(
config: AgentConfig,
provider: Arc<dyn LLMProvider>,
state_db: Option<Arc<SessionDb>>,
tool_registry: Option<Arc<ToolRegistry>>,
context_engine: Option<Arc<dyn crate::context_engine::ContextEngine>>,
goal_store: Option<Arc<dyn crate::goals::GoalStore>>,
) -> 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());
let (steer_tx, steer_rx) = crate::steering::steering_channel();
let goal_store =
goal_store.unwrap_or_else(|| crate::goals::goal_store_for_db(state_db.clone()));
Self {
config: RwLock::new(config),
provider: RwLock::new(provider),
state_db,
tool_registry: RwLock::new(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()),
goal_store,
context_engine,
steer_tx,
steer_rx: std::sync::Mutex::new(Some(steer_rx)),
steer_event_tx: std::sync::Mutex::new(None),
steer_pending: std::sync::atomic::AtomicUsize::new(0),
}
}
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 set_shadow_judge_enabled(&self, enabled: bool) {
self.config.write().await.shadow_judge.enabled = enabled;
}
pub async fn shadow_judge_enabled(&self) -> bool {
self.config.read().await.shadow_judge.enabled
}
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(OriginChat::new(platform, chat_id));
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(OriginChat::new(platform, chat_id));
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> {
tracing::info!(message_len = message.len(), "chat_streaming: entered");
let streaming_enabled = {
let config = self.config.read().await;
tracing::info!(streaming = config.streaming, model = %config.model, "chat_streaming: config read");
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, 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,
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),
None,
)
.await
}
pub async fn fork_isolated(&self, options: IsolatedAgentOptions) -> Result<Self, AgentError> {
let profile = options
.profile
.as_deref()
.map(crate::kanban_profiles::normalize_profile_name)
.filter(|s| !s.is_empty());
let (mut config, provider, state_db) = if let Some(profile_name) = profile {
let root = options
.install_root
.clone()
.unwrap_or_else(crate::kanban_profiles::install_root);
let profile_cfg =
crate::kanban_profiles::load_config_for_profile(&root, &profile_name)?;
let model = profile_cfg.model.default_model.clone();
let provider = if let Some((provider_name, model_name)) = model.split_once('/') {
let canonical = lingshu_tools::vision_models::normalize_provider_name(provider_name);
lingshu_tools::create_provider_for_model(&canonical, model_name)
.map_err(|e| AgentError::Config(format!("kanban worker provider: {e}")))?
} else {
return Err(AgentError::Config(format!(
"kanban worker profile '{profile_name}' model must be provider/model form, got '{model}'"
)));
};
let _profile_home =
crate::kanban_profiles::profile_effective_home(&root, &profile_name);
let db = Arc::new(
lingshu_state::SessionDb::open_default()
.map_err(|e| AgentError::Database(format!("profile state db: {e}")))?,
);
let agent_cfg = AgentBuilder::from_config(&profile_cfg).into_agent_config();
(agent_cfg, provider, Some(db))
} else {
let config = self.config.read().await.clone();
let provider = self.provider.read().await.clone();
let state_db = self.state_db.clone();
(config, provider, state_db)
};
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;
config.kanban_task_id = options.kanban_task_id;
if config.kanban_task_id.is_some()
&& !config
.enabled_toolsets
.iter()
.any(|t| t.eq_ignore_ascii_case("kanban"))
{
config.enabled_toolsets.push("kanban".into());
}
let child = Self::build_runtime_clone(
config,
provider,
state_db,
self.tool_registry.read().await.clone(),
self.context_engine.clone(),
None,
);
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 process_table(&self) -> Arc<ProcessTable> {
self.process_table.clone()
}
pub fn is_cancelled(&self) -> bool {
self.cancel
.lock()
.expect("cancel mutex not poisoned")
.is_cancelled()
}
pub fn steer_sender(&self) -> crate::steering::SteeringSender {
self.steer_tx.clone()
}
pub fn send_steering(
&self,
event: crate::steering::SteeringEvent,
) -> Result<(), tokio::sync::mpsc::error::SendError<crate::steering::SteeringEvent>> {
use std::sync::atomic::Ordering;
let count = self.steer_pending.fetch_add(1, Ordering::Relaxed) + 1;
if let Some(tx) = self
.steer_event_tx
.lock()
.expect("steer_event_tx mutex not poisoned")
.clone()
{
let _ = tx.send(crate::StreamEvent::SteerPending { count });
}
self.steer_tx.send(event)
}
pub async fn new_session(&self) {
let previous_session_id = self.session.read().await.session_id.clone();
let config = self.config.read().await.clone();
if let Some(session_id) = previous_session_id.as_deref() {
self.fire_session_boundary_hooks(
&config,
"on_session_finalize",
session_id,
config.platform,
)
.await;
}
let passthrough = { config.terminal_env_passthrough.clone() };
lingshu_tools::tools::backends::local::clear_env_passthrough();
if !passthrough.is_empty() {
lingshu_tools::tools::backends::local::register_env_passthrough(&passthrough);
}
let mut session = self.session.write().await;
*session = SessionState::default();
let new_session_id = uuid::Uuid::new_v4().to_string();
session.session_id = Some(new_session_id.clone());
self.budget.reset();
drop(session);
self.fire_session_boundary_hooks(
&config,
"on_session_reset",
&new_session_id,
config.platform,
)
.await;
}
pub async fn finalize_session(&self) {
let session_id = self.session.read().await.session_id.clone();
let config = self.config.read().await.clone();
if let Some(session_id) = session_id.as_deref() {
self.fire_session_boundary_hooks(
&config,
"on_session_finalize",
session_id,
config.platform,
)
.await;
}
}
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 switch_model_fast(
&self,
target_spec: &str,
) -> Result<crate::model_transfer::ModelSwitchOutcome, crate::model_transfer::ModelTransferError>
{
use crate::model_catalog::ModelCatalog;
use crate::model_transfer::{
ModelSwitchOutcome, create_model_transfer_provider, resolve_model_transfer_target,
};
let target = resolve_model_transfer_target(target_spec)?;
let current_model = self.config.read().await.model.clone();
if ModelCatalog::equivalent_model_specs(&target.display, ¤t_model) {
return Err(crate::model_transfer::ModelTransferError::SameModel);
}
let provider = create_model_transfer_provider(&target)?;
self.swap_model(target.display.clone(), provider).await;
if let Some(db) = &self.state_db
&& let Some(session_id) = self.session.read().await.session_id.as_deref()
&& let Err(err) = db.update_session_model(session_id, &target.display).await
{
tracing::warn!(error = %err, "fast model switch session DB update failed");
}
Ok(ModelSwitchOutcome {
from_model: current_model,
to_model: target.display,
})
}
pub async fn perform_model_transfer(
&self,
target_model: &str,
events: Option<tokio::sync::mpsc::UnboundedSender<StreamEvent>>,
) -> Result<
crate::model_transfer::ModelTransferOutcome,
crate::model_transfer::ModelTransferError,
> {
let (compression_cfg, auxiliary_model, current_model, main_provider) = {
let cfg = self.config.read().await;
(
cfg.compression.clone(),
cfg.auxiliary.model.clone(),
cfg.model.clone(),
self.provider.read().await.clone(),
)
};
let (outcome, new_provider) = {
let mut session = self.session.write().await;
let system_prompt = session.cached_system_prompt.clone();
let mut ctx = crate::model_transfer::ModelTransferContext {
current_model: ¤t_model,
messages: &mut session.messages,
system_prompt: system_prompt.as_deref(),
compression_cfg: &compression_cfg,
main_provider,
auxiliary_model: auxiliary_model.as_deref(),
};
let (outcome, new_provider) =
crate::model_transfer::ModelTransferOrchestrator::execute(&mut ctx, target_model)
.await?;
session.apply_model_transfer_outcome(&outcome);
(outcome, new_provider)
};
self.swap_model(outcome.to_model.clone(), new_provider)
.await;
if outcome.compressed
&& let Some(session_id) = self.session.read().await.session_id.as_deref()
{
lingshu_tools::read_tracker::reset_read_dedup(session_id);
}
if let Some(db) = &self.state_db
&& let Some(session_id) = self.session.read().await.session_id.as_deref()
{
if let Err(err) = db.update_session_model(session_id, &outcome.to_model).await {
tracing::warn!(error = %err, "handoff session model DB update failed");
}
if let Err(err) = db.record_model_transfer(
session_id,
&outcome.from_model,
&outcome.to_model,
&outcome.brief,
).await {
tracing::warn!(error = %err, "handoff DB record failed");
}
}
if let Some(tx) = events {
let _ = tx.send(StreamEvent::ModelTransferComplete {
from: outcome.from_model.clone(),
to: outcome.to_model.clone(),
brief: outcome.brief.clone(),
compressed: outcome.compressed,
});
}
Ok(outcome)
}
pub async fn model(&self) -> String {
self.config.read().await.model.clone()
}
async fn fire_session_boundary_hooks(
&self,
config: &AgentConfig,
hook_name: &str,
session_id: &str,
platform: Platform,
) {
let Ok(discovery) = discover_plugins(&config.plugins_config, platform) else {
return;
};
for plugin in discovery
.plugins
.iter()
.filter(|plugin| hermes_supports_hook(plugin, hook_name))
{
if let Err(error) = invoke_hermes_hook(
plugin,
hook_name,
serde_json::json!({
"session_id": session_id,
"platform": platform.to_string(),
}),
)
.await
{
tracing::warn!(plugin = %plugin.name, ?error, hook = hook_name, "Hermes session-boundary hook failed");
}
}
}
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 async fn last_run_outcome(&self) -> Option<RunOutcome> {
self.session.read().await.last_run_outcome.clone()
}
pub async fn custom_system_prompt(&self) -> Option<String> {
self.config.read().await.custom_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;
session.cached_stable_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 set_custom_system_prompt(&self, prompt: Option<String>) {
{
let mut config = self.config.write().await;
config.custom_system_prompt = prompt.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then_some(trimmed.to_string())
});
}
self.invalidate_system_prompt().await;
}
pub async fn preloaded_skills(&self) -> Vec<String> {
self.config.read().await.skills_config.preloaded.clone()
}
pub fn preloaded_skills_snapshot(&self) -> Vec<String> {
self.config
.try_read()
.map(|config| config.skills_config.preloaded.clone())
.unwrap_or_default()
}
pub async fn set_preloaded_skills(&self, skills: Vec<String>) {
{
let mut config = self.config.write().await;
config.skills_config.preloaded = skills;
}
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,
);
let (compressed, _llm_succeeded) =
crate::compression::compress_with_llm(&session.messages, ¶ms, &provider, None)
.await;
session.messages = compressed;
}
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).await;
}
}
async fn ensure_session_id(&self) -> Result<String, AgentError> {
let session_id = {
let mut session = self.session.write().await;
if session.session_id.is_none() {
session.session_id = Some(uuid::Uuid::new_v4().to_string());
}
session
.session_id
.clone()
.expect("session_id initialized above")
};
if let Some(db) = &self.state_db {
let config = self.config.read().await;
let (source, user_id) = match &config.origin_chat {
Some(origin) => (origin.platform.clone(), Some(origin.chat_id.clone())),
None => (config.platform.to_string(), None),
};
db.ensure_session_row(
&session_id,
&source,
user_id.as_deref(),
Some(config.model.as_str()),
).await?;
}
Ok(session_id)
}
pub async fn ensure_persisted_session(&self) -> Result<String, AgentError> {
self.ensure_session_id().await
}
pub async fn goal_set(&self, text: &str) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let max_turns = self.config.read().await.goals.max_turns.max(1);
self.goal_store.set_goal(&session_id, text, max_turns)?;
Ok(format!(
"⊙ Goal set ({max_turns}-turn budget): {}\n\
After each turn, a judge model checks if the goal is done. \
Lingshu keeps working until it is, you pause/clear it, or the budget is exhausted.",
text.trim()
))
}
pub async fn goal_status(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if state.is_empty() {
return Ok(crate::goals::status_line(&state));
}
if state.subgoals.is_empty() {
return Ok(crate::goals::status_line(&state));
}
Ok(format!(
"{}\n{}",
crate::goals::status_line(&state),
crate::goals::render_subgoals_list(&state)
))
}
pub async fn goal_state(&self) -> Result<crate::goals::GoalState, AgentError> {
let session_id = self.ensure_session_id().await?;
self.goal_store.active(&session_id)
}
pub async fn goal_resume_with_kickoff(&self) -> Result<(String, Option<String>), AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if state.is_empty() {
return Ok(("No goal to resume.".into(), None));
}
self.goal_store.resume(&session_id, true)?;
let refreshed = self.goal_store.active(&session_id)?;
let continuation = crate::goals::next_continuation_prompt(&refreshed);
let msg = format!(
"▶ Goal resumed: {}",
refreshed.goal_text.as_deref().unwrap_or("")
);
Ok((msg, continuation))
}
pub async fn goal_show(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if state.is_empty() {
return Ok("No active goal. Use /goal <text> to set one.".into());
}
Ok(crate::goals::render_goal_block(&state))
}
pub async fn goal_clear(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let had = !self.goal_store.active(&session_id)?.is_empty();
self.goal_store.clear(&session_id)?;
Ok(if had {
"✓ Goal cleared.".into()
} else {
"No active goal.".into()
})
}
pub async fn goal_pause(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if state.is_empty() {
return Ok("No goal set.".into());
}
self.goal_store.pause(&session_id, "user-paused")?;
Ok(format!(
"⏸ Goal paused: {}",
state.goal_text.as_deref().unwrap_or("")
))
}
pub async fn goal_resume(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if state.is_empty() {
return Ok("No goal to resume.".into());
}
self.goal_store.resume(&session_id, true)?;
Ok(format!(
"▶ Goal resumed: {}",
state.goal_text.as_deref().unwrap_or("")
))
}
pub async fn subgoal_push(&self, text: &str) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
self.goal_store.push_subgoal(&session_id, text)?;
Ok(format!("📌 Subgoal added: {}", text.trim()))
}
pub async fn subgoal_list(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
if !state.has_goal() {
return Ok("No active goal. Set one with /goal <text>.".into());
}
Ok(format!(
"{}\n{}",
crate::goals::status_line(&state),
crate::goals::render_subgoals_list(&state)
))
}
pub async fn subgoal_remove(&self, index_1based: usize) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let removed = self.goal_store.remove_subgoal(&session_id, index_1based)?;
Ok(format!("Removed subgoal #{index_1based}: {removed}"))
}
pub async fn subgoal_clear(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
let prev = self.goal_store.clear_subgoals(&session_id)?;
Ok(format!("Cleared {prev} subgoal(s)."))
}
pub async fn subgoal_done(&self) -> Result<String, AgentError> {
let session_id = self.ensure_session_id().await?;
match self.goal_store.complete_subgoal(&session_id)? {
Some(sub) => Ok(format!("✅ Subgoal done: {}", sub.text)),
None => Ok("No incomplete subgoals to mark done.".into()),
}
}
pub async fn goal_is_active(&self) -> Result<bool, AgentError> {
let session_id = self.ensure_session_id().await?;
let state = self.goal_store.active(&session_id)?;
Ok(state.status == crate::goals::GoalStatus::Active && state.has_goal())
}
pub async fn goal_evaluate_after_turn(
&self,
last_response: &str,
interrupted: bool,
) -> Result<crate::goals::GoalContinuationDecision, AgentError> {
let session_id = self.ensure_session_id().await?;
let cfg = self.config.read().await.clone();
let provider = self.provider.read().await.clone();
let model = cfg.model.clone();
crate::goals::evaluate_goal_after_turn(
self.goal_store.clone(),
&session_id,
last_response,
interrupted,
&cfg.goals,
&cfg.auxiliary.goal_judge,
cfg.auxiliary.model.as_deref(),
provider,
&model,
)
.await
}
pub fn goal_store(&self) -> Arc<dyn crate::goals::GoalStore> {
self.goal_store.clone()
}
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).await?
.ok_or_else(|| AgentError::Config(format!("Session not found: {id}")))?;
let messages = db.get_messages(id).await?;
{
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, lingshu_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 async fn list_sessions(
&self,
limit: usize,
) -> Result<Vec<lingshu_state::SessionSummary>, AgentError> {
match &self.state_db {
Some(db) => db.list_sessions(limit).await,
None => Ok(Vec::new()),
}
}
pub async fn delete_session(&self, id: &str) -> Result<(), AgentError> {
match &self.state_db {
Some(db) => db.delete_session(id).await,
None => Err(AgentError::Config("No state database configured".into())),
}
}
pub async fn rename_session(&self, id: &str, title: &str) -> Result<(), AgentError> {
match &self.state_db {
Some(db) => db.update_session_title(id, title).await,
None => Err(AgentError::Config("No state database configured".into())),
}
}
pub async 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).await,
None => Ok(0),
}
}
pub fn has_state_db(&self) -> bool {
self.state_db.is_some()
}
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.read().await.as_ref() {
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.read().await.as_ref() {
Some(reg) => reg.toolset_summary(),
None => Vec::new(),
}
}
pub async fn tool_inventory(&self) -> Vec<ToolInventoryEntry> {
let registry = self.tool_registry.read().await.clone();
let Some(registry) = 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 = lingshu_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: Some(registry.clone()),
delegate_depth: 0,
delegate_agent_id: None,
delegate_parent_id: None,
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,
injected_messages: None,
tool_progress_tx: None,
watch_notification_tx: None,
mutation_turn: None,
lsp_gate: None,
kanban_task_id: config.kanban_task_id.clone(),
};
registry.tool_inventory(&ctx)
}
pub async fn set_tool_registry(&self, registry: Arc<ToolRegistry>) {
*self.tool_registry.write().await = Some(registry);
}
pub async fn set_plugins_config(&self, plugins_config: lingshu_plugins::PluginsConfig) {
let mut config = self.config.write().await;
config.plugins_config = plugins_config;
}
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_approvals_config(&self, approvals: crate::config::ApprovalsConfig) {
let mut config = self.config.write().await;
config.approvals_config = approvals;
}
pub async fn set_kanban_config(&self, kanban: crate::config::KanbanConfig) {
let mut config = self.config.write().await;
config.kanban_config = kanban;
}
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 reload_web_search_from_disk(&self) {
let path = crate::config::lingshu_home().join("config.yaml");
let loaded = if path.is_file() {
crate::config::AppConfig::load_from(&path)
.map(|cfg| cfg.web_search)
.unwrap_or_default()
} else {
crate::config::WebSearchConfig::default()
};
let mut config = self.config.write().await;
config.web_search = loaded;
}
pub async fn lsp_config(&self) -> crate::config::LspConfig {
self.config.read().await.lsp.clone()
}
pub async fn set_lsp_enabled(&self, enabled: bool) {
let mut config = self.config.write().await;
config.lsp.enabled = enabled;
}
pub async fn set_computer_use_enabled(&self, enabled: bool) {
let mut config = self.config.write().await;
config.computer_use_enabled = enabled;
if enabled
&& !config
.enabled_toolsets
.iter()
.any(|name| name.eq_ignore_ascii_case("computer_use"))
{
config.enabled_toolsets.push("computer_use".to_string());
}
}
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 {
self.last_prompt_tokens
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApprovalChoice {
Once,
Session,
Always,
Deny,
}
pub enum StreamEvent {
Token(String),
Reasoning(String),
ToolGenerating {
tool_call_id: String,
name: String,
partial_args: 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,
depth: u32,
agent_id: String,
parent_id: Option<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>,
},
RunFinished { outcome: RunOutcome },
Footer(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,
},
ActivityNotice(String),
LlmWaitProgress {
provider: String,
elapsed_secs: u64,
has_tools: bool,
prompt_tokens_estimated: Option<u64>,
context_length: Option<u64>,
prefill_pct: Option<f32>,
},
BackgroundProcessTail {
process_id: String,
command_preview: String,
tail: String,
},
BackgroundProcessFinished {
process_id: String,
exit_code: Option<i32>,
},
SteerPending { count: usize },
SteerApplied {
message: String,
},
ModelTransferComplete {
from: String,
to: String,
brief: String,
compressed: bool,
},
}
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::ToolGenerating { name, .. } => write!(f, "ToolGenerating({name:?})"),
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,
depth,
agent_id,
parent_id,
} => write!(
f,
"SubAgentStart({}/{}, d{depth}, id={agent_id}, parent={parent_id:?}, {:?})",
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::RunFinished { outcome } => write!(
f,
"RunFinished(state={:?}, exit_reason={:?})",
outcome.state, outcome.exit_reason
),
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})"
),
Self::ActivityNotice(text) => {
let preview = &text[..text.len().min(60)];
write!(f, "ActivityNotice({preview:?}…)")
}
Self::LlmWaitProgress {
provider,
elapsed_secs,
has_tools,
prompt_tokens_estimated,
context_length,
prefill_pct,
} => write!(
f,
"LlmWaitProgress({provider}, {elapsed_secs}s, tools={has_tools}, \
prompt={prompt_tokens_estimated:?}, ctx={context_length:?}, prefill={prefill_pct:?})"
),
Self::BackgroundProcessTail {
process_id, tail, ..
} => {
let preview = &tail[..tail.len().min(40)];
write!(f, "BackgroundProcessTail({process_id}, {preview:?}…)")
}
Self::BackgroundProcessFinished {
process_id,
exit_code,
} => write!(f, "BackgroundProcessFinished({process_id}, {exit_code:?})"),
Self::SteerPending { count } => write!(f, "SteerPending({count})"),
Self::SteerApplied { message } => {
let preview = &message[..message.len().min(60)];
write!(f, "SteerApplied({preview:?}…)")
}
Self::ModelTransferComplete {
from,
to,
compressed,
..
} => {
write!(
f,
"ModelTransferComplete({from:?} → {to:?}, compressed={compressed})"
)
}
Self::Footer(text) => {
let preview = &text[..text.len().min(60)];
write!(f, "Footer({preview:?}…)")
}
}
}
}
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>>,
context_engine: Option<Arc<dyn crate::context_engine::ContextEngine>>,
goal_store: Option<Arc<dyn crate::goals::GoalStore>>,
}
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,
context_engine: None,
goal_store: None,
}
}
pub fn from_config(config: &AppConfig) -> Self {
let personality_addon =
crate::config::resolve_personality(config, &config.display.personality);
Self {
config: AgentConfig {
custom_system_prompt: (!config.agent.system_prompt.trim().is_empty())
.then_some(config.agent.system_prompt.clone()),
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(),
memory_config: config.memory.clone(),
approvals_config: config.approvals.clone(),
kanban_config: config.kanban.clone(),
kanban_task_id: None,
plugins_config: config.plugins.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,
checkpoints_max_total_size_mb: config.checkpoints.max_total_size_mb,
checkpoints_max_file_size_mb: config.checkpoints.max_file_size_mb,
computer_use_enabled: config.computer_use.enabled,
computer_use_keep_last_n_screenshots: config.computer_use.keep_last_n_screenshots,
computer_use_confirm_destructive: config.computer_use.confirm_destructive,
computer_use_cua_cmd: config.computer_use.cua_driver_cmd.clone(),
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(),
shadow_judge: config.shadow_judge.clone(),
goals: config.goals.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(),
result_spill: config.tools.result_spill,
result_spill_threshold: config.tools.result_spill_threshold,
result_spill_preview_lines: config.tools.result_spill_preview_lines,
result_turn_budget_chars: config.tools.result_turn_budget_chars,
max_write_payload_kib: config.tools.file.max_write_payload_kib,
local_write_create_dirs: config.local_inference.write_create_dirs,
local_max_tool_turn_tokens: config.local_inference.max_tool_turn_tokens,
file_mutation_verifier: config.display.file_mutation_verifier,
cache: config.cache.clone(),
web_search: config.web_search.clone(),
web: config.web.clone(),
..Default::default()
},
provider: None,
state_db: None,
tool_registry: None,
context_engine: None,
goal_store: 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 context_engine(mut self, engine: Arc<dyn crate::context_engine::ContextEngine>) -> Self {
self.context_engine = Some(engine);
self
}
pub fn goal_store(mut self, store: Arc<dyn crate::goals::GoalStore>) -> Self {
self.goal_store = Some(store);
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(OriginChat::new(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 enabled_toolsets(mut self, toolsets: Vec<String>) -> Self {
self.config.enabled_toolsets = toolsets;
self
}
pub fn disabled_toolsets(mut self, toolsets: Vec<String>) -> Self {
self.config.disabled_toolsets = toolsets;
self
}
pub fn enabled_tools(mut self, tools: Vec<String>) -> Self {
self.config.enabled_tools = tools;
self
}
pub fn disabled_tools(mut self, tools: Vec<String>) -> Self {
self.config.disabled_tools = tools;
self
}
pub fn custom_system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.config.custom_system_prompt = Some(prompt.into());
self
}
pub fn skip_context_files(mut self, skip: bool) -> Self {
self.config.skip_context_files = skip;
self
}
pub fn skip_memory(mut self, skip: bool) -> Self {
self.config.skip_memory = skip;
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,
self.context_engine,
self.goal_store,
))
}
pub fn into_agent_config(self) -> AgentConfig {
self.config
}
}
impl Drop for Agent {
fn drop(&mut self) {
self.gc_cancel.cancel();
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use lingshu_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()])
}
}
fn write_session_hook_plugin(dir: &Path) {
std::fs::write(
dir.join("plugin.yaml"),
r#"
name: session-hooks
version: "1.0.0"
description: Session boundary recorder
provides_hooks:
- on_session_finalize
- on_session_reset
"#,
)
.expect("manifest");
std::fs::write(
dir.join("__init__.py"),
r#"
import json
from pathlib import Path
def _append(event_name, session_id, platform, **kwargs):
target = Path(__file__).with_name("session-hooks.jsonl")
with target.open("a", encoding="utf-8") as handle:
handle.write(json.dumps({"event": event_name, "session_id": session_id, "platform": platform}) + "\n")
def register(ctx):
ctx.register_hook("on_session_finalize", lambda **kwargs: _append("on_session_finalize", **kwargs))
ctx.register_hook("on_session_reset", lambda **kwargs: _append("on_session_reset", **kwargs))
"#,
)
.expect("plugin");
}
#[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(flavor = "multi_thread")]
async fn session_boundary_hooks_fire_on_new_and_finalize() {
let temp = tempfile::TempDir::new().expect("tempdir");
let plugin_dir = temp.path().join("session-hooks");
std::fs::create_dir_all(&plugin_dir).expect("plugin dir");
write_session_hook_plugin(&plugin_dir);
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let mut config = AppConfig::default();
config.plugins.external_dirs = vec![temp.path().to_string_lossy().to_string()];
let agent = AgentBuilder::from_config(&config)
.provider(provider)
.build()
.expect("build agent");
let old_session_id = "existing-session".to_string();
agent.session.write().await.session_id = Some(old_session_id.clone());
agent.new_session().await;
let new_session_id = agent
.session
.read()
.await
.session_id
.clone()
.expect("new session id");
agent.finalize_session().await;
let log = std::fs::read_to_string(plugin_dir.join("session-hooks.jsonl")).expect("log");
assert!(log.contains("on_session_finalize"));
assert!(log.contains("on_session_reset"));
assert!(log.contains(&old_session_id));
assert!(log.contains(&new_session_id));
}
#[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_default().expect("db"));
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let session = lingshu_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()),
handoff_state: None,
handoff_platform: None,
handoff_error: None,
};
db.save_session(&session).await.expect("save session");
db.save_message("restore-me", &Message::user("hello"))
.await
.expect("save user");
db.save_message("restore-me", &Message::assistant("hi"))
.await
.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")
);
}
#[test]
fn llm_response_cache_write_tokens_defaults_to_none() {
let r = edgequake_llm::LLMResponse::new("hello", "test-model");
assert!(
r.cache_write_tokens.is_none(),
"cache_write_tokens must default to None"
);
}
#[test]
fn llm_response_cache_write_tokens_builder() {
let r = edgequake_llm::LLMResponse::new("hello", "test-model").with_cache_write_tokens(42);
assert_eq!(r.cache_write_tokens, Some(42));
}
#[test]
fn stream_usage_cache_write_tokens_builder() {
use edgequake_llm::traits::StreamUsage;
let u = StreamUsage::new(10, 5).with_cache_write_tokens(99);
assert_eq!(u.cache_write_tokens, Some(99));
}
#[test]
fn session_state_first_compression_done_defaults_false() {
let s = SessionState::default();
assert!(
!s.first_compression_done,
"first_compression_done must start as false"
);
}
#[test]
fn session_state_first_compression_done_can_be_set() {
let s = SessionState {
first_compression_done: true,
..SessionState::default()
};
assert!(s.first_compression_done);
}
#[test]
fn apply_model_transfer_outcome_clears_cache_and_injects_brief() {
let outcome = crate::model_transfer::ModelTransferOutcome {
from_model: "a/m1".into(),
to_model: "b/m2".into(),
brief: "Continue tests.".into(),
compressed: true,
from_context_window: 200_000,
target_context_window: 128_000,
};
let mut session = SessionState {
cached_system_prompt: Some("stable".into()),
cached_stable_prompt: Some("stable".into()),
messages: vec![Message::user("hello")],
..SessionState::default()
};
session.apply_model_transfer_outcome(&outcome);
assert!(session.cached_system_prompt.is_none());
assert!(session.cached_stable_prompt.is_none());
assert!(session.first_compression_done);
assert_eq!(session.messages.len(), 2);
assert!(
session.messages[1]
.text_content()
.contains("Continue tests.")
);
}
#[tokio::test]
async fn goal_set_before_first_chat_avoids_foreign_key_error() {
let db = Arc::new(lingshu_state::SessionDb::open_default().expect("open db"));
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.state_db(db)
.build()
.expect("build");
let msg = agent
.goal_set("Refactor demo module in ./demo")
.await
.expect("goal set before any chat turn");
assert!(msg.contains("Goal set"));
let shown = agent.goal_show().await.expect("show");
assert!(shown.contains("Refactor demo module"));
}
#[tokio::test]
async fn goal_set_does_not_mutate_cached_system_prompt() {
use crate::goals::InMemoryGoalStore;
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let goal_store = Arc::new(InMemoryGoalStore::new());
let agent = AgentBuilder::new("mock")
.provider(provider)
.goal_store(goal_store)
.build()
.expect("build");
agent.chat("hello").await.expect("chat");
let before = agent.session.read().await.cached_system_prompt.clone();
agent
.goal_set("Refactor payment service")
.await
.expect("goal set");
let after = agent.session.read().await.cached_system_prompt.clone();
assert_eq!(before, after, "cached system prompt must remain unchanged");
}
#[tokio::test]
async fn perform_model_transfer_clears_cached_prompt_and_injects_brief() {
use crate::goals::{GoalStore, InMemoryGoalStore};
use crate::model_transfer::{
ModelTransferContext, ModelTransferOrchestrator, resolve_model_transfer_target,
};
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let goal_store = Arc::new(InMemoryGoalStore::new());
let db = Arc::new(lingshu_state::SessionDb::open_default().expect("open db"));
let agent = AgentBuilder::new("anthropic/claude-opus-4.6")
.provider(provider.clone())
.goal_store(goal_store.clone())
.state_db(db)
.build()
.expect("build");
agent
.goal_set("finish session handoff tests")
.await
.expect("goal");
agent.chat("implement session handoff").await.expect("chat");
let session_id = agent.session.read().await.session_id.clone().expect("sid");
{
let mut session = agent.session.write().await;
session.cached_system_prompt = Some("stable block".into());
session.cached_stable_prompt = Some("stable".into());
}
let target = resolve_model_transfer_target("anthropic/claude-haiku-4.5").expect("catalog");
let mut messages = agent.messages().await;
let cfg = agent.config.read().await.compression.clone();
let mut ctx = ModelTransferContext {
current_model: "anthropic/claude-opus-4.6",
messages: &mut messages,
system_prompt: Some("stable block"),
compression_cfg: &cfg,
main_provider: provider.clone(),
auxiliary_model: None,
};
let (outcome, new_provider) =
ModelTransferOrchestrator::execute_with_provider(&mut ctx, &target, provider)
.await
.expect("handoff");
{
let mut session = agent.session.write().await;
session.messages = messages;
session.apply_model_transfer_outcome(&outcome);
}
agent
.swap_model(outcome.to_model.clone(), new_provider)
.await;
let session = agent.session.read().await;
assert!(session.cached_system_prompt.is_none());
assert!(session.cached_stable_prompt.is_none());
assert!(
session.messages.last().is_some_and(|m| {
m.role == Role::User
&& m.text_content()
.starts_with("Continuing from previous session")
}),
"handoff brief should be injected as user message"
);
assert_eq!(agent.model().await, "anthropic/claude-haiku-4.5");
let goal = goal_store.as_ref().active(&session_id).expect("goal load");
assert!(
goal.goal_text
.as_deref()
.is_some_and(|g| g.contains("handoff")),
"standing goal must survive model handoff"
);
}
#[tokio::test]
async fn perform_model_transfer_persists_audit_trail() {
use crate::model_transfer::{
ModelTransferContext, ModelTransferOrchestrator, resolve_model_transfer_target,
};
let provider: Arc<dyn LLMProvider> = Arc::new(edgequake_llm::MockProvider::new());
let db = Arc::new(lingshu_state::SessionDb::open_default().expect("open db"));
let agent = AgentBuilder::new("anthropic/claude-opus-4.6")
.provider(provider.clone())
.state_db(db.clone())
.build()
.expect("build");
agent
.chat("work on model transfer unification")
.await
.expect("chat");
let session_id = agent.session.read().await.session_id.clone().expect("sid");
let target = resolve_model_transfer_target("anthropic/claude-haiku-4.5").expect("catalog");
let mut messages = agent.messages().await;
let cfg = agent.config.read().await.compression.clone();
let mut ctx = ModelTransferContext {
current_model: "anthropic/claude-opus-4.6",
messages: &mut messages,
system_prompt: None,
compression_cfg: &cfg,
main_provider: provider.clone(),
auxiliary_model: None,
};
let (outcome, new_provider) =
ModelTransferOrchestrator::execute_with_provider(&mut ctx, &target, provider)
.await
.expect("handoff");
{
let mut session = agent.session.write().await;
session.messages = messages;
session.apply_model_transfer_outcome(&outcome);
}
agent
.swap_model(outcome.to_model.clone(), new_provider)
.await;
db.update_session_model(&session_id, &outcome.to_model)
.await
.expect("session model update");
db.record_model_transfer(
&session_id,
&outcome.from_model,
&outcome.to_model,
&outcome.brief,
)
.await
.expect("record transfer");
assert_eq!(outcome.to_model, "anthropic/claude-haiku-4.5");
assert_eq!(agent.model().await, "anthropic/claude-haiku-4.5");
let records = db
.list_model_transfers(&session_id)
.await
.expect("list transfers");
assert_eq!(records.len(), 1);
assert_eq!(records[0].from_model, "anthropic/claude-opus-4.6");
assert_eq!(records[0].to_model, "anthropic/claude-haiku-4.5");
assert!(!records[0].brief.trim().is_empty());
}
}