//! Per-request mutable state for a single Loki interaction. //! //! `RequestContext` holds everything that was formerly the runtime //! (`#[serde(skip)]`) half of [`Config`](super::Config): the active role, //! session, agent, RAG, supervisor state, inbox/escalation queues, and //! the conversation's "last message" cursor. //! //! Each frontend constructs and owns a `RequestContext`: //! //! * **CLI** — one `RequestContext` per invocation, dropped at exit. //! * **REPL** — one long-lived `RequestContext` mutated across turns. //! * **API** — one `RequestContext` per HTTP request, hydrated from a //! persisted session and written back at the end. //! //! # Tool scope and agent runtime (planned) //! //! # Flat fields and sub-struct fields coexist during the bridge //! //! The flat fields (`functions`, `tool_call_tracker`, `supervisor`, //! `inbox`, `root_escalation_queue`, `self_agent_id`, `current_depth`, //! `parent_supervisor`) mirror the runtime half of today's `Config` //! and are populated by //! [`Config::to_request_context`](super::Config::to_request_context) //! during the bridge. //! //! Step 6.5 added two **sub-struct fields** alongside the flat ones: //! //! * [`tool_scope: ToolScope`](super::tool_scope::ToolScope) — the //! planned home for `functions`, `mcp_runtime`, and `tool_tracker` //! * [`agent_runtime: Option`](super::agent_runtime::AgentRuntime) — //! the planned home for `supervisor`, `inbox`, `escalation_queue`, //! `todo_list`, `self_agent_id`, `parent_supervisor`, and `depth` //! //! During the bridge window the sub-struct fields are **additive //! scaffolding**: they're initialized to defaults in //! [`RequestContext::new`] and stay empty until Step 8 rewrites the //! entry points to build them explicitly. The flat fields continue //! to carry the actual state from `Config` during the bridge. //! //! # Phase 1 refactor history //! //! * **Step 0** — struct introduced alongside `Config` //! * **Step 5** — added 13 request-read methods //! * **Step 6** — added 12 request-write methods //! * **Step 6.5** — added `tool_scope` and `agent_runtime` sub-struct //! fields as additive scaffolding //! * **Step 7** — added 14 mixed-method splits that take `&AppConfig` //! as a parameter for the serialized half //! //! See `docs/PHASE-1-IMPLEMENTATION-PLAN.md` for the full migration //! plan. use super::MessageContentToolCalls; use super::rag_cache::{RagCache, RagKey}; use super::session::Session; use super::todo::TodoList; use super::tool_scope::{McpRuntime, ToolScope}; use super::{ AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, CREATE_TITLE_ROLE, Input, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role, RoleLike, SESSIONS_DIR_NAME, SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags, TEMP_ROLE_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, paths, }; use crate::client::{Model, ModelType, list_models}; use crate::function::{ FunctionDeclaration, Functions, ToolCallTracker, ToolResult, user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, }; use crate::rag::Rag; use crate::supervisor::Supervisor; use crate::supervisor::escalation::EscalationQueue; use crate::supervisor::mailbox::Inbox; use crate::utils::{ AbortSignal, abortable_run_with_spinner, edit_file, fuzzy_filter, get_env_name, list_file_names, now, render_prompt, temp_file, }; use anyhow::{Context, Error, Result, bail}; use indoc::formatdoc; use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation}; use parking_lot::RwLock; use std::collections::{HashMap, HashSet}; use std::env; use std::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file}; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; pub struct RequestContext { pub app: Arc, pub macro_flag: bool, pub info_flag: bool, pub working_mode: WorkingMode, pub model: Model, pub agent_variables: Option, pub role: Option, pub session: Option, pub rag: Option>, pub agent: Option, pub last_message: Option, pub tool_scope: ToolScope, pub supervisor: Option>>, pub parent_supervisor: Option>>, pub self_agent_id: Option, pub inbox: Option>, pub escalation_queue: Option>, pub current_depth: usize, pub auto_continue_count: usize, pub todo_list: TodoList, pub last_continuation_response: Option, } impl RequestContext { pub fn new(app: Arc, working_mode: WorkingMode) -> Self { Self { app, macro_flag: false, info_flag: false, working_mode, model: Default::default(), agent_variables: None, role: None, session: None, rag: None, agent: None, last_message: None, tool_scope: ToolScope::default(), supervisor: None, parent_supervisor: None, self_agent_id: None, inbox: None, escalation_queue: None, current_depth: 0, auto_continue_count: 0, todo_list: TodoList::default(), last_continuation_response: None, } } pub fn new_for_child( app: Arc, parent: &Self, current_depth: usize, inbox: Arc, self_agent_id: String, ) -> Self { let tool_call_tracker = ToolCallTracker::new(4, 10); Self { app, macro_flag: parent.macro_flag, info_flag: parent.info_flag, working_mode: WorkingMode::Cmd, model: parent.model.clone(), agent_variables: parent.agent_variables.clone(), role: None, session: None, rag: None, agent: None, last_message: None, tool_scope: ToolScope { functions: Functions::default(), mcp_runtime: McpRuntime::default(), tool_tracker: tool_call_tracker, }, supervisor: None, parent_supervisor: parent.supervisor.clone(), self_agent_id: Some(self_agent_id), inbox: Some(inbox), escalation_queue: parent.escalation_queue.clone(), current_depth, auto_continue_count: 0, todo_list: TodoList::default(), last_continuation_response: None, } } fn update_app_config(&mut self, update: impl FnOnce(&mut AppConfig)) { let mut app_config = (*self.app.config).clone(); update(&mut app_config); let mut app_state = (*self.app).clone(); app_state.config = Arc::new(app_config); self.app = Arc::new(app_state); } pub fn root_escalation_queue(&self) -> Option<&Arc> { self.escalation_queue.as_ref() } pub fn ensure_root_escalation_queue(&mut self) -> Arc { self.escalation_queue .get_or_insert_with(|| Arc::new(EscalationQueue::new())) .clone() } pub fn rag_cache(&self) -> &Arc { &self.app.rag_cache } 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.auto_continue_count = 0; } pub fn increment_auto_continue_count(&mut self) { self.auto_continue_count += 1; } pub fn reset_continuation_count(&mut self) { self.auto_continue_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 state(&self) -> StateFlags { let mut flags = StateFlags::empty(); if let Some(session) = &self.session { if session.is_empty() { flags |= StateFlags::SESSION_EMPTY; } else { flags |= StateFlags::SESSION; } if session.role_name().is_some() { flags |= StateFlags::ROLE; } } else if self.role.is_some() { flags |= StateFlags::ROLE; } if self.agent.is_some() { flags |= StateFlags::AGENT; } if self.rag.is_some() { flags |= StateFlags::RAG; } flags } pub fn messages_file(&self) -> PathBuf { match &self.agent { None => match env::var(get_env_name("messages_file")) { Ok(value) => PathBuf::from(value), Err(_) => paths::cache_path().join(MESSAGES_FILE_NAME), }, Some(agent) => paths::cache_path() .join(AGENTS_DIR_NAME) .join(agent.name()) .join(MESSAGES_FILE_NAME), } } pub fn sessions_dir(&self) -> PathBuf { match &self.agent { None => match env::var(get_env_name("sessions_dir")) { Ok(value) => PathBuf::from(value), Err(_) => paths::local_path(SESSIONS_DIR_NAME), }, Some(agent) => paths::agent_data_dir(agent.name()).join(SESSIONS_DIR_NAME), } } pub fn session_file(&self, name: &str) -> PathBuf { match name.split_once("/") { Some((dir, name)) => self.sessions_dir().join(dir).join(format!("{name}.yaml")), None => self.sessions_dir().join(format!("{name}.yaml")), } } pub fn rag_file(&self, name: &str) -> PathBuf { match &self.agent { Some(agent) => paths::agent_rag_file(agent.name(), name), None => paths::rags_dir().join(format!("{name}.yaml")), } } pub fn role_info(&self) -> Result { if let Some(session) = &self.session { if session.role_name().is_some() { let role = session.to_role(); Ok(role.export()) } else { bail!("No session role") } } else if let Some(role) = &self.role { Ok(role.export()) } else { bail!("No role") } } pub fn agent_info(&self) -> Result { if let Some(agent) = &self.agent { agent.export() } else { bail!("No agent") } } pub fn agent_banner(&self) -> Result { if let Some(agent) = &self.agent { Ok(agent.banner()) } else { bail!("No agent") } } pub fn rag_info(&self) -> Result { if let Some(rag) = &self.rag { rag.export() } else { bail!("No RAG") } } pub fn list_sessions(&self) -> Vec { list_file_names(self.sessions_dir(), ".yaml") } pub fn list_autoname_sessions(&self) -> Vec { list_file_names(self.sessions_dir().join("_"), ".yaml") } pub fn is_compressing_session(&self) -> bool { self.session .as_ref() .map(|v| v.compressing()) .unwrap_or_default() } pub fn role_like_mut(&mut self) -> Option<&mut dyn RoleLike> { if let Some(session) = self.session.as_mut() { Some(session) } else if let Some(agent) = self.agent.as_mut() { Some(agent) } else if let Some(role) = self.role.as_mut() { Some(role) } else { None } } pub fn use_role_obj(&mut self, role: Role) -> Result<()> { if self.agent.is_some() { bail!("Cannot perform this operation because you are using a agent") } if let Some(session) = self.session.as_mut() { session.guard_empty()?; session.set_role(role); } else { self.role = Some(role); } Ok(()) } pub fn exit_role(&mut self) -> Result<()> { if let Some(session) = self.session.as_mut() { session.guard_empty()?; session.clear_role(); } else if self.role.is_some() { self.role = None; } Ok(()) } pub fn exit_session(&mut self) -> Result<()> { if let Some(mut session) = self.session.take() { let sessions_dir = self.sessions_dir(); session.exit(&sessions_dir, self.working_mode.is_repl())?; self.discontinuous_last_message(); } Ok(()) } pub fn save_session(&mut self, name: Option<&str>) -> Result<()> { let session_name = match &self.session { Some(session) => match name { Some(v) => v.to_string(), None => session .autoname() .unwrap_or_else(|| session.name()) .to_string(), }, None => bail!("No session"), }; let session_path = self.session_file(&session_name); if let Some(session) = self.session.as_mut() { session.save(&session_name, &session_path, self.working_mode.is_repl())?; } Ok(()) } pub fn empty_session(&mut self) -> Result<()> { if let Some(session) = self.session.as_mut() { if let Some(agent) = self.agent.as_ref() { session.sync_agent(agent); } session.clear_messages(); } else { bail!("No session") } self.discontinuous_last_message(); Ok(()) } pub fn set_save_session_this_time(&mut self) -> Result<()> { if let Some(session) = self.session.as_mut() { session.set_save_session_this_time(); } else { bail!("No session") } Ok(()) } pub fn exit_rag(&mut self) -> Result<()> { self.rag.take(); Ok(()) } pub fn exit_agent_session(&mut self) -> Result<()> { self.exit_session()?; if let Some(agent) = self.agent.as_mut() { agent.exit_session(); if self.working_mode.is_repl() { self.init_agent_shared_variables()?; } } Ok(()) } pub fn before_chat_completion(&mut self, input: &Input) -> Result<()> { self.last_message = Some(LastMessage::new(input.clone(), String::new())); Ok(()) } pub fn discontinuous_last_message(&mut self) { if let Some(last_message) = self.last_message.as_mut() { last_message.continuous = false; } } pub fn init_agent_shared_variables(&mut self) -> Result<()> { let agent = match self.agent.as_mut() { Some(v) => v, None => return Ok(()), }; if !agent.defined_variables().is_empty() && agent.shared_variables().is_empty() { let new_variables = Agent::init_agent_variables( agent.defined_variables(), self.agent_variables.as_ref(), self.info_flag, )?; agent.set_shared_variables(new_variables); } if !self.info_flag { agent.update_shared_dynamic_instructions(false)?; } Ok(()) } pub fn init_agent_session_variables(&mut self, new_session: bool) -> Result<()> { let (agent, session) = match (self.agent.as_mut(), self.session.as_mut()) { (Some(agent), Some(session)) => (agent, session), _ => return Ok(()), }; if new_session { let shared_variables = agent.shared_variables().clone(); let session_variables = if !agent.defined_variables().is_empty() && shared_variables.is_empty() { let new_variables = Agent::init_agent_variables( agent.defined_variables(), self.agent_variables.as_ref(), self.info_flag, )?; agent.set_shared_variables(new_variables.clone()); new_variables } else { shared_variables }; agent.set_session_variables(session_variables); if !self.info_flag { agent.update_session_dynamic_instructions(None)?; } session.sync_agent(agent); } else { let variables = session.agent_variables(); agent.set_session_variables(variables.clone()); agent.update_session_dynamic_instructions(Some( session.agent_instructions().to_string(), ))?; } Ok(()) } pub fn current_model(&self) -> &Model { if let Some(session) = self.session.as_ref() { session.model() } else if let Some(agent) = self.agent.as_ref() { agent.model() } else if let Some(role) = self.role.as_ref() { role.model() } else { &self.model } } pub fn extract_role(&self, app: &AppConfig) -> Role { if let Some(session) = self.session.as_ref() { session.to_role() } else if let Some(agent) = self.agent.as_ref() { agent.to_role() } else if let Some(role) = self.role.as_ref() { role.clone() } else { let mut role = Role::default(); role.batch_set( &self.model, app.temperature, app.top_p, app.enabled_tools.clone(), app.enabled_mcp_servers.clone(), ); role } } pub fn set_temperature_on_role_like(&mut self, value: Option) -> bool { match self.role_like_mut() { Some(role_like) => { role_like.set_temperature(value); true } None => false, } } pub fn set_top_p_on_role_like(&mut self, value: Option) -> bool { match self.role_like_mut() { Some(role_like) => { role_like.set_top_p(value); true } None => false, } } pub fn set_enabled_tools_on_role_like(&mut self, value: Option) -> bool { match self.role_like_mut() { Some(role_like) => { role_like.set_enabled_tools(value); true } None => false, } } pub fn set_enabled_mcp_servers_on_role_like(&mut self, value: Option) -> bool { match self.role_like_mut() { Some(role_like) => { role_like.set_enabled_mcp_servers(value); true } None => false, } } pub fn set_save_session_on_session(&mut self, value: Option) -> bool { match self.session.as_mut() { Some(session) => { session.set_save_session(value); true } None => false, } } pub fn set_compression_threshold_on_session(&mut self, value: Option) -> bool { match self.session.as_mut() { Some(session) => { session.set_compression_threshold(value); true } None => false, } } pub fn set_max_output_tokens_on_role_like(&mut self, value: Option) -> bool { match self.role_like_mut() { Some(role_like) => { let mut model = role_like.model().clone(); model.set_max_tokens(value, true); role_like.set_model(model); true } None => false, } } pub fn save_message(&mut self, app: &AppConfig, input: &Input, output: &str) -> Result<()> { let mut input = input.clone(); input.clear_patch(); if let Some(session) = input.session_mut(&mut self.session) { session.add_message(&input, output)?; return Ok(()); } if !app.save { return Ok(()); } let mut file = self.open_message_file()?; if output.is_empty() && input.tool_calls().is_none() { return Ok(()); } let now = now(); let summary = input.summary(); let raw_input = input.raw(); let scope = if self.agent.is_none() { let role_name = if input.role().is_derived() { None } else { Some(input.role().name()) }; match (role_name, input.rag_name()) { (Some(role), Some(rag_name)) => format!(" ({role}#{rag_name})"), (Some(role), _) => format!(" ({role})"), (None, Some(rag_name)) => format!(" (#{rag_name})"), _ => String::new(), } } else { String::new() }; let tool_calls = match input.tool_calls() { Some(MessageContentToolCalls { tool_results, text, .. }) => { let mut lines = vec!["".to_string()]; if !text.is_empty() { lines.push(text.clone()); } lines.push(serde_json::to_string(&tool_results).unwrap_or_default()); lines.push("\n".to_string()); lines.join("\n") } None => String::new(), }; let output = format!( "# CHAT: {summary} [{now}]{scope}\n{raw_input}\n--------\n{tool_calls}{output}\n--------\n\n", ); file.write_all(output.as_bytes()) .with_context(|| "Failed to save message") } fn open_message_file(&self) -> Result { let path = self.messages_file(); ensure_parent_exists(&path)?; OpenOptions::new() .create(true) .append(true) .open(&path) .with_context(|| format!("Failed to create/append {}", path.display())) } pub fn after_chat_completion( &mut self, app: &AppConfig, input: &Input, output: &str, tool_results: &[ToolResult], ) -> Result<()> { if !tool_results.is_empty() { return Ok(()); } self.last_message = Some(LastMessage::new(input.clone(), output.to_string())); if !app.dry_run { self.save_message(app, input, output)?; } Ok(()) } pub fn sysinfo(&self, app: &AppConfig) -> Result { let display_path = |path: &Path| path.display().to_string(); let wrap = app .wrap .clone() .map_or_else(|| String::from("no"), |v| v.to_string()); let (rag_reranker_model, rag_top_k) = match &self.rag { Some(rag) => rag.get_config(), None => (app.rag_reranker_model.clone(), app.rag_top_k), }; let role = self.extract_role(app); let mut items = vec![ ("model", role.model().id()), ( "temperature", super::format_option_value(&role.temperature()), ), ("top_p", super::format_option_value(&role.top_p())), ( "enabled_tools", super::format_option_value(&role.enabled_tools()), ), ( "enabled_mcp_servers", super::format_option_value(&role.enabled_mcp_servers()), ), ( "max_output_tokens", role.model() .max_tokens_param() .map(|v| format!("{v} (current model)")) .unwrap_or_else(|| "null".into()), ), ( "save_session", super::format_option_value(&app.save_session), ), ( "compression_threshold", app.compression_threshold.to_string(), ), ( "rag_reranker_model", super::format_option_value(&rag_reranker_model), ), ("rag_top_k", rag_top_k.to_string()), ("dry_run", app.dry_run.to_string()), ( "function_calling_support", app.function_calling_support.to_string(), ), ("mcp_server_support", app.mcp_server_support.to_string()), ("stream", app.stream.to_string()), ("save", app.save.to_string()), ("keybindings", app.keybindings.clone()), ("wrap", wrap), ("wrap_code", app.wrap_code.to_string()), ("highlight", app.highlight.to_string()), ("theme", super::format_option_value(&app.theme)), ("config_file", display_path(&paths::config_file())), ("env_file", display_path(&paths::env_file())), ("agents_dir", display_path(&paths::agents_data_dir())), ("roles_dir", display_path(&paths::roles_dir())), ("sessions_dir", display_path(&self.sessions_dir())), ("rags_dir", display_path(&paths::rags_dir())), ("macros_dir", display_path(&paths::macros_dir())), ("functions_dir", display_path(&paths::functions_dir())), ("messages_file", display_path(&self.messages_file())), ( "vault_password_file", display_path(&app.vault_password_file()), ), ]; if let Ok((_, Some(log_path))) = paths::log_config() { items.push(("log_path", display_path(&log_path))); } let output = items .iter() .map(|(name, value)| format!("{name:<30}{value}\n")) .collect::>() .join(""); Ok(output) } pub fn info(&self, app: &AppConfig) -> Result { if let Some(agent) = &self.agent { let output = agent.export()?; if let Some(session) = &self.session { let session = session .export()? .split('\n') .map(|v| format!(" {v}")) .collect::>() .join("\n"); Ok(format!("{output}session:\n{session}")) } else { Ok(output) } } else if let Some(session) = &self.session { session.export() } else if let Some(role) = &self.role { Ok(role.export()) } else if let Some(rag) = &self.rag { rag.export() } else { self.sysinfo(app) } } pub fn session_info(&self, app: &AppConfig) -> Result { if let Some(session) = &self.session { let render_options = app.render_options()?; let mut markdown_render = crate::render::MarkdownRender::init(render_options)?; let agent_info: Option<(String, Vec)> = self.agent.as_ref().map(|agent| { let functions = agent .functions() .declarations() .iter() .filter_map(|v| if v.agent { Some(v.name.clone()) } else { None }) .collect(); (agent.name().to_string(), functions) }); session.render(&mut markdown_render, &agent_info) } else { bail!("No session") } } pub fn generate_prompt_context(&self, app: &AppConfig) -> HashMap<&str, String> { let mut output = HashMap::new(); let role = self.extract_role(app); output.insert("model", role.model().id()); output.insert("client_name", role.model().client_name().to_string()); output.insert("model_name", role.model().name().to_string()); output.insert( "max_input_tokens", role.model() .max_input_tokens() .unwrap_or_default() .to_string(), ); if let Some(temperature) = role.temperature() && temperature != 0.0 { output.insert("temperature", temperature.to_string()); } if let Some(top_p) = role.top_p() && top_p != 0.0 { output.insert("top_p", top_p.to_string()); } if app.dry_run { output.insert("dry_run", "true".to_string()); } if app.stream { output.insert("stream", "true".to_string()); } if app.save { output.insert("save", "true".to_string()); } if let Some(wrap) = &app.wrap && wrap != "no" { output.insert("wrap", wrap.clone()); } if !role.is_derived() { output.insert("role", role.name().to_string()); } if let Some(session) = &self.session { output.insert("session", session.name().to_string()); if let Some(autoname) = session.autoname() { output.insert("session_autoname", autoname.to_string()); } output.insert("dirty", session.dirty().to_string()); let (tokens, percent) = session.tokens_usage(); output.insert("consume_tokens", tokens.to_string()); output.insert("consume_percent", percent.to_string()); output.insert("user_messages_len", session.user_messages_len().to_string()); } if let Some(rag) = &self.rag { output.insert("rag", rag.name().to_string()); } if let Some(agent) = &self.agent { output.insert("agent", agent.name().to_string()); } if app.highlight { output.insert("color.reset", "\u{1b}[0m".to_string()); output.insert("color.black", "\u{1b}[30m".to_string()); output.insert("color.dark_gray", "\u{1b}[90m".to_string()); output.insert("color.red", "\u{1b}[31m".to_string()); output.insert("color.light_red", "\u{1b}[91m".to_string()); output.insert("color.green", "\u{1b}[32m".to_string()); output.insert("color.light_green", "\u{1b}[92m".to_string()); output.insert("color.yellow", "\u{1b}[33m".to_string()); output.insert("color.light_yellow", "\u{1b}[93m".to_string()); output.insert("color.blue", "\u{1b}[34m".to_string()); output.insert("color.light_blue", "\u{1b}[94m".to_string()); output.insert("color.purple", "\u{1b}[35m".to_string()); output.insert("color.light_purple", "\u{1b}[95m".to_string()); output.insert("color.magenta", "\u{1b}[35m".to_string()); output.insert("color.light_magenta", "\u{1b}[95m".to_string()); output.insert("color.cyan", "\u{1b}[36m".to_string()); output.insert("color.light_cyan", "\u{1b}[96m".to_string()); output.insert("color.white", "\u{1b}[37m".to_string()); output.insert("color.light_gray", "\u{1b}[97m".to_string()); } output } pub fn render_prompt_left(&self, app: &AppConfig) -> String { let variables = self.generate_prompt_context(app); let left_prompt = app.left_prompt.as_deref().unwrap_or(LEFT_PROMPT); render_prompt(left_prompt, &variables) } pub fn render_prompt_right(&self, app: &AppConfig) -> String { let variables = self.generate_prompt_context(app); let right_prompt = app.right_prompt.as_deref().unwrap_or(RIGHT_PROMPT); render_prompt(right_prompt, &variables) } pub fn select_enabled_functions(&self, role: &Role) -> Vec { let app = self.app.config.as_ref(); let mut functions = vec![]; if app.function_calling_support { if let Some(enabled_tools) = role.enabled_tools() { let mut tool_names: HashSet = Default::default(); let declaration_names: HashSet = self .tool_scope .functions .declarations() .iter() .filter(|v| { !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) }) .map(|v| v.name.to_string()) .collect(); if enabled_tools == "all" { tool_names.extend(declaration_names); } else { for item in enabled_tools.split(',') { let item = item.trim(); if let Some(values) = app.mapping_tools.get(item) { tool_names.extend( values .split(',') .map(|v| v.to_string()) .filter(|v| declaration_names.contains(v)), ) } else if declaration_names.contains(item) { tool_names.insert(item.to_string()); } } } functions = self .tool_scope .functions .declarations() .iter() .filter_map(|v| { if tool_names.contains(&v.name) { Some(v.clone()) } else { None } }) .collect(); } if self.agent.is_none() { let existing: HashSet = functions.iter().map(|f| f.name.clone()).collect(); let builtin_functions: Vec = self .tool_scope .functions .declarations() .iter() .filter(|v| { v.name.starts_with(USER_FUNCTION_PREFIX) && !existing.contains(&v.name) }) .cloned() .collect(); functions.extend(builtin_functions); } if let Some(agent) = &self.agent { let mut agent_functions: Vec = agent .functions() .declarations() .to_vec() .into_iter() .filter(|v| { !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) }) .collect(); let tool_names: HashSet = agent_functions .iter() .filter_map(|v| { if v.agent { None } else { Some(v.name.to_string()) } }) .collect(); agent_functions.extend( functions .into_iter() .filter(|v| !tool_names.contains(&v.name)), ); functions = agent_functions; } } functions } pub fn select_enabled_mcp_servers(&self, role: &Role) -> Vec { let app = self.app.config.as_ref(); let mut mcp_functions = vec![]; if app.mcp_server_support { if let Some(enabled_mcp_servers) = role.enabled_mcp_servers() { let mut server_names: HashSet = Default::default(); let mcp_declaration_names: HashSet = self .tool_scope .functions .declarations() .iter() .filter(|v| { v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) }) .map(|v| v.name.to_string()) .collect(); if enabled_mcp_servers == "all" { server_names.extend(mcp_declaration_names); } else { for item in enabled_mcp_servers.split(',') { let item = item.trim(); let item_invoke_name = format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); let item_search_name = format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); let item_describe_name = format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); if let Some(values) = app.mapping_mcp_servers.get(item) { server_names.extend( values .split(',') .flat_map(|v| { vec![ format!( "{}_{}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX, v.to_string() ), format!( "{}_{}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX, v.to_string() ), format!( "{}_{}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, v.to_string() ), ] }) .filter(|v| mcp_declaration_names.contains(v)), ) } else if mcp_declaration_names.contains(&item_invoke_name) { server_names.insert(item_invoke_name); server_names.insert(item_search_name); server_names.insert(item_describe_name); } } } mcp_functions = self .tool_scope .functions .declarations() .iter() .filter_map(|v| { if server_names.contains(&v.name) { Some(v.clone()) } else { None } }) .collect(); } if let Some(agent) = &self.agent { let mut agent_functions: Vec = agent .functions() .declarations() .to_vec() .into_iter() .filter(|v| { v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) }) .collect(); let tool_names: HashSet = agent_functions .iter() .filter_map(|v| { if v.agent { None } else { Some(v.name.to_string()) } }) .collect(); agent_functions.extend( mcp_functions .into_iter() .filter(|v| !tool_names.contains(&v.name)), ); mcp_functions = agent_functions; } } mcp_functions } pub fn select_functions(&self, role: &Role) -> Option> { let mut functions = vec![]; functions.extend(self.select_enabled_functions(role)); functions.extend(self.select_enabled_mcp_servers(role)); if functions.is_empty() { None } else { Some(functions) } } pub fn retrieve_role(&self, app: &AppConfig, name: &str) -> Result { let names = paths::list_roles(false); let mut role = if names.contains(&name.to_string()) { let path = paths::role_file(name); let content = read_to_string(&path)?; Role::new(name, &content) } else { Role::builtin(name)? }; let current_model = self.current_model().clone(); match role.model_id() { Some(model_id) => { if current_model.id() != model_id { let model = Model::retrieve_model(app, model_id, ModelType::Chat)?; role.set_model(model); } else { role.set_model(current_model); } } None => { role.set_model(current_model); if role.temperature().is_none() { role.set_temperature(app.temperature); } if role.top_p().is_none() { role.set_top_p(app.top_p); } } } Ok(role) } /// Returns `Ok(true)` if a role-like was mutated, `Ok(false)` if /// the model was set on `ctx.model` directly (no role-like active). pub fn set_model_on_role_like(&mut self, app: &AppConfig, model_id: &str) -> Result { let model = Model::retrieve_model(app, model_id, ModelType::Chat)?; match self.role_like_mut() { Some(role_like) => { role_like.set_model(model); Ok(true) } None => { self.model = model; Ok(false) } } } #[allow(dead_code)] pub fn reload_current_model(&mut self, app: &AppConfig, model_id: &str) -> Result<()> { let model = Model::retrieve_model(app, model_id, ModelType::Chat)?; self.model = model; Ok(()) } pub fn use_prompt(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> { let mut role = Role::new(TEMP_ROLE_NAME, prompt); role.set_model(self.current_model().clone()); self.use_role_obj(role) } pub fn edit_config(&self) -> Result<()> { let config_path = paths::config_file(); let editor = self.app.config.editor()?; edit_file(&editor, &config_path)?; println!( "NOTE: Remember to restart {} if there are changes made to '{}'", env!("CARGO_CRATE_NAME"), config_path.display(), ); Ok(()) } pub fn new_role(&self, app: &AppConfig, name: &str) -> Result<()> { if self.macro_flag { bail!("No role"); } let ans = Confirm::new("Create a new role?") .with_default(true) .prompt()?; if ans { self.upsert_role(app, name)?; } else { bail!("No role"); } Ok(()) } pub fn save_role(&mut self, name: Option<&str>) -> Result<()> { let mut role_name = match &self.role { Some(role) => { if role.has_args() { bail!("Unable to save the role with arguments (whose name contains '#')") } match name { Some(v) => v.to_string(), None => role.name().to_string(), } } None => bail!("No role"), }; if role_name == TEMP_ROLE_NAME { role_name = Text::new("Role name:") .with_validator(|input: &str| { let input = input.trim(); if input.is_empty() { Ok(Validation::Invalid("This name is required".into())) } else if input == TEMP_ROLE_NAME { Ok(Validation::Invalid("This name is reserved".into())) } else { Ok(Validation::Valid) } }) .prompt()?; } let role_path = paths::role_file(&role_name); if let Some(role) = self.role.as_mut() { role.save(&role_name, &role_path, self.working_mode.is_repl())?; } Ok(()) } pub fn edit_session(&mut self, app: &AppConfig) -> Result<()> { let name = match &self.session { Some(session) => session.name().to_string(), None => bail!("No session"), }; let session_path = self.session_file(&name); self.save_session(Some(&name))?; let editor = app.editor()?; edit_file(&editor, &session_path).with_context(|| { format!( "Failed to edit '{}' with '{}'", session_path.display(), editor ) })?; self.session = Some(Session::load_from_ctx(self, app, &name, &session_path)?); self.discontinuous_last_message(); Ok(()) } pub fn edit_agent_config(&self, app: &AppConfig) -> Result<()> { let agent_name = match &self.agent { Some(agent) => agent.name(), None => bail!("No agent"), }; let agent_config_path = paths::agent_config_file(agent_name); ensure_parent_exists(&agent_config_path)?; if !agent_config_path.exists() { std::fs::write( &agent_config_path, "# see https://github.com/Dark-Alex-17/loki/blob/main/config.agent.example.yaml\n", ) .with_context(|| format!("Failed to write to '{}'", agent_config_path.display()))?; } let editor = app.editor()?; edit_file(&editor, &agent_config_path)?; println!( "NOTE: Remember to reload the agent if there are changes made to '{}'", agent_config_path.display() ); Ok(()) } pub fn new_macro(&self, app: &AppConfig, name: &str) -> Result<()> { if self.macro_flag { bail!("No macro"); } let ans = Confirm::new("Create a new macro?") .with_default(true) .prompt()?; if ans { let macro_path = paths::macro_file(name); ensure_parent_exists(¯o_path)?; let editor = app.editor()?; edit_file(&editor, ¯o_path)?; } else { bail!("No macro"); } Ok(()) } pub fn delete(&self, kind: &str) -> Result<()> { let (dir, file_ext) = match kind { "role" => (paths::roles_dir(), Some(".md")), "session" => (self.sessions_dir(), Some(".yaml")), "rag" => (paths::rags_dir(), Some(".yaml")), "macro" => (paths::macros_dir(), Some(".yaml")), "agent-data" => (paths::agents_data_dir(), None), _ => bail!("Unknown kind '{kind}'"), }; let names = match read_dir(&dir) { Ok(rd) => { let mut names = vec![]; for entry in rd.flatten() { let name = entry.file_name(); match file_ext { Some(file_ext) => { if let Some(name) = name.to_string_lossy().strip_suffix(file_ext) { names.push(name.to_string()); } } None => { if entry.path().is_dir() { names.push(name.to_string_lossy().to_string()); } } } } names.sort_unstable(); names } Err(_) => vec![], }; if names.is_empty() { bail!("No {kind} to delete") } let select_names = MultiSelect::new(&format!("Select {kind} to delete:"), names) .with_validator(|list: &[ListOption<&String>]| { if list.is_empty() { Ok(Validation::Invalid( "At least one item must be selected".into(), )) } else { Ok(Validation::Valid) } }) .prompt()?; for name in select_names { match file_ext { Some(ext) => { let path = dir.join(format!("{name}{ext}")); remove_file(&path).with_context(|| { format!("Failed to delete {kind} at '{}'", path.display()) })?; } None => { let path = dir.join(name); remove_dir_all(&path).with_context(|| { format!("Failed to delete {kind} at '{}'", path.display()) })?; } } } println!("✓ Successfully deleted {kind}."); Ok(()) } pub fn rag_sources(&self) -> Result { match self.rag.as_ref() { Some(rag) => match rag.get_last_sources() { Some(v) => Ok(v), None => bail!("No sources"), }, None => bail!("No RAG"), } } pub async fn update(&mut self, data: &str, abort_signal: AbortSignal) -> Result<()> { let parts: Vec<&str> = data.split_whitespace().collect(); if parts.len() != 2 { bail!("Usage: .set . If value is null, unset key."); } let key = parts[0]; let value = parts[1]; match key { "temperature" => { let value = super::parse_value(value)?; if !self.set_temperature_on_role_like(value) { self.update_app_config(|app| app.temperature = value); } } "top_p" => { let value = super::parse_value(value)?; if !self.set_top_p_on_role_like(value) { self.update_app_config(|app| app.top_p = value); } } "enabled_tools" => { let value = super::parse_value(value)?; if !self.set_enabled_tools_on_role_like(value.clone()) { self.update_app_config(|app| app.enabled_tools = value); } } "enabled_mcp_servers" => { let value: Option = super::parse_value(value)?; if let Some(servers) = value.as_ref() { let Some(mcp_config) = &self.app.mcp_config else { bail!( "No MCP servers are configured. Please configure MCP servers first before setting 'enabled_mcp_servers'." ); }; if mcp_config.mcp_servers.is_empty() { bail!( "No MCP servers are configured. Please configure MCP servers first before setting 'enabled_mcp_servers'." ); } if !servers.split(',').all(|s| { let server = s.trim(); server == "all" || mcp_config.mcp_servers.contains_key(server) }) { bail!( "Some of the specified MCP servers in 'enabled_mcp_servers' are not fully configured. Please check your MCP server configuration." ); } } if !self.set_enabled_mcp_servers_on_role_like(value.clone()) { self.update_app_config(|app| app.enabled_mcp_servers = value.clone()); } if self.app.config.mcp_server_support { let app = Arc::clone(&self.app.config); self.bootstrap_tools(app.as_ref(), true, abort_signal.clone()) .await?; } } "max_output_tokens" => { let value = super::parse_value(value)?; if !self.set_max_output_tokens_on_role_like(value) { self.model.set_max_tokens(value, true); } } "save_session" => { let value = super::parse_value(value)?; if !self.set_save_session_on_session(value) { self.update_app_config(|app| app.save_session = value); } } "compression_threshold" => { let value = super::parse_value(value)?; if !self.set_compression_threshold_on_session(value) { self.update_app_config(|app| { app.compression_threshold = value.unwrap_or_default(); }); } } "rag_reranker_model" => { let value = super::parse_value(value)?; let app = Arc::clone(&self.app.config); if !self.set_rag_reranker_model(app.as_ref(), value.clone())? { self.update_app_config(|app| app.rag_reranker_model = value); } } "rag_top_k" => { let value = value.parse().with_context(|| "Invalid value")?; if !self.set_rag_top_k(value)? { self.update_app_config(|app| app.rag_top_k = value); } } "dry_run" => { let value = value.parse().with_context(|| "Invalid value")?; self.update_app_config(|app| app.dry_run = value); } "function_calling_support" => { let value = value.parse().with_context(|| "Invalid value")?; if value && self.tool_scope.functions.is_empty() { bail!("Function calling cannot be enabled because no functions are installed.") } self.update_app_config(|app| app.function_calling_support = value); } "mcp_server_support" => { let value = value.parse().with_context(|| "Invalid value")?; self.update_app_config(|app| app.mcp_server_support = value); let app = Arc::clone(&self.app.config); self.bootstrap_tools(app.as_ref(), value, abort_signal.clone()) .await?; } "stream" => { let value = value.parse().with_context(|| "Invalid value")?; self.update_app_config(|app| app.stream = value); } "save" => { let value = value.parse().with_context(|| "Invalid value")?; self.update_app_config(|app| app.save = value); } "highlight" => { let value = value.parse().with_context(|| "Invalid value")?; self.update_app_config(|app| app.highlight = value); } _ => bail!("Unknown key '{key}'"), } Ok(()) } /// Returns `Ok(true)` if the active RAG was mutated, `Ok(false)` if /// no RAG is active (caller should fall back to the `AppConfig` default). pub fn set_rag_reranker_model( &mut self, app: &AppConfig, value: Option, ) -> Result { if let Some(id) = &value { Model::retrieve_model(app, id, ModelType::Reranker)?; } match &self.rag { Some(_) => { let mut rag = self.rag.as_ref().expect("checked above").as_ref().clone(); rag.set_reranker_model(value)?; self.rag = Some(Arc::new(rag)); Ok(true) } None => Ok(false), } } pub fn set_rag_top_k(&mut self, value: usize) -> Result { match &self.rag { Some(_) => { let mut rag = self.rag.as_ref().expect("checked above").as_ref().clone(); rag.set_top_k(value)?; self.rag = Some(Arc::new(rag)); Ok(true) } None => Ok(false), } } pub fn repl_complete( &self, cmd: &str, args: &[&str], _line: &str, ) -> Vec<(String, Option)> { let app = self.app.config.as_ref(); let mut values: Vec<(String, Option)> = vec![]; let filter = args.last().unwrap_or(&""); if args.len() == 1 { values = match cmd { ".role" => super::map_completion_values(paths::list_roles(true)), ".model" => list_models(app, ModelType::Chat) .into_iter() .map(|v| (v.id(), Some(v.description()))) .collect(), ".session" => { if args[0].starts_with("_/") { super::map_completion_values( self.list_autoname_sessions() .iter() .rev() .map(|v| format!("_/{v}")) .collect::>(), ) } else { super::map_completion_values(self.list_sessions()) } } ".rag" => super::map_completion_values(paths::list_rags()), ".agent" => super::map_completion_values(list_agents()), ".macro" => super::map_completion_values(paths::list_macros()), ".starter" => match &self.agent { Some(agent) => agent .conversation_starters() .iter() .enumerate() .map(|(i, v)| ((i + 1).to_string(), Some(v.to_string()))) .collect(), None => vec![], }, ".set" => { let mut values = vec![ "temperature", "top_p", "enabled_tools", "enabled_mcp_servers", "save_session", "compression_threshold", "rag_reranker_model", "rag_top_k", "max_output_tokens", "dry_run", "function_calling_support", "mcp_server_support", "stream", "save", "highlight", ]; values.sort_unstable(); values .into_iter() .map(|v| (format!("{v} "), None)) .collect() } ".delete" => super::map_completion_values(vec![ "role", "session", "rag", "macro", "agent-data", ]), ".vault" => { let mut values = vec!["add", "get", "update", "delete", "list"]; values.sort_unstable(); values .into_iter() .map(|v| (format!("{v} "), None)) .collect() } _ => vec![], }; } else if cmd == ".set" && args.len() == 2 { let candidates = match args[0] { "max_output_tokens" => match self.current_model().max_output_tokens() { Some(v) => vec![v.to_string()], None => vec![], }, "dry_run" => super::complete_bool(app.dry_run), "stream" => super::complete_bool(app.stream), "save" => super::complete_bool(app.save), "function_calling_support" => super::complete_bool(app.function_calling_support), "enabled_tools" => { let mut prefix = String::new(); let mut ignores = HashSet::new(); if let Some((v, _)) = args[1].rsplit_once(',') { ignores = v.split(',').collect(); prefix = format!("{v},"); } let mut values = vec![]; if prefix.is_empty() { values.push("all".to_string()); } values.extend( self.tool_scope .functions .declarations() .iter() .filter(|v| { !v.name.starts_with("user__") && !v.name.starts_with("mcp_") && !v.name.starts_with("todo__") && !v.name.starts_with("agent__") }) .map(|v| v.name.clone()), ); values.extend(app.mapping_tools.keys().map(|v| v.to_string())); values .into_iter() .filter(|v| !ignores.contains(v.as_str())) .map(|v| format!("{prefix}{v}")) .collect() } "mcp_server_support" => super::complete_bool(app.mcp_server_support), "enabled_mcp_servers" => { let mut prefix = String::new(); let mut ignores = HashSet::new(); if let Some((v, _)) = args[1].rsplit_once(',') { ignores = v.split(',').collect(); prefix = format!("{v},"); } let mut values = vec![]; if prefix.is_empty() { values.push("all".to_string()); } if let Some(mcp_config) = &self.app.mcp_config { values.extend(mcp_config.mcp_servers.keys().map(|v| v.to_string())); } values.extend(app.mapping_mcp_servers.keys().map(|v| v.to_string())); values.sort(); values.dedup(); values .into_iter() .filter(|v| !ignores.contains(v.as_str())) .map(|v| format!("{prefix}{v}")) .collect() } "save_session" => { let save_session = if let Some(session) = &self.session { session.save_session() } else { app.save_session }; super::complete_option_bool(save_session) } "rag_reranker_model" => list_models(app, ModelType::Reranker) .iter() .map(|v| v.id()) .collect(), "highlight" => super::complete_bool(app.highlight), _ => vec![], }; values = candidates.into_iter().map(|v| (v, None)).collect(); } else if cmd == ".vault" && args.len() == 2 { values = self .app .vault .list_secrets(false) .unwrap_or_default() .into_iter() .map(|v| (v, None)) .collect(); } else if cmd == ".agent" { if args.len() == 2 { let dir = paths::agent_data_dir(args[0]).join(super::SESSIONS_DIR_NAME); values = list_file_names(dir, ".yaml") .into_iter() .map(|v| (v, None)) .collect(); } values.extend(super::complete_agent_variables(args[0])); }; fuzzy_filter(values, |v| v.0.as_str(), filter) } async fn rebuild_tool_scope( &mut self, app: &AppConfig, enabled_mcp_servers: Option, abort_signal: AbortSignal, ) -> Result<()> { let mut mcp_runtime = McpRuntime::new(); if app.mcp_server_support && let Some(mcp_config) = &self.app.mcp_config { let server_ids: Vec = match &enabled_mcp_servers { Some(servers) if servers == "all" => { mcp_config.mcp_servers.keys().cloned().collect() } Some(servers) => { let mut ids = Vec::new(); for item in servers.split(',').map(|s| s.trim()) { if mcp_config.mcp_servers.contains_key(item) { ids.push(item.to_string()); } else if let Some(mapped) = app.mapping_mcp_servers.get(item) { for mapped_id in mapped.split(',').map(|s| s.trim()) { if mcp_config.mcp_servers.contains_key(mapped_id) { ids.push(mapped_id.to_string()); } } } } ids } None => vec![], }; if !server_ids.is_empty() { let app_ref = &self.app; let acquire_all = async { let mut handles = Vec::new(); for id in &server_ids { if let Some(spec) = mcp_config.mcp_servers.get(id) { let handle = app_ref .mcp_factory .acquire(id, spec, app_ref.mcp_log_path.as_deref()) .await?; handles.push((id.clone(), handle)); } } Ok::<_, Error>(handles) }; let handles = abortable_run_with_spinner( acquire_all, "Loading MCP servers", abort_signal.clone(), ) .await?; for (id, handle) in handles { mcp_runtime.insert(id, handle); } } } let mut functions = Functions::init(app.visible_tools.as_ref().unwrap_or(&Vec::new()))?; if self.working_mode.is_repl() { functions.append_user_interaction_functions(); } if !mcp_runtime.is_empty() { functions.append_mcp_meta_functions(mcp_runtime.server_names()); } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { functions, mcp_runtime, tool_tracker, }; Ok(()) } pub async fn use_role( &mut self, app: &AppConfig, name: &str, abort_signal: AbortSignal, ) -> Result<()> { let role = self.retrieve_role(app, name)?; let mcp_servers = if app.mcp_server_support { role.enabled_mcp_servers() } else { if role.enabled_mcp_servers().is_some() { eprintln!( "{}", formatdoc!( " This role uses MCP servers, but MCP support is disabled. To enable it, exit the role and set 'mcp_server_support: true', then try again " ) ); } None }; self.rebuild_tool_scope(app, mcp_servers, abort_signal) .await?; self.use_role_obj(role) } pub async fn use_session( &mut self, app: &AppConfig, session_name: Option<&str>, abort_signal: AbortSignal, ) -> Result<()> { if self.session.is_some() { bail!( "Already in a session, please run '.exit session' first to exit the current session." ); } let mut session; match session_name { None | Some(TEMP_SESSION_NAME) => { let session_file = self.session_file(TEMP_SESSION_NAME); if session_file.exists() { remove_file(session_file).with_context(|| { format!("Failed to cleanup previous '{TEMP_SESSION_NAME}' session") })?; } session = Some(Session::new_from_ctx(self, app, TEMP_SESSION_NAME)); } Some(name) => { let session_path = self.session_file(name); if !session_path.exists() { session = Some(Session::new_from_ctx(self, app, name)); } else { session = Some(Session::load_from_ctx(self, app, name, &session_path)?); } } } let mut new_session = false; if let Some(session) = session.as_mut() { let mcp_servers = if app.mcp_server_support { session.enabled_mcp_servers() } else { if session.enabled_mcp_servers().is_some() { eprintln!( "{}", formatdoc!( " This session uses MCP servers, but MCP support is disabled. To enable it, exit the session and set 'mcp_server_support: true', then try again " ) ); } None }; self.rebuild_tool_scope(app, mcp_servers, abort_signal.clone()) .await?; if session.is_empty() { new_session = true; if let Some(LastMessage { input, output, continuous, }) = &self.last_message && (*continuous && !output.is_empty()) && self.agent.is_some() == input.with_agent() { let ans = Confirm::new( "Start a session that incorporates the last question and answer?", ) .with_default(false) .prompt()?; if ans { session.add_message(input, output)?; } } } } self.session = session; self.init_agent_session_variables(new_session)?; Ok(()) } pub async fn use_agent( &mut self, app: &AppConfig, agent_name: &str, session_name: Option<&str>, abort_signal: AbortSignal, ) -> Result<()> { if !app.function_calling_support { bail!("Please enable function calling support before using the agent."); } if self.agent.is_some() { bail!("Already in an agent, please run '.exit agent' first to exit the current agent."); } let current_model = self.current_model().clone(); let agent = Agent::init( app, &self.app, ¤t_model, self.info_flag, agent_name, abort_signal.clone(), ) .await?; let mcp_servers = if app.mcp_server_support { (!agent.mcp_server_names().is_empty()).then(|| agent.mcp_server_names().join(",")) } else { if !agent.mcp_server_names().is_empty() { bail!( "This agent uses MCP servers, but MCP support is disabled.\nTo enable it, set 'mcp_server_support: true', then try again." ); } None }; self.rebuild_tool_scope(app, mcp_servers, abort_signal.clone()) .await?; if !agent.model().supports_function_calling() { eprintln!( "Warning: The model '{}' does not support function calling. Agent tools (including todo, spawning, and user interaction) will not be available.", agent.model().id() ); } let session_name = session_name.map(|v| v.to_string()).or_else(|| { if self.macro_flag { None } else { agent.agent_session().map(|v| v.to_string()) } }); let should_init_supervisor = agent.can_spawn_agents(); let max_concurrent = agent.max_concurrent_agents(); let max_depth = agent.max_agent_depth(); let supervisor = should_init_supervisor .then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth)))); self.rag = agent.rag(); self.agent = Some(agent); self.supervisor = supervisor; self.inbox = None; self.escalation_queue = None; self.self_agent_id = None; self.parent_supervisor = None; self.current_depth = 0; self.auto_continue_count = 0; self.todo_list = TodoList::default(); if let Some(session_name) = session_name.as_deref() { self.use_session(app, Some(session_name), abort_signal) .await?; } else { self.init_agent_shared_variables()?; } self.agent_variables = None; Ok(()) } pub fn exit_agent(&mut self, app: &AppConfig) -> Result<()> { self.exit_session()?; let mut functions = Functions::init(app.visible_tools.as_ref().unwrap_or(&Vec::new()))?; if self.working_mode.is_repl() { functions.append_user_interaction_functions(); } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { functions, mcp_runtime: McpRuntime::default(), tool_tracker, }; if self.agent.take().is_some() { if let Some(supervisor) = self.supervisor.clone() { supervisor.read().cancel_all(); } self.supervisor = None; self.parent_supervisor = None; self.self_agent_id = None; self.inbox = None; self.escalation_queue = None; self.current_depth = 0; self.auto_continue_count = 0; self.todo_list = TodoList::default(); self.rag.take(); self.discontinuous_last_message(); } Ok(()) } pub async fn edit_role(&mut self, app: &AppConfig, abort_signal: AbortSignal) -> Result<()> { let role_name; if let Some(session) = self.session.as_ref() { if let Some(name) = session.role_name().map(|v| v.to_string()) { if session.is_empty() { role_name = Some(name); } else { bail!("Cannot perform this operation because you are in a non-empty session") } } else { bail!("No role") } } else { role_name = self.role.as_ref().map(|v| v.name().to_string()); } let name = role_name.ok_or_else(|| anyhow::anyhow!("No role"))?; self.upsert_role(app, &name)?; self.use_role(app, &name, abort_signal).await } fn upsert_role(&self, app: &AppConfig, name: &str) -> Result<()> { let role_path = paths::role_file(name); ensure_parent_exists(&role_path)?; let editor = app.editor()?; edit_file(&editor, &role_path)?; if self.working_mode.is_repl() { println!("✓ Saved the role to '{}'.", role_path.display()); } Ok(()) } pub async fn apply_prelude( &mut self, app: &AppConfig, abort_signal: AbortSignal, ) -> Result<()> { if self.macro_flag || !self.state().is_empty() { return Ok(()); } let prelude = match self.working_mode { WorkingMode::Repl => app.repl_prelude.as_ref(), WorkingMode::Cmd => app.cmd_prelude.as_ref(), }; let prelude = match prelude { Some(v) => { if v.is_empty() { return Ok(()); } v.to_string() } None => return Ok(()), }; let err_msg = || format!("Invalid prelude '{prelude}"); match prelude.split_once(':') { Some(("role", name)) => { self.use_role(app, name, abort_signal) .await .with_context(err_msg)?; } Some(("session", name)) => { self.use_session(app, Some(name), abort_signal) .await .with_context(err_msg)?; } Some((session_name, role_name)) => { self.use_session(app, Some(session_name), abort_signal.clone()) .await .with_context(err_msg)?; if let Some(true) = self.session.as_ref().map(|v| v.is_empty()) { self.use_role(app, role_name, abort_signal) .await .with_context(err_msg)?; } } _ => { bail!("{}", err_msg()) } } Ok(()) } pub fn maybe_autoname_session(&mut self) -> bool { if let Some(session) = self.session.as_mut() && session.need_autoname() { session.set_autonaming(true); true } else { false } } fn enabled_mcp_servers_for_current_scope( &self, app: &AppConfig, start_mcp_servers: bool, ) -> Option { if !start_mcp_servers || !app.mcp_server_support { return None; } if let Some(agent) = self.agent.as_ref() { return (!agent.mcp_server_names().is_empty()) .then(|| agent.mcp_server_names().join(",")); } if let Some(session) = self.session.as_ref() { return session.enabled_mcp_servers(); } if let Some(role) = self.role.as_ref() { return role.enabled_mcp_servers(); } app.enabled_mcp_servers.clone() } pub async fn bootstrap_tools( &mut self, app: &AppConfig, start_mcp_servers: bool, abort_signal: AbortSignal, ) -> Result<()> { let enabled_mcp_servers = self.enabled_mcp_servers_for_current_scope(app, start_mcp_servers); self.rebuild_tool_scope(app, enabled_mcp_servers, abort_signal) .await } pub async fn compress_session(&mut self) -> Result<()> { match self.session.as_ref() { Some(session) => { if !session.has_user_messages() { bail!("No need to compress since there are no messages in the session") } } None => bail!("No session"), } let prompt = self .app .config .summarization_prompt .clone() .unwrap_or_else(|| SUMMARIZATION_PROMPT.into()); let input = Input::from_str(self, &prompt, None); let summary = input.fetch_chat_text().await?; let summary_context_prompt = self .app .config .summary_context_prompt .clone() .unwrap_or_else(|| SUMMARY_CONTEXT_PROMPT.into()); let todo_prefix = if self.agent.is_some() && !self.todo_list.is_empty() { format!( "[ACTIVE TODO LIST]\n{}\n\n", self.todo_list.render_for_model() ) } else { String::new() }; if let Some(session) = self.session.as_mut() { session.compress(format!("{todo_prefix}{summary_context_prompt}{summary}")); } self.discontinuous_last_message(); Ok(()) } pub async fn autoname_session(&mut self, app: &AppConfig) -> Result<()> { let text = match self .session .as_ref() .and_then(|session| session.chat_history_for_autonaming()) { Some(v) => v, None => bail!("No chat history"), }; let role = self.retrieve_role(app, CREATE_TITLE_ROLE)?; let input = Input::from_str(self, &text, Some(role)); let text = input.fetch_chat_text().await?; if let Some(session) = self.session.as_mut() { session.set_autoname(&text); } Ok(()) } pub async fn use_rag(&mut self, rag: Option<&str>, abort_signal: AbortSignal) -> Result<()> { if self.agent.is_some() { bail!("Cannot perform this operation because you are using a agent") } let app = self.app.config.clone(); let rag_cache = self.rag_cache(); let working_mode = self.working_mode; let rag: Arc = match rag { None => { let rag_path = self.rag_file(super::TEMP_RAG_NAME); if rag_path.exists() { remove_file(&rag_path).with_context(|| { format!("Failed to cleanup previous '{}' rag", super::TEMP_RAG_NAME) })?; } Arc::new(Rag::init(&app, super::TEMP_RAG_NAME, &rag_path, &[], abort_signal).await?) } Some(name) => { let rag_path = self.rag_file(name); let key = RagKey::Named(name.to_string()); rag_cache .load_with(key, || { let app = app.clone(); let rag_path = rag_path.clone(); let abort_signal = abort_signal.clone(); async move { if !rag_path.exists() { if working_mode.is_cmd() { bail!("Unknown RAG '{name}'"); } Rag::init(&app, name, &rag_path, &[], abort_signal.clone()).await } else { Rag::load(&app, name, &rag_path) } } }) .await? } }; self.rag = Some(rag); Ok(()) } pub async fn edit_rag_docs(&mut self, abort_signal: AbortSignal) -> Result<()> { let mut rag = match self.rag.clone() { Some(v) => v.as_ref().clone(), None => bail!("No RAG"), }; let document_paths = rag.document_paths(); let temp_file = temp_file(&format!("-rag-{}", rag.name()), ".txt"); tokio::fs::write(&temp_file, &document_paths.join("\n")) .await .with_context(|| format!("Failed to write to '{}'", temp_file.display()))?; let editor = self.app.config.editor()?; edit_file(&editor, &temp_file)?; let new_document_paths = tokio::fs::read_to_string(&temp_file) .await .with_context(|| format!("Failed to read '{}'", temp_file.display()))?; let new_document_paths = new_document_paths .split('\n') .filter_map(|v| { let v = v.trim(); if v.is_empty() { None } else { Some(v.to_string()) } }) .collect::>(); if new_document_paths.is_empty() || new_document_paths == document_paths { bail!("No changes") } let key = if self.agent.is_some() { RagKey::Agent(rag.name().to_string()) } else { RagKey::Named(rag.name().to_string()) }; self.rag_cache().invalidate(&key); rag.refresh_document_paths(&new_document_paths, false, &self.app.config, abort_signal) .await?; self.rag = Some(Arc::new(rag)); Ok(()) } pub async fn rebuild_rag(&mut self, abort_signal: AbortSignal) -> Result<()> { let mut rag = match self.rag.clone() { Some(v) => v.as_ref().clone(), None => bail!("No RAG"), }; let key = if self.agent.is_some() { RagKey::Agent(rag.name().to_string()) } else { RagKey::Named(rag.name().to_string()) }; self.rag_cache().invalidate(&key); let document_paths = rag.document_paths().to_vec(); rag.refresh_document_paths(&document_paths, true, &self.app.config, abort_signal) .await?; self.rag = Some(Arc::new(rag)); Ok(()) } } #[cfg(test)] mod tests { use super::*; use crate::config::{AppState, Config}; use crate::utils::get_env_name; use std::env; use std::fs::{create_dir_all, remove_dir_all, write}; use std::path::PathBuf; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; struct TestConfigDirGuard { key: String, previous: Option, path: PathBuf, } impl TestConfigDirGuard { fn new() -> Self { let key = get_env_name("config_dir"); let previous = env::var_os(&key); let unique = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_nanos(); let path = env::temp_dir().join(format!("loki-request-context-tests-{unique}")); create_dir_all(&path).unwrap(); unsafe { env::set_var(&key, &path); } Self { key, previous, path, } } } impl Drop for TestConfigDirGuard { fn drop(&mut self) { if let Some(previous) = &self.previous { unsafe { env::set_var(&self.key, previous); } } else { unsafe { env::remove_var(&self.key); } } let _ = remove_dir_all(&self.path); } } fn app_state_from_config(cfg: &Config) -> Arc { Arc::new(AppState { config: Arc::new(cfg.to_app_config()), vault: cfg.vault.clone(), mcp_factory: Arc::new(super::super::mcp_factory::McpFactory::default()), rag_cache: Arc::new(RagCache::default()), mcp_config: None, mcp_log_path: None, }) } fn create_test_ctx() -> RequestContext { let cfg = Config::default(); let app_state = app_state_from_config(&cfg); cfg.to_request_context(app_state) } #[test] fn to_request_context_creates_clean_state() { let cfg = Config::default(); let app_state = app_state_from_config(&cfg); let ctx = cfg.to_request_context(app_state); assert!(ctx.role.is_none()); assert!(ctx.session.is_none()); assert!(ctx.agent.is_none()); assert!(ctx.rag.is_none()); assert!(ctx.supervisor.is_none()); assert!(ctx.tool_scope.mcp_runtime.is_empty()); assert_eq!(ctx.current_depth, 0); } #[test] fn update_app_config_persists_changes() { let cfg = Config::default(); let app_state = app_state_from_config(&cfg); let mut ctx = RequestContext::new(app_state, WorkingMode::Cmd); let previous = Arc::clone(&ctx.app.config); ctx.update_app_config(|app| { app.save = true; app.compression_threshold = 1234; }); assert!(ctx.app.config.save); assert_eq!(ctx.app.config.compression_threshold, 1234); assert!(!Arc::ptr_eq(&ctx.app.config, &previous)); } #[test] fn use_role_obj_sets_role() { let mut ctx = create_test_ctx(); let role = Role::new("test", "test prompt"); ctx.use_role_obj(role).unwrap(); assert!(ctx.role.is_some()); assert_eq!(ctx.role.as_ref().unwrap().name(), "test"); } #[test] fn exit_role_clears_role() { let mut ctx = create_test_ctx(); let role = Role::new("test", "prompt"); ctx.use_role_obj(role).unwrap(); assert!(ctx.role.is_some()); ctx.exit_role().unwrap(); assert!(ctx.role.is_none()); } #[test] fn use_prompt_creates_temp_role() { let mut ctx = create_test_ctx(); let app = ctx.app.config.clone(); ctx.use_prompt(&app, "you are a pirate").unwrap(); assert!(ctx.role.is_some()); assert_eq!(ctx.role.as_ref().unwrap().name(), "temp"); assert!( ctx.role .as_ref() .unwrap() .prompt() .contains("you are a pirate") ); } #[test] fn extract_role_returns_standalone_role() { let mut ctx = create_test_ctx(); let app = ctx.app.config.clone(); let role = Role::new("myrole", "my prompt"); ctx.use_role_obj(role).unwrap(); let extracted = ctx.extract_role(&app); assert_eq!(extracted.name(), "myrole"); } #[test] fn extract_role_returns_default_when_nothing_active() { let ctx = create_test_ctx(); let app = ctx.app.config.clone(); let extracted = ctx.extract_role(&app); assert_eq!(extracted.name(), ""); } #[test] fn exit_session_clears_session() { let mut ctx = create_test_ctx(); ctx.session = Some(Session::default()); assert!(ctx.session.is_some()); ctx.exit_session().unwrap(); assert!(ctx.session.is_none()); } #[test] fn empty_session_clears_messages() { let mut ctx = create_test_ctx(); ctx.session = Some(Session::default()); ctx.empty_session().unwrap(); assert!(ctx.session.is_some()); assert!(ctx.session.as_ref().unwrap().is_empty()); } #[test] fn maybe_autoname_session_returns_false_when_no_session() { let mut ctx = create_test_ctx(); assert!(!ctx.maybe_autoname_session()); } #[test] fn exit_agent_clears_all_agent_state() { let _guard = TestConfigDirGuard::new(); let mut ctx = create_test_ctx(); let app = ctx.app.config.clone(); let agent_name = format!( "test_agent_{}", SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_nanos() ); let agent_dir = paths::agent_data_dir(&agent_name); create_dir_all(&agent_dir).unwrap(); write( agent_dir.join("config.yaml"), format!("name: {agent_name}\ninstructions: hi\n"), ) .unwrap(); tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { ctx.use_agent(&app, &agent_name, None, crate::utils::create_abort_signal()) .await .unwrap(); }); assert!(ctx.agent.is_some()); ctx.exit_agent(&app).unwrap(); assert!(ctx.agent.is_none()); assert!(ctx.rag.is_none()); } #[test] fn current_depth_default_is_zero() { let ctx = create_test_ctx(); assert_eq!(ctx.current_depth, 0); } #[test] fn current_depth_can_be_set() { let mut ctx = create_test_ctx(); ctx.current_depth = 3; assert_eq!(ctx.current_depth, 3); } #[test] fn supervisor_defaults_to_none() { let ctx = create_test_ctx(); assert!(ctx.supervisor.is_none()); } #[test] fn inbox_defaults_to_none() { let ctx = create_test_ctx(); assert!(ctx.inbox.is_none()); } #[test] fn escalation_queue_defaults_to_none() { let ctx = create_test_ctx(); assert!(ctx.root_escalation_queue().is_none()); } }