feat: 99% complete migration to new state structs to get away from God-Config struct; i.e. AppConfig, AppState, and RequestContext

This commit is contained in:
2026-04-19 17:05:27 -06:00
parent c85adfd00e
commit b32bcf8fbc
90 changed files with 18983 additions and 3448 deletions
+126 -106
View File
@@ -1,4 +1,3 @@
use super::todo::TodoList;
use super::*;
use crate::{
@@ -6,6 +5,7 @@ use crate::{
function::{Functions, run_llm_function},
};
use crate::config::paths;
use crate::config::prompts::{
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
@@ -38,16 +38,13 @@ pub struct Agent {
rag: Option<Arc<Rag>>,
model: Model,
vault: GlobalVault,
todo_list: TodoList,
continuation_count: usize,
last_continuation_response: Option<String>,
}
impl Agent {
pub fn install_builtin_agents() -> Result<()> {
info!(
"Installing built-in agents in {}",
Config::agents_data_dir().display()
paths::agents_data_dir().display()
);
for file in AgentAssets::iter() {
@@ -56,7 +53,7 @@ impl Agent {
let embedded_file = AgentAssets::get(&file)
.ok_or_else(|| anyhow!("Failed to load embedded agent file: {}", file.as_ref()))?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = Config::agents_data_dir().join(file.as_ref());
let file_path = paths::agents_data_dir().join(file.as_ref());
let file_extension = file_path
.extension()
.and_then(OsStr::to_str)
@@ -88,14 +85,17 @@ impl Agent {
}
pub async fn init(
config: &GlobalConfig,
app: &AppConfig,
app_state: &AppState,
current_model: &Model,
info_flag: bool,
name: &str,
abort_signal: AbortSignal,
) -> Result<Self> {
let agent_data_dir = Config::agent_data_dir(name);
let loaders = config.read().document_loaders.clone();
let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME);
let config_path = Config::agent_config_file(name);
let agent_data_dir = paths::agent_data_dir(name);
let loaders = app.document_loaders.clone();
let rag_path = paths::agent_rag_file(name, DEFAULT_AGENT_NAME);
let config_path = paths::agent_config_file(name);
let mut agent_config = if config_path.exists() {
AgentConfig::load(&config_path)?
} else {
@@ -103,57 +103,24 @@ impl Agent {
};
let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;
config.write().functions.clear_mcp_meta_functions();
let mcp_servers = if config.read().mcp_server_support {
(!agent_config.mcp_servers.is_empty()).then(|| agent_config.mcp_servers.join(","))
} else {
eprintln!(
"{}",
formatdoc!(
"
This agent uses MCP servers, but MCP support is disabled.
To enable it, exit the agent and set 'mcp_server_support: true', then try again
"
)
);
None
};
agent_config.load_envs(app);
let registry = config
.write()
.mcp_registry
.take()
.with_context(|| "MCP registry should be populated")?;
let new_mcp_registry =
McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;
if !new_mcp_registry.is_empty() {
functions.append_mcp_meta_functions(new_mcp_registry.list_started_servers());
}
config.write().mcp_registry = Some(new_mcp_registry);
agent_config.load_envs(&config.read());
let model = {
let config = config.read();
match agent_config.model_id.as_ref() {
Some(model_id) => Model::retrieve_model(&config, model_id, ModelType::Chat)?,
None => {
if agent_config.temperature.is_none() {
agent_config.temperature = config.temperature;
}
if agent_config.top_p.is_none() {
agent_config.top_p = config.top_p;
}
config.current_model().clone()
let model = match agent_config.model_id.as_ref() {
Some(model_id) => Model::retrieve_model(app, model_id, ModelType::Chat)?,
None => {
if agent_config.temperature.is_none() {
agent_config.temperature = app.temperature;
}
if agent_config.top_p.is_none() {
agent_config.top_p = app.top_p;
}
current_model.clone()
}
};
let rag = if rag_path.exists() {
Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?))
} else if !agent_config.documents.is_empty() && !config.read().info_flag {
Some(Arc::new(Rag::load(app, DEFAULT_AGENT_NAME, &rag_path)?))
} else if !agent_config.documents.is_empty() && !info_flag {
let mut ans = false;
if *IS_STDOUT_TERMINAL {
ans = Confirm::new("The agent has documents attached, init RAG?")
@@ -185,8 +152,7 @@ impl Agent {
document_paths.push(path.to_string())
}
}
let rag =
Rag::init(config, "rag", &rag_path, &document_paths, abort_signal).await?;
let rag = Rag::init(app, "rag", &rag_path, &document_paths, abort_signal).await?;
Some(Arc::new(rag))
} else {
None
@@ -218,10 +184,7 @@ impl Agent {
functions,
rag,
model,
vault: Arc::clone(&config.read().vault),
todo_list: TodoList::default(),
continuation_count: 0,
last_continuation_response: None,
vault: app_state.vault.clone(),
})
}
@@ -295,11 +258,11 @@ impl Agent {
let mut config = self.config.clone();
config.instructions = self.interpolated_instructions();
value["definition"] = json!(config);
value["data_dir"] = Config::agent_data_dir(&self.name)
value["data_dir"] = paths::agent_data_dir(&self.name)
.display()
.to_string()
.into();
value["config_file"] = Config::agent_config_file(&self.name)
value["config_file"] = paths::agent_config_file(&self.name)
.display()
.to_string()
.into();
@@ -323,6 +286,14 @@ impl Agent {
self.rag.clone()
}
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
self.functions.append_mcp_meta_functions(mcp_servers);
}
pub fn mcp_server_names(&self) -> &[String] {
&self.config.mcp_servers
}
pub fn conversation_starters(&self) -> Vec<String> {
self.config
.conversation_starters
@@ -443,44 +414,6 @@ impl Agent {
self.config.escalation_timeout
}
pub fn continuation_count(&self) -> usize {
self.continuation_count
}
pub fn increment_continuation(&mut self) {
self.continuation_count += 1;
}
pub fn reset_continuation(&mut self) {
self.continuation_count = 0;
self.last_continuation_response = None;
}
pub fn set_last_continuation_response(&mut self, response: String) {
self.last_continuation_response = Some(response);
}
pub fn todo_list(&self) -> &TodoList {
&self.todo_list
}
pub fn init_todo_list(&mut self, goal: &str) {
self.todo_list = TodoList::new(goal);
}
pub fn add_todo(&mut self, task: &str) -> usize {
self.todo_list.add(task)
}
pub fn mark_todo_done(&mut self, id: usize) -> bool {
self.todo_list.mark_done(id)
}
pub fn clear_todo_list(&mut self) {
self.todo_list.clear();
self.reset_continuation();
}
pub fn continuation_prompt(&self) -> String {
self.config.continuation_prompt.clone().unwrap_or_else(|| {
formatdoc! {"
@@ -696,12 +629,12 @@ impl AgentConfig {
Ok(agent_config)
}
fn load_envs(&mut self, config: &Config) {
fn load_envs(&mut self, app: &AppConfig) {
let name = &self.name;
let with_prefix = |v: &str| normalize_env_name(&format!("{name}_{v}"));
if self.agent_session.is_none() {
self.agent_session = config.agent_session.clone();
self.agent_session = app.agent_session.clone();
}
if let Some(v) = read_env_value::<String>(&with_prefix("model")) {
@@ -793,7 +726,7 @@ pub struct AgentVariable {
}
pub fn list_agents() -> Vec<String> {
let agents_data_dir = Config::agents_data_dir();
let agents_data_dir = paths::agents_data_dir();
if !agents_data_dir.exists() {
return vec![];
}
@@ -803,6 +736,7 @@ pub fn list_agents() -> Vec<String> {
for entry in entries.flatten() {
if entry.path().is_dir()
&& let Some(name) = entry.file_name().to_str()
&& !name.starts_with('.')
{
agents.push(name.to_string());
}
@@ -813,7 +747,7 @@ pub fn list_agents() -> Vec<String> {
}
pub fn complete_agent_variables(agent_name: &str) -> Vec<(String, Option<String>)> {
let config_path = Config::agent_config_file(agent_name);
let config_path = paths::agent_config_file(agent_name);
if !config_path.exists() {
return vec![];
}
@@ -832,3 +766,89 @@ pub fn complete_agent_variables(agent_name: &str) -> Vec<(String, Option<String>
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn agent_config_parses_from_yaml() {
let yaml = r#"
name: test-agent
description: A test agent
instructions: You are helpful
auto_continue: true
max_auto_continues: 5
can_spawn_agents: true
max_concurrent_agents: 8
max_agent_depth: 2
mcp_servers:
- github
- jira
global_tools:
- execute_command.sh
- fs_read.sh
conversation_starters:
- "Hello!"
- "How are you?"
variables:
- name: username
description: Your name
"#;
let config: AgentConfig = serde_yaml::from_str(yaml).unwrap();
assert_eq!(config.name, "test-agent");
assert_eq!(config.description, "A test agent");
assert!(config.auto_continue);
assert_eq!(config.max_auto_continues, 5);
assert!(config.can_spawn_agents);
assert_eq!(config.max_concurrent_agents, 8);
assert_eq!(config.max_agent_depth, 2);
assert_eq!(config.mcp_servers, vec!["github", "jira"]);
assert_eq!(config.global_tools.len(), 2);
assert_eq!(config.conversation_starters.len(), 2);
assert_eq!(config.variables.len(), 1);
assert_eq!(config.variables[0].name, "username");
}
#[test]
fn agent_config_defaults() {
let yaml = "name: minimal\ninstructions: hi\n";
let config: AgentConfig = serde_yaml::from_str(yaml).unwrap();
assert_eq!(config.name, "minimal");
assert!(!config.auto_continue);
assert!(!config.can_spawn_agents);
assert_eq!(config.max_concurrent_agents, 4);
assert_eq!(config.max_agent_depth, 3);
assert_eq!(config.max_auto_continues, 10);
assert!(config.mcp_servers.is_empty());
assert!(config.global_tools.is_empty());
assert!(config.conversation_starters.is_empty());
assert!(config.variables.is_empty());
assert!(config.model_id.is_none());
assert!(config.temperature.is_none());
assert!(config.top_p.is_none());
}
#[test]
fn agent_config_with_model() {
let yaml =
"name: test\nmodel: openai:gpt-4\ntemperature: 0.7\ntop_p: 0.9\ninstructions: hi\n";
let config: AgentConfig = serde_yaml::from_str(yaml).unwrap();
assert_eq!(config.model_id, Some("openai:gpt-4".to_string()));
assert_eq!(config.temperature, Some(0.7));
assert_eq!(config.top_p, Some(0.9));
}
#[test]
fn agent_config_inject_defaults_true() {
let yaml = "name: test\ninstructions: hi\n";
let config: AgentConfig = serde_yaml::from_str(yaml).unwrap();
assert!(config.inject_todo_instructions);
assert!(config.inject_spawn_instructions);
}
}
+586
View File
@@ -0,0 +1,586 @@
//! Immutable, server-wide application configuration.
//!
//! `AppConfig` contains the settings loaded from `config.yaml` that are
//! global to the Loki process: LLM provider configs, UI preferences, tool
//! and MCP settings, RAG defaults, etc.
//!
//! This is Phase 1, Step 0 of the REST API refactor: the struct is
//! introduced alongside the existing [`Config`](super::Config) and is not
//! yet wired into the runtime. See `docs/PHASE-1-IMPLEMENTATION-PLAN.md`
//! for the full migration plan.
//!
//! # Relationship to `Config`
//!
//! `AppConfig` mirrors the **serialized** fields of [`Config`] — that is,
//! every field that is NOT marked `#[serde(skip)]`. The deserialization
//! shape is identical so an existing `config.yaml` can be loaded into
//! either type without modification.
//!
//! Runtime-only state (current role, session, agent, supervisor, etc.)
//! lives on [`RequestContext`](super::request_context::RequestContext).
use crate::client::ClientConfig;
use crate::render::{MarkdownRender, RenderOptions};
use crate::utils::{IS_STDOUT_TERMINAL, NO_COLOR, decode_bin, get_env_name};
use super::paths;
use anyhow::{Context, Result, anyhow};
use indexmap::IndexMap;
use serde::Deserialize;
use std::collections::HashMap;
use std::env;
use std::path::PathBuf;
use syntect::highlighting::ThemeSet;
use terminal_colorsaurus::{ColorScheme, QueryOptions, color_scheme};
#[derive(Debug, Clone, Deserialize)]
#[serde(default)]
pub struct AppConfig {
#[serde(rename(serialize = "model", deserialize = "model"))]
#[serde(default)]
#[allow(dead_code)]
pub model_id: String,
pub temperature: Option<f64>,
pub top_p: Option<f64>,
pub dry_run: bool,
pub stream: bool,
pub save: bool,
pub keybindings: String,
pub editor: Option<String>,
pub wrap: Option<String>,
pub wrap_code: bool,
pub(crate) vault_password_file: Option<PathBuf>,
pub function_calling_support: bool,
pub mapping_tools: IndexMap<String, String>,
pub enabled_tools: Option<String>,
pub visible_tools: Option<Vec<String>>,
pub mcp_server_support: bool,
pub mapping_mcp_servers: IndexMap<String, String>,
pub enabled_mcp_servers: Option<String>,
pub repl_prelude: Option<String>,
pub cmd_prelude: Option<String>,
pub agent_session: Option<String>,
pub save_session: Option<bool>,
pub compression_threshold: usize,
pub summarization_prompt: Option<String>,
pub summary_context_prompt: Option<String>,
pub rag_embedding_model: Option<String>,
pub rag_reranker_model: Option<String>,
pub rag_top_k: usize,
pub rag_chunk_size: Option<usize>,
pub rag_chunk_overlap: Option<usize>,
pub rag_template: Option<String>,
#[serde(default)]
pub document_loaders: HashMap<String, String>,
pub highlight: bool,
pub theme: Option<String>,
pub left_prompt: Option<String>,
pub right_prompt: Option<String>,
pub user_agent: Option<String>,
pub save_shell_history: bool,
pub sync_models_url: Option<String>,
pub clients: Vec<ClientConfig>,
}
impl Default for AppConfig {
fn default() -> Self {
Self {
model_id: Default::default(),
temperature: None,
top_p: None,
dry_run: false,
stream: true,
save: false,
keybindings: "emacs".into(),
editor: None,
wrap: None,
wrap_code: false,
vault_password_file: None,
function_calling_support: true,
mapping_tools: Default::default(),
enabled_tools: None,
visible_tools: None,
mcp_server_support: true,
mapping_mcp_servers: Default::default(),
enabled_mcp_servers: None,
repl_prelude: None,
cmd_prelude: None,
agent_session: None,
save_session: None,
compression_threshold: 4000,
summarization_prompt: None,
summary_context_prompt: None,
rag_embedding_model: None,
rag_reranker_model: None,
rag_top_k: 5,
rag_chunk_size: None,
rag_chunk_overlap: None,
rag_template: None,
document_loaders: Default::default(),
highlight: true,
theme: None,
left_prompt: None,
right_prompt: None,
user_agent: None,
save_shell_history: true,
sync_models_url: None,
clients: vec![],
}
}
}
impl AppConfig {
pub fn vault_password_file(&self) -> PathBuf {
match &self.vault_password_file {
Some(path) => match path.exists() {
true => path.clone(),
false => gman::config::Config::local_provider_password_file(),
},
None => gman::config::Config::local_provider_password_file(),
}
}
pub fn editor(&self) -> Result<String> {
super::EDITOR.get_or_init(move || {
let editor = self.editor.clone()
.or_else(|| env::var("VISUAL").ok().or_else(|| env::var("EDITOR").ok()))
.unwrap_or_else(|| {
if cfg!(windows) {
"notepad".to_string()
} else {
"nano".to_string()
}
});
which::which(&editor).ok().map(|_| editor)
})
.clone()
.ok_or_else(|| anyhow!("Editor not found. Please add the `editor` configuration or set the $EDITOR or $VISUAL environment variable."))
}
pub fn sync_models_url(&self) -> String {
self.sync_models_url
.clone()
.unwrap_or_else(|| super::SYNC_MODELS_URL.into())
}
pub fn light_theme(&self) -> bool {
matches!(self.theme.as_deref(), Some("light"))
}
pub fn render_options(&self) -> Result<RenderOptions> {
let theme = if self.highlight {
let theme_mode = if self.light_theme() { "light" } else { "dark" };
let theme_filename = format!("{theme_mode}.tmTheme");
let theme_path = paths::local_path(&theme_filename);
if theme_path.exists() {
let theme = ThemeSet::get_theme(&theme_path)
.with_context(|| format!("Invalid theme at '{}'", theme_path.display()))?;
Some(theme)
} else {
let theme = if self.light_theme() {
decode_bin(super::LIGHT_THEME).context("Invalid builtin light theme")?
} else {
decode_bin(super::DARK_THEME).context("Invalid builtin dark theme")?
};
Some(theme)
}
} else {
None
};
let wrap = if *IS_STDOUT_TERMINAL {
self.wrap.clone()
} else {
None
};
let truecolor = matches!(
env::var("COLORTERM").as_ref().map(|v| v.as_str()),
Ok("truecolor")
);
Ok(RenderOptions::new(theme, wrap, self.wrap_code, truecolor))
}
pub fn print_markdown(&self, text: &str) -> Result<()> {
if *IS_STDOUT_TERMINAL {
let render_options = self.render_options()?;
let mut markdown_render = MarkdownRender::init(render_options)?;
println!("{}", markdown_render.render(text));
} else {
println!("{text}");
}
Ok(())
}
}
impl AppConfig {
#[allow(dead_code)]
pub fn set_wrap(&mut self, value: &str) -> Result<()> {
if value == "no" {
self.wrap = None;
} else if value == "auto" {
self.wrap = Some(value.into());
} else {
value
.parse::<u16>()
.map_err(|_| anyhow!("Invalid wrap value"))?;
self.wrap = Some(value.into())
}
Ok(())
}
#[allow(dead_code)]
pub fn setup_document_loaders(&mut self) {
[("pdf", "pdftotext $1 -"), ("docx", "pandoc --to plain $1")]
.into_iter()
.for_each(|(k, v)| {
let (k, v) = (k.to_string(), v.to_string());
self.document_loaders.entry(k).or_insert(v);
});
}
#[allow(dead_code)]
pub fn setup_user_agent(&mut self) {
if let Some("auto") = self.user_agent.as_deref() {
self.user_agent = Some(format!(
"{}/{}",
env!("CARGO_CRATE_NAME"),
env!("CARGO_PKG_VERSION")
));
}
}
#[allow(dead_code)]
pub fn load_envs(&mut self) {
if let Ok(v) = env::var(get_env_name("model")) {
self.model_id = v;
}
if let Some(v) = super::read_env_value::<f64>(&get_env_name("temperature")) {
self.temperature = v;
}
if let Some(v) = super::read_env_value::<f64>(&get_env_name("top_p")) {
self.top_p = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("dry_run")) {
self.dry_run = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("stream")) {
self.stream = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("save")) {
self.save = v;
}
if let Ok(v) = env::var(get_env_name("keybindings"))
&& v == "vi"
{
self.keybindings = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("editor")) {
self.editor = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("wrap")) {
self.wrap = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("wrap_code")) {
self.wrap_code = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("function_calling_support")) {
self.function_calling_support = v;
}
if let Ok(v) = env::var(get_env_name("mapping_tools"))
&& let Ok(v) = serde_json::from_str(&v)
{
self.mapping_tools = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("enabled_tools")) {
self.enabled_tools = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("mcp_server_support")) {
self.mcp_server_support = v;
}
if let Ok(v) = env::var(get_env_name("mapping_mcp_servers"))
&& let Ok(v) = serde_json::from_str(&v)
{
self.mapping_mcp_servers = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("enabled_mcp_servers")) {
self.enabled_mcp_servers = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("repl_prelude")) {
self.repl_prelude = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("cmd_prelude")) {
self.cmd_prelude = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("agent_session")) {
self.agent_session = v;
}
if let Some(v) = super::read_env_bool(&get_env_name("save_session")) {
self.save_session = v;
}
if let Some(Some(v)) =
super::read_env_value::<usize>(&get_env_name("compression_threshold"))
{
self.compression_threshold = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("summarization_prompt")) {
self.summarization_prompt = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("summary_context_prompt")) {
self.summary_context_prompt = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_embedding_model")) {
self.rag_embedding_model = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_reranker_model")) {
self.rag_reranker_model = v;
}
if let Some(Some(v)) = super::read_env_value::<usize>(&get_env_name("rag_top_k")) {
self.rag_top_k = v;
}
if let Some(v) = super::read_env_value::<usize>(&get_env_name("rag_chunk_size")) {
self.rag_chunk_size = v;
}
if let Some(v) = super::read_env_value::<usize>(&get_env_name("rag_chunk_overlap")) {
self.rag_chunk_overlap = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_template")) {
self.rag_template = v;
}
if let Ok(v) = env::var(get_env_name("document_loaders"))
&& let Ok(v) = serde_json::from_str(&v)
{
self.document_loaders = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("highlight")) {
self.highlight = v;
}
if *NO_COLOR {
self.highlight = false;
}
if self.highlight && self.theme.is_none() {
if let Some(v) = super::read_env_value::<String>(&get_env_name("theme")) {
self.theme = v;
} else if *IS_STDOUT_TERMINAL
&& let Ok(color_scheme) = color_scheme(QueryOptions::default())
{
let theme = match color_scheme {
ColorScheme::Dark => "dark",
ColorScheme::Light => "light",
};
self.theme = Some(theme.into());
}
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("left_prompt")) {
self.left_prompt = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("right_prompt")) {
self.right_prompt = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("user_agent")) {
self.user_agent = v;
}
if let Some(Some(v)) = super::read_env_bool(&get_env_name("save_shell_history")) {
self.save_shell_history = v;
}
if let Some(v) = super::read_env_value::<String>(&get_env_name("sync_models_url")) {
self.sync_models_url = v;
}
}
}
impl AppConfig {
#[allow(dead_code)]
pub fn set_temperature_default(&mut self, value: Option<f64>) {
self.temperature = value;
}
#[allow(dead_code)]
pub fn set_top_p_default(&mut self, value: Option<f64>) {
self.top_p = value;
}
#[allow(dead_code)]
pub fn set_enabled_tools_default(&mut self, value: Option<String>) {
self.enabled_tools = value;
}
#[allow(dead_code)]
pub fn set_enabled_mcp_servers_default(&mut self, value: Option<String>) {
self.enabled_mcp_servers = value;
}
#[allow(dead_code)]
pub fn set_save_session_default(&mut self, value: Option<bool>) {
self.save_session = value;
}
#[allow(dead_code)]
pub fn set_compression_threshold_default(&mut self, value: Option<usize>) {
self.compression_threshold = value.unwrap_or_default();
}
#[allow(dead_code)]
pub fn set_rag_reranker_model_default(&mut self, value: Option<String>) {
self.rag_reranker_model = value;
}
#[allow(dead_code)]
pub fn set_rag_top_k_default(&mut self, value: usize) {
self.rag_top_k = value;
}
#[allow(dead_code)]
pub fn set_model_id_default(&mut self, model_id: String) {
self.model_id = model_id;
}
#[allow(dead_code)]
pub fn ensure_default_model_id(&mut self) -> Result<String> {
if self.model_id.is_empty() {
let models = crate::client::list_models(self, crate::client::ModelType::Chat);
if models.is_empty() {
anyhow::bail!("No available model");
}
self.model_id = models[0].id();
}
Ok(self.model_id.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
fn cached_editor() -> Option<String> {
super::super::EDITOR.get().cloned().flatten()
}
#[test]
fn to_app_config_copies_serialized_fields() {
let cfg = Config {
model_id: "test-model".to_string(),
temperature: Some(0.7),
top_p: Some(0.9),
dry_run: true,
stream: false,
save: true,
highlight: false,
compression_threshold: 2000,
rag_top_k: 10,
..Config::default()
};
let app = cfg.to_app_config();
assert_eq!(app.model_id, "test-model");
assert_eq!(app.temperature, Some(0.7));
assert_eq!(app.top_p, Some(0.9));
assert!(app.dry_run);
assert!(!app.stream);
assert!(app.save);
assert!(!app.highlight);
assert_eq!(app.compression_threshold, 2000);
assert_eq!(app.rag_top_k, 10);
}
#[test]
fn to_app_config_copies_clients() {
let cfg = Config::default();
let app = cfg.to_app_config();
assert!(app.clients.is_empty());
}
#[test]
fn to_app_config_copies_mapping_fields() {
let mut cfg = Config::default();
cfg.mapping_tools
.insert("alias".to_string(), "real_tool".to_string());
cfg.mapping_mcp_servers
.insert("gh".to_string(), "github-mcp".to_string());
let app = cfg.to_app_config();
assert_eq!(
app.mapping_tools.get("alias"),
Some(&"real_tool".to_string())
);
assert_eq!(
app.mapping_mcp_servers.get("gh"),
Some(&"github-mcp".to_string())
);
}
#[test]
fn editor_returns_configured_value() {
let configured = cached_editor()
.unwrap_or_else(|| std::env::current_exe().unwrap().display().to_string());
let app = AppConfig {
editor: Some(configured.clone()),
..AppConfig::default()
};
assert_eq!(app.editor().unwrap(), configured);
}
#[test]
fn editor_falls_back_to_env() {
if let Some(expected) = cached_editor() {
let app = AppConfig::default();
assert_eq!(app.editor().unwrap(), expected);
return;
}
let expected = std::env::current_exe().unwrap().display().to_string();
unsafe {
std::env::set_var("VISUAL", &expected);
}
let app = AppConfig::default();
let result = app.editor();
assert!(result.is_ok());
assert_eq!(result.unwrap(), expected);
}
#[test]
fn light_theme_default_is_false() {
let app = AppConfig::default();
assert!(!app.light_theme());
}
#[test]
fn sync_models_url_has_default() {
let app = AppConfig::default();
let url = app.sync_models_url();
assert!(!url.is_empty());
}
}
+49
View File
@@ -0,0 +1,49 @@
//! Shared global services for a running Loki process.
//!
//! `AppState` holds the services that are genuinely process-wide and
//! immutable during request handling: the frozen [`AppConfig`], the
//! credential [`Vault`](GlobalVault), the [`McpFactory`](super::mcp_factory::McpFactory)
//! for MCP subprocess sharing, and the [`RagCache`](super::rag_cache::RagCache)
//! for shared RAG instances. It is intended to be wrapped in `Arc`
//! and shared across every [`RequestContext`] that a frontend (CLI,
//! REPL, API) creates.
//!
//! This struct deliberately does **not** hold a live `McpRegistry`.
//! MCP server processes are scoped to whichever `RoleLike`
//! (role/session/agent) is currently active, because each scope may
//! demand a different enabled server set. Live MCP processes are
//! owned by per-scope
//! [`ToolScope`](super::tool_scope::ToolScope)s on the
//! [`RequestContext`] and acquired through `McpFactory`.
//!
//! # Phase 1 scope
//!
//! This is Phase 1 of the REST API refactor:
//!
//! * **Step 0** introduced this struct alongside the existing
//! [`Config`](super::Config)
//! * **Step 6.5** added the `mcp_factory` and `rag_cache` fields
//!
//! Neither field is wired into the runtime yet — they exist as
//! additive scaffolding that Step 8+ will connect when the entry
//! points migrate. See `docs/PHASE-1-IMPLEMENTATION-PLAN.md`.
use super::mcp_factory::McpFactory;
use super::rag_cache::RagCache;
use crate::config::AppConfig;
use crate::mcp::McpServersConfig;
use crate::vault::GlobalVault;
use std::path::PathBuf;
use std::sync::Arc;
#[derive(Clone)]
pub struct AppState {
pub config: Arc<AppConfig>,
pub vault: GlobalVault,
pub mcp_factory: Arc<McpFactory>,
#[allow(dead_code)]
pub rag_cache: Arc<RagCache>,
pub mcp_config: Option<McpServersConfig>,
pub mcp_log_path: Option<PathBuf>,
}
+103
View File
@@ -0,0 +1,103 @@
//! Transitional conversions between the legacy [`Config`] struct and the
//! new [`AppConfig`] + [`RequestContext`] split.
use crate::config::todo::TodoList;
use super::{AppConfig, AppState, Config, RequestContext};
use std::sync::Arc;
impl Config {
pub fn to_app_config(&self) -> AppConfig {
AppConfig {
model_id: self.model_id.clone(),
temperature: self.temperature,
top_p: self.top_p,
dry_run: self.dry_run,
stream: self.stream,
save: self.save,
keybindings: self.keybindings.clone(),
editor: self.editor.clone(),
wrap: self.wrap.clone(),
wrap_code: self.wrap_code,
vault_password_file: self.vault_password_file.clone(),
function_calling_support: self.function_calling_support,
mapping_tools: self.mapping_tools.clone(),
enabled_tools: self.enabled_tools.clone(),
visible_tools: self.visible_tools.clone(),
mcp_server_support: self.mcp_server_support,
mapping_mcp_servers: self.mapping_mcp_servers.clone(),
enabled_mcp_servers: self.enabled_mcp_servers.clone(),
repl_prelude: self.repl_prelude.clone(),
cmd_prelude: self.cmd_prelude.clone(),
agent_session: self.agent_session.clone(),
save_session: self.save_session,
compression_threshold: self.compression_threshold,
summarization_prompt: self.summarization_prompt.clone(),
summary_context_prompt: self.summary_context_prompt.clone(),
rag_embedding_model: self.rag_embedding_model.clone(),
rag_reranker_model: self.rag_reranker_model.clone(),
rag_top_k: self.rag_top_k,
rag_chunk_size: self.rag_chunk_size,
rag_chunk_overlap: self.rag_chunk_overlap,
rag_template: self.rag_template.clone(),
document_loaders: self.document_loaders.clone(),
highlight: self.highlight,
theme: self.theme.clone(),
left_prompt: self.left_prompt.clone(),
right_prompt: self.right_prompt.clone(),
user_agent: self.user_agent.clone(),
save_shell_history: self.save_shell_history,
sync_models_url: self.sync_models_url.clone(),
clients: self.clients.clone(),
}
}
pub fn to_request_context(&self, app: Arc<AppState>) -> RequestContext {
let mut mcp_runtime = super::tool_scope::McpRuntime::default();
if let Some(registry) = &self.mcp_registry {
mcp_runtime.sync_from_registry(registry);
}
let tool_tracker = self
.tool_call_tracker
.clone()
.unwrap_or_else(crate::function::ToolCallTracker::default);
RequestContext {
app,
macro_flag: self.macro_flag,
info_flag: self.info_flag,
working_mode: self.working_mode,
model: self.model.clone(),
agent_variables: self.agent_variables.clone(),
role: self.role.clone(),
session: self.session.clone(),
rag: self.rag.clone(),
agent: self.agent.clone(),
last_message: self.last_message.clone(),
tool_scope: super::tool_scope::ToolScope {
functions: self.functions.clone(),
mcp_runtime,
tool_tracker,
},
supervisor: self.supervisor.clone(),
parent_supervisor: self.parent_supervisor.clone(),
self_agent_id: self.self_agent_id.clone(),
inbox: self.inbox.clone(),
escalation_queue: self.root_escalation_queue.clone(),
current_depth: self.current_depth,
auto_continue_count: 0,
todo_list: TodoList::default(),
last_continuation_response: None,
}
}
}
+57 -28
View File
@@ -9,7 +9,7 @@ use crate::utils::{AbortSignal, base64_encode, is_loader_protocol, sha256};
use anyhow::{Context, Result, bail};
use indexmap::IndexSet;
use std::{collections::HashMap, fs::File, io::Read};
use std::{collections::HashMap, fs::File, io::Read, sync::Arc};
use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
const IMAGE_EXTS: [&str; 5] = ["png", "jpeg", "jpg", "webp", "gif"];
@@ -17,7 +17,11 @@ const SUMMARY_MAX_WIDTH: usize = 80;
#[derive(Debug, Clone)]
pub struct Input {
config: GlobalConfig,
app_config: Arc<AppConfig>,
stream_enabled: bool,
session: Option<Session>,
rag: Option<Arc<Rag>>,
functions: Option<Vec<FunctionDeclaration>>,
text: String,
raw: (String, Vec<String>),
patched_text: Option<String>,
@@ -34,10 +38,15 @@ pub struct Input {
}
impl Input {
pub fn from_str(config: &GlobalConfig, text: &str, role: Option<Role>) -> Self {
let (role, with_session, with_agent) = resolve_role(&config.read(), role);
pub fn from_str(ctx: &RequestContext, text: &str, role: Option<Role>) -> Self {
let (role, with_session, with_agent) = resolve_role(ctx, role);
let captured = capture_input_config(ctx, &role);
Self {
config: config.clone(),
app_config: Arc::clone(&ctx.app.config),
stream_enabled: captured.stream_enabled,
session: captured.session,
rag: captured.rag,
functions: captured.functions,
text: text.to_string(),
raw: (text.to_string(), vec![]),
patched_text: None,
@@ -55,12 +64,12 @@ impl Input {
}
pub async fn from_files(
config: &GlobalConfig,
ctx: &RequestContext,
raw_text: &str,
paths: Vec<String>,
role: Option<Role>,
) -> Result<Self> {
let loaders = config.read().document_loaders.clone();
let loaders = ctx.app.config.document_loaders.clone();
let (raw_paths, local_paths, remote_urls, external_cmds, protocol_paths, with_last_reply) =
resolve_paths(&loaders, paths)?;
let mut last_reply = None;
@@ -78,7 +87,7 @@ impl Input {
texts.push(raw_text.to_string());
};
if with_last_reply {
if let Some(LastMessage { input, output, .. }) = config.read().last_message.as_ref() {
if let Some(LastMessage { input, output, .. }) = ctx.last_message.as_ref() {
if !output.is_empty() {
last_reply = Some(output.clone())
} else if let Some(v) = input.last_reply.as_ref() {
@@ -102,9 +111,14 @@ impl Input {
));
}
}
let (role, with_session, with_agent) = resolve_role(&config.read(), role);
let (role, with_session, with_agent) = resolve_role(ctx, role);
let captured = capture_input_config(ctx, &role);
Ok(Self {
config: config.clone(),
app_config: Arc::clone(&ctx.app.config),
stream_enabled: captured.stream_enabled,
session: captured.session,
rag: captured.rag,
functions: captured.functions,
text: texts.join("\n"),
raw: (raw_text.to_string(), raw_paths),
patched_text: None,
@@ -122,14 +136,14 @@ impl Input {
}
pub async fn from_files_with_spinner(
config: &GlobalConfig,
ctx: &RequestContext,
raw_text: &str,
paths: Vec<String>,
role: Option<Role>,
abort_signal: AbortSignal,
) -> Result<Self> {
abortable_run_with_spinner(
Input::from_files(config, raw_text, paths, role),
Input::from_files(ctx, raw_text, paths, role),
"Loading files",
abort_signal,
)
@@ -164,7 +178,7 @@ impl Input {
}
pub fn stream(&self) -> bool {
self.config.read().stream && !self.role().model().no_stream()
self.stream_enabled && !self.role().model().no_stream()
}
pub fn continue_output(&self) -> Option<&str> {
@@ -183,10 +197,9 @@ impl Input {
self.regenerate
}
pub fn set_regenerate(&mut self) {
let role = self.config.read().extract_role();
if role.name() == self.role().name() {
self.role = role;
pub fn set_regenerate(&mut self, current_role: Role) {
if current_role.name() == self.role().name() {
self.role = current_role;
}
self.regenerate = true;
self.tool_calls = None;
@@ -196,9 +209,9 @@ impl Input {
if self.text.is_empty() {
return Ok(());
}
let rag = self.config.read().rag.clone();
if let Some(rag) = rag {
let result = Config::search_rag(&self.config, &rag, &self.text, abort_signal).await?;
if let Some(rag) = &self.rag {
let result =
Config::search_rag(&self.app_config, rag, &self.text, abort_signal).await?;
self.patched_text = Some(result);
self.rag_name = Some(rag.name().to_string());
}
@@ -220,7 +233,7 @@ impl Input {
}
pub fn create_client(&self) -> Result<Box<dyn Client>> {
init_client(&self.config, Some(self.role().model().clone()))
init_client(&self.app_config, self.role().model().clone())
}
pub async fn fetch_chat_text(&self) -> Result<String> {
@@ -240,7 +253,7 @@ impl Input {
model.guard_max_input_tokens(&messages)?;
let (temperature, top_p) = (self.role().temperature(), self.role().top_p());
let functions = if model.supports_function_calling() {
let fns = self.config.read().select_functions(self.role());
let fns = self.functions.clone();
if let Some(vec) = &fns {
for def in vec {
debug!("Function definition: {:?}", def.name);
@@ -260,7 +273,7 @@ impl Input {
}
pub fn build_messages(&self) -> Result<Vec<Message>> {
let mut messages = if let Some(session) = self.session(&self.config.read().session) {
let mut messages = if let Some(session) = self.session(&self.session) {
session.build_messages(self)
} else {
self.role().build_messages(self)
@@ -275,7 +288,7 @@ impl Input {
}
pub fn echo_messages(&self) -> String {
if let Some(session) = self.session(&self.config.read().session) {
if let Some(session) = self.session(&self.session) {
session.echo_messages(self)
} else {
self.role().echo_messages(self)
@@ -384,17 +397,33 @@ impl Input {
}
}
fn resolve_role(config: &Config, role: Option<Role>) -> (Role, bool, bool) {
fn resolve_role(ctx: &RequestContext, role: Option<Role>) -> (Role, bool, bool) {
match role {
Some(v) => (v, false, false),
None => (
config.extract_role(),
config.session.is_some(),
config.agent.is_some(),
ctx.extract_role(ctx.app.config.as_ref()),
ctx.session.is_some(),
ctx.agent.is_some(),
),
}
}
struct CapturedInputConfig {
stream_enabled: bool,
session: Option<Session>,
rag: Option<Arc<Rag>>,
functions: Option<Vec<FunctionDeclaration>>,
}
fn capture_input_config(ctx: &RequestContext, role: &Role) -> CapturedInputConfig {
CapturedInputConfig {
stream_enabled: ctx.app.config.stream,
session: ctx.session.clone(),
rag: ctx.rag.clone(),
functions: ctx.select_functions(role),
}
}
type ResolvePathsOutput = (
Vec<String>,
Vec<String>,
+38 -22
View File
@@ -1,14 +1,13 @@
use crate::config::{Config, GlobalConfig, RoleLike, ensure_parent_exists};
use crate::config::paths;
use crate::config::{Config, RequestContext, RoleLike, ensure_parent_exists};
use crate::repl::{run_repl_command, split_args_text};
use crate::utils::{AbortSignal, multiline_text};
use anyhow::{Result, anyhow};
use indexmap::IndexMap;
use parking_lot::RwLock;
use rust_embed::Embed;
use serde::Deserialize;
use std::fs::File;
use std::io::Write;
use std::sync::Arc;
#[derive(Embed)]
#[folder = "assets/macros"]
@@ -16,7 +15,7 @@ struct MacroAssets;
#[async_recursion::async_recursion]
pub async fn macro_execute(
config: &GlobalConfig,
ctx: &mut RequestContext,
name: &str,
args: Option<&str>,
abort_signal: AbortSignal,
@@ -29,25 +28,42 @@ pub async fn macro_execute(
let variables = macro_value
.resolve_variables(&new_args)
.map_err(|err| anyhow!("{err}. Usage: {}", macro_value.usage(name)))?;
let role = config.read().extract_role();
let mut config = config.read().clone();
config.temperature = role.temperature();
config.top_p = role.top_p();
config.enabled_tools = role.enabled_tools().clone();
config.enabled_mcp_servers = role.enabled_mcp_servers().clone();
config.macro_flag = true;
config.model = role.model().clone();
config.role = None;
config.session = None;
config.rag = None;
config.agent = None;
config.discontinuous_last_message();
let config = Arc::new(RwLock::new(config));
config.write().macro_flag = true;
let role = ctx.extract_role(ctx.app.config.as_ref());
let mut app_config = (*ctx.app.config).clone();
app_config.temperature = role.temperature();
app_config.top_p = role.top_p();
app_config.enabled_tools = role.enabled_tools().clone();
app_config.enabled_mcp_servers = role.enabled_mcp_servers().clone();
let mut app_state = (*ctx.app).clone();
app_state.config = std::sync::Arc::new(app_config);
let mut macro_ctx = RequestContext::new(std::sync::Arc::new(app_state), ctx.working_mode);
macro_ctx.macro_flag = true;
macro_ctx.info_flag = ctx.info_flag;
macro_ctx.model = role.model().clone();
macro_ctx.agent_variables = ctx.agent_variables.clone();
macro_ctx.last_message = ctx.last_message.clone();
macro_ctx.supervisor = ctx.supervisor.clone();
macro_ctx.parent_supervisor = ctx.parent_supervisor.clone();
macro_ctx.self_agent_id = ctx.self_agent_id.clone();
macro_ctx.inbox = ctx.inbox.clone();
macro_ctx.escalation_queue = ctx.escalation_queue.clone();
macro_ctx.current_depth = ctx.current_depth;
macro_ctx.auto_continue_count = ctx.auto_continue_count;
macro_ctx.todo_list = ctx.todo_list.clone();
macro_ctx.tool_scope.tool_tracker = ctx.tool_scope.tool_tracker.clone();
macro_ctx.discontinuous_last_message();
let app = macro_ctx.app.config.clone();
macro_ctx
.bootstrap_tools(app.as_ref(), true, abort_signal.clone())
.await?;
for step in &macro_value.steps {
let command = Macro::interpolate_command(step, &variables);
println!(">> {}", multiline_text(&command));
run_repl_command(&config, abort_signal.clone(), &command).await?;
run_repl_command(&mut macro_ctx, abort_signal.clone(), &command).await?;
}
Ok(())
}
@@ -63,7 +79,7 @@ impl Macro {
pub fn install_macros() -> Result<()> {
info!(
"Installing built-in macros in {}",
Config::macros_dir().display()
paths::macros_dir().display()
);
for file in MacroAssets::iter() {
@@ -71,7 +87,7 @@ impl Macro {
let embedded_file = MacroAssets::get(&file)
.ok_or_else(|| anyhow!("Failed to load embedded macro file: {}", file.as_ref()))?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = Config::macros_dir().join(file.as_ref());
let file_path = paths::macros_dir().join(file.as_ref());
if file_path.exists() {
debug!(
+122
View File
@@ -0,0 +1,122 @@
//! Per-process factory for MCP subprocess handles.
//!
//! `McpFactory` lives on [`AppState`](super::AppState) and is the
//! single entrypoint that scopes use to obtain `Arc<ConnectedServer>`
//! handles for MCP tool servers. Multiple scopes requesting the same
//! server can (eventually) share a single subprocess via `Arc`
//! reference counting.
//!
//! # Phase 1 Step 6.5 scope
//!
//! This file introduces the factory scaffolding with a trivial
//! implementation:
//!
//! * `active` — `Mutex<HashMap<McpServerKey, Weak<ConnectedServer>>>`
//! for future Arc-based sharing across scopes
//! * `acquire` — unimplemented stub for now; will be filled in when
//! Step 8 rewrites `use_role` / `use_session` / `use_agent` to
//! actually build `ToolScope`s
//!
//! The full design (idle pool, reaper task, per-server TTL, health
//! checks, graceful shutdown) lands in **Phase 5** per
//! `docs/PHASE-5-IMPLEMENTATION-PLAN.md`. Phase 1 Step 6.5 ships just
//! enough for the type to exist on `AppState` and participate in
//! construction / test round-trips.
//!
//! The key type `McpServerKey` hashes the server name plus its full
//! command/args/env so that two scopes requesting an identically-
//! configured server share an `Arc`, while two scopes requesting
//! differently-configured servers (e.g., different API tokens) get
//! independent subprocesses. This is the sharing-vs-isolation property
//! described in `docs/REST-API-ARCHITECTURE.md` section 5.
use crate::mcp::{ConnectedServer, JsonField, McpServer, spawn_mcp_server};
use anyhow::Result;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Weak};
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct McpServerKey {
pub name: String,
pub command: String,
pub args: Vec<String>,
pub env: Vec<(String, String)>,
}
impl McpServerKey {
pub fn new(
name: impl Into<String>,
command: impl Into<String>,
args: impl IntoIterator<Item = String>,
env: impl IntoIterator<Item = (String, String)>,
) -> Self {
let mut args: Vec<String> = args.into_iter().collect();
args.sort();
let mut env: Vec<(String, String)> = env.into_iter().collect();
env.sort();
Self {
name: name.into(),
command: command.into(),
args,
env,
}
}
pub fn from_spec(name: &str, spec: &McpServer) -> Self {
let args = spec.args.clone().unwrap_or_default();
let env: Vec<(String, String)> = spec
.env
.as_ref()
.map(|e| {
e.iter()
.map(|(k, v)| {
let v_str = match v {
JsonField::Str(s) => s.clone(),
JsonField::Bool(b) => b.to_string(),
JsonField::Int(i) => i.to_string(),
};
(k.clone(), v_str)
})
.collect()
})
.unwrap_or_default();
Self::new(name, &spec.command, args, env)
}
}
#[derive(Default)]
pub struct McpFactory {
active: Mutex<HashMap<McpServerKey, Weak<ConnectedServer>>>,
}
impl McpFactory {
pub fn try_get_active(&self, key: &McpServerKey) -> Option<Arc<ConnectedServer>> {
let map = self.active.lock();
map.get(key).and_then(|weak| weak.upgrade())
}
pub fn insert_active(&self, key: McpServerKey, handle: &Arc<ConnectedServer>) {
let mut map = self.active.lock();
map.insert(key, Arc::downgrade(handle));
}
pub async fn acquire(
&self,
name: &str,
spec: &McpServer,
log_path: Option<&Path>,
) -> Result<Arc<ConnectedServer>> {
let key = McpServerKey::from_spec(name, spec);
if let Some(existing) = self.try_get_active(&key) {
return Ok(existing);
}
let handle = spawn_mcp_server(spec, log_path).await?;
self.insert_active(key, &handle);
Ok(handle)
}
}
+75 -2283
View File
File diff suppressed because it is too large Load Diff
+265
View File
@@ -0,0 +1,265 @@
//! Static path and filesystem-lookup helpers that used to live as
//! associated functions on [`Config`](super::Config).
//!
//! None of these functions depend on any `Config` instance data — they
//! compute paths from environment variables, XDG directories, or the
//! crate constant for the config root. Moving them here is Phase 1
//! Step 2 of the REST API refactor: the `Config` struct is shedding
//! anything that doesn't actually need per-instance state so the
//! eventual split into `AppConfig` + `RequestContext` has a clean
//! division line.
//!
//! # Compatibility shim during migration
//!
//! The existing associated functions on `Config` (e.g.,
//! `Config::config_dir()`) are kept as `#[deprecated]` forwarders that
//! call into this module. Callers are migrated module-by-module; when
//! the last caller is updated, the forwarders are deleted in a later
//! sub-step of Step 2.
use super::role::Role;
use super::{
AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME, ENV_FILE_NAME,
FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME, GLOBAL_TOOLS_UTILS_DIR_NAME,
MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME,
};
use crate::client::ProviderModels;
use crate::utils::{get_env_name, list_file_names, normalize_env_name};
use anyhow::{Context, Result, anyhow, bail};
use log::LevelFilter;
use std::collections::HashSet;
use std::env;
use std::fs::{read_dir, read_to_string};
use std::path::PathBuf;
pub fn config_dir() -> PathBuf {
if let Ok(v) = env::var(get_env_name("config_dir")) {
PathBuf::from(v)
} else if let Ok(v) = env::var("XDG_CONFIG_HOME") {
PathBuf::from(v).join(env!("CARGO_CRATE_NAME"))
} else {
let dir = dirs::config_dir().expect("No user's config directory");
dir.join(env!("CARGO_CRATE_NAME"))
}
}
pub fn local_path(name: &str) -> PathBuf {
config_dir().join(name)
}
pub fn cache_path() -> PathBuf {
let base_dir = dirs::cache_dir().unwrap_or_else(env::temp_dir);
base_dir.join(env!("CARGO_CRATE_NAME"))
}
pub fn oauth_tokens_path() -> PathBuf {
cache_path().join("oauth")
}
pub fn token_file(client_name: &str) -> PathBuf {
oauth_tokens_path().join(format!("{client_name}_oauth_tokens.json"))
}
pub fn log_path() -> PathBuf {
cache_path().join(format!("{}.log", env!("CARGO_CRATE_NAME")))
}
pub fn config_file() -> PathBuf {
match env::var(get_env_name("config_file")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(CONFIG_FILE_NAME),
}
}
pub fn roles_dir() -> PathBuf {
match env::var(get_env_name("roles_dir")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(ROLES_DIR_NAME),
}
}
pub fn role_file(name: &str) -> PathBuf {
roles_dir().join(format!("{name}.md"))
}
pub fn macros_dir() -> PathBuf {
match env::var(get_env_name("macros_dir")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(MACROS_DIR_NAME),
}
}
pub fn macro_file(name: &str) -> PathBuf {
macros_dir().join(format!("{name}.yaml"))
}
pub fn env_file() -> PathBuf {
match env::var(get_env_name("env_file")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(ENV_FILE_NAME),
}
}
pub fn rags_dir() -> PathBuf {
match env::var(get_env_name("rags_dir")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(RAGS_DIR_NAME),
}
}
pub fn functions_dir() -> PathBuf {
match env::var(get_env_name("functions_dir")) {
Ok(value) => PathBuf::from(value),
Err(_) => local_path(FUNCTIONS_DIR_NAME),
}
}
pub fn functions_bin_dir() -> PathBuf {
functions_dir().join(FUNCTIONS_BIN_DIR_NAME)
}
pub fn mcp_config_file() -> PathBuf {
functions_dir().join(MCP_FILE_NAME)
}
pub fn global_tools_dir() -> PathBuf {
functions_dir().join(GLOBAL_TOOLS_DIR_NAME)
}
pub fn global_utils_dir() -> PathBuf {
functions_dir().join(GLOBAL_TOOLS_UTILS_DIR_NAME)
}
pub fn bash_prompt_utils_file() -> PathBuf {
global_utils_dir().join(BASH_PROMPT_UTILS_FILE_NAME)
}
pub fn agents_data_dir() -> PathBuf {
local_path(AGENTS_DIR_NAME)
}
pub fn agent_data_dir(name: &str) -> PathBuf {
match env::var(format!("{}_DATA_DIR", normalize_env_name(name))) {
Ok(value) => PathBuf::from(value),
Err(_) => agents_data_dir().join(name),
}
}
pub fn agent_config_file(name: &str) -> PathBuf {
match env::var(format!("{}_CONFIG_FILE", normalize_env_name(name))) {
Ok(value) => PathBuf::from(value),
Err(_) => agent_data_dir(name).join(CONFIG_FILE_NAME),
}
}
pub fn agent_bin_dir(name: &str) -> PathBuf {
agent_data_dir(name).join(FUNCTIONS_BIN_DIR_NAME)
}
pub fn agent_rag_file(agent_name: &str, rag_name: &str) -> PathBuf {
agent_data_dir(agent_name).join(format!("{rag_name}.yaml"))
}
pub fn agent_functions_file(name: &str) -> Result<PathBuf> {
let priority = ["tools.sh", "tools.py", "tools.ts", "tools.js"];
let dir = agent_data_dir(name);
for filename in priority {
let path = dir.join(filename);
if path.exists() {
return Ok(path);
}
}
Err(anyhow!(
"No tools script found in agent functions directory"
))
}
pub fn models_override_file() -> PathBuf {
local_path("models-override.yaml")
}
pub fn log_config() -> Result<(LevelFilter, Option<PathBuf>)> {
let log_level = env::var(get_env_name("log_level"))
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(match cfg!(debug_assertions) {
true => LevelFilter::Debug,
false => LevelFilter::Info,
});
let resolved_log_path = match env::var(get_env_name("log_path")) {
Ok(v) => Some(PathBuf::from(v)),
Err(_) => Some(log_path()),
};
Ok((log_level, resolved_log_path))
}
pub fn list_roles(with_builtin: bool) -> Vec<String> {
let mut names = HashSet::new();
if let Ok(rd) = read_dir(roles_dir()) {
for entry in rd.flatten() {
if let Some(name) = entry
.file_name()
.to_str()
.and_then(|v| v.strip_suffix(".md"))
{
names.insert(name.to_string());
}
}
}
if with_builtin {
names.extend(Role::list_builtin_role_names());
}
let mut names: Vec<_> = names.into_iter().collect();
names.sort_unstable();
names
}
pub fn has_role(name: &str) -> bool {
let names = list_roles(true);
names.contains(&name.to_string())
}
pub fn list_rags() -> Vec<String> {
match read_dir(rags_dir()) {
Ok(rd) => {
let mut names = vec![];
for entry in rd.flatten() {
let name = entry.file_name();
if let Some(name) = name.to_string_lossy().strip_suffix(".yaml") {
names.push(name.to_string());
}
}
names.sort_unstable();
names
}
Err(_) => vec![],
}
}
pub fn list_macros() -> Vec<String> {
list_file_names(macros_dir(), ".yaml")
}
pub fn has_macro(name: &str) -> bool {
let names = list_macros();
names.contains(&name.to_string())
}
pub fn local_models_override() -> Result<Vec<ProviderModels>> {
let model_override_path = models_override_file();
let err = || {
format!(
"Failed to load models at '{}'",
model_override_path.display()
)
};
let content = read_to_string(&model_override_path).with_context(err)?;
let models_override: ModelsOverride = serde_yaml::from_str(&content).with_context(err)?;
if models_override.version != env!("CARGO_PKG_VERSION") {
bail!("Incompatible version")
}
Ok(models_override.list)
}
+74
View File
@@ -0,0 +1,74 @@
//! Per-process RAG instance cache with weak-reference sharing.
//!
//! `RagCache` lives on [`AppState`](super::AppState) and serves both
//! standalone RAGs (attached via `.rag <name>`) and agent-owned RAGs
//! (loaded from an agent's `documents:` field). The cache keys with
//! [`RagKey`] so that agent RAGs and standalone RAGs occupy distinct
//! namespaces even if they share a name.
//!
//! Entries are held as `Weak<Rag>` so the cache never keeps a RAG
//! alive on its own — once all active scopes drop their `Arc<Rag>`,
//! the cache entry becomes unupgradable and the next `load()` falls
//! through to a fresh disk read.
//!
//! # Phase 1 Step 6.5 scope
//!
//! This file introduces the type scaffolding. Actual cache population
//! (i.e., routing `use_rag`, `use_agent`, and sub-agent spawning
//! through the cache) is deferred to Step 8 when the entry points get
//! rewritten. During the bridge window, `Config.rag` keeps serving
//! today's callers via direct `Rag::load` / `Rag::init` calls and
//! `RagCache` sits on `AppState` as an unused-but-ready service.
//!
//! See `docs/REST-API-ARCHITECTURE.md` section 5 ("RAG Cache") for
//! the full design including concurrent first-load serialization and
//! invalidation semantics.
use crate::rag::Rag;
use anyhow::Result;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::{Arc, Weak};
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum RagKey {
Named(String),
Agent(String),
}
#[derive(Default)]
pub struct RagCache {
entries: RwLock<HashMap<RagKey, Weak<Rag>>>,
}
impl RagCache {
pub fn try_get(&self, key: &RagKey) -> Option<Arc<Rag>> {
let map = self.entries.read();
map.get(key).and_then(|weak| weak.upgrade())
}
pub fn insert(&self, key: RagKey, rag: &Arc<Rag>) {
let mut map = self.entries.write();
map.insert(key, Arc::downgrade(rag));
}
pub fn invalidate(&self, key: &RagKey) {
let mut map = self.entries.write();
map.remove(key);
}
pub async fn load_with<F, Fut>(&self, key: RagKey, loader: F) -> Result<Arc<Rag>>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Rag>>,
{
if let Some(existing) = self.try_get(&key) {
return Ok(existing);
}
let rag = loader().await?;
let arc = Arc::new(rag);
self.insert(key, &arc);
Ok(arc)
}
}
File diff suppressed because it is too large Load Diff
+94
View File
@@ -374,6 +374,100 @@ fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) {
mod tests {
use super::*;
#[test]
fn role_new_parses_prompt() {
let role = Role::new("test", "You are a helpful assistant");
assert_eq!(role.name(), "test");
assert_eq!(role.prompt(), "You are a helpful assistant");
}
#[test]
fn role_new_parses_metadata() {
let content =
"---\nmodel: openai:gpt-4\ntemperature: 0.7\ntop_p: 0.9\n---\nYou are helpful";
let role = Role::new("test", content);
assert_eq!(role.model_id(), Some("openai:gpt-4"));
assert_eq!(role.temperature(), Some(0.7));
assert_eq!(role.top_p(), Some(0.9));
assert_eq!(role.prompt(), "You are helpful");
}
#[test]
fn role_new_parses_enabled_tools() {
let content = "---\nenabled_tools: tool1,tool2\n---\nPrompt";
let role = Role::new("test", content);
assert_eq!(role.enabled_tools(), Some("tool1,tool2".to_string()));
}
#[test]
fn role_new_parses_enabled_mcp_servers() {
let content = "---\nenabled_mcp_servers: github,jira\n---\nPrompt";
let role = Role::new("test", content);
assert_eq!(role.enabled_mcp_servers(), Some("github,jira".to_string()));
}
#[test]
fn role_new_no_metadata_has_none_fields() {
let role = Role::new("test", "Just a prompt");
assert_eq!(role.model_id(), None);
assert_eq!(role.temperature(), None);
assert_eq!(role.top_p(), None);
assert_eq!(role.enabled_tools(), None);
assert_eq!(role.enabled_mcp_servers(), None);
}
#[test]
fn role_builtin_shell_loads() {
let role = Role::builtin("shell").unwrap();
assert_eq!(role.name(), "shell");
assert!(!role.prompt().is_empty());
}
#[test]
fn role_builtin_code_loads() {
let role = Role::builtin("code").unwrap();
assert_eq!(role.name(), "code");
assert!(!role.prompt().is_empty());
}
#[test]
fn role_builtin_nonexistent_errors() {
let result = Role::builtin("nonexistent_role_xyz");
assert!(result.is_err());
}
#[test]
fn role_default_has_empty_fields() {
let role = Role::default();
assert_eq!(role.name(), "");
assert_eq!(role.prompt(), "");
assert_eq!(role.model_id(), None);
}
#[test]
fn role_set_model_updates_model() {
let mut role = Role::new("test", "prompt");
let model = Model::default();
role.set_model(model.clone());
assert_eq!(role.model().id(), model.id());
}
#[test]
fn role_set_temperature_works() {
let mut role = Role::new("test", "prompt");
role.set_temperature(Some(0.5));
assert_eq!(role.temperature(), Some(0.5));
}
#[test]
fn role_export_includes_metadata() {
let content = "---\ntemperature: 0.8\n---\nMy prompt";
let role = Role::new("test", content);
let exported = role.export();
assert!(exported.contains("temperature"));
assert!(exported.contains("My prompt"));
}
#[test]
fn test_parse_structure_prompt1() {
let prompt = r#"
+180 -6
View File
@@ -67,11 +67,11 @@ pub struct Session {
}
impl Session {
pub fn new(config: &Config, name: &str) -> Self {
let role = config.extract_role();
pub fn new_from_ctx(ctx: &RequestContext, app: &AppConfig, name: &str) -> Self {
let role = ctx.extract_role(app);
let mut session = Self {
name: name.to_string(),
save_session: config.save_session,
save_session: app.save_session,
..Default::default()
};
session.set_role(role);
@@ -79,13 +79,18 @@ impl Session {
session
}
pub fn load(config: &Config, name: &str, path: &Path) -> Result<Self> {
pub fn load_from_ctx(
ctx: &RequestContext,
app: &AppConfig,
name: &str,
path: &Path,
) -> Result<Self> {
let content = read_to_string(path)
.with_context(|| format!("Failed to load session {} at {}", name, path.display()))?;
let mut session: Self =
serde_yaml::from_str(&content).with_context(|| format!("Invalid session {name}"))?;
session.model = Model::retrieve_model(config, &session.model_id, ModelType::Chat)?;
session.model = Model::retrieve_model(app, &session.model_id, ModelType::Chat)?;
if let Some(autoname) = name.strip_prefix("_/") {
session.name = TEMP_SESSION_NAME.to_string();
@@ -99,7 +104,7 @@ impl Session {
}
if let Some(role_name) = &session.role_name
&& let Ok(role) = config.retrieve_role(role_name)
&& let Ok(role) = ctx.retrieve_role(app, role_name)
{
session.role_prompt = role.prompt().to_string();
}
@@ -664,3 +669,172 @@ impl AutoName {
!self.naming && self.chat_history.is_some() && self.name.is_none()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::{Message, MessageContent, MessageRole, Model};
use crate::config::{AppState, Config};
use std::sync::Arc;
#[test]
fn session_default_is_empty() {
let session = Session::default();
assert!(session.is_empty());
assert_eq!(session.name(), "");
assert_eq!(session.role_name(), None);
assert!(!session.dirty());
}
#[test]
fn session_new_from_ctx_captures_save_session() {
let cfg = Config::default();
let app_config = Arc::new(cfg.to_app_config());
let app_state = Arc::new(AppState {
config: app_config.clone(),
vault: cfg.vault.clone(),
mcp_factory: Arc::new(mcp_factory::McpFactory::default()),
rag_cache: Arc::new(rag_cache::RagCache::default()),
mcp_config: None,
mcp_log_path: None,
});
let ctx = cfg.to_request_context(app_state);
let session = Session::new_from_ctx(&ctx, &app_config, "test-session");
assert_eq!(session.name(), "test-session");
assert_eq!(session.save_session(), app_config.save_session);
assert!(session.is_empty());
assert!(!session.dirty());
}
#[test]
fn session_set_role_captures_role_info() {
let mut session = Session::default();
let content = "---\ntemperature: 0.5\n---\nYou are a coder";
let mut role = Role::new("coder", content);
role.set_model(Model::default());
session.set_role(role);
assert_eq!(session.role_name(), Some("coder"));
assert_eq!(session.temperature(), Some(0.5));
assert!(session.dirty());
}
#[test]
fn session_clear_role() {
let mut session = Session::default();
let mut role = Role::new("test", "prompt");
role.set_model(Model::default());
session.set_role(role);
assert_eq!(session.role_name(), Some("test"));
session.clear_role();
assert_eq!(session.role_name(), None);
}
#[test]
fn session_guard_empty_passes_when_empty() {
let session = Session::default();
assert!(session.guard_empty().is_ok());
}
#[test]
fn session_needs_compression_threshold() {
let session = Session::default();
assert!(!session.needs_compression(4000));
}
#[test]
fn session_needs_compression_returns_false_when_compressing() {
let mut session = Session::default();
session.set_compressing(true);
assert!(!session.needs_compression(0));
}
#[test]
fn session_needs_compression_returns_false_when_threshold_zero() {
let session = Session::default();
assert!(!session.needs_compression(0));
}
#[test]
fn session_set_compressing_flag() {
let mut session = Session::default();
assert!(!session.compressing());
session.set_compressing(true);
assert!(session.compressing());
session.set_compressing(false);
assert!(!session.compressing());
}
#[test]
fn session_set_save_session_this_time() {
let mut session = Session::default();
assert!(!session.save_session_this_time);
session.set_save_session_this_time();
assert!(session.save_session_this_time);
}
#[test]
fn session_save_session_returns_configured_value() {
let mut session = Session::default();
assert_eq!(session.save_session(), None);
session.set_save_session(Some(true));
assert_eq!(session.save_session(), Some(true));
session.set_save_session(Some(false));
assert_eq!(session.save_session(), Some(false));
session.set_save_session(None);
assert_eq!(session.save_session(), None);
}
#[test]
fn session_compress_moves_messages() {
let mut session = Session::default();
session.messages.push(Message::new(
MessageRole::System,
MessageContent::Text("system prompt".to_string()),
));
session.messages.push(Message::new(
MessageRole::User,
MessageContent::Text("hello".to_string()),
));
assert_eq!(session.messages.len(), 2);
assert!(session.compressed_messages.is_empty());
session.compress("Summary of conversation".to_string());
assert!(!session.compressed_messages.is_empty());
assert_eq!(session.messages.len(), 1);
assert!(session.dirty());
}
#[test]
fn session_is_not_empty_after_compress() {
let mut session = Session::default();
session.messages.push(Message::new(
MessageRole::User,
MessageContent::Text("hello".to_string()),
));
session.compress("Summary".to_string());
assert!(!session.is_empty());
}
#[test]
fn session_need_autoname_default_false() {
let session = Session::default();
assert!(!session.need_autoname());
}
#[test]
fn session_set_autonaming_doesnt_panic_without_autoname() {
let mut session = Session::default();
session.set_autonaming(true);
assert!(!session.need_autoname());
}
}
+180
View File
@@ -0,0 +1,180 @@
//! Per-scope tool runtime: resolved functions + live MCP handles +
//! call tracker.
//!
//! `ToolScope` is the unit of tool availability for a single request.
//! Every active `RoleLike` (role, session, agent) conceptually owns one.
//! The contents are:
//!
//! * `functions` — the `Functions` declarations visible to the LLM for
//! this scope (global tools + role/session/agent filters applied)
//! * `mcp_runtime` — live MCP subprocess handles for the servers this
//! scope has enabled, keyed by server name
//! * `tool_tracker` — per-scope tool call history for auto-continuation
//! and looping detection
//!
//! # Phase 1 Step 6.5 scope
//!
//! This file introduces the type scaffolding. Scope transitions
//! (`use_role`, `use_session`, `use_agent`, `exit_*`) that actually
//! build and swap `ToolScope` instances are deferred to Step 8 when
//! the entry points (`main.rs`, `repl/mod.rs`) get rewritten to thread
//! `RequestContext` through the pipeline. During the bridge window,
//! `Config.functions` / `Config.mcp_registry` keep serving today's
//! callers and `ToolScope` sits alongside them on `RequestContext` as
//! an unused (but compiling and tested) parallel structure.
//!
//! The fields mirror the plan in `docs/REST-API-ARCHITECTURE.md`
//! section 5 and `docs/PHASE-1-IMPLEMENTATION-PLAN.md` Step 6.5.
use crate::function::{Functions, ToolCallTracker};
use crate::mcp::{CatalogItem, ConnectedServer, McpRegistry};
use anyhow::{Context, Result, anyhow};
use bm25::{Document, Language, SearchEngineBuilder};
use rmcp::model::{CallToolRequestParams, CallToolResult};
use serde_json::{Value, json};
use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
pub struct ToolScope {
pub functions: Functions,
pub mcp_runtime: McpRuntime,
pub tool_tracker: ToolCallTracker,
}
impl Default for ToolScope {
fn default() -> Self {
Self {
functions: Functions::default(),
mcp_runtime: McpRuntime::default(),
tool_tracker: ToolCallTracker::default(),
}
}
}
#[derive(Default)]
pub struct McpRuntime {
pub servers: HashMap<String, Arc<ConnectedServer>>,
}
impl McpRuntime {
pub fn new() -> Self {
Self::default()
}
pub fn is_empty(&self) -> bool {
self.servers.is_empty()
}
pub fn insert(&mut self, name: String, handle: Arc<ConnectedServer>) {
self.servers.insert(name, handle);
}
pub fn get(&self, name: &str) -> Option<&Arc<ConnectedServer>> {
self.servers.get(name)
}
pub fn server_names(&self) -> Vec<String> {
self.servers.keys().cloned().collect()
}
pub fn sync_from_registry(&mut self, registry: &McpRegistry) {
self.servers.clear();
for (name, handle) in registry.running_servers() {
self.servers.insert(name.clone(), Arc::clone(handle));
}
}
async fn catalog_items(&self, server: &str) -> Result<HashMap<String, CatalogItem>> {
let server_handle = self
.get(server)
.cloned()
.with_context(|| format!("{server} MCP server not found in runtime"))?;
let tools = server_handle.list_tools(None).await?;
let mut items = HashMap::new();
for tool in tools.tools {
let item = CatalogItem {
name: tool.name.to_string(),
server: server.to_string(),
description: tool.description.unwrap_or_default().to_string(),
};
items.insert(item.name.clone(), item);
}
Ok(items)
}
pub async fn search(
&self,
server: &str,
query: &str,
top_k: usize,
) -> Result<Vec<CatalogItem>> {
let items = self.catalog_items(server).await?;
let docs = items.values().map(|item| Document {
id: item.name.clone(),
contents: format!(
"{}\n{}\nserver:{}",
item.name, item.description, item.server
),
});
let engine = SearchEngineBuilder::<String>::with_documents(Language::English, docs).build();
Ok(engine
.search(query, top_k.min(20))
.into_iter()
.filter_map(|result| items.get(&result.document.id))
.take(top_k)
.cloned()
.collect())
}
pub async fn describe(&self, server: &str, tool: &str) -> Result<Value> {
let server_handle = self
.get(server)
.cloned()
.with_context(|| format!("{server} MCP server not found in runtime"))?;
let tool_schema = server_handle
.list_tools(None)
.await?
.tools
.into_iter()
.find(|item| item.name == tool)
.ok_or_else(|| anyhow!("{tool} not found in {server} MCP server catalog"))?
.input_schema;
Ok(json!({
"type": "object",
"properties": {
"tool": {
"type": "string",
},
"arguments": tool_schema
}
}))
}
pub async fn invoke(
&self,
server: &str,
tool: &str,
arguments: Value,
) -> Result<CallToolResult> {
let server_handle = self
.get(server)
.cloned()
.with_context(|| format!("Invoked MCP server does not exist: {server}"))?;
let request = CallToolRequestParams {
name: Cow::Owned(tool.to_owned()),
arguments: arguments.as_object().cloned(),
meta: None,
task: None,
};
server_handle.call_tool(request).await.map_err(Into::into)
}
}