#[cfg(test)]
mod embed_roles_tests;
pub mod attachment;
pub mod calc;
pub mod chats;
pub mod code;
pub mod confirm;
pub mod control;
pub mod datetime;
pub mod dialogue;
pub mod fetch;
pub mod fs;
pub mod history;
pub mod introspection;
pub mod llm;
pub mod mcp;
pub mod meta;
pub mod notes;
pub mod present;
pub mod python;
pub mod rag;
mod reach;
pub mod self_model;
pub mod subagent;
pub mod web;
pub mod youtube;
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Arc, LazyLock};
use std::time::Duration;
use anyhow::Result;
use chrono::{DateTime, Utc};
use uuid::Uuid;
use crate::entities::profile::ToolId;
use crate::entities::sampling::SamplingConfig;
use crate::entities::self_model::SelfModelParams;
use crate::shared::api::{Embedder, EngineBackend, ToolSchema};
use crate::shared::config::{AppConfig, CloudProvider, PythonMode};
use crate::shared::sandbox::WasmerSandbox;
use crate::shared::storage::Storage;
pub use introspection::{GET_SAMPLING_ID, SET_SAMPLING_ID};
#[derive(Clone)]
pub struct ToolContext {
pub profile_id: Uuid,
pub chat_id: Uuid,
pub system_message: String,
pub effective_sampling: SamplingConfig,
pub last_user_message_at: Option<DateTime<Utc>>,
pub storage: Arc<Storage>,
pub engine: Arc<dyn EngineBackend>,
pub embedder: Arc<dyn Embedder>,
pub chunk_params: rag::ChunkParams,
pub self_model_params: SelfModelParams,
pub recall_includes_self: bool,
pub loc: &'static crate::shared::i18n::Locale,
pub file_hint: Option<&'static str>,
pub attachments: std::sync::Arc<[crate::entities::attachment::Attachment]>,
pub attachment_cfg: crate::shared::config::AttachmentSettings,
pub history: Option<Arc<crate::features::compaction::HistoryView>>,
pub history_page_tokens: usize,
pub other_chats: std::sync::Arc<[chats::ChatRef]>,
pub mcp_images: bool,
pub python_net: bool,
pub python_mode: crate::shared::config::PythonMode,
pub workspace_journal: Option<std::path::PathBuf>,
pub files_dir: Option<std::path::PathBuf>,
pub inputs: std::sync::Arc<[crate::features::chat_inputs::ChatInput]>,
pub files: std::sync::Arc<[crate::entities::chat_file::ChatFile]>,
pub images: std::sync::Arc<[crate::entities::message_image::MessageImage]>,
pub stages_files: bool,
pub workspace: Option<crate::entities::workspace::Workspace>,
pub workspace_cfg: crate::shared::config::WorkspaceSettings,
pub named_secrets: std::sync::Arc<[String]>,
pub cancel: tokio_util::sync::CancellationToken,
pub model_name: Option<String>,
pub engine_mode: crate::shared::config::ServerMode,
pub sessions: Option<Arc<crate::shared::session_budget::SessionBudget>>,
pub silent_lane: bool,
}
#[derive(Clone)]
pub struct ToolDeps {
pub storage: Arc<Storage>,
pub engine: Arc<dyn EngineBackend>,
pub embedder: Arc<dyn Embedder>,
}
#[derive(Clone)]
pub struct ToolParams {
pub chunk_params: rag::ChunkParams,
pub self_model_params: SelfModelParams,
pub recall_includes_self: bool,
pub attachments: crate::shared::config::AttachmentSettings,
pub history_page_tokens: usize,
pub mcp_images: bool,
pub python_net: bool,
pub python_mode: crate::shared::config::PythonMode,
pub workspace: crate::shared::config::WorkspaceSettings,
pub file_hint: Option<&'static str>,
pub named_secrets: std::sync::Arc<[String]>,
}
impl ToolParams {
pub fn from_config(cfg: &AppConfig) -> Self {
Self {
chunk_params: rag::ChunkParams::from_settings(&cfg.rag),
self_model_params: SelfModelParams::from_settings(&cfg.self_model),
recall_includes_self: cfg.notes.recall_includes_self,
attachments: cfg.attachments,
history_page_tokens: cfg.compaction.page_tokens,
mcp_images: cfg.tools.mcp_images,
python_net: cfg.tools.python_net_enabled
|| matches!(
cfg.tools.python_mode,
crate::shared::config::PythonMode::Local
),
python_mode: cfg.tools.python_mode,
workspace: cfg.workspace,
file_hint: crate::shared::text_decode::tld_hint(cfg.interface.language),
named_secrets: crate::shared::config::named_key_env_vars(cfg).into(),
}
}
}
pub struct TurnInfo {
pub profile_id: Uuid,
pub chat_id: Uuid,
pub system_message: String,
pub effective_sampling: SamplingConfig,
pub last_user_message_at: Option<DateTime<Utc>>,
pub attachments: std::sync::Arc<[crate::entities::attachment::Attachment]>,
pub history: Option<Arc<crate::features::compaction::HistoryView>>,
pub other_chats: std::sync::Arc<[chats::ChatRef]>,
pub workspace: Option<crate::entities::workspace::Workspace>,
pub workspace_journal: Option<std::path::PathBuf>,
pub files_dir: Option<std::path::PathBuf>,
pub files: std::sync::Arc<[crate::entities::chat_file::ChatFile]>,
pub inputs: std::sync::Arc<[crate::features::chat_inputs::ChatInput]>,
pub images: std::sync::Arc<[crate::entities::message_image::MessageImage]>,
pub stages_files: bool,
pub lang: crate::shared::i18n::Lang,
pub cancel: tokio_util::sync::CancellationToken,
pub model_name: Option<String>,
pub engine_mode: crate::shared::config::ServerMode,
pub sessions: Option<Arc<crate::shared::session_budget::SessionBudget>>,
pub silent_lane: bool,
}
impl ToolContext {
pub fn new(deps: ToolDeps, params: ToolParams, turn: TurnInfo) -> Self {
Self {
profile_id: turn.profile_id,
chat_id: turn.chat_id,
system_message: turn.system_message,
effective_sampling: turn.effective_sampling,
last_user_message_at: turn.last_user_message_at,
attachments: turn.attachments,
attachment_cfg: params.attachments,
history: turn.history,
history_page_tokens: params.history_page_tokens,
other_chats: turn.other_chats,
workspace: turn.workspace,
workspace_journal: turn.workspace_journal,
files_dir: turn.files_dir,
files: turn.files,
inputs: turn.inputs,
images: turn.images,
stages_files: turn.stages_files,
workspace_cfg: params.workspace,
named_secrets: params.named_secrets,
mcp_images: params.mcp_images,
python_net: params.python_net,
python_mode: params.python_mode,
storage: deps.storage,
engine: deps.engine,
embedder: deps.embedder,
chunk_params: params.chunk_params,
self_model_params: params.self_model_params,
recall_includes_self: params.recall_includes_self,
loc: crate::shared::i18n::locale(turn.lang),
file_hint: params.file_hint,
cancel: turn.cancel,
model_name: turn.model_name,
engine_mode: turn.engine_mode,
sessions: turn.sessions,
silent_lane: turn.silent_lane,
}
}
pub fn sync_inputs(&mut self) {
let dir = self.files_dir.clone().unwrap_or_default();
self.inputs = crate::features::chat_inputs::reconcile(
&self.inputs,
&self.attachments,
&self.files,
&self.images.iter().collect::<Vec<_>>(),
&dir,
)
.into();
}
}
pub(crate) fn search_parameters(
loc: &crate::shared::i18n::Locale,
query_key: &str,
top_k_key: &str,
) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": loc.t(query_key)
},
"top_k": {
"type": "integer",
"minimum": 1,
"description": loc.t(top_k_key)
}
},
"required": ["query"]
})
}
pub(crate) fn paged_read_parameters(
loc: &crate::shared::i18n::Locale,
target: &str,
target_key: &str,
page_key: &str,
) -> serde_json::Value {
let mut properties = serde_json::Map::new();
properties.insert(
target.to_string(),
serde_json::json!({
"type": "string",
"description": loc.t(target_key)
}),
);
properties.insert(
"page".to_string(),
serde_json::json!({
"type": "integer",
"minimum": 1,
"description": loc.t(page_key)
}),
);
serde_json::json!({
"type": "object",
"properties": properties,
"required": [target]
})
}
pub(crate) fn search_args<'a>(
args: &'a serde_json::Value,
loc: &crate::shared::i18n::Locale,
err_key: &str,
default_k: usize,
) -> Result<(&'a str, usize)> {
let query = args
.get("query")
.and_then(|v| v.as_str())
.map(str::trim)
.filter(|s| !s.is_empty())
.ok_or_else(|| anyhow::anyhow!(loc.t(err_key).to_string()))?;
let k = args
.get("top_k")
.and_then(|v| v.as_u64())
.map(|n| n as usize)
.unwrap_or(default_k);
Ok((query, k))
}
#[derive(Debug, Clone, PartialEq)]
pub enum ChatEffect {
SetSystemMessage(String),
SetSamplingOverride(Box<SamplingConfig>),
AddAttachment(Box<crate::entities::attachment::Attachment>),
AddChatFile(Box<crate::entities::chat_file::ChatFile>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct ToolOutcome {
pub result: String,
pub effects: Vec<ChatEffect>,
pub images: Vec<ToolImage>,
pub wrote: bool,
pub prefill: Option<crate::shared::api::contract::Prefill>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolImage {
pub mime: String,
pub data: String,
pub entry: Option<String>,
}
impl ToolOutcome {
pub fn text(result: impl Into<String>) -> Self {
Self {
result: result.into(),
effects: Vec::new(),
images: Vec::new(),
wrote: false,
prefill: None,
}
}
pub fn with_effects(result: impl Into<String>, effects: Vec<ChatEffect>) -> Self {
Self {
result: result.into(),
effects,
images: Vec::new(),
wrote: false,
prefill: None,
}
}
pub fn with_images(mut self, images: Vec<ToolImage>) -> Self {
self.images = images;
self
}
pub fn wrote(self) -> Self {
self.wrote_if(true)
}
pub fn wrote_if(mut self, wrote: bool) -> Self {
self.wrote = wrote;
self
}
pub fn with_prefill(mut self, prefill: Option<crate::shared::api::contract::Prefill>) -> Self {
self.prefill = prefill;
self
}
}
#[async_trait::async_trait]
pub trait Tool: Send + Sync {
fn id(&self) -> ToolId;
fn description(&self, loc: &crate::shared::i18n::Locale) -> String;
fn parameters(&self, loc: &crate::shared::i18n::Locale) -> serde_json::Value;
async fn invoke(&self, ctx: &ToolContext, args: serde_json::Value) -> Result<ToolOutcome>;
fn schema(&self, loc: &crate::shared::i18n::Locale) -> ToolSchema {
ToolSchema {
name: self.id(),
description: self.description(loc),
parameters: self.parameters(loc),
}
}
fn group(&self) -> meta::ToolGroup;
fn ui_label(&self) -> &'static str;
fn gate(&self) -> Option<meta::ToolGate> {
None
}
fn danger(&self) -> bool {
false
}
fn enabled_by_default(&self) -> bool {
true
}
fn counts_toward_round_limit(&self) -> bool {
true
}
fn concurrent(&self) -> bool {
false
}
}
pub const WEB_SEARCH_ID: &str = "web_search";
pub const FETCH_URL_ID: &str = "fetch_url";
pub const PYTHON_EXEC_ID: &str = "python_exec";
pub const YOUTUBE_WATCH_ID: &str = "youtube_watch";
static CATALOG: LazyLock<Vec<meta::ToolInfo>> =
LazyLock::new(|| standard_registry(&ToolConfig::default()).infos());
pub fn tool_catalog() -> Vec<meta::ToolInfo> {
CATALOG.clone()
}
pub fn default_tool_ids() -> Vec<ToolId> {
CATALOG
.iter()
.filter(|i| i.enabled_by_default)
.map(|i| i.id.clone())
.collect()
}
#[allow(dead_code)]
pub fn all_tool_ids() -> Vec<ToolId> {
CATALOG.iter().map(|i| i.id.clone()).collect()
}
#[derive(Debug, Clone, Default)]
pub struct ToolGates {
pub web: bool,
pub python: bool,
pub fs: bool,
pub mcp: bool,
pub background: bool,
pub history: bool,
pub workspace: bool,
pub workspace_commands: code::WorkspaceCommands,
pub sampling_provider: Option<CloudProvider>,
pub sampling_endpoint: Option<std::sync::Arc<[String]>>,
}
pub fn effective_tool_ids(enabled: &[ToolId], gates: &ToolGates) -> Vec<ToolId> {
let sampling_available = !crate::entities::sampling::available_sampling_fields(
gates.sampling_provider,
gates.sampling_endpoint.as_deref(),
)
.is_empty();
let gate_of = |id: &str| CATALOG.iter().find(|i| i.id == id).and_then(|i| i.gate);
enabled
.iter()
.filter(|id| {
if id.as_str() == GET_SAMPLING_ID || id.as_str() == SET_SAMPLING_ID {
return sampling_available;
}
if id.as_str() == history::HISTORY_READ_ID || id.as_str() == history::HISTORY_SEARCH_ID
{
return gates.history;
}
if let Some(offered) = code::offered(id, gates.workspace, gates.workspace_commands) {
return offered;
}
if id.starts_with(mcp::MCP_TOOL_PREFIX) {
return gates.mcp;
}
match gate_of(id) {
Some(meta::ToolGate::Web) => gates.web,
Some(meta::ToolGate::Python) => gates.python,
Some(meta::ToolGate::Fs) => gates.fs,
Some(meta::ToolGate::Mcp) => gates.mcp,
Some(meta::ToolGate::Background) => gates.background,
None => true,
}
})
.cloned()
.collect()
}
#[derive(Debug, Clone)]
pub struct ToolConfig {
pub python_mode: PythonMode,
pub python_path: Option<String>,
pub python_net: bool,
pub python_wasm_timeout: Duration,
pub python_wasm_memory_mb: Option<u64>,
pub python_local_memory_mb: Option<u64>,
pub python_images: bool,
pub sandbox_dir: Option<PathBuf>,
pub web_fetch_content: bool,
pub web_allow_private: bool,
pub web_provider: crate::shared::config::WebProvider,
pub web_search_keys: Vec<(crate::shared::secrets::SearchSlot, String)>,
pub fs_root: Option<String>,
pub named_secrets: Vec<String>,
pub subagent_parallel: u32,
pub video: Option<crate::shared::video::VideoConfig>,
pub sampling_provider: Option<CloudProvider>,
pub sampling_endpoint: Option<std::sync::Arc<[String]>>,
}
impl Default for ToolConfig {
fn default() -> Self {
Self {
python_mode: PythonMode::default(),
python_path: None,
python_net: true,
python_wasm_timeout: Duration::from_secs(
crate::shared::config::DEFAULT_PYTHON_WASM_TIMEOUT_SECS,
),
python_wasm_memory_mb: None,
python_local_memory_mb: None,
python_images: true,
sandbox_dir: None,
web_fetch_content: true,
web_allow_private: false,
web_provider: crate::shared::config::WebProvider::default(),
web_search_keys: Vec::new(),
fs_root: None,
named_secrets: Vec::new(),
subagent_parallel: 1,
video: None,
sampling_provider: None,
sampling_endpoint: None,
}
}
}
pub fn standard_registry(cfg: &ToolConfig) -> ToolRegistry {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(introspection::GetSampling::new(
cfg.sampling_provider,
cfg.sampling_endpoint.clone(),
)));
reg.register(Arc::new(introspection::SetSampling::new(
cfg.sampling_provider,
cfg.sampling_endpoint.clone(),
)));
reg.register(Arc::new(introspection::GetSystemMessage));
reg.register(Arc::new(introspection::SetSystemMessage));
reg.register(Arc::new(introspection::GetLastUserMessageTime));
reg.register(Arc::new(llm::GetLlmName));
reg.register(Arc::new(llm::GetLlmHistory));
reg.register(Arc::new(notes::NoteSave));
reg.register(Arc::new(notes::NoteRecall));
reg.register(Arc::new(notes::NoteRevise));
reg.register(Arc::new(notes::NoteLink));
reg.register(Arc::new(notes::NoteNeighbors));
reg.register(Arc::new(notes::NoteSupersede));
reg.register(Arc::new(notes::NoteMerge));
reg.register(Arc::new(notes::ConsolidateNotes));
reg.register(Arc::new(notes::NoteCiteSource));
reg.register(Arc::new(rag::RagAdd));
reg.register(Arc::new(rag::RagSearch));
reg.register(Arc::new(subagent::CallSubagent {
parallel: cfg.subagent_parallel,
}));
reg.register(Arc::new(subagent::StartSubagent));
reg.register(Arc::new(dialogue::RunDialogue));
reg.register(Arc::new(dialogue::StartDialogue));
let policy = crate::shared::net::AddressPolicy::from_allow_private(cfg.web_allow_private);
reg.register(Arc::new(web::WebSearch::new(
cfg.web_fetch_content,
policy,
web::keyed_backends(cfg.web_provider, &cfg.web_search_keys),
)));
reg.register(Arc::new(fetch::FetchUrl::new(policy)));
reg.register(Arc::new(youtube::YoutubeWatch::new(
cfg.video.clone().map(|c| {
Arc::new(crate::shared::video::gemini::GeminiVideo::new(c))
as Arc<dyn crate::shared::video::VideoUnderstanding>
}),
cfg.video
.as_ref()
.map_or(crate::shared::config::DEFAULT_VIDEO_MAX_MINUTES, |c| {
c.max_minutes
}),
)));
let runner: Arc<dyn crate::shared::sandbox::SandboxRunner> = match cfg.python_mode {
crate::shared::config::PythonMode::Wasmer => Arc::new(
WasmerSandbox::new(cfg.sandbox_dir.clone())
.with_memory_limit(cfg.python_wasm_memory_mb)
.with_private_network(cfg.web_allow_private),
),
crate::shared::config::PythonMode::Local => Arc::new(
crate::shared::sandbox::LocalSandbox::new(cfg.python_path.clone())
.with_memory_limit(cfg.python_local_memory_mb)
.with_named_secrets(cfg.named_secrets.clone()),
),
};
reg.register(Arc::new(
python::PythonExec::new(
cfg.python_mode,
runner,
cfg.python_net,
cfg.python_wasm_timeout,
)
.with_images(cfg.python_images)
.with_private_network(cfg.web_allow_private),
));
reg.register(Arc::new(calc::Calculate));
reg.register(Arc::new(datetime::CurrentTime));
reg.register(Arc::new(fs::FsRead::new(cfg.fs_root.clone())));
reg.register(Arc::new(fs::FsWrite::new(cfg.fs_root.clone())));
reg.register(Arc::new(fs::FsList::new(cfg.fs_root.clone())));
for tool in code::ALL {
reg.register(Arc::new(tool));
}
reg.register(Arc::new(attachment::AttachmentRead));
reg.register(Arc::new(attachment::AttachmentSearch));
reg.register(Arc::new(history::HistoryRead));
reg.register(Arc::new(history::HistorySearch));
reg.register(Arc::new(chats::ChatSearch));
reg.register(Arc::new(chats::ChatRead));
reg.register(Arc::new(control::SendFollowupMessage));
reg.register(Arc::new(control::RewriteCurrentMessage));
reg.register(Arc::new(self_model::GetSelfModel));
reg.register(Arc::new(self_model::Reflect));
reg.register(Arc::new(self_model::UpdateSelfModel));
reg.register(Arc::new(self_model::UpdateUserModel));
reg.register(Arc::new(self_model::AddInsight));
reg
}
#[derive(Default)]
pub struct ToolRegistry {
tools: BTreeMap<ToolId, Arc<dyn Tool>>,
}
impl ToolRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, tool: Arc<dyn Tool>) {
self.tools.insert(tool.id(), tool);
}
#[allow(dead_code)]
pub fn get(&self, id: &str) -> Option<&Arc<dyn Tool>> {
self.tools.get(id)
}
pub fn infos(&self) -> Vec<meta::ToolInfo> {
self.tools
.values()
.map(|t| meta::ToolInfo {
id: t.id(),
group: t.group(),
label: t.ui_label(),
gate: t.gate(),
enabled_by_default: t.enabled_by_default(),
concurrent: t.concurrent(),
description: None,
})
.collect()
}
pub fn is_concurrent(&self, id: &str) -> bool {
self.tools.get(id).is_some_and(|t| t.concurrent())
}
pub fn schemas_for(
&self,
enabled: &[ToolId],
loc: &crate::shared::i18n::Locale,
) -> Vec<ToolSchema> {
enabled
.iter()
.filter_map(|id| self.tools.get(id))
.map(|t| t.schema(loc))
.collect()
}
pub async fn invoke(
&self,
id: &str,
ctx: &ToolContext,
args: serde_json::Value,
) -> Result<ToolOutcome> {
match self.tools.get(id) {
Some(tool) => tool.invoke(ctx, args).await,
None => anyhow::bail!("unknown tool: {id}"),
}
}
}
#[cfg(test)]
pub(crate) mod testkit {
use super::*;
use crate::shared::api::EmbedRole;
use crate::shared::api::mock::{MockBackend, MockEmbedder};
use crate::shared::paths::Paths;
fn test_params() -> ToolParams {
ToolParams {
chunk_params: rag::ChunkParams::default(),
self_model_params: SelfModelParams::default(),
recall_includes_self: false,
history_page_tokens: crate::shared::config::DEFAULT_COMPACTION_PAGE_TOKENS,
attachments: crate::shared::config::AttachmentSettings::default(),
mcp_images: true,
python_net: false,
python_mode: crate::shared::config::PythonMode::Wasmer,
workspace: crate::shared::config::WorkspaceSettings::default(),
file_hint: None,
named_secrets: Vec::new().into(),
}
}
fn test_turn(profile_id: Uuid) -> TurnInfo {
TurnInfo {
profile_id,
chat_id: Uuid::new_v4(),
system_message: "системное сообщение".into(),
effective_sampling: SamplingConfig::default(),
last_user_message_at: None,
attachments: std::sync::Arc::from(Vec::new()),
history: None,
other_chats: std::sync::Arc::from(Vec::new()),
workspace: None,
workspace_journal: None,
files_dir: None,
files: std::sync::Arc::from(Vec::new()),
inputs: std::sync::Arc::from(Vec::new()),
images: std::sync::Arc::from(Vec::new()),
stages_files: false,
lang: crate::shared::i18n::Lang::Ru,
cancel: tokio_util::sync::CancellationToken::new(),
model_name: None,
engine_mode: crate::shared::config::ServerMode::default(),
sessions: None,
silent_lane: false,
}
}
pub fn ctx_with_storage(profile_id: Uuid) -> (tempfile::TempDir, Arc<Storage>, ToolContext) {
ctx_with_storage_lang(profile_id, crate::shared::i18n::Lang::Ru)
}
pub fn ctx_with_storage_lang(
profile_id: Uuid,
lang: crate::shared::i18n::Lang,
) -> (tempfile::TempDir, Arc<Storage>, ToolContext) {
let dir = tempfile::tempdir().unwrap();
let storage = Arc::new(Storage::open_in_memory(Paths::with_root(dir.path())).unwrap());
let deps = ToolDeps {
storage: storage.clone(),
engine: Arc::new(MockBackend::scripted(vec![])),
embedder: Arc::new(MockEmbedder::new(16)),
};
let mut turn = test_turn(profile_id);
turn.lang = lang;
let ctx = ToolContext::new(deps, test_params(), turn);
(dir, storage, ctx)
}
pub fn ctx_with_deps(profile_id: Uuid, deps: ToolDeps) -> ToolContext {
ToolContext::new(deps, test_params(), test_turn(profile_id))
}
pub struct RoleRecorder {
inner: MockEmbedder,
pub calls: std::sync::Mutex<Vec<(Vec<String>, EmbedRole)>>,
}
impl RoleRecorder {
pub fn new() -> Self {
Self {
inner: MockEmbedder::new(16),
calls: std::sync::Mutex::new(Vec::new()),
}
}
pub fn roles(&self) -> Vec<EmbedRole> {
self.calls.lock().unwrap().iter().map(|(_, r)| *r).collect()
}
pub fn all_were(&self, role: EmbedRole) -> bool {
let roles = self.roles();
!roles.is_empty() && roles.iter().all(|r| *r == role)
}
}
#[async_trait::async_trait]
impl Embedder for RoleRecorder {
async fn embed(
&self,
texts: Vec<String>,
role: EmbedRole,
) -> anyhow::Result<Vec<Vec<f32>>> {
self.calls.lock().unwrap().push((texts.clone(), role));
self.inner.embed(texts, role).await
}
}
pub fn ctx_with_backends(
profile_id: Uuid,
engine: Arc<dyn EngineBackend>,
embedder: Arc<dyn Embedder>,
) -> (tempfile::TempDir, Arc<Storage>, ToolContext) {
let dir = tempfile::tempdir().unwrap();
let storage = Arc::new(Storage::open_in_memory(Paths::with_root(dir.path())).unwrap());
let deps = ToolDeps {
storage: storage.clone(),
engine,
embedder,
};
let ctx = ToolContext::new(deps, test_params(), test_turn(profile_id));
(dir, storage, ctx)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tool_params_take_the_file_hint_from_the_interface_language() {
let mut cfg = AppConfig::default();
cfg.interface.language = crate::shared::i18n::Lang::Ru;
assert_eq!(ToolParams::from_config(&cfg).file_hint, Some("ru"));
cfg.interface.language = crate::shared::i18n::Lang::En;
assert_eq!(ToolParams::from_config(&cfg).file_hint, None);
}
#[test]
fn local_mode_reports_the_network_as_reachable_whatever_the_switch_says() {
let mut cfg = AppConfig::default();
cfg.tools.python_net_enabled = false;
cfg.tools.python_mode = crate::shared::config::PythonMode::Wasmer;
assert!(!ToolParams::from_config(&cfg).python_net);
cfg.tools.python_mode = crate::shared::config::PythonMode::Local;
let params = ToolParams::from_config(&cfg);
assert!(params.python_net);
assert_eq!(params.python_mode, crate::shared::config::PythonMode::Local);
}
#[test]
fn concurrent_tools_are_the_documented_set_and_never_dangerous() {
let reg = standard_registry(&ToolConfig::default());
let infos = reg.infos();
let mut marked: Vec<&str> = infos
.iter()
.filter(|i| i.concurrent)
.map(|i| i.id.as_str())
.collect();
marked.sort_unstable();
let mut documented = vec![
fs::FS_READ_ID,
fs::FS_LIST_ID,
code::CODE_READ_ID,
code::CODE_GREP_ID,
code::CODE_LIST_ID,
attachment::ATTACHMENT_READ_ID,
attachment::ATTACHMENT_SEARCH_ID,
chats::CHAT_SEARCH_ID,
chats::CHAT_READ_ID,
history::HISTORY_READ_ID,
history::HISTORY_SEARCH_ID,
self_model::GET_SELF_MODEL_ID,
GET_SAMPLING_ID,
llm::GET_LLM_NAME_ID,
llm::GET_LLM_HISTORY_ID,
FETCH_URL_ID,
];
documented.sort_unstable();
assert_eq!(marked, documented);
for info in infos.iter().filter(|i| i.concurrent) {
let tool = reg.get(&info.id).unwrap();
assert!(
!tool.danger(),
"{} is marked concurrent and dangerous",
info.id
);
}
assert!(reg.is_concurrent(FETCH_URL_ID));
assert!(!reg.is_concurrent(fs::FS_WRITE_ID));
assert!(
!reg.is_concurrent(WEB_SEARCH_ID),
"out until measured (fork F4)"
);
assert!(
!reg.is_concurrent(subagent::CALL_SUBAGENT_ID),
"the group's own path"
);
assert!(!reg.is_concurrent("no_such_tool"));
}
#[test]
fn tool_context_storage_touches_no_disk() {
let (dir, storage, _ctx) = testkit::ctx_with_storage(Uuid::new_v4());
storage
.db()
.note_insert(&crate::entities::note::Note::new(
Uuid::new_v4(),
"заметка",
vec![],
))
.unwrap();
let paths = crate::shared::paths::Paths::with_root(dir.path());
assert!(
!paths.data_db().exists(),
"the tool testkit must not create data.db — see the doc comment above"
);
assert!(!paths.cache_db().exists(), "nor cache.db");
}
struct Echo;
#[async_trait::async_trait]
impl Tool for Echo {
fn id(&self) -> ToolId {
"echo".into()
}
fn description(&self, _loc: &crate::shared::i18n::Locale) -> String {
"Возвращает аргумент text".into()
}
fn parameters(&self, _loc: &crate::shared::i18n::Locale) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
})
}
async fn invoke(&self, _ctx: &ToolContext, args: serde_json::Value) -> Result<ToolOutcome> {
let text = args["text"].as_str().unwrap_or_default();
Ok(ToolOutcome::text(text))
}
fn group(&self) -> meta::ToolGroup {
meta::ToolGroup::Utils
}
fn ui_label(&self) -> &'static str {
"echo"
}
}
#[tokio::test]
async fn registry_invokes_registered_tool() {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(Echo));
let (_d, _s, ctx) = testkit::ctx_with_storage(Uuid::new_v4());
let out = reg
.invoke("echo", &ctx, serde_json::json!({"text": "hi"}))
.await
.unwrap();
assert_eq!(out.result, "hi");
assert!(out.effects.is_empty());
}
#[tokio::test]
async fn registry_unknown_tool_errors() {
let reg = ToolRegistry::new();
let (_d, _s, ctx) = testkit::ctx_with_storage(Uuid::new_v4());
assert!(
reg.invoke("nope", &ctx, serde_json::json!({}))
.await
.is_err()
);
}
#[test]
fn standard_registry_has_all_default_tools() {
let reg = standard_registry(&ToolConfig::default());
for id in all_tool_ids() {
assert!(reg.get(&id).is_some(), "tool {id} not registered");
}
assert_eq!(
reg.schemas_for(
&all_tool_ids(),
crate::shared::i18n::locale(crate::shared::i18n::Lang::Ru)
)
.len(),
all_tool_ids().len()
);
}
#[test]
fn all_tool_descriptions_localized_to_en() {
use crate::shared::i18n::{Lang, locale};
let reg = standard_registry(&ToolConfig::default());
let (ru, en) = (locale(Lang::Ru), locale(Lang::En));
for id in all_tool_ids() {
let t = reg.get(&id).expect("tool in registry");
let d_en = t.description(en);
assert!(
!d_en
.chars()
.any(|c| ('а'..='я').contains(&c) || ('А'..='Я').contains(&c)),
"{id}: Cyrillic in en description: {d_en}"
);
assert_ne!(t.description(ru), d_en, "{id}: description not localized");
}
}
#[test]
fn note_revise_is_default_tool() {
assert!(
default_tool_ids()
.iter()
.any(|t| t == notes::NOTE_REVISE_ID)
);
}
#[test]
fn control_tools_optional_not_in_defaults() {
assert!(
!default_tool_ids()
.iter()
.any(|t| t == control::SEND_FOLLOWUP_ID)
);
assert!(
!default_tool_ids()
.iter()
.any(|t| t == control::REWRITE_CURRENT_ID)
);
assert!(
all_tool_ids()
.iter()
.any(|t| t == control::SEND_FOLLOWUP_ID)
);
assert!(
all_tool_ids()
.iter()
.any(|t| t == control::REWRITE_CURRENT_ID)
);
}
#[test]
fn self_model_tools_optional_not_in_defaults() {
for id in [
self_model::GET_SELF_MODEL_ID,
self_model::REFLECT_ID,
self_model::UPDATE_SELF_MODEL_ID,
self_model::UPDATE_USER_MODEL_ID,
self_model::ADD_INSIGHT_ID,
] {
assert!(
!default_tool_ids().iter().any(|t| t == id),
"{id} in defaults"
);
assert!(
all_tool_ids().iter().any(|t| t == id),
"{id} not in catalog"
);
}
let eff = effective_tool_ids(
&all_tool_ids(),
&ToolGates {
background: false,
history: true,
..Default::default()
},
);
assert!(eff.iter().any(|t| t == self_model::GET_SELF_MODEL_ID));
assert!(eff.iter().any(|t| t == self_model::UPDATE_SELF_MODEL_ID));
}
#[test]
fn effective_tool_ids_gates_external_tools() {
let enabled = default_tool_ids();
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
web: true,
history: true,
..Default::default()
},
);
assert!(eff.iter().any(|t| t == WEB_SEARCH_ID));
assert!(eff.iter().any(|t| t == FETCH_URL_ID));
assert!(!eff.iter().any(|t| t == PYTHON_EXEC_ID));
assert!(!eff.iter().any(|t| t == fs::FS_READ_ID));
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
..Default::default()
},
);
assert!(!eff.iter().any(|t| t == WEB_SEARCH_ID || t == FETCH_URL_ID));
assert!(
!eff.iter()
.any(|t| t == fs::FS_READ_ID || t == fs::FS_WRITE_ID || t == fs::FS_LIST_ID)
);
assert!(eff.iter().any(|t| t == "note_save"));
assert!(eff.iter().any(|t| t == "calculate"));
assert!(eff.iter().any(|t| t == "current_time"));
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
fs: true,
history: true,
..Default::default()
},
);
assert!(eff.iter().any(|t| t == fs::FS_READ_ID));
assert!(eff.iter().any(|t| t == fs::FS_WRITE_ID));
assert!(eff.iter().any(|t| t == fs::FS_LIST_ID));
}
#[test]
fn effective_tool_ids_keeps_sampling_tools_when_params_available() {
let enabled = default_tool_ids();
for provider in [
None,
Some(CloudProvider::OpenAi),
Some(CloudProvider::Gemini),
Some(CloudProvider::Claude),
] {
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
sampling_provider: provider,
..Default::default()
},
);
assert!(
eff.iter().any(|t| t == GET_SAMPLING_ID),
"get_sampling must be available for {provider:?}"
);
assert!(eff.iter().any(|t| t == SET_SAMPLING_ID));
}
}
#[test]
fn effective_tool_ids_gates_history_tools_by_the_chat() {
let enabled: Vec<ToolId> = vec![
"note_save".into(),
history::HISTORY_READ_ID.into(),
history::HISTORY_SEARCH_ID.into(),
];
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
..Default::default()
},
);
assert!(
!eff.iter()
.any(|t| t == history::HISTORY_READ_ID || t == history::HISTORY_SEARCH_ID),
"nothing folded → the tools must not be offered"
);
assert!(eff.iter().any(|t| t == "note_save"), "unrelated tools stay");
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
..Default::default()
},
);
assert!(eff.iter().any(|t| t == history::HISTORY_READ_ID));
assert!(eff.iter().any(|t| t == history::HISTORY_SEARCH_ID));
}
#[test]
fn effective_tool_ids_gates_mcp_tools_by_prefix() {
let enabled: Vec<ToolId> = vec!["note_save".into(), "mcp__fs__read_text_file".into()];
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
..Default::default()
},
);
assert!(!eff.iter().any(|t| t.starts_with(mcp::MCP_TOOL_PREFIX)));
assert!(eff.iter().any(|t| t == "note_save"));
let eff = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
mcp: true,
history: true,
..Default::default()
},
);
assert!(eff.iter().any(|t| t == "mcp__fs__read_text_file"));
}
#[test]
fn workspace_tools_need_an_attached_project() {
let enabled: Vec<ToolId> = code::WORKSPACE_TOOL_IDS
.iter()
.map(|id| ToolId::from(*id))
.collect();
let detached = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
..Default::default()
},
);
assert!(
detached.is_empty(),
"with no project the tools must not be offered: {detached:?}"
);
let attached = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
workspace: true,
..Default::default()
},
);
assert_eq!(
attached.len(),
code::WORKSPACE_TOOL_IDS.len() - 3,
"a slot with no line must not be offered: {attached:?}"
);
for id in [code::CODE_BUILD_ID, code::CODE_RUN_ID, code::CODE_TEST_ID] {
assert!(!attached.iter().any(|t| t == id), "{id} without a line");
}
for slot in crate::entities::workspace::CommandSlot::ALL {
let mut ws = crate::entities::workspace::Workspace::new("/p");
ws.set_command(slot, Some("cargo build".into()));
let offered = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
workspace: true,
workspace_commands: code::WorkspaceCommands::of(&ws),
..Default::default()
},
);
let wanted = crate::features::tools::code::CodeTool::Command(slot).id();
assert!(
offered.iter().any(|t| t == wanted),
"{wanted} must be offered once its slot has a line: {offered:?}"
);
assert_eq!(
offered.len(),
code::WORKSPACE_TOOL_IDS.len() - 2,
"only this slot's tool joins: {offered:?}"
);
}
let one: Vec<ToolId> = vec![code::CODE_READ_ID.into()];
let narrow = effective_tool_ids(
&one,
&ToolGates {
background: false,
history: true,
workspace: true,
..Default::default()
},
);
assert_eq!(narrow, one);
}
#[test]
fn workspace_tools_do_not_ride_the_fs_switch() {
let enabled: Vec<ToolId> = vec![code::CODE_LIST_ID.into()];
let fs_off = effective_tool_ids(
&enabled,
&ToolGates {
background: false,
history: true,
workspace: true,
..Default::default()
},
);
assert_eq!(fs_off, enabled, "fs_enabled must not gate the code tools");
}
#[test]
fn schemas_for_filters_and_orders() {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(Echo));
let schemas = reg.schemas_for(
&["echo".into(), "missing".into()],
crate::shared::i18n::locale(crate::shared::i18n::Lang::Ru),
);
assert_eq!(schemas.len(), 1);
assert_eq!(schemas[0].name, "echo");
assert!(
reg.schemas_for(
&[],
crate::shared::i18n::locale(crate::shared::i18n::Lang::Ru)
)
.is_empty()
);
}
}