diff --git a/docs/PHASE-1-IMPLEMENTATION-PLAN.md b/docs/PHASE-1-IMPLEMENTATION-PLAN.md index 7b48a3f..318a655 100644 --- a/docs/PHASE-1-IMPLEMENTATION-PLAN.md +++ b/docs/PHASE-1-IMPLEMENTATION-PLAN.md @@ -1058,6 +1058,45 @@ At this point no code references `GlobalConfig`. **Blocked by:** Step 14. +--- + +### Step 16: Complete Config → AppConfig Migration (Post-QA) + +**Status:** PENDING — to be completed after QA testing phase + +The current bridge has a bug: `Config::init` mutates Config during startup (env vars, model resolution, etc.), but `to_app_config()` only copies serialized fields, losing those mutations. + +Current startup flow (broken): +``` +YAML → Config (serde deserialize) + → config.load_envs() ← mutates Config + → config.setup_model() ← resolves model + → config.load_mcp_servers() ← starts MCP + → cfg.to_app_config() ← COPIES ONLY serialized fields! + → AppConfig loses mutations +``` + +**Problem:** Mutations in Config are lost when building AppConfig. + +**Solution:** Move mutations AFTER the bridge: + +1. Move `load_envs()`, `set_wrap()`, `setup_model()`, `load_mcp_servers()`, `setup_document_loaders()`, `setup_user_agent()` from Config to AppConfig +2. In `main.rs`, apply these mutations AFTER `to_app_config()` +3. Delete duplicated methods from AppConfig (they become reachable) +4. Simplify Config to pure serde deserialization only +5. Remove bridge if Config becomes just a deserialization target (or keep for backwards compat) + +**Files to modify:** +- `src/config/mod.rs` — remove init mutations, keep only serde + deserialization +- `src/config/app_config.rs` — enable mutations, remove duplication +- `src/main.rs` — reorder bridge + mutations + +**Goal:** Config becomes a simple POJO. All runtime configuration lives in AppConfig/AppState. + +**Blocked by:** QA testing (Step 16 can begin after tests pass) + +--- + Phase 1 complete. --- diff --git a/docs/implementation/PHASE-1-STEP-16-NOTES.md b/docs/implementation/PHASE-1-STEP-16-NOTES.md new file mode 100644 index 0000000..489f2d9 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16-NOTES.md @@ -0,0 +1,443 @@ +# Phase 1 Step 16 — Implementation Notes + +## Status + +Pending. Architecture plan approved; ready for sub-phase execution. + +## Plan reference + +- Plan: `docs/PHASE-1-IMPLEMENTATION-PLAN.md` +- Section: "Step 16: Complete Config → AppConfig Migration (Post-QA)" + +## Problem + +The current startup flow mutates `Config` during `Config::init()`, +then converts it to `AppConfig` via `bridge.rs::to_app_config()`. This +design was transitional — it let us build the new structs alongside +the old one without a big-bang migration. + +Now that the transition is nearly done, we want `Config` to be a +genuine serde POJO: no runtime state, no init logic, nothing that +couldn't round-trip through YAML. The structs that actually represent +runtime state (`AppConfig`, `AppState`, `RequestContext`) should own +their own initialization logic. + +## Target architecture + +Instead of migrating mutations incrementally through the bridge, we +**pivot the initialization direction**. Each struct owns its own init. + +``` +YAML file + ↓ Config::load_from_file (serde only — no init logic) +Config (pure POJO) + ↓ AppConfig::from_config(config) → AppConfig +AppConfig (immutable app-wide settings, self-initializing) + ↓ AppState::init(Arc, ...).await → AppState +AppState (shared process state: vault, mcp_factory, rag_cache, mcp_registry, functions) + ↓ RequestContext::new(Arc, working_mode) +RequestContext (per-request mutable state, unchanged) +``` + +### Struct responsibilities (post-16) + +**`Config`** — trivial serde POJO: +- Only `#[serde(...)]` fields (no `#[serde(skip)]`) +- Only method: `load_from_file(path) -> Result<(Config, String)>` + (returns parsed Config + raw YAML content for secret interpolation) +- Can be round-tripped via YAML + +**`AppConfig::from_config(config) -> Result`** absorbs: +- Field-copy from `Config` (same as today's `to_app_config`) +- `load_envs()` — env var overrides +- `set_wrap()` — wrap string validation +- `setup_document_loaders()` — default pdf/docx loaders +- `setup_user_agent()` — resolve `"auto"` user agent +- `resolve_model()` — logic from `Config::setup_model` (picks default + model if `model_id` is empty) + +**`AppState::init(app_config, log_path, start_mcp_servers, abort_signal)`** absorbs: +- `Vault::init(&app_config)` (vault moves from Config to AppState) +- `McpRegistry::init(...)` (currently `Config::load_mcp_servers`) +- `Functions::init(...)` (currently `Config::load_functions`) +- Returns fully-wired `AppState` + +**`install_builtins()`** — top-level free function (replaces +`Agent::install_builtin_agents()` + `Macro::install_macros()` being +called inside `Config::init`). Called once from `main.rs` before any +config loading. Config-independent — just copies embedded assets. + +**`bridge.rs` — deleted.** No more `to_app_config()` / +`to_request_context()`. + +## Sub-phase layout + +| Sub-phase | Scope | +|-----------|-------| +| 16a | Build `AppConfig::from_config` absorbing env/wrap/docs/user-agent/model-resolution logic | +| 16b | Extract `install_builtins()` as top-level function | +| 16c | Migrate `Vault` onto `AppState` | +| 16d | Build `AppState::init` absorbing MCP-registry/functions logic | +| 16e | Update `main.rs` + audit all 15 `Config::init()` callers, switch to new flow | +| 16f | Delete Config runtime fields, bridge.rs, `Config::init`, duplicated methods | + +Sub-phases 16a–16d can largely proceed in parallel (each adds new +entry points without removing the old ones). 16e switches callers. +16f is the final cleanup. + +## 16a — AppConfig::from_config + +**Target signature:** +```rust +impl AppConfig { + pub fn from_config(config: Config) -> Result { + let mut app_config = Self { + // Copy all serde fields from config + }; + app_config.load_envs(); + if let Some(wrap) = app_config.wrap.clone() { + app_config.set_wrap(&wrap)?; + } + app_config.setup_document_loaders(); + app_config.setup_user_agent(); + app_config.resolve_model()?; + Ok(app_config) + } + + fn resolve_model(&mut self) -> Result<()> { + if self.model_id.is_empty() { + let models = crate::client::list_models(self, ModelType::Chat); + if models.is_empty() { + bail!("No available model"); + } + self.model_id = models[0].id(); + } + Ok(()) + } +} +``` + +**New method: `AppConfig::resolve_model()`** — moves logic from +`Config::setup_model`. Ensures `model_id` is a valid, non-empty +concrete model reference. + +**Note on `Model` vs `model_id`:** `Model` (the resolved runtime +handle) stays on `RequestContext`. AppConfig owns `model_id: String` +(the config default). RequestContext.model is built by calling +`Model::retrieve_model(&app_config, &model_id, ModelType::Chat)` +during context construction. They're different types for a reason. + +**Files modified (16a):** +- `src/config/app_config.rs` — add `from_config`, `resolve_model` +- Also remove `#[allow(dead_code)]` from `load_envs`, `set_wrap`, + `setup_document_loaders`, `setup_user_agent`, `set_*_default`, + `ensure_default_model_id` (they all become reachable) + +**Bridge still exists after 16a.** `Config::init` still calls its own +mutations for now. 16a just introduces the new entry point. + +## 16b — install_builtins() + +**Target signature:** +```rust +// In src/config/mod.rs or a new module +pub fn install_builtins() -> Result<()> { + Agent::install_builtin_agents()?; + Macro::install_macros()?; + Ok(()) +} +``` + +**Changes:** +- Remove `Agent::install_builtin_agents()?;` and + `Macro::install_macros()?;` calls from inside `Config::init` +- Add `install_builtins()?;` to `main.rs` as the first step before + any config loading + +Both functions are Config-independent (they just copy embedded +assets to the config directory), so this is a straightforward +extraction. + +**Files modified (16b):** +- `src/config/mod.rs` — remove calls from `Config::init`, expose + `install_builtins` as a module-level pub fn +- `src/main.rs` — call `install_builtins()?;` at startup + +## 16c — Vault → AppState + +Today `Config.vault: Arc` is a `#[serde(skip)]` runtime +field populated by `Vault::init(config)`. Post-16c, the vault lives +natively on `AppState`. + +**Current:** +```rust +pub struct AppState { + pub config: Arc, + pub vault: GlobalVault, // Already here, sourced from config.vault + ... +} +``` + +Wait — `AppState.vault` already exists. The work in 16c is just: + +1. Change `Vault::init(config: &Config)` → `Vault::init(config: &AppConfig)` + - `Vault::init` only reads `config.vault_password_file()`, which + is already a serde field on AppConfig. Rename the param. +2. Delete `Config.vault` field (no longer needed once 16e routes + through AppState) +3. Update `main.rs` to call `Vault::init(&app_config)` instead of + `cfg.vault.clone()` + +**Files modified (16c):** +- `src/vault/mod.rs` — `Vault::init` takes `&AppConfig` +- `src/config/mod.rs` — delete `Config.vault` field (after callers + switch) + +## 16d — AppState::init + +**Target signature:** +```rust +impl AppState { + pub async fn init( + config: Arc, + log_path: Option, + start_mcp_servers: bool, + abort_signal: AbortSignal, + ) -> Result { + let vault = Vault::init(&config); + let functions = { + let mut fns = Functions::init( + config.visible_tools.as_ref().unwrap_or(&Vec::new()) + )?; + // REPL-specific fns appended by RequestContext, not here + fns + }; + let mcp_registry = McpRegistry::init( + log_path.clone(), + start_mcp_servers, + config.enabled_mcp_servers.clone(), + abort_signal, + &config, // new signature: &AppConfig + ).await?; + let (mcp_config, mcp_log_path) = ( + mcp_registry.mcp_config().cloned(), + mcp_registry.log_path().cloned(), + ); + Ok(Self { + config, + vault, + mcp_factory: Default::default(), + rag_cache: Default::default(), + mcp_config, + mcp_log_path, + mcp_registry: Some(mcp_registry), // NEW field + functions, // NEW field + }) + } +} +``` + +**New AppState fields:** +- `mcp_registry: Option` — the live registry of started + MCP servers (currently on Config) +- `functions: Functions` — the global function declarations (currently + on Config) + +These become the "source of truth" that `ToolScope` copies from when +a scope transition happens. + +**`McpRegistry::init` signature change:** today takes `&Config`, +needs to take `&AppConfig`. Only reads serialized fields. + +**Files modified (16d):** +- `src/config/app_state.rs` — add `init`, add `mcp_registry` + + `functions` fields +- `src/mcp/mod.rs` — `McpRegistry::init` takes `&AppConfig` + +**Important:** `Functions.append_user_interaction_functions()` is +currently called inside `Config::load_functions` when in REPL mode. +That logic is working-mode-dependent and belongs on `RequestContext` +(which knows its mode), not `AppState`. The migration moves that +append step to `RequestContext::new` or similar. + +## 16e — Switch main.rs and 15 callers + +**New `main.rs` flow:** +```rust +install_builtins()?; +let (config, raw_yaml) = Config::load_from_file(&paths::config_file())?; + +// Secret interpolation (two-pass) +let bootstrap_vault = Vault::init_from_password_file(&config.vault_password_file()); +let interpolated = interpolate_secrets_or_err(&raw_yaml, &bootstrap_vault, info_flag)?; +let final_config = if interpolated != raw_yaml { + Config::load_from_str(&interpolated)? +} else { + config +}; + +let app_config = Arc::new(AppConfig::from_config(final_config)?); +let app_state = Arc::new( + AppState::init( + app_config.clone(), + log_path, + start_mcp_servers, + abort_signal.clone(), + ).await? +); +let ctx = RequestContext::new(app_state.clone(), working_mode); +``` + +**Secret interpolation complication:** Today's `Config::init` does a +two-pass YAML parse — load, init vault, interpolate secrets into raw +content, re-parse if content changed. In the new flow: +1. Load Config from YAML (also returns raw content) +2. Bootstrap Vault using Config's `vault_password_file` serde field +3. Interpolate secrets in raw content +4. If content changed, re-parse Config +5. Build AppConfig from final Config +6. Build AppState (creates the full Vault via `Vault::init(&app_config)`) + +Step 2 and step 6 create the vault twice — once bootstrap (to decrypt +secrets in raw YAML), once full (for AppState). This matches current +behavior, just made explicit. + +**15 callers of `Config::init()`** — audit required. Discovery +happens during 16e execution. Open questions flagged for user input +as discovered. + +| File | Expected Action | +|------|-----------------| +| `main.rs` | Use new flow | +| `client/common.rs` | Probably needs AppConfig only | +| `vault/mod.rs` | Already uses `Config::vault_password_file`; switch to AppConfig | +| `config/request_context.rs` | Test helper — use `AppState::test_default()` or build directly | +| `config/session.rs` | Test helper — same | +| `rag/mod.rs` | Probably AppConfig | +| `function/supervisor.rs` | Test helper | +| `utils/request.rs` | Probably AppConfig | +| `config/role.rs` | Test helper | +| `utils/clipboard.rs` | Probably AppConfig | +| `supervisor/mod.rs` | Test helper | +| `repl/mod.rs` | Test helper | +| `parsers/common.rs` | Probably AppConfig | +| `utils/abort_signal.rs` | Probably AppConfig | +| `utils/spinner.rs` | Probably AppConfig | + +**Files modified (16e):** +- `src/main.rs` +- Any of the 15 files above that aren't trivial — may need + `test_default()` helpers added + +## 16f — Final cleanup + +**Delete from `Config`:** +- All `#[serde(skip)]` fields: `vault`, `macro_flag`, `info_flag`, + `agent_variables`, `model`, `functions`, `mcp_registry`, + `working_mode`, `last_message`, `role`, `session`, `rag`, `agent`, + `tool_call_tracker`, `supervisor`, `parent_supervisor`, + `self_agent_id`, `current_depth`, `inbox`, + `root_escalation_queue` +- `Config::init` (whole function) +- `Config::load_envs`, `Config::load_functions`, + `Config::load_mcp_servers`, `Config::setup_model`, + `Config::set_model`, `Config::role_like_mut`, + `Config::sessions_dir`, `Config::set_wrap`, + `Config::setup_document_loaders`, `Config::setup_user_agent`, + `Config::vault_password_file` (redundant with AppConfig's) + +**Delete:** +- `src/config/bridge.rs` (entire file) +- `mod bridge;` declaration in `config/mod.rs` + +**Resulting `Config`:** +```rust +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(default)] +pub struct Config { + // Only serde-annotated fields — the YAML shape + pub model_id: String, + pub temperature: Option, + // ... all the other serde fields +} + +impl Config { + pub fn load_from_file(path: &Path) -> Result<(Config, String)> { ... } + pub fn load_from_str(content: &str) -> Result { ... } +} +``` + +A genuine POJO. No runtime state. No init logic. Just shape. + +## Open questions (for execution) + +1. **`Vault::init_bare`** — currently used as a fallback when no + Config exists. Does it still need to exist? The default + `vault_password_file` location is static. Might need + `AppConfig::default().vault_password_file()` or a free function. +2. **Secret interpolation ownership** — does `AppConfig::from_config` + handle it internally (takes raw YAML string and interpolates), or + does `main.rs` orchestrate the two-pass explicitly? Leaning toward + `main.rs` orchestration (cleaner separation). +3. **REPL-only `append_user_interaction_functions`** — moves to + `RequestContext::new`? Or stays as a post-init append called + explicitly? +4. **`Functions::init` + MCP meta functions** — today + `load_mcp_servers` calls `self.functions.append_mcp_meta_functions(...)` + after starting servers. In the new flow, `AppState::init` does + this. Verify ordering is preserved. +5. **Testing strategy** — User said don't worry unless trivial. If + test helpers need refactoring to work with new flow, prefer + adding `test_default()` methods gated by `#[cfg(test)]` over + rewriting tests. + +## Dependencies between sub-phases + +``` +16a ──┐ +16b ──┤ +16c ──┼──→ 16d ──→ 16e ──→ 16f + │ +16b, 16c, 16a independent and can run in any order +16d depends on 16c (vault on AppConfig) +16e depends on 16a, 16d (needs the new entry points) +16f depends on 16e (needs all callers switched) +``` + +## Rationale for this architecture + +The original Step 16 plan migrated mutations piecewise through the +existing `to_app_config()` bridge. That works but: + +- Leaves the bridge in place indefinitely +- Keeps `Config` burdened with both YAML shape AND runtime state +- Requires careful ordering to avoid breaking downstream consumers + like `load_functions`/`load_mcp_servers`/`setup_model` +- Creates transitional states where some mutations live on Config, + some on AppConfig + +The new approach: + +- Eliminates `bridge.rs` entirely +- Makes `Config` a true POJO +- Makes `AppConfig`/`AppState` self-contained (initialize from YAML + directly) +- REST API path is trivial: `AppConfig::from_config(yaml_string)` +- Test helpers can build `AppConfig`/`AppState` without Config +- Each struct owns exactly its concerns + +## Verification criteria for each sub-phase + +- 16a: `cargo check` + `cargo test` clean. `AppConfig::from_config` + produces the same state as `Config::init` + `to_app_config()` for + the same YAML input. +- 16b: `install_builtins()` called once from `main.rs`; agents and + macros still install on first startup. +- 16c: `Vault::init` takes `&AppConfig`; `Config.vault` field deleted. +- 16d: `AppState::init` builds a fully-wired `AppState` from + `Arc` + startup context. MCP servers start; functions + load. +- 16e: REPL starts, all CLI flags work, all env vars honored, all + existing tests pass. +- 16f: Grep for `Config {` / `Config::init(` / `bridge::to_` shows + zero non-test hits. `Config` has only serde fields. diff --git a/docs/implementation/PHASE-1-STEP-16a-NOTES.md b/docs/implementation/PHASE-1-STEP-16a-NOTES.md new file mode 100644 index 0000000..49b1589 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16a-NOTES.md @@ -0,0 +1,179 @@ +# Phase 1 Step 16a — Implementation Notes + +## Status + +Done. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Sub-phase goal: "Build `AppConfig::from_config` absorbing + env/wrap/docs/user-agent/model-resolution logic" + +## Summary + +Introduced `AppConfig::from_config(config: Config) -> Result` +as the canonical constructor for a fully-initialized `AppConfig`. The +new constructor chains the four mutation methods (`load_envs`, +`set_wrap`, `setup_document_loaders`, `setup_user_agent`) plus a new +`resolve_model()` method that picks a default model when `model_id` +is empty. + +The existing bridge (`Config::init` + `Config::to_app_config`) is +untouched. `from_config` is currently dead code (gated with +`#[allow(dead_code)]`) — it becomes the entry point in Step 16e when +`main.rs` switches over. The methods it calls (`load_envs`, etc.) are +no longer dead code because they're reachable via `from_config`, so +their `#[allow(dead_code)]` gates were removed. + +## What was changed + +### `src/config/app_config.rs` + +**Added:** +- `AppConfig::from_config(config) -> Result` — canonical + constructor that copies serde fields, applies env overrides, + validates wrap, installs doc loaders, resolves user agent, and + ensures a default model. +- `AppConfig::resolve_model(&mut self) -> Result<()>` — if + `model_id` is empty, picks the first available chat model. Errors + if no models are available. Replaces the logic from + `Config::setup_model` that belongs on `AppConfig` (the + `Model`-resolution half of `Config::setup_model` stays in Config + for now — that moves in 16e). +- 8 unit tests covering field copying, doc loader insertion, user + agent resolution, wrap validation (valid + invalid), and + `resolve_model` error/happy paths. + +**Removed `#[allow(dead_code)]` from:** +- `set_wrap` +- `setup_document_loaders` +- `setup_user_agent` +- `load_envs` + +These are now reachable via `from_config`. They remain `pub` because +REPL-mode mutations (via `.set wrap ` or similar) will go +through them once `RequestContext` stops mutating `Config`. + +**Removed entirely:** +- `AppConfig::ensure_default_model_id` — redundant with the new + `resolve_model`. Had no callers outside itself (confirmed via + grep). + +### Behavior parity notes + +1. **`from_config` is non-destructive:** it consumes `Config` by + value (not `&Config`) since post-bridge, Config is no longer + needed. This matches the long-term design. + +2. **`from_config` vs `to_app_config` + mutations:** The methods + called inside `from_config` are identical bodies to the ones + currently called on `Config` inside `Config::init`. Env var + reads, wrap validation, doc loader defaults, and user agent + resolution all produce the same state. + +3. **`resolve_model` vs `Config::setup_model`:** + - `Config::setup_model` does TWO things: + (a) ensure `model_id` is non-empty (pick default if empty) + (b) resolve the `Model` struct via `Model::retrieve_model` and + store it in `self.model` + - `AppConfig::resolve_model` only does (a). + - (b) happens today in `cfg.set_model(&model_id)` inside + `Config::setup_model`. In the new architecture, the `Model` + struct lives on `RequestContext.model`, and + `Model::retrieve_model(&app_config, &app_config.model_id, ...)` + will be called inside `RequestContext::new` (or equivalent) + once the bridge is removed in 16e. + +## Files modified + +- `src/config/app_config.rs` — 2 new methods, 4 + `#[allow(dead_code)]` gates removed, 1 method deleted, 8 new + tests. + +## Files NOT modified + +- `src/config/mod.rs` — `Config::init` still runs all mutations on + Config; bridge still copies to AppConfig. Unchanged in 16a. +- `src/config/bridge.rs` — Untouched. Used by `from_config` + internally (`config.to_app_config()`). +- `src/main.rs` — Still uses the bridge flow. Switch happens in 16e. + +## Assumptions made + +1. **`from_config` consumes `Config` by value** (not `&Config`) + — aligns with the long-term design where `Config` is discarded + after conversion. No current caller would benefit from keeping + the Config around after conversion. + +2. **`resolve_model` narrow scope**: only responsible for ensuring + `model_id` is non-empty. Does NOT resolve a `Model` struct — + that's RequestContext's job. This matches the split between + `AppConfig` (the configuration) and `RequestContext` (the + resolved runtime handle). + +3. **`#[allow(dead_code)]` on `from_config` and `resolve_model`**: + they're unused until 16e. The gate is explicit so grep-hunts can + find them when 16e switches over. + +4. **User agent prefix in tests**: I assumed the user agent prefix + is not critical to test literally (it depends on the crate name). + The test checks for a non-"auto" value containing `/` rather + than matching `loki-ai/`. Safer against crate rename. + +## Open questions (parked for later sub-phases) + +1. **Should `from_config` also run secret interpolation?** Currently + `Config::init` does a two-pass YAML parse where the raw content + gets secrets injected from the vault, then the Config is + re-parsed. In the new architecture this belongs in `main.rs` or + a separate helper (the Config comes in already-interpolated). + Not a 16a concern. + +2. **Test naming convention**: Existing tests use + `fn test_name_returns_value_when_condition`. New tests use + `fn from_config_does_thing`. Both styles present in the file; + kept consistent with new code. + +3. **`ensure_default_model_id` deletion**: confirmed via grep that + it had no callers outside itself. Deleted cleanly. If a future + sub-phase needs the Option return variant, it can be + re-added. + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy` — clean, zero warnings +- `cargo test` — 122 passing (114 pre-16a + 8 new), zero failures + +## Remaining work for Step 16 + +- **16b**: Extract `install_builtins()` as top-level free function +- **16c**: Migrate `Vault::init(&Config)` → `Vault::init(&AppConfig)` +- **16d**: Build `AppState::init(app_config, ...).await` +- **16e**: Switch `main.rs` and all 15 `Config::init()` callers to + the new flow +- **16f**: Delete Config runtime fields, bridge.rs, `Config::init`, + duplicated methods + +## Migration direction preserved + +After 16a, no runtime behavior has changed. The new entry point +exists but isn't wired in. The bridge flow continues as before: + +``` +YAML → Config::load_from_file + → Config::init (unchanged, does all current mutations) + - load_envs, set_wrap, setup_document_loaders, ... + - setup_model, load_functions, load_mcp_servers + → cfg.to_app_config() → AppConfig (via bridge) + → cfg.to_request_context(AppState) → RequestContext +``` + +New entry point ready for 16e: + +``` +AppConfig::from_config(config) → AppConfig + (internally: to_app_config, load_envs, set_wrap, + setup_document_loaders, setup_user_agent, resolve_model) +``` diff --git a/docs/implementation/PHASE-1-STEP-16b-NOTES.md b/docs/implementation/PHASE-1-STEP-16b-NOTES.md new file mode 100644 index 0000000..8ed1f49 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16b-NOTES.md @@ -0,0 +1,170 @@ +# Phase 1 Step 16b — Implementation Notes + +## Status + +Done. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Sub-phase goal: "Extract `install_builtins()` as top-level function" + +## Summary + +Extracted `Agent::install_builtin_agents()` and `Macro::install_macros()` +from inside `Config::init` into a new top-level free function +`config::install_builtins()`. Called once from `main.rs` before any +config-loading path. + +Both functions are Config-independent — they just copy embedded +agent/macro assets from the binary into the user's config directory. +Extracting them clears the way for `Config::init`'s eventual +deletion in Step 16f. + +## What was changed + +### `src/config/mod.rs` + +**Added:** +```rust +pub fn install_builtins() -> Result<()> { + Agent::install_builtin_agents()?; + Macro::install_macros()?; + Ok(()) +} +``` + +Placed after the `Config::Default` impl, before the `impl Config` +block. Module-level `pub fn` (not a method on any type). + +**Removed from `Config::init` (inside the async `setup` closure):** +- `Agent::install_builtin_agents()?;` (was at top of setup block) +- `Macro::install_macros()?;` (was at bottom of setup block) + +### `src/main.rs` + +**Added:** +- `install_builtins` to the `use crate::config::{...}` import list +- `install_builtins()?;` call after `setup_logger()?` and before any + of the three config-loading paths (oauth, vault flags, main config) + +### Placement rationale + +The early-return paths (`cli.completions`, `cli.tail_logs`) +legitimately don't need builtins — they return before touching any +config. Those skip the install. + +The three config paths (oauth via `Config::init_bare`, vault flags +via `Config::init_bare`, main via `Config::init`) all benefit from +builtins being installed once at startup. `install_builtins()` is +idempotent — it checks file existence and skips if already present — +so calling it unconditionally in the common path is safe. + +## Behavior parity + +- `install_builtin_agents` and `install_macros` are static methods + with no `&self` or Config arguments. Nothing observable changes + about their execution. +- The two functions ran on every `Config::init` call before. Now + they run once per `main.rs` invocation, which is equivalent for + the REPL and CLI paths. +- `Config::init_bare()` no longer triggers the installs + transitively. The oauth and vault-flag paths now rely on `main.rs` + having called `install_builtins()?` first. This is a minor + behavior shift — those paths previously installed builtins as a + side effect of calling `Config::init_bare`. Since we now call + `install_builtins()` unconditionally in `main.rs` before those + paths, the observable behavior is identical. + +## Files modified + +- `src/config/mod.rs` — added `install_builtins()` free function; + removed 2 calls from `Config::init`. +- `src/main.rs` — added import; added `install_builtins()?` call + after logger setup. + +## Assumptions made + +1. **`install_builtins` should always run unconditionally.** Even + if the user is only running `--completions` or `--tail-logs` + (early-return paths), those return before the install call. + The three config-using paths all benefit from it. No downside to + running it early. + +2. **Module-level `pub fn` is the right API surface.** Could have + made it a method on `AppState` or `AppConfig`, but: + - It's called before any config/state exists + - It has no `self` parameter + - It's a static side-effectful operation (filesystem) + A free function at the module level is the honest signature. + +3. **No tests added.** `install_builtins` is a thin wrapper around + two side-effectful functions that write files. Testing would + require filesystem mocking or temp dirs, which is + disproportionate for a 3-line function. The underlying + `install_builtin_agents` and `install_macros` functions have + existing behavior in the codebase; the extraction doesn't change + their contracts. + +## Open questions + +1. **Should `install_builtins` accept a "skip install" flag?** + Currently it always runs. For a server/REST API deployment, you + might want to skip this to avoid writing to the user's config + dir at startup. Deferring this question until REST API path + exists — can add a flag or a `_skip_install()` variant later. + +2. **Do CI/test environments break because of the filesystem write?** + The install functions already existed in the codebase and ran on + every Config::init. No new risk introduced. Watch for flaky + tests after this change, but expected clean. + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy` — clean, zero warnings +- `cargo test` — 122 passing, zero failures +- Grep confirmation: + - `install_builtin_agents` only defined in `src/config/agent.rs` + and called only via `install_builtins` + - `install_macros` only defined in `src/config/macros.rs` and + called only via `install_builtins` + - `install_builtins` has one caller (`main.rs`) + +## Remaining work for Step 16 + +- **16c**: Migrate `Vault::init(&Config)` → `Vault::init(&AppConfig)`; + eventually move vault ownership to `AppState`. +- **16d**: Build `AppState::init(app_config, ...).await`. +- **16e**: Switch `main.rs` and all 15 `Config::init()` callers to + the new flow. +- **16f**: Delete `Config` runtime fields, `bridge.rs`, `Config::init`, + duplicated methods. + +## Migration direction preserved + +After 16b, `Config::init` no longer handles builtin-asset installation. +This is a small but meaningful piece of responsibility removal — when +`Config::init` is eventually deleted in 16f, we don't need to worry +about orphaning the install logic. + +Startup flow now: +``` +main() + → install_builtins()? [NEW: extracted from Config::init] + → if oauth: Config::init_bare → oauth flow + → if vault flags: Config::init_bare → vault handler + → else: Config::init → to_app_config → AppState → ctx → run +``` + +The `Config::init` calls still do: +- load_envs +- set_wrap +- load_functions +- load_mcp_servers +- setup_model +- setup_document_loaders +- setup_user_agent + +Those move to `AppConfig::from_config` (already built in 16a but +unused) and `AppState::init` (16d) in later sub-phases. diff --git a/docs/implementation/PHASE-1-STEP-16c-NOTES.md b/docs/implementation/PHASE-1-STEP-16c-NOTES.md new file mode 100644 index 0000000..f59cefb --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16c-NOTES.md @@ -0,0 +1,196 @@ +# Phase 1 Step 16c — Implementation Notes + +## Status + +Done. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Sub-phase goal: "Migrate Vault onto AppState; Vault::init takes + `&AppConfig`" + +## Summary + +Changed `Vault::init(&Config)` and `Vault::init_bare()` to operate on +`AppConfig` instead of `Config`. Simplified `Vault::handle_vault_flags` +to accept a `&Vault` directly instead of extracting one from a Config +argument. Deleted the duplicate `Config::vault_password_file` method +(the canonical version lives on `AppConfig`). + +Vault is now fully Config-independent at the signature level. The +`Config.vault` runtime field still exists because it's populated +inside `Config::init` (for the current bridge-era flow), but nothing +about `Vault`'s API references `Config` anymore. The field itself +gets deleted in Step 16f when Config becomes a pure POJO. + +## What was changed + +### `src/vault/mod.rs` + +**Signature change:** +```rust +// Before +pub fn init(config: &Config) -> Self +// After +pub fn init(config: &AppConfig) -> Self +``` + +**`Vault::init_bare` now uses `AppConfig::default()`** instead of +`Config::default()` to get the default vault password file path. +Behavioral parity — both `.default().vault_password_file()` calls +resolve to the same path (the fallback from `gman::config`). + +**`Vault::handle_vault_flags` simplified:** +```rust +// Before +pub fn handle_vault_flags(cli: Cli, config: Config) -> Result<()> +// After +pub fn handle_vault_flags(cli: Cli, vault: &Vault) -> Result<()> +``` + +The old signature took a `Config` by value just to extract +`config.vault`. The new signature takes the Vault directly, which +decouples this function from Config entirely. Callers pass +`&config.vault` or equivalent. + +**Import updated:** +- `use crate::config::Config` removed +- `use crate::config::AppConfig` added + +### `src/config/mod.rs` + +**In `Config::init`:** +```rust +// Before +let vault = Vault::init(config); +// After +let vault = Vault::init(&config.to_app_config()); +``` + +This is a transitional call — `Config::init` builds a temporary +`AppConfig` via the bridge just to satisfy the new signature. This +temporary conversion disappears in Step 16e when `main.rs` stops +calling `Config::init` entirely and `AppConfig` is built first. + +**Deleted `Config::vault_password_file`** method. The identical body +lives on `AppConfig`. All callers go through AppConfig now. + +### `src/main.rs` + +**Vault-flags path:** +```rust +// Before +return Vault::handle_vault_flags(cli, Config::init_bare()?); +// After +let cfg = Config::init_bare()?; +return Vault::handle_vault_flags(cli, &cfg.vault); +``` + +This is a minor restructure — same observable behavior, but the +Vault is extracted from the Config and passed directly. Makes the +vault-flags path obvious about what it actually needs (a Vault, not +a Config). + +## Behavior parity + +- `Vault::init` reads `config.vault_password_file()` — identical + method on both Config and AppConfig (removed from Config in this + step, kept on AppConfig). +- Password file initialization (`ensure_password_file_initialized`) + still runs in `Vault::init` as before. +- `Vault::init_bare` fallback path resolves to the same default + password file location. +- `handle_vault_flags` operates on the same Vault instance either + way — just receives it directly instead of indirectly via Config. + +## Files modified + +- `src/vault/mod.rs` — imports, `init` signature, `init_bare` + fallback source, `handle_vault_flags` signature. +- `src/config/mod.rs` — transitional `to_app_config()` call in + `Config::init`; deleted duplicate `vault_password_file` method. +- `src/main.rs` — vault-flags path takes `&cfg.vault` directly. + +## Assumptions made + +1. **`AppConfig::default().vault_password_file()` behaves identically + to `Config::default().vault_password_file()`.** Verified by + comparing method bodies — identical logic, same fallback via + `gman::config::Config::local_provider_password_file()`. Tests + confirm 122 passing, no regressions. + +2. **Transitional `&config.to_app_config()` in `Config::init` is + acceptable.** The conversion happens once per Config::init call + — trivial cost for a non-hot path. Disappears entirely in 16e. + +3. **`handle_vault_flags` taking `&Vault` is a strict improvement.** + The old signature took `Config` by value (wasteful for a function + that only needed one field). The new signature is honest about + its dependency. + +4. **`Config.vault` field stays for now.** The `#[serde(skip)]` field + on Config still exists because `Config::init` populates it for + downstream Bridge flow consumers. Deletion deferred to 16f. + +## Open questions + +1. **Should `Vault::init` return `Result` instead of panicking?** + Currently uses `.expect("Failed to initialize password file")`. + The vault flags path can't do anything useful without a vault, + so panic vs early return is pragmatically equivalent. Leaving + as-is to minimize change surface in 16c. + +2. **Is `Config::init_bare` still needed after 16e?** It's called + from the oauth path in `main.rs`. In 16e we'll audit whether + those paths really need full Config init or just an AppConfig. + Deferred to 16e. + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy` — clean, zero warnings +- `cargo test` — 122 passing, zero failures +- Grep confirmation: + - `Vault::init(` has one caller (in `Config::init`); one + remaining via test path (none found — deferred to 16d/16e where + AppState::init will own vault construction) + - `Vault::init_bare(` has one caller (via `interpolate_secrets` + flow); no other references + - `Config::vault_password_file` — zero references anywhere + - `Vault::handle_vault_flags` — single caller in `main.rs`, + signature verified + +## Remaining work for Step 16 + +- **16d**: Build `AppState::init(app_config, ...).await` that takes + ownership of vault construction (replacing the current + `AppState { vault: cfg.vault.clone(), ... }` pattern). +- **16e**: Switch `main.rs` and all 15 `Config::init()` callers to + the new flow. +- **16f**: Delete `Config.vault` field, `Config::init`, bridge.rs, + and all remaining Config runtime fields. + +## Migration direction preserved + +Before 16c: +``` +Vault::init(&Config) ← tight coupling to Config +Vault::init_bare() ← uses Config::default() internally +handle_vault_flags(Cli, Config) ← takes Config, extracts vault +Config::vault_password_file() ← duplicate with AppConfig's +``` + +After 16c: +``` +Vault::init(&AppConfig) ← depends only on AppConfig +Vault::init_bare() ← uses AppConfig::default() internally +handle_vault_flags(Cli, &Vault) ← takes Vault directly +Config::vault_password_file() ← DELETED +``` + +The Vault module now has no Config dependency in its public API. +This means Step 16d can build `AppState::init` that calls +`Vault::init(&app_config)` without touching Config at all. It also +means `Config` is one step closer to being a pure POJO — one fewer +method on its surface, one fewer implicit dependency. diff --git a/docs/implementation/PHASE-1-STEP-16d-NOTES.md b/docs/implementation/PHASE-1-STEP-16d-NOTES.md new file mode 100644 index 0000000..3b4aef5 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16d-NOTES.md @@ -0,0 +1,271 @@ +# Phase 1 Step 16d — Implementation Notes + +## Status + +Done. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Sub-phase goal: "Build `AppState::init(app_config, ...).await` + absorbing MCP registry startup and global functions loading" + +## Summary + +Added `AppState::init()` async constructor that self-initializes all +process-wide shared state from an `Arc` and startup +context. Two new fields on `AppState`: `mcp_registry` (holds initial +MCP server Arcs alive) and `functions` (the global base +`Functions`). Changed `McpRegistry::init` to take `&AppConfig + +&Vault` instead of `&Config`. + +The constructor is dead-code-gated (`#[allow(dead_code)]`) until +Step 16e switches `main.rs` over to call it. The bridge flow +continues to populate the new fields via Config's existing +`functions` and `mcp_registry` so nothing breaks. + +## What was changed + +### `src/mcp/mod.rs` + +**`McpRegistry::init` signature change:** +```rust +// Before +pub async fn init( + log_path: Option, + start_mcp_servers: bool, + enabled_mcp_servers: Option, + abort_signal: AbortSignal, + config: &Config, +) -> Result + +// After +pub async fn init( + log_path: Option, + start_mcp_servers: bool, + enabled_mcp_servers: Option, + abort_signal: AbortSignal, + app_config: &AppConfig, + vault: &Vault, +) -> Result +``` + +The function reads two things from its config argument: +- `config.vault` for secret interpolation → now takes `&Vault` + directly +- `config.mcp_server_support` for the start-servers gate → a serde + field already present on `AppConfig` + +Both dependencies are now explicit. No Config reference anywhere in +the MCP module. + +**Imports updated:** +- Removed `use crate::config::Config` +- Added `use crate::config::AppConfig` +- Added `use crate::vault::Vault` + +### `src/config/mod.rs` + +**`Config::load_mcp_servers` updated** to build a temporary AppConfig +via `self.to_app_config()` and pass it plus `&self.vault` to +`McpRegistry::init`. This is a transitional bridge — disappears in +16e when `main.rs` stops calling `Config::init` and +`AppState::init` becomes the sole entry point. + +### `src/config/app_state.rs` + +**New fields:** +```rust +pub mcp_registry: Option>, +pub functions: Functions, +``` + +**New `AppState::init()` async constructor** absorbs: +- `Vault::init(&config)` (replaces the old + `AppState { vault: cfg.vault.clone() }` pattern) +- `McpRegistry::init(...)` (previously inside `Config::init`) +- Registers initial MCP servers with `McpFactory` via + `insert_active(McpServerKey, &handle)` — this is NEW behavior, see + below. +- `Functions::init(config.visible_tools)` (previously inside + `Config::init`) +- `functions.append_mcp_meta_functions(...)` when MCP support is on + and servers started +- Wraps `McpRegistry` in `Arc` to keep initial server handles alive + across scope transitions (see below) + +**Imports expanded:** +- `McpServerKey` from `super::mcp_factory` +- `Functions` from `crate::function` +- `McpRegistry` from `crate::mcp` +- `AbortSignal` from `crate::utils` +- `Vault` from `crate::vault` +- `anyhow::Result` + +### `src/main.rs` + +**AppState struct literal extended** to populate the two new fields +from `cfg.mcp_registry` and `cfg.functions`. This keeps the bridge +flow working unchanged. When 16e replaces this struct literal with +`AppState::init(...)`, these field references go away entirely. + +### `src/function/supervisor.rs` + +**Child AppState construction extended** to propagate the new +fields from parent: `mcp_registry: ctx.app.mcp_registry.clone()` and +`functions: ctx.app.functions.clone()`. This maintains parent-child +sharing of the MCP factory cache (which was already fixed earlier in +this work stream). + +### `src/config/request_context.rs` and `src/config/session.rs` + +**Test helper AppState construction extended** to include the two +new fields with safe defaults (`None`, and `cfg.functions.clone()` +respectively). + +## New behavior: McpFactory pre-registration + +`AppState::init` registers every initial server with `McpFactory` via +`insert_active`. This fixes a latent issue in the current bridge +flow: + +- Before: initial servers were held on `Config.mcp_registry.servers`; + when the first scope transition (e.g., `.role coder`) ran + `rebuild_tool_scope`, it called `McpFactory::acquire(name, spec, + log_path)` which saw an empty cache and **spawned duplicate + servers**. The original servers died when the initial ToolScope + was replaced. +- After (via `AppState::init`): the factory's Weak map is seeded + with the initial server Arcs. The registry itself is wrapped in + Arc and held on AppState so the Arcs stay alive. Scope + transitions now hit the factory cache and reuse the same + subprocesses. + +This is a real improvement that shows up once `main.rs` switches to +`AppState::init` in 16e. During the 16d bridge window, nothing +reads the factory pre-registration yet. + +## Behavior parity (16d window) + +- `main.rs` still calls `Config::init`, still uses the bridge; all + new fields populated from Config's own state. +- `AppState::init` is present but unused in production code paths. +- Test helpers still use struct literals; they pass `None` for + `mcp_registry` and clone `cfg.functions` which is the same as + what the bridge was doing. +- No observable runtime change for users. + +## Files modified + +- `src/mcp/mod.rs` — `McpRegistry::init` signature; imports. +- `src/config/mod.rs` — `Config::load_mcp_servers` bridges to new + signature. +- `src/config/app_state.rs` — added 2 fields, added `init` + constructor, expanded imports. +- `src/main.rs` — struct literal populates 2 new fields. +- `src/function/supervisor.rs` — child struct literal populates 2 + new fields. +- `src/config/request_context.rs` — test helper populates 2 new + fields. +- `src/config/session.rs` — test helper populates 2 new fields. + +## Assumptions made + +1. **`McpFactory::insert_active` with Weak is sufficient to seed the + cache.** The initial ServerArcs live on `AppState.mcp_registry` + (wrapped in Arc to enable clone across child states). Scope + transitions call `McpFactory::acquire` which does + `try_get_active(key).unwrap_or_else(spawn_new)`. The Weak in + factory upgrades because Arc is alive in `mcp_registry`. Verified + by reading code paths; not yet verified at runtime since bridge + still drives today's flow. + +2. **`functions: Functions` is Clone-safe.** The struct contains + `Vec` and related fields; cloning is cheap + enough at startup and child-agent spawn. Inspected the + definition; no references to check. + +3. **`mcp_server_support` gate still applies.** `Config::init` used + to gate the MCP meta function append with both `is_empty()` and + `mcp_server_support`; `AppState::init` preserves both checks. + Parity confirmed. + +4. **Mode-specific function additions (REPL-only + `append_user_interaction_functions`) do NOT live on AppState.** + They are added per-scope in `rebuild_tool_scope` and in the + initial `RequestContext::new` path. `AppState.functions` is the + mode-agnostic base. This matches the long-term design + (RequestContext owns mode-aware additions). + +5. **`mcp_registry: Option>` vs `McpRegistry`.** + Went with `Option>` so: + - `None` when no MCP config exists (can skip work) + - `Arc` so parent/child AppStates share the same registry + (keeping initial server handles alive across the tree) + +## Open questions + +1. **Registry on AppState vs factory-owned lifecycle**. The factory + holds Weak; the registry holds Arc. Keeping the registry alive + on AppState extends server lifetime to the process lifetime. + This differs from the current "servers die on scope transition" + behavior. In practice this is what users expect — start the + servers once, keep them alive. But it means long-running REPL + sessions retain all server subprocesses even if the user switches + away from them. Acceptable trade-off for Phase 1. + +2. **Should `AppState::init` return `Arc` directly?** + Currently returns `Self`. Caller wraps in Arc. Symmetric with + other init functions; caller has full flexibility. Keep as-is. + +3. **Unit tests for `AppState::init`.** Didn't add any because the + function is heavily async, touches filesystem (paths), + subprocess startup (MCP), and the vault. A meaningful unit test + would require mocking. Integration-level validation happens in + 16e when main.rs switches over. Deferred. + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy` — clean, zero warnings +- `cargo test` — 122 passing, zero failures +- New `AppState::init` gated with `#[allow(dead_code)]` — no + warnings for being unused + +## Remaining work for Step 16 + +- **16e**: Switch `main.rs` to call + `AppState::init(app_config, log_path, start_mcp_servers, + abort_signal).await?` instead of the bridge pattern. Audit the + 15 `Config::init()` callers. Remove `#[allow(dead_code)]` from + `AppConfig::from_config`, `AppConfig::resolve_model`, and + `AppState::init`. +- **16f**: Delete `Config.vault`, `Config.functions`, + `Config.mcp_registry`, all other `#[serde(skip)]` runtime fields. + Delete `Config::init`, `Config::load_envs`, `Config::load_functions`, + `Config::load_mcp_servers`, `Config::setup_model`, + `Config::set_model`, etc. Delete `bridge.rs`. + +## Migration direction preserved + +Before 16d: +``` +AppState { + config, vault, mcp_factory, rag_cache, mcp_config, mcp_log_path +} +``` +Constructed only via struct literal from Config fields via bridge. + +After 16d: +``` +AppState { + config, vault, mcp_factory, rag_cache, mcp_config, mcp_log_path, + mcp_registry, functions +} +impl AppState { + pub async fn init(config, log_path, start_mcp_servers, abort_signal) + -> Result +} +``` +New fields present on all code paths. New self-initializing +constructor ready for 16e's switchover. diff --git a/docs/implementation/PHASE-1-STEP-16e-NOTES.md b/docs/implementation/PHASE-1-STEP-16e-NOTES.md new file mode 100644 index 0000000..3240441 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16e-NOTES.md @@ -0,0 +1,228 @@ +# Phase 1 Step 16e — Implementation Notes + +## Status + +Done. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Sub-phase goal: "Switch main.rs and all Config::init() callers + to the new flow" + +## Summary + +`main.rs` and `cli/completer.rs` no longer call `Config::init` or +`Config::init_bare` — they use the new flow: +`Config::load_with_interpolation` → `AppConfig::from_config` → +`AppState::init` → `RequestContext::bootstrap`. + +The bridge `Config::to_request_context` and the old `Config::init` +are now dead code, gated with `#[allow(dead_code)]` pending deletion +in 16f. + +## What was changed + +### New helpers + +**`Config::load_with_interpolation(info_flag: bool) -> Result`** in +`src/config/mod.rs` — absorbs the two-pass YAML parse with secret +interpolation. Handles: +1. Missing config file (creates via `create_config_file` if TTY, or + `load_dynamic` from env vars) +2. Reading the raw YAML content +3. Bootstrapping a Vault from the freshly-parsed Config +4. Interpolating secrets +5. Re-parsing Config if interpolation changed anything +6. Sets `config.vault` (legacy field — deleted in 16f) + +**`config::default_sessions_dir() -> PathBuf`** and +**`config::list_sessions() -> Vec`** free functions — +provide session listing without needing a Config instance. Used by +the session completer. + +**`RequestContext::bootstrap(app: Arc, working_mode, +info_flag) -> Result`** in `src/config/request_context.rs` — +the new entry point for creating the initial RequestContext. Builds: +- Resolved `Model` from `app.config.model_id` +- `ToolScope.functions` cloned from `app.functions` with + `append_user_interaction_functions` added in REPL mode +- `ToolScope.mcp_runtime` synced from `app.mcp_registry` + +### Made public in Config for new flow + +- `Config::load_from_file` (was `fn`) +- `Config::load_from_str` (was `fn`) +- `Config::load_dynamic` (was `fn`) +- `config::create_config_file` (was `async fn`) + +### src/main.rs + +Three startup paths rewired: + +```rust +// Path 1: --authenticate +let cfg = Config::load_with_interpolation(true).await?; +let app_config = AppConfig::from_config(cfg)?; +let (client_name, provider) = + resolve_oauth_client(client_arg.as_deref(), &app_config.clients)?; +oauth::run_oauth_flow(&*provider, &client_name).await?; + +// Path 2: vault flags +let cfg = Config::load_with_interpolation(true).await?; +let app_config = AppConfig::from_config(cfg)?; +let vault = Vault::init(&app_config); +return Vault::handle_vault_flags(cli, &vault); + +// Path 3: main +let cfg = Config::load_with_interpolation(info_flag).await?; +let app_config: Arc = Arc::new(AppConfig::from_config(cfg)?); +let app_state: Arc = Arc::new( + AppState::init(app_config, log_path, start_mcp_servers, abort_signal.clone()).await? +); +let ctx = RequestContext::bootstrap(app_state, working_mode, info_flag)?; +``` + +No more `Config::init`, `Config::to_app_config`, `cfg.mcp_registry`, +or `cfg.to_request_context` references in `main.rs`. + +### src/cli/completer.rs + +Three completers that needed config access updated: + +- `model_completer` → uses new `load_app_config_for_completion()` + helper (runs `Config::load_with_interpolation` synchronously from + the completion context; async via `Handle::try_current` or a fresh + runtime) +- `session_completer` → uses the new free function + `list_sessions()` (no Config needed) +- `secrets_completer` → uses `Vault::init(&app_config)` directly + +### #[allow(dead_code)] removed + +- `AppConfig::from_config` +- `AppConfig::resolve_model` +- `AppState::init` +- `AppState.rag_cache` (was flagged dead; now wired in) + +### #[allow(dead_code)] added (temporary, deleted in 16f) + +- `Config::init_bare` — no longer called +- `Config::sessions_dir` — replaced by free function +- `Config::list_sessions` — replaced by free function +- `Config::to_request_context` — replaced by `RequestContext::bootstrap` + +## Behavior parity + +- `main.rs` startup now invokes: + - `install_builtins()` (installs builtin global tools, agents, + macros — same files get copied as before, Step 16b) + - `Config::load_with_interpolation` (same YAML loading + secret + interpolation as old `Config::init`) + - `AppConfig::from_config` (same env/wrap/docs/user-agent/model + resolution as old Config mutations) + - `AppState::init` (same vault init + MCP registry startup + + global Functions loading as old Config methods, now owned by + AppState; also pre-registers initial servers with McpFactory — + new behavior that fixes a latent cache miss bug) + - `RequestContext::bootstrap` (same initial state as old bridge + `to_request_context`: resolved Model, Functions with REPL + extensions, MCP runtime from registry) + +- Completer paths now use a lighter-weight config load (no MCP + startup) which is appropriate since shell completion isn't + supposed to start subprocesses. + +## Files modified + +- `src/config/mod.rs` — added `load_with_interpolation`, + `default_sessions_dir`, `list_sessions`; made 3 methods public; + added `#[allow(dead_code)]` to `Config::init_bare`, + `sessions_dir`, `list_sessions`. +- `src/config/request_context.rs` — added `bootstrap`. +- `src/config/app_config.rs` — removed 2 `#[allow(dead_code)]` + gates. +- `src/config/app_state.rs` — removed 2 `#[allow(dead_code)]` + gates. +- `src/config/bridge.rs` — added `#[allow(dead_code)]` to + `to_request_context`. +- `src/main.rs` — rewired three startup paths. +- `src/cli/completer.rs` — rewired three completers. + +## Assumptions made + +1. **Completer helper runtime handling**: The three completers run + in a sync context (clap completion). The new + `load_app_config_for_completion` uses + `Handle::try_current().ok()` to detect if a tokio runtime + exists; if so, uses `block_in_place`; otherwise creates a + fresh runtime. This matches the old `Config::init_bare` pattern + (which also used `block_in_place` + `block_on`). + +2. **`Config::to_request_context` kept with `#[allow(dead_code)]`**: + It's unused now but 16f deletes it cleanly. Leaving it in place + keeps 16e a non-destructive switchover. + +3. **`RequestContext::bootstrap` returns `Result` not + `Arc`**: caller decides wrapping. main.rs doesn't wrap; + the REPL wraps `Arc>` a few lines later. + +4. **`install_builtin_global_tools` added to `install_builtins`**: + A function added in user's 16b commit extracted builtin tool + installation out of `Functions::init` into a standalone function. + My Step 16b commit that extracted `install_builtins` missed + including this function — fixed in this step. + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy` — clean, zero warnings +- `cargo test` — 122 passing, zero failures +- Grep confirmation: + - `Config::init(` — only called from `Config::init_bare` (which + is now dead) + - `Config::init_bare` — no external callers (test helper uses + `#[allow(dead_code)]`) + - `to_request_context` — zero callers outside bridge.rs + - `cfg.mcp_registry` / `cfg.functions` / `cfg.vault` references + in main.rs — zero + +## Remaining work for Step 16 + +- **16f**: Delete all `#[allow(dead_code)]` scaffolding: + - `Config::init`, `Config::init_bare` + - `Config::sessions_dir`, `Config::list_sessions` + - `Config::set_wrap`, `Config::setup_document_loaders`, + `Config::setup_user_agent`, `Config::load_envs`, + `Config::load_functions`, `Config::load_mcp_servers`, + `Config::setup_model`, `Config::set_model`, + `Config::role_like_mut`, `Config::vault_password_file` + - `bridge.rs` — delete entirely + - All `#[serde(skip)]` runtime fields on `Config` + - `mod bridge;` declaration + +After 16f, `Config` will be a pure serde POJO with only serialized +fields and `load_from_file` / `load_from_str` / `load_dynamic` / +`load_with_interpolation` methods. + +## Migration direction achieved + +Before 16e: +``` +main.rs: Config::init → to_app_config → AppState {...} → to_request_context +``` + +After 16e: +``` +main.rs: + install_builtins() + Config::load_with_interpolation → AppConfig::from_config + AppState::init(app_config, ...).await + RequestContext::bootstrap(app_state, working_mode, info_flag) +``` + +No more god-init. Each struct owns its initialization. The REST +API path is now trivial: skip `install_builtins()` if not desired, +call `AppConfig::from_config(yaml_string)`, call +`AppState::init(...)`, create per-request `RequestContext` as +needed. diff --git a/docs/implementation/PHASE-1-STEP-16f-NOTES.md b/docs/implementation/PHASE-1-STEP-16f-NOTES.md new file mode 100644 index 0000000..b31fff5 --- /dev/null +++ b/docs/implementation/PHASE-1-STEP-16f-NOTES.md @@ -0,0 +1,355 @@ +# Phase 1 Step 16f — Implementation Notes + +## Status + +Done. Phase 1 Step 16 (Config → AppConfig migration) complete. + +## Plan reference + +- Parent plan: `docs/implementation/PHASE-1-STEP-16-NOTES.md` +- Predecessor: `docs/implementation/PHASE-1-STEP-16e-NOTES.md` +- Sub-phase goal: "Delete all #[allow(dead_code)] scaffolding from + Config and bridge.rs, delete runtime fields from Config, delete + bridge.rs entirely." + +## Summary + +`Config` is now a pure serde POJO. `bridge.rs` is gone. Every +runtime field, every `Config::init*` flavor, and every Config method +that was scaffolding for the old god-init has been deleted. The +project compiles clean, clippy clean, and all 122 tests pass. + +## What was changed + +### Deleted: `src/config/bridge.rs` + +Whole file removed. `mod bridge;` declaration in `config/mod.rs` +removed. The two methods (`Config::to_app_config` and +`Config::to_request_context`) had no remaining callers after 16e. + +### `src/config/mod.rs` — Config slimmed to a POJO + +**Deleted runtime (`#[serde(skip)]`) fields from `Config`:** +- `vault`, `macro_flag`, `info_flag`, `agent_variables` +- `model`, `functions`, `mcp_registry`, `working_mode`, + `last_message` +- `role`, `session`, `rag`, `agent`, `tool_call_tracker` +- `supervisor`, `parent_supervisor`, `self_agent_id`, + `current_depth`, `inbox`, `root_escalation_queue` + +**Deleted methods on `Config`:** +- `init`, `init_bare` (god-init replaced by + `load_with_interpolation` + `AppConfig::from_config` + + `AppState::init` + `RequestContext::bootstrap`) +- `sessions_dir`, `list_sessions` (replaced by + `config::default_sessions_dir` / `config::list_sessions` free + functions for use without a Config; per-context paths live on + `RequestContext::sessions_dir` / `RequestContext::list_sessions`) +- `role_like_mut` (lives on `RequestContext` post-migration) +- `set_wrap`, `setup_document_loaders`, `setup_user_agent`, + `load_envs` (lives on `AppConfig` post-migration) +- `set_model`, `setup_model` (model resolution now in + `AppConfig::resolve_model`; per-scope model selection lives on + `RequestContext`) +- `load_functions`, `load_mcp_servers` (absorbed by + `AppState::init`) + +**Default impl entries** for the deleted runtime fields removed. + +**Imports cleaned up:** removed unused `ToolCallTracker`, +`McpRegistry`, `Supervisor`, `EscalationQueue`, `Inbox`, `RwLock`, +`ColorScheme`, `QueryOptions`, `color_scheme`, `Handle`. Kept +`Model`, `ModelType`, `GlobalVault` because sibling modules +(`role.rs`, `input.rs`, `agent.rs`, `session.rs`) use +`use super::*;` and depend on those re-exports. + +**Removed assertions** for the deleted runtime fields from +`config_defaults_match_expected` test. + +### `src/config/mod.rs` — `load_with_interpolation` no longer touches AppConfig::to_app_config + +Previously called `config.to_app_config()` to build a Vault for +secret interpolation. Now constructs a minimal `AppConfig` inline +with only `vault_password_file` populated, since that's all +`Vault::init` reads. Also removed the `config.vault = Arc::new(vault)` +assignment that was the last write to the deleted runtime field. + +### `src/config/mod.rs` — `vault_password_file` made `pub(super)` + +Previously private. Now `pub(super)` so `AppConfig::from_config` (a +sibling module under `config/`) can read it during the field-copy. + +### `src/config/app_config.rs` — `AppConfig::from_config` self-contained + +Previously delegated to `Config::to_app_config()` (lived on bridge) +for the field-copy. Now inlines the field-copy directly in +`from_config`, then runs `load_envs`, `set_wrap`, +`setup_document_loaders`, `setup_user_agent`, and `resolve_model` +as before. + +**Removed `#[allow(dead_code)]` from `AppConfig.model_id`** — it's +read from `app.config.model_id` in `RequestContext::bootstrap` so +the lint exemption was stale. + +**Test refactor:** the three `to_app_config_*` tests rewritten as +`from_config_*` tests using `AppConfig::from_config(cfg).unwrap()`. +A `ClientConfig::default()` and non-empty `model_id: "test-model"` +were added so `resolve_model()` doesn't bail with "No available +model" during the runtime initialization. + +### `src/config/session.rs` — Test helper rewired + +`session_new_from_ctx_captures_save_session` rewritten to build the +test `AppState` directly with `AppConfig::default()`, +`Vault::default()`, `Functions::default()` instead of going through +`cfg.to_app_config()` / `cfg.vault` / `cfg.functions`. Then uses +`RequestContext::new(app_state, WorkingMode::Cmd)` instead of the +deleted `cfg.to_request_context(app_state)`. + +### `src/config/request_context.rs` — Test helpers rewired + +The `app_state_from_config(&Config)` helper rewritten as +`default_app_state()` — no longer takes a Config, builds AppState +from `AppConfig::default()` + `Vault::default()` + `Functions::default()` +directly. The two callers (`create_test_ctx`, +`update_app_config_persists_changes`) updated. + +The `to_request_context_creates_clean_state` test renamed to +`new_creates_clean_state` and rewritten to use `RequestContext::new` +directly. + +### Doc comment refresh + +Three module docstrings rewritten to reflect the post-16f world: + +- `app_config.rs` — was "Phase 1 Step 0 ... not yet wired into the + runtime." Now describes `AppConfig` as the runtime-resolved + view of YAML, built via `AppConfig::from_config`. +- `app_state.rs` — was "Step 6.5 added mcp_factory and rag_cache + ... neither wired in yet ... Step 8+ will connect." Now + describes `AppState::init` as the wiring point. +- `request_context.rs` — was an extensive description of the + bridge window with flat fields vs sub-struct fields, citing + `Config::to_request_context`. Now describes the type's actual + ownership/lifecycle without referring to deleted entry points. +- `tool_scope.rs` — was "Step 6.5 scope ... unused parallel + structure ... Step 8 will rewrite." Now describes `ToolScope` + as the live per-scope tool runtime. + +(Other phase-era comments in `paths.rs`, `mcp_factory.rs`, +`rag_cache.rs` not touched. They reference Step 2 / Step 6.5 / +Step 8 but the affected types still exist and the descriptions +aren't actively misleading — those files weren't part of 16f +scope. Future cleanup if desired.) + +## What Config looks like now + +```rust +#[derive(Debug, Clone, Deserialize)] +#[serde(default)] +pub struct Config { + pub model_id: String, + pub temperature: Option, + pub top_p: Option, + pub dry_run: bool, + pub stream: bool, + pub save: bool, + pub keybindings: String, + pub editor: Option, + pub wrap: Option, + pub wrap_code: bool, + pub(super) vault_password_file: Option, + pub function_calling_support: bool, + pub mapping_tools: IndexMap, + pub enabled_tools: Option, + pub visible_tools: Option>, + pub mcp_server_support: bool, + pub mapping_mcp_servers: IndexMap, + pub enabled_mcp_servers: Option, + pub repl_prelude: Option, + pub cmd_prelude: Option, + pub agent_session: Option, + pub save_session: Option, + pub compression_threshold: usize, + pub summarization_prompt: Option, + pub summary_context_prompt: Option, + pub rag_embedding_model: Option, + pub rag_reranker_model: Option, + pub rag_top_k: usize, + pub rag_chunk_size: Option, + pub rag_chunk_overlap: Option, + pub rag_template: Option, + pub document_loaders: HashMap, + pub highlight: bool, + pub theme: Option, + pub left_prompt: Option, + pub right_prompt: Option, + pub user_agent: Option, + pub save_shell_history: bool, + pub sync_models_url: Option, + pub clients: Vec, +} + +impl Config { + pub async fn load_with_interpolation(info_flag: bool) -> Result { ... } + pub fn load_from_file(config_path: &Path) -> Result<(Self, String)> { ... } + pub fn load_from_str(content: &str) -> Result { ... } + pub fn load_dynamic(model_id: &str) -> Result { ... } +} +``` + +Just shape + four loaders. The three associated functions that +used to live here (`search_rag`, `load_macro`, `sync_models`) +were relocated in the 16f cleanup pass below — none of them +touched Config state, they were squatters from the god-object +era. + +## Assumptions made + +1. **Doc cleanup scope**: The user asked to "delete the dead-code + scaffolding from Config and bridge.rs." Doc comments in + `paths.rs`, `mcp_factory.rs`, `rag_cache.rs` still reference + "Phase 1 Step 6.5 / Step 8" but the types they describe are + still real and the descriptions aren't actively wrong (just + historically dated). Left them alone. Updated only the docs in + `app_config.rs`, `app_state.rs`, `request_context.rs`, and + `tool_scope.rs` because those were either pointing at deleted + types (`Config::to_request_context`) or making explicitly + false claims ("not wired into the runtime yet"). + +2. **`set_*_default` helpers on AppConfig**: Lines 485–528 of + `app_config.rs` define nine `#[allow(dead_code)]` + `set_*_default` methods. These were added in earlier sub-phases + as planned setters for runtime overrides. They're still unused. + The 16-NOTES plan flagged them ("set_*_default ... become + reachable") but reachability never happened. Since the user's + directive was specifically "Config and bridge.rs scaffolding," + I left these untouched. Removing them is independent cleanup + that doesn't block 16f. + +3. **`reload_current_model` on RequestContext**: Same situation — + one `#[allow(dead_code)]` left on a RequestContext method. + Belongs to a different cleanup task; not Config or bridge + scaffolding. + +4. **`vault_password_file` visibility**: `Config.vault_password_file` + was a private field. Made it `pub(super)` so + `AppConfig::from_config` (sibling under `config/`) can read it + for the field-copy. This is the minimum viable visibility — + no code outside `config/` can touch it, matching the previous + intent. + +5. **Bootstrap vault construction in `load_with_interpolation`**: + Used `AppConfig { vault_password_file: ..., ..AppConfig::default() }` + instead of e.g. a dedicated helper. The vault only reads + `vault_password_file` so this is sufficient. A comment explains + the dual-vault pattern (bootstrap for secret interpolation vs + canonical from `AppState::init`). + +## Verification + +- `cargo check` — clean, zero warnings +- `cargo clippy --all-targets` — clean, zero warnings +- `cargo test` — 122 passing, zero failures (same count as 16e) +- Grep confirmation: + - `to_app_config` — zero hits in `src/` + - `to_request_context` — zero hits in `src/` + - `Config::init` / `Config::init_bare` — zero hits in `src/` + - `bridge::` / `config::bridge` / `mod bridge` — zero hits in `src/` + - `src/config/bridge.rs` — file deleted +- Config now contains only serde fields and load/helper + functions; no runtime state. + +## Phase 1 Step 16 — overall outcome + +The full migration is complete: + +| Sub-phase | Outcome | +|-----------|---------| +| 16a | `AppConfig::from_config` built | +| 16b | `install_builtins()` extracted | +| 16c | Vault on AppState (already-existing field, `Vault::init` rewired to `&AppConfig`) | +| 16d | `AppState::init` built | +| 16e | `main.rs` + completers + `RequestContext::bootstrap` switched to new flow | +| 16f | Bridge + Config runtime fields + dead methods deleted | + +`Config` is a serde POJO. `AppConfig` is the runtime-resolved +process-wide settings. `AppState` owns process-wide services +(vault, MCP registry, base functions, MCP factory, RAG cache). +`RequestContext` owns per-request mutable state. Each struct +owns its initialization. The REST API surface is now trivial: +parse YAML → `AppConfig::from_config` → `AppState::init` → +per-request `RequestContext`. + +## Files modified (16f) + +- `src/config/mod.rs` — runtime fields/methods/Default entries + deleted, imports cleaned up, `vault_password_file` made + `pub(super)`, `load_with_interpolation` decoupled from + `to_app_config`, default-test simplified +- `src/config/app_config.rs` — `from_config` inlines field-copy, + `#[allow(dead_code)]` on `model_id` removed, three tests + rewritten, module docstring refreshed +- `src/config/session.rs` — test helper rewired, imports updated +- `src/config/request_context.rs` — test helpers rewired, + imports updated, module docstring refreshed +- `src/config/app_state.rs` — module docstring refreshed +- `src/config/tool_scope.rs` — module docstring refreshed + +## Files deleted (16f) + +- `src/config/bridge.rs` + +## 16f cleanup pass — Config straggler relocation + +After the main 16f deletions landed, three associated functions +remained on `impl Config` that took no `&self` and didn't touch +any Config field — they were holdovers from the god-object era, +attached to `Config` only because Config used to be the +namespace for everything. Relocated each to its rightful owner: + +| Method | New home | Why | +|--------|----------|-----| +| `Config::load_macro(name)` | `Macro::load(name)` in `src/config/macros.rs` | Sibling of `Macro::install_macros` already there. The function loads a macro from disk and parses it into a `Macro` — pure macro concern. | +| `Config::search_rag(app, rag, text, signal)` | `Rag::search_with_template(&self, app, text, signal)` in `src/rag/mod.rs` | Operates on a `Rag` instance and one field of `AppConfig`. Pulled `RAG_TEMPLATE` constant along with it. | +| `Config::sync_models(url, signal)` | Free function `config::sync_models(url, signal)` in `src/config/mod.rs` | Fetches a URL, parses YAML, writes to `paths::models_override_file()`. No Config state involved. Sibling pattern to `install_builtins`, `default_sessions_dir`, `list_sessions`. | + +### Caller updates + +- `src/config/macros.rs:23` — `Config::load_macro(name)` → `Macro::load(name)` +- `src/config/input.rs:214` — `Config::search_rag(&self.app_config, rag, &self.text, abort_signal)` → `rag.search_with_template(&self.app_config, &self.text, abort_signal)` +- `src/main.rs:149` — `Config::sync_models(&url, abort_signal.clone())` → `sync_models(&url, abort_signal.clone())` (added `sync_models` to the `crate::config::{...}` import list) + +### Constants relocated + +- `RAG_TEMPLATE` moved from `src/config/mod.rs` to `src/rag/mod.rs` alongside the new `search_with_template` method that uses it. + +### Final shape of `impl Config` + +```rust +impl Config { + pub async fn load_with_interpolation(info_flag: bool) -> Result { ... } + pub fn load_from_file(config_path: &Path) -> Result<(Self, String)> { ... } + pub fn load_from_str(content: &str) -> Result { ... } + pub fn load_dynamic(model_id: &str) -> Result { ... } +} +``` + +Four loaders, all returning `Self` or `(Self, String)`. Nothing +else. The `Config` type is now genuinely what its docstring +claims: a serde POJO with constructors. No squatters. + +### Verification (cleanup pass) + +- `cargo check` — clean +- `cargo clippy --all-targets` — clean +- `cargo test` — 122 passing, zero failures +- `Config::sync_models` / `Config::load_macro` / `Config::search_rag` — zero hits in `src/` + +### Files modified (cleanup pass) + +- `src/config/mod.rs` — deleted `Config::load_macro`, `Config::search_rag`, `Config::sync_models`, and `RAG_TEMPLATE` const; added free `sync_models` function +- `src/config/macros.rs` — added `Macro::load`, updated import (added `Context`, `read_to_string`; removed `Config`) +- `src/rag/mod.rs` — added `RAG_TEMPLATE` const and `Rag::search_with_template` method +- `src/config/input.rs` — updated caller to `rag.search_with_template` +- `src/main.rs` — added `sync_models` to import list, updated caller diff --git a/src/cli/completer.rs b/src/cli/completer.rs index bc0695b..1809bbc 100644 --- a/src/cli/completer.rs +++ b/src/cli/completer.rs @@ -1,6 +1,7 @@ use crate::client::{ModelType, list_models}; use crate::config::paths; -use crate::config::{Config, list_agents}; +use crate::config::{AppConfig, Config, list_agents, list_sessions}; +use crate::vault::Vault; use clap_complete::{CompletionCandidate, Shell, generate}; use clap_complete_nushell::Nushell; use std::ffi::OsStr; @@ -33,8 +34,8 @@ impl ShellCompletion { pub(super) fn model_completer(current: &OsStr) -> Vec { let cur = current.to_string_lossy(); - match Config::init_bare() { - Ok(config) => list_models(&config.to_app_config(), ModelType::Chat) + match load_app_config_for_completion() { + Ok(app_config) => list_models(&app_config, ModelType::Chat) .into_iter() .filter(|&m| m.id().starts_with(&*cur)) .map(|m| CompletionCandidate::new(m.id())) @@ -43,6 +44,20 @@ pub(super) fn model_completer(current: &OsStr) -> Vec { } } +fn load_app_config_for_completion() -> anyhow::Result { + let h = tokio::runtime::Handle::try_current().ok(); + let cfg = match h { + Some(handle) => { + tokio::task::block_in_place(|| handle.block_on(Config::load_with_interpolation(true)))? + } + None => { + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(Config::load_with_interpolation(true))? + } + }; + AppConfig::from_config(cfg) +} + pub(super) fn role_completer(current: &OsStr) -> Vec { let cur = current.to_string_lossy(); paths::list_roles(true) @@ -81,22 +96,17 @@ pub(super) fn macro_completer(current: &OsStr) -> Vec { pub(super) fn session_completer(current: &OsStr) -> Vec { let cur = current.to_string_lossy(); - match Config::init_bare() { - Ok(config) => config - .list_sessions() - .into_iter() - .filter(|s| s.starts_with(&*cur)) - .map(CompletionCandidate::new) - .collect(), - Err(_) => vec![], - } + list_sessions() + .into_iter() + .filter(|s| s.starts_with(&*cur)) + .map(CompletionCandidate::new) + .collect() } pub(super) fn secrets_completer(current: &OsStr) -> Vec { let cur = current.to_string_lossy(); - match Config::init_bare() { - Ok(config) => config - .vault + match load_app_config_for_completion() { + Ok(app_config) => Vault::init(&app_config) .list_secrets(false) .unwrap_or_default() .into_iter() diff --git a/src/client/common.rs b/src/client/common.rs index 401335b..1b84cb4 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -417,7 +417,7 @@ pub async fn call_chat_completions( ctx: &mut RequestContext, abort_signal: AbortSignal, ) -> Result<(String, Vec)> { - let is_child_agent = ctx.current_depth() > 0; + let is_child_agent = ctx.current_depth > 0; let spinner_message = if is_child_agent { "" } else { "Generating" }; let ret = abortable_run_with_spinner( client.chat_completions(input.clone()), diff --git a/src/config/agent.rs b/src/config/agent.rs index 2abaa77..6f0f302 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -1,4 +1,3 @@ -use super::todo::TodoList; use super::*; use crate::{ @@ -6,6 +5,7 @@ use crate::{ function::{Functions, run_llm_function}, }; +use super::rag_cache::RagKey; use crate::config::paths; use crate::config::prompts::{ DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS, @@ -39,9 +39,6 @@ pub struct Agent { rag: Option>, model: Model, vault: GlobalVault, - todo_list: TodoList, - continuation_count: usize, - last_continuation_response: Option, } impl Agent { @@ -123,12 +120,16 @@ impl Agent { }; let rag = if rag_path.exists() { - Some(Arc::new(Rag::load( - app, - &app.clients, - DEFAULT_AGENT_NAME, - &rag_path, - )?)) + let key = RagKey::Agent(name.to_string()); + let app_clone = app.clone(); + let rag_path_clone = rag_path.clone(); + let rag = app_state + .rag_cache + .load_with(key, || async move { + Rag::load(&app_clone, DEFAULT_AGENT_NAME, &rag_path_clone) + }) + .await?; + Some(rag) } else if !agent_config.documents.is_empty() && !info_flag { let mut ans = false; if *IS_STDOUT_TERMINAL { @@ -161,16 +162,23 @@ impl Agent { document_paths.push(path.to_string()) } } - let rag = Rag::init( - app, - &app.clients, - "rag", - &rag_path, - &document_paths, - abort_signal, - ) - .await?; - Some(Arc::new(rag)) + let key = RagKey::Agent(name.to_string()); + let app_clone = app.clone(); + let rag_path_clone = rag_path.clone(); + let rag = app_state + .rag_cache + .load_with(key, || async move { + Rag::init( + &app_clone, + "rag", + &rag_path_clone, + &document_paths, + abort_signal, + ) + .await + }) + .await?; + Some(rag) } else { None } @@ -202,9 +210,6 @@ impl Agent { rag, model, vault: app_state.vault.clone(), - todo_list: TodoList::default(), - continuation_count: 0, - last_continuation_response: None, }) } @@ -434,44 +439,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! {" diff --git a/src/config/app_config.rs b/src/config/app_config.rs new file mode 100644 index 0000000..fdd1da8 --- /dev/null +++ b/src/config/app_config.rs @@ -0,0 +1,757 @@ +//! Immutable, process-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. +//! +//! `AppConfig` mirrors the field shape of [`Config`](super::Config) (the +//! serde POJO loaded from YAML) but is the runtime-resolved form: env +//! var overrides applied, wrap validated, default document loaders +//! installed, user agent resolved, default model picked. Build it via +//! [`AppConfig::from_config`]. +//! +//! Runtime-only state (current role, session, agent, supervisor, etc.) +//! lives on [`RequestContext`](super::request_context::RequestContext). +//! Process-wide services (vault, MCP registry, function registry) live +//! on [`AppState`](super::app_state::AppState). + +use crate::client::{ClientConfig, list_models}; +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)] + pub model_id: String, + pub temperature: Option, + pub top_p: Option, + + pub dry_run: bool, + pub stream: bool, + pub save: bool, + pub keybindings: String, + pub editor: Option, + pub wrap: Option, + pub wrap_code: bool, + pub(crate) vault_password_file: Option, + + pub function_calling_support: bool, + pub mapping_tools: IndexMap, + pub enabled_tools: Option, + pub visible_tools: Option>, + + pub mcp_server_support: bool, + pub mapping_mcp_servers: IndexMap, + pub enabled_mcp_servers: Option, + + pub repl_prelude: Option, + pub cmd_prelude: Option, + pub agent_session: Option, + + pub save_session: Option, + pub compression_threshold: usize, + pub summarization_prompt: Option, + pub summary_context_prompt: Option, + + pub rag_embedding_model: Option, + pub rag_reranker_model: Option, + pub rag_top_k: usize, + pub rag_chunk_size: Option, + pub rag_chunk_overlap: Option, + pub rag_template: Option, + + #[serde(default)] + pub document_loaders: HashMap, + + pub highlight: bool, + pub theme: Option, + pub left_prompt: Option, + pub right_prompt: Option, + + pub user_agent: Option, + pub save_shell_history: bool, + pub sync_models_url: Option, + + pub clients: Vec, +} + +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 from_config(config: super::Config) -> Result { + let mut app_config = Self { + model_id: config.model_id, + temperature: config.temperature, + top_p: config.top_p, + + dry_run: config.dry_run, + stream: config.stream, + save: config.save, + keybindings: config.keybindings, + editor: config.editor, + wrap: config.wrap, + wrap_code: config.wrap_code, + vault_password_file: config.vault_password_file, + + function_calling_support: config.function_calling_support, + mapping_tools: config.mapping_tools, + enabled_tools: config.enabled_tools, + visible_tools: config.visible_tools, + + mcp_server_support: config.mcp_server_support, + mapping_mcp_servers: config.mapping_mcp_servers, + enabled_mcp_servers: config.enabled_mcp_servers, + + repl_prelude: config.repl_prelude, + cmd_prelude: config.cmd_prelude, + agent_session: config.agent_session, + + save_session: config.save_session, + compression_threshold: config.compression_threshold, + summarization_prompt: config.summarization_prompt, + summary_context_prompt: config.summary_context_prompt, + + rag_embedding_model: config.rag_embedding_model, + rag_reranker_model: config.rag_reranker_model, + rag_top_k: config.rag_top_k, + rag_chunk_size: config.rag_chunk_size, + rag_chunk_overlap: config.rag_chunk_overlap, + rag_template: config.rag_template, + + document_loaders: config.document_loaders, + + highlight: config.highlight, + theme: config.theme, + left_prompt: config.left_prompt, + right_prompt: config.right_prompt, + + user_agent: config.user_agent, + save_shell_history: config.save_shell_history, + sync_models_url: config.sync_models_url, + + clients: config.clients, + }; + app_config.load_envs(); + if let Some(wrap) = app_config.wrap.clone() { + app_config.set_wrap(&wrap)?; + } + app_config.setup_document_loaders(); + app_config.setup_user_agent(); + app_config.resolve_model()?; + Ok(app_config) + } + + pub fn resolve_model(&mut self) -> Result<()> { + if self.model_id.is_empty() { + let models = list_models(self, crate::client::ModelType::Chat); + if models.is_empty() { + anyhow::bail!("No available model"); + } + self.model_id = models[0].id(); + } + Ok(()) + } + + 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 { + 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 { + 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 { + 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::() + .map_err(|_| anyhow!("Invalid wrap value"))?; + self.wrap = Some(value.into()) + } + Ok(()) + } + + 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); + }); + } + + 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") + )); + } + } + + 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::(&get_env_name("temperature")) { + self.temperature = v; + } + if let Some(v) = super::read_env_value::(&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::(&get_env_name("editor")) { + self.editor = v; + } + if let Some(v) = super::read_env_value::(&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::(&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::(&get_env_name("enabled_mcp_servers")) { + self.enabled_mcp_servers = v; + } + + if let Some(v) = super::read_env_value::(&get_env_name("repl_prelude")) { + self.repl_prelude = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("cmd_prelude")) { + self.cmd_prelude = v; + } + if let Some(v) = super::read_env_value::(&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::(&get_env_name("compression_threshold")) + { + self.compression_threshold = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("summarization_prompt")) { + self.summarization_prompt = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("summary_context_prompt")) { + self.summary_context_prompt = v; + } + + if let Some(v) = super::read_env_value::(&get_env_name("rag_embedding_model")) { + self.rag_embedding_model = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("rag_reranker_model")) { + self.rag_reranker_model = v; + } + if let Some(Some(v)) = super::read_env_value::(&get_env_name("rag_top_k")) { + self.rag_top_k = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("rag_chunk_size")) { + self.rag_chunk_size = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("rag_chunk_overlap")) { + self.rag_chunk_overlap = v; + } + if let Some(v) = super::read_env_value::(&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::(&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::(&get_env_name("left_prompt")) { + self.left_prompt = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("right_prompt")) { + self.right_prompt = v; + } + if let Some(v) = super::read_env_value::(&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::(&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) { + self.temperature = value; + } + + #[allow(dead_code)] + pub fn set_top_p_default(&mut self, value: Option) { + self.top_p = value; + } + + #[allow(dead_code)] + pub fn set_enabled_tools_default(&mut self, value: Option) { + self.enabled_tools = value; + } + + #[allow(dead_code)] + pub fn set_enabled_mcp_servers_default(&mut self, value: Option) { + self.enabled_mcp_servers = value; + } + + #[allow(dead_code)] + pub fn set_save_session_default(&mut self, value: Option) { + self.save_session = value; + } + + #[allow(dead_code)] + pub fn set_compression_threshold_default(&mut self, value: Option) { + self.compression_threshold = value.unwrap_or_default(); + } + + #[allow(dead_code)] + pub fn set_rag_reranker_model_default(&mut self, value: Option) { + 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; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + + fn cached_editor() -> Option { + super::super::EDITOR.get().cloned().flatten() + } + + #[test] + fn from_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, + clients: vec![ClientConfig::default()], + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + 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 from_config_copies_clients() { + let cfg = Config { + model_id: "test-model".to_string(), + clients: vec![ClientConfig::default()], + ..Config::default() + }; + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!(app.clients.len(), 1); + } + + #[test] + fn from_config_copies_mapping_fields() { + let mut cfg = Config { + model_id: "test-model".to_string(), + clients: vec![ClientConfig::default()], + ..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 = AppConfig::from_config(cfg).unwrap(); + + 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()); + } + + #[test] + fn from_config_copies_serde_fields() { + let cfg = Config { + model_id: "provider:model-x".to_string(), + temperature: Some(0.42), + compression_threshold: 1234, + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!(app.model_id, "provider:model-x"); + assert_eq!(app.temperature, Some(0.42)); + assert_eq!(app.compression_threshold, 1234); + } + + #[test] + fn from_config_installs_default_document_loaders() { + let cfg = Config { + model_id: "provider:test".to_string(), + ..Config::default() + }; + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!( + app.document_loaders.get("pdf"), + Some(&"pdftotext $1 -".to_string()) + ); + assert_eq!( + app.document_loaders.get("docx"), + Some(&"pandoc --to plain $1".to_string()) + ); + } + + #[test] + fn from_config_resolves_auto_user_agent() { + let cfg = Config { + model_id: "provider:test".to_string(), + user_agent: Some("auto".to_string()), + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + let ua = app.user_agent.as_deref().unwrap(); + assert!(ua != "auto", "user_agent should have been resolved"); + assert!(ua.contains('/'), "user_agent should be '/'"); + } + + #[test] + fn from_config_preserves_explicit_user_agent() { + let cfg = Config { + model_id: "provider:test".to_string(), + user_agent: Some("custom/1.0".to_string()), + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!(app.user_agent.as_deref(), Some("custom/1.0")); + } + + #[test] + fn from_config_validates_wrap_value() { + let cfg = Config { + model_id: "provider:test".to_string(), + wrap: Some("invalid".to_string()), + ..Config::default() + }; + + let result = AppConfig::from_config(cfg); + assert!(result.is_err()); + } + + #[test] + fn from_config_accepts_wrap_auto() { + let cfg = Config { + model_id: "provider:test".to_string(), + wrap: Some("auto".to_string()), + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + assert_eq!(app.wrap.as_deref(), Some("auto")); + } + + #[test] + fn resolve_model_errors_when_no_models_available() { + let mut app = AppConfig { + model_id: String::new(), + clients: vec![], + ..AppConfig::default() + }; + + let result = app.resolve_model(); + assert!(result.is_err()); + } + + #[test] + fn resolve_model_keeps_explicit_model_id() { + let mut app = AppConfig { + model_id: "provider:explicit".to_string(), + ..AppConfig::default() + }; + + app.resolve_model().unwrap(); + assert_eq!(app.model_id, "provider:explicit"); + } +} diff --git a/src/config/app_state.rs b/src/config/app_state.rs new file mode 100644 index 0000000..fdc18b2 --- /dev/null +++ b/src/config/app_state.rs @@ -0,0 +1,95 @@ +//! 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, the [`RagCache`](super::rag_cache::RagCache) +//! for shared RAG instances, the global MCP registry, and the base +//! [`Functions`] declarations seeded into per-request `ToolScope`s. It +//! is wrapped in `Arc` and shared across every [`RequestContext`] that +//! a frontend (CLI, REPL, API) creates. +//! +//! Built via [`AppState::init`] from an `Arc` plus +//! startup context (log path, MCP-start flag, abort signal). The +//! `init` call is the single place that wires the vault, MCP +//! registry, and global functions together. + +use super::mcp_factory::{McpFactory, McpServerKey}; +use super::rag_cache::RagCache; +use crate::config::AppConfig; +use crate::function::Functions; +use crate::mcp::{McpRegistry, McpServersConfig}; +use crate::utils::AbortSignal; +use crate::vault::{GlobalVault, Vault}; + +use anyhow::Result; +use std::path::PathBuf; +use std::sync::Arc; + +#[derive(Clone)] +pub struct AppState { + pub config: Arc, + pub vault: GlobalVault, + pub mcp_factory: Arc, + pub rag_cache: Arc, + pub mcp_config: Option, + pub mcp_log_path: Option, + pub mcp_registry: Option>, + pub functions: Functions, +} + +impl AppState { + pub async fn init( + config: Arc, + log_path: Option, + start_mcp_servers: bool, + abort_signal: AbortSignal, + ) -> Result { + let vault = Arc::new(Vault::init(&config)); + + let mcp_registry = McpRegistry::init( + log_path, + start_mcp_servers, + config.enabled_mcp_servers.clone(), + abort_signal, + &config, + &vault, + ) + .await?; + + let mcp_config = mcp_registry.mcp_config().cloned(); + let mcp_log_path = mcp_registry.log_path().cloned(); + + let mcp_factory = Arc::new(McpFactory::default()); + if let Some(mcp_servers_config) = &mcp_config { + for (id, handle) in mcp_registry.running_servers() { + if let Some(spec) = mcp_servers_config.mcp_servers.get(id) { + let key = McpServerKey::from_spec(id, spec); + mcp_factory.insert_active(key, handle); + } + } + } + + let mut functions = Functions::init(config.visible_tools.as_ref().unwrap_or(&Vec::new()))?; + if !mcp_registry.is_empty() && config.mcp_server_support { + functions.append_mcp_meta_functions(mcp_registry.list_started_servers()); + } + + let mcp_registry = if mcp_registry.is_empty() { + None + } else { + Some(Arc::new(mcp_registry)) + }; + + Ok(Self { + config, + vault, + mcp_factory, + rag_cache: Arc::new(RagCache::default()), + mcp_config, + mcp_log_path, + mcp_registry, + functions, + }) + } +} diff --git a/src/config/input.rs b/src/config/input.rs index 71cb683..b133bec 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -20,7 +20,7 @@ pub struct Input { app_config: Arc, stream_enabled: bool, session: Option, - rag: Option>, + rag: Option>, functions: Option>, text: String, raw: (String, Vec), @@ -210,8 +210,9 @@ impl Input { return Ok(()); } if let Some(rag) = &self.rag { - let result = - Config::search_rag(&self.app_config, rag, &self.text, abort_signal).await?; + let result = rag + .search_with_template(&self.app_config, &self.text, abort_signal) + .await?; self.patched_text = Some(result); self.rag_name = Some(rag.name().to_string()); } @@ -411,7 +412,7 @@ fn resolve_role(ctx: &RequestContext, role: Option) -> (Role, bool, bool) struct CapturedInputConfig { stream_enabled: bool, session: Option, - rag: Option>, + rag: Option>, functions: Option>, } diff --git a/src/config/macros.rs b/src/config/macros.rs index f0825c4..86dd34e 100644 --- a/src/config/macros.rs +++ b/src/config/macros.rs @@ -1,12 +1,12 @@ use crate::config::paths; -use crate::config::{Config, RequestContext, RoleLike, ensure_parent_exists}; +use crate::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 anyhow::{Context, Result, anyhow}; use indexmap::IndexMap; use rust_embed::Embed; use serde::Deserialize; -use std::fs::File; +use std::fs::{File, read_to_string}; use std::io::Write; #[derive(Embed)] @@ -20,7 +20,7 @@ pub async fn macro_execute( args: Option<&str>, abort_signal: AbortSignal, ) -> Result<()> { - let macro_value = Config::load_macro(name)?; + let macro_value = Macro::load(name)?; let (mut new_args, text) = split_args_text(args.unwrap_or_default(), cfg!(windows)); if !text.is_empty() { new_args.push(text.to_string()); @@ -44,7 +44,14 @@ pub async fn macro_execute( macro_ctx.model = role.model().clone(); macro_ctx.agent_variables = ctx.agent_variables.clone(); macro_ctx.last_message = ctx.last_message.clone(); - macro_ctx.agent_runtime = ctx.agent_runtime.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(); @@ -69,6 +76,14 @@ pub struct Macro { } impl Macro { + pub fn load(name: &str) -> Result { + let path = paths::macro_file(name); + let err = || format!("Failed to load macro '{name}' at '{}'", path.display()); + let content = read_to_string(&path).with_context(err)?; + let value: Macro = serde_yaml::from_str(&content).with_context(err)?; + Ok(value) + } + pub fn install_macros() -> Result<()> { info!( "Installing built-in macros in {}", diff --git a/src/config/mcp_factory.rs b/src/config/mcp_factory.rs new file mode 100644 index 0000000..74dcd57 --- /dev/null +++ b/src/config/mcp_factory.rs @@ -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` +//! 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>>` +//! 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, + pub env: Vec<(String, String)>, +} + +impl McpServerKey { + pub fn new( + name: impl Into, + command: impl Into, + args: impl IntoIterator, + env: impl IntoIterator, + ) -> Self { + let mut args: Vec = 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>>, +} + +impl McpFactory { + pub fn try_get_active(&self, key: &McpServerKey) -> Option> { + let map = self.active.lock(); + map.get(key).and_then(|weak| weak.upgrade()) + } + + pub fn insert_active(&self, key: McpServerKey, handle: &Arc) { + 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> { + 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) + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index 2d69ee0..37466c3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,8 +1,6 @@ mod agent; -mod agent_runtime; mod app_config; mod app_state; -mod bridge; mod input; mod macros; mod mcp_factory; @@ -29,25 +27,20 @@ pub use self::role::{ use self::session::Session; use crate::client::{ ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS, - ProviderModels, create_client_config, list_client_types, list_models, + ProviderModels, create_client_config, list_client_types, }; -use crate::function::{FunctionDeclaration, Functions, ToolCallTracker}; +use crate::function::{FunctionDeclaration, Functions}; use crate::rag::Rag; use crate::utils::*; pub use macros::macro_execute; use crate::config::macros::Macro; -use crate::mcp::McpRegistry; -use crate::supervisor::Supervisor; -use crate::supervisor::escalation::EscalationQueue; -use crate::supervisor::mailbox::Inbox; use crate::vault::{GlobalVault, Vault, create_vault_password_file, interpolate_secrets}; use anyhow::{Context, Result, anyhow, bail}; use fancy_regex::Regex; use indexmap::IndexMap; use indoc::formatdoc; use inquire::{Confirm, Select}; -use parking_lot::RwLock; use serde::{Deserialize, Serialize}; use serde_json::json; use std::collections::HashMap; @@ -60,8 +53,6 @@ use std::{ process, sync::{Arc, OnceLock}, }; -use terminal_colorsaurus::{ColorScheme, QueryOptions, color_scheme}; -use tokio::runtime::Handle; pub const TEMP_ROLE_NAME: &str = "temp"; pub const TEMP_RAG_NAME: &str = "temp"; @@ -98,30 +89,6 @@ const SUMMARIZATION_PROMPT: &str = "Summarize the discussion briefly in 200 words or less to use as a prompt for future context."; const SUMMARY_CONTEXT_PROMPT: &str = "This is a summary of the chat history as a recap: "; -const RAG_TEMPLATE: &str = r#"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags) - - -__CONTEXT__ - - - -__SOURCES__ - - - -- If you don't know, just say so. -- If you are not sure, ask for clarification. -- Answer in the same language as the user query. -- If the context appears unreadable or of poor quality, tell the user then answer as best as you can. -- If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge. -- Answer directly and without using xml tags. -- When using information from the context, cite the relevant source from the section. - - - -__INPUT__ -"#; - const LEFT_PROMPT: &str = "{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} "; const RIGHT_PROMPT: &str = "{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}"; @@ -143,7 +110,7 @@ pub struct Config { pub editor: Option, pub wrap: Option, pub wrap_code: bool, - vault_password_file: Option, + pub(super) vault_password_file: Option, pub function_calling_support: bool, pub mapping_tools: IndexMap, @@ -183,50 +150,6 @@ pub struct Config { pub sync_models_url: Option, pub clients: Vec, - - #[serde(skip)] - pub vault: GlobalVault, - - #[serde(skip)] - pub macro_flag: bool, - #[serde(skip)] - pub info_flag: bool, - #[serde(skip)] - pub agent_variables: Option, - - #[serde(skip)] - pub model: Model, - #[serde(skip)] - pub functions: Functions, - #[serde(skip)] - pub mcp_registry: Option, - #[serde(skip)] - pub working_mode: WorkingMode, - #[serde(skip)] - pub last_message: Option, - - #[serde(skip)] - pub role: Option, - #[serde(skip)] - pub session: Option, - #[serde(skip)] - pub rag: Option>, - #[serde(skip)] - pub agent: Option, - #[serde(skip)] - pub(crate) tool_call_tracker: Option, - #[serde(skip)] - pub supervisor: Option>>, - #[serde(skip)] - pub parent_supervisor: Option>>, - #[serde(skip)] - pub self_agent_id: Option, - #[serde(skip)] - pub current_depth: usize, - #[serde(skip)] - pub inbox: Option>, - #[serde(skip)] - pub root_escalation_queue: Option>, } impl Default for Config { @@ -282,55 +205,52 @@ impl Default for Config { sync_models_url: None, clients: vec![], - - vault: Default::default(), - - macro_flag: false, - info_flag: false, - agent_variables: None, - - model: Default::default(), - functions: Default::default(), - mcp_registry: Default::default(), - working_mode: WorkingMode::Cmd, - last_message: None, - - role: None, - session: None, - rag: None, - agent: None, - tool_call_tracker: Some(ToolCallTracker::default()), - supervisor: None, - parent_supervisor: None, - self_agent_id: None, - current_depth: 0, - inbox: None, - root_escalation_queue: None, } } } -impl Config { - pub fn init_bare() -> Result { - let h = Handle::current(); - tokio::task::block_in_place(|| { - h.block_on(Self::init( - WorkingMode::Cmd, - true, - false, - None, - create_abort_signal(), - )) - }) - } +pub fn install_builtins() -> Result<()> { + Functions::install_builtin_global_tools()?; + Agent::install_builtin_agents()?; + Macro::install_macros()?; + Ok(()) +} - pub async fn init( - working_mode: WorkingMode, - info_flag: bool, - start_mcp_servers: bool, - log_path: Option, - abort_signal: AbortSignal, - ) -> Result { +pub fn default_sessions_dir() -> PathBuf { + match env::var(get_env_name("sessions_dir")) { + Ok(value) => PathBuf::from(value), + Err(_) => paths::local_path(SESSIONS_DIR_NAME), + } +} + +pub fn list_sessions() -> Vec { + list_file_names(default_sessions_dir(), ".yaml") +} + +pub async fn sync_models(url: &str, abort_signal: AbortSignal) -> Result<()> { + let content = abortable_run_with_spinner(fetch(url), "Fetching models.yaml", abort_signal) + .await + .with_context(|| format!("Failed to fetch '{url}'"))?; + println!("✓ Fetched '{url}'"); + let list = serde_yaml::from_str::>(&content) + .with_context(|| "Failed to parse models.yaml")?; + let models_override = ModelsOverride { + version: env!("CARGO_PKG_VERSION").to_string(), + list, + }; + let models_override_data = + serde_yaml::to_string(&models_override).with_context(|| "Failed to serde {}")?; + + let model_override_path = paths::models_override_file(); + ensure_parent_exists(&model_override_path)?; + std::fs::write(&model_override_path, models_override_data) + .with_context(|| format!("Failed to write to '{}'", model_override_path.display()))?; + println!("✓ Updated '{}'", model_override_path.display()); + Ok(()) +} + +impl Config { + pub async fn load_with_interpolation(info_flag: bool) -> Result { let config_path = paths::config_file(); let (mut config, content) = if !config_path.exists() { match env::var(get_env_name("provider")) @@ -349,180 +269,39 @@ impl Config { Self::load_from_file(&config_path)? }; - let setup = async |config: &mut Self| -> Result<()> { - let vault = Vault::init(config); - - let (parsed_config, missing_secrets) = interpolate_secrets(&content, &vault); - if !missing_secrets.is_empty() && !info_flag { - debug!( - "Global config references secrets that are missing from the vault: {missing_secrets:?}" - ); - return Err(anyhow!(formatdoc!( - " - Global config file references secrets that are missing from the vault: {:?} - Please add these secrets to the vault and try again.", - missing_secrets - ))); - } - - if !parsed_config.is_empty() && !info_flag { - debug!("Global config is invalid once secrets are injected: {parsed_config}"); - let new_config = Self::load_from_str(&parsed_config).with_context(|| { - formatdoc!( - " - Global config is invalid once secrets are injected. - Double check the secret values and file syntax, then try again. - " - ) - })?; - *config = new_config.clone(); - } - - config.working_mode = working_mode; - config.info_flag = info_flag; - config.vault = Arc::new(vault); - - Agent::install_builtin_agents()?; - - config.load_envs(); - - if let Some(wrap) = config.wrap.clone() { - config.set_wrap(&wrap)?; - } - - config.load_functions()?; - config - .load_mcp_servers(log_path, start_mcp_servers, abort_signal) - .await?; - - config.setup_model()?; - config.setup_document_loaders(); - config.setup_user_agent(); - Macro::install_macros()?; - Ok(()) + let bootstrap_app = AppConfig { + vault_password_file: config.vault_password_file.clone(), + ..AppConfig::default() }; - let ret = setup(&mut config).await; - if !info_flag { - ret?; + let vault = Vault::init(&bootstrap_app); + let (parsed_config, missing_secrets) = interpolate_secrets(&content, &vault); + if !missing_secrets.is_empty() && !info_flag { + debug!( + "Global config references secrets that are missing from the vault: {missing_secrets:?}" + ); + return Err(anyhow!(formatdoc!( + " + Global config file references secrets that are missing from the vault: {:?} + Please add these secrets to the vault and try again.", + missing_secrets + ))); + } + if !parsed_config.is_empty() && !info_flag { + debug!("Global config is invalid once secrets are injected: {parsed_config}"); + let new_config = Self::load_from_str(&parsed_config).with_context(|| { + formatdoc!( + " + Global config is invalid once secrets are injected. + Double check the secret values and file syntax, then try again. + " + ) + })?; + config = new_config; } Ok(config) } - 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 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 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 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::() - .map_err(|_| anyhow!("Invalid wrap value"))?; - self.wrap = Some(value.into()) - } - Ok(()) - } - - pub fn set_model(&mut self, model_id: &str) -> Result<()> { - let model = Model::retrieve_model(&self.to_app_config(), model_id, ModelType::Chat)?; - match self.role_like_mut() { - Some(role_like) => role_like.set_model(model), - None => { - self.model = model; - } - } - Ok(()) - } - - pub fn list_sessions(&self) -> Vec { - list_file_names(self.sessions_dir(), ".yaml") - } - - pub async fn search_rag( - app: &AppConfig, - rag: &Rag, - text: &str, - abort_signal: AbortSignal, - ) -> Result { - let (reranker_model, top_k) = rag.get_config(); - let (embeddings, sources, ids) = rag - .search(text, top_k, reranker_model.as_deref(), abort_signal) - .await?; - let rag_template = app.rag_template.as_deref().unwrap_or(RAG_TEMPLATE); - let text = if embeddings.is_empty() { - text.to_string() - } else { - rag_template - .replace("__CONTEXT__", &embeddings) - .replace("__SOURCES__", &sources) - .replace("__INPUT__", text) - }; - rag.set_last_sources(&ids); - Ok(text) - } - - pub fn load_macro(name: &str) -> Result { - let path = paths::macro_file(name); - let err = || format!("Failed to load macro '{name}' at '{}'", path.display()); - let content = read_to_string(&path).with_context(err)?; - let value: Macro = serde_yaml::from_str(&content).with_context(err)?; - Ok(value) - } - - pub async fn sync_models(url: &str, abort_signal: AbortSignal) -> Result<()> { - let content = abortable_run_with_spinner(fetch(url), "Fetching models.yaml", abort_signal) - .await - .with_context(|| format!("Failed to fetch '{url}'"))?; - println!("✓ Fetched '{url}'"); - let list = serde_yaml::from_str::>(&content) - .with_context(|| "Failed to parse models.yaml")?; - let models_override = ModelsOverride { - version: env!("CARGO_PKG_VERSION").to_string(), - list, - }; - let models_override_data = - serde_yaml::to_string(&models_override).with_context(|| "Failed to serde {}")?; - - let model_override_path = paths::models_override_file(); - ensure_parent_exists(&model_override_path)?; - std::fs::write(&model_override_path, models_override_data) - .with_context(|| format!("Failed to write to '{}'", model_override_path.display()))?; - println!("✓ Updated '{}'", model_override_path.display()); - Ok(()) - } - - fn load_from_file(config_path: &Path) -> Result<(Self, String)> { + pub fn load_from_file(config_path: &Path) -> Result<(Self, String)> { let err = || format!("Failed to load config at '{}'", config_path.display()); let content = read_to_string(config_path).with_context(err)?; let config = Self::load_from_str(&content).with_context(err)?; @@ -530,7 +309,7 @@ impl Config { Ok((config, content)) } - fn load_from_str(content: &str) -> Result { + pub fn load_from_str(content: &str) -> Result { if PASSWORD_FILE_SECRET_RE.is_match(content)? { bail!("secret injection cannot be done on the vault_password_file property"); } @@ -556,7 +335,7 @@ impl Config { Ok(config) } - fn load_dynamic(model_id: &str) -> Result { + pub fn load_dynamic(model_id: &str) -> Result { let provider = match model_id.split_once(':') { Some((v, _)) => v, _ => model_id, @@ -578,225 +357,6 @@ impl Config { serde_json::from_value(config).with_context(|| "Failed to load config from env")?; Ok(config) } - - fn load_envs(&mut self) { - if let Ok(v) = env::var(get_env_name("model")) { - self.model_id = v; - } - if let Some(v) = read_env_value::(&get_env_name("temperature")) { - self.temperature = v; - } - if let Some(v) = read_env_value::(&get_env_name("top_p")) { - self.top_p = v; - } - - if let Some(Some(v)) = read_env_bool(&get_env_name("dry_run")) { - self.dry_run = v; - } - if let Some(Some(v)) = read_env_bool(&get_env_name("stream")) { - self.stream = v; - } - if let Some(Some(v)) = 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) = read_env_value::(&get_env_name("editor")) { - self.editor = v; - } - if let Some(v) = read_env_value::(&get_env_name("wrap")) { - self.wrap = v; - } - if let Some(Some(v)) = read_env_bool(&get_env_name("wrap_code")) { - self.wrap_code = v; - } - - if let Some(Some(v)) = 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) = read_env_value::(&get_env_name("enabled_tools")) { - self.enabled_tools = v; - } - - if let Some(Some(v)) = 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) = read_env_value::(&get_env_name("enabled_mcp_servers")) { - self.enabled_mcp_servers = v; - } - - if let Some(v) = read_env_value::(&get_env_name("repl_prelude")) { - self.repl_prelude = v; - } - if let Some(v) = read_env_value::(&get_env_name("cmd_prelude")) { - self.cmd_prelude = v; - } - if let Some(v) = read_env_value::(&get_env_name("agent_session")) { - self.agent_session = v; - } - - if let Some(v) = read_env_bool(&get_env_name("save_session")) { - self.save_session = v; - } - if let Some(Some(v)) = read_env_value::(&get_env_name("compression_threshold")) { - self.compression_threshold = v; - } - if let Some(v) = read_env_value::(&get_env_name("summarization_prompt")) { - self.summarization_prompt = v; - } - if let Some(v) = read_env_value::(&get_env_name("summary_context_prompt")) { - self.summary_context_prompt = v; - } - - if let Some(v) = read_env_value::(&get_env_name("rag_embedding_model")) { - self.rag_embedding_model = v; - } - if let Some(v) = read_env_value::(&get_env_name("rag_reranker_model")) { - self.rag_reranker_model = v; - } - if let Some(Some(v)) = read_env_value::(&get_env_name("rag_top_k")) { - self.rag_top_k = v; - } - if let Some(v) = read_env_value::(&get_env_name("rag_chunk_size")) { - self.rag_chunk_size = v; - } - if let Some(v) = read_env_value::(&get_env_name("rag_chunk_overlap")) { - self.rag_chunk_overlap = v; - } - if let Some(v) = read_env_value::(&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)) = 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) = read_env_value::(&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) = read_env_value::(&get_env_name("left_prompt")) { - self.left_prompt = v; - } - if let Some(v) = read_env_value::(&get_env_name("right_prompt")) { - self.right_prompt = v; - } - if let Some(v) = read_env_value::(&get_env_name("user_agent")) { - self.user_agent = v; - } - if let Some(Some(v)) = read_env_bool(&get_env_name("save_shell_history")) { - self.save_shell_history = v; - } - if let Some(v) = read_env_value::(&get_env_name("sync_models_url")) { - self.sync_models_url = v; - } - } - - fn load_functions(&mut self) -> Result<()> { - self.functions = Functions::init(self.visible_tools.as_ref().unwrap_or(&Vec::new()))?; - if self.working_mode.is_repl() { - self.functions.append_user_interaction_functions(); - } - Ok(()) - } - - async fn load_mcp_servers( - &mut self, - log_path: Option, - start_mcp_servers: bool, - abort_signal: AbortSignal, - ) -> Result<()> { - let mcp_registry = McpRegistry::init( - log_path, - start_mcp_servers, - self.enabled_mcp_servers.clone(), - abort_signal.clone(), - self, - ) - .await?; - match mcp_registry.is_empty() { - false => { - if self.mcp_server_support { - self.functions - .append_mcp_meta_functions(mcp_registry.list_started_servers()); - } else { - debug!( - "Skipping global MCP functions registration since 'mcp_server_support' was 'false'" - ); - } - } - _ => debug!( - "Skipping global MCP functions registration since 'start_mcp_servers' was 'false'" - ), - } - self.mcp_registry = Some(mcp_registry); - - Ok(()) - } - - fn setup_model(&mut self) -> Result<()> { - let mut model_id = self.model_id.clone(); - if model_id.is_empty() { - let models = list_models(&self.to_app_config(), ModelType::Chat); - if models.is_empty() { - bail!("No available model"); - } - model_id = models[0].id() - } - self.set_model(&model_id)?; - self.model_id = model_id; - - Ok(()) - } - - 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); - }); - } - - 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") - )); - } - } } pub fn load_env_file() -> Result<()> { @@ -897,7 +457,7 @@ impl AssertState { } } -async fn create_config_file(config_path: &Path) -> Result<()> { +pub async fn create_config_file(config_path: &Path) -> Result<()> { let ans = Confirm::new("No config file, create a new one?") .with_default(true) .prompt()?; @@ -1021,21 +581,17 @@ mod tests { assert_eq!(cfg.model_id, ""); assert_eq!(cfg.temperature, None); assert_eq!(cfg.top_p, None); - assert_eq!(cfg.dry_run, false); - assert_eq!(cfg.stream, true); - assert_eq!(cfg.save, false); - assert_eq!(cfg.highlight, true); - assert_eq!(cfg.function_calling_support, true); - assert_eq!(cfg.mcp_server_support, true); + assert!(!cfg.dry_run); + assert!(cfg.stream); + assert!(!cfg.save); + assert!(cfg.highlight); + assert!(cfg.function_calling_support); + assert!(cfg.mcp_server_support); assert_eq!(cfg.compression_threshold, 4000); assert_eq!(cfg.rag_top_k, 5); - assert_eq!(cfg.save_shell_history, true); + assert!(cfg.save_shell_history); assert_eq!(cfg.keybindings, "emacs"); assert!(cfg.clients.is_empty()); - assert!(cfg.role.is_none()); - assert!(cfg.session.is_none()); - assert!(cfg.agent.is_none()); - assert!(cfg.rag.is_none()); assert!(cfg.save_session.is_none()); assert!(cfg.enabled_tools.is_none()); assert!(cfg.enabled_mcp_servers.is_none()); diff --git a/src/config/paths.rs b/src/config/paths.rs new file mode 100644 index 0000000..3f778d9 --- /dev/null +++ b/src/config/paths.rs @@ -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 { + 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)> { + 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 { + 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 { + 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 { + 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> { + 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) +} diff --git a/src/config/rag_cache.rs b/src/config/rag_cache.rs new file mode 100644 index 0000000..e2658f4 --- /dev/null +++ b/src/config/rag_cache.rs @@ -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 `) 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` so the cache never keeps a RAG +//! alive on its own — once all active scopes drop their `Arc`, +//! 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>>, +} + +impl RagCache { + pub fn try_get(&self, key: &RagKey) -> Option> { + let map = self.entries.read(); + map.get(key).and_then(|weak| weak.upgrade()) + } + + pub fn insert(&self, key: RagKey, rag: &Arc) { + 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(&self, key: RagKey, loader: F) -> Result> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + 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) + } +} diff --git a/src/config/request_context.rs b/src/config/request_context.rs new file mode 100644 index 0000000..f4dfce3 --- /dev/null +++ b/src/config/request_context.rs @@ -0,0 +1,2609 @@ +//! Per-request mutable state for a single Loki interaction. +//! +//! `RequestContext` owns the runtime state that was previously stored +//! on `Config` as `#[serde(skip)]` fields: the active role, session, +//! agent, RAG, supervisor state, inbox/escalation queues, the +//! conversation's "last message" cursor, and the per-scope +//! [`ToolScope`](super::tool_scope::ToolScope) carrying functions and +//! live MCP handles. +//! +//! 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. +//! +//! `RequestContext` is built via [`RequestContext::bootstrap`] (CLI/REPL +//! entry point) or [`RequestContext::new`] (test/child-agent helper). +//! It holds an `Arc` for shared, immutable services +//! (config, vault, MCP factory, RAG cache, MCP registry, base +//! functions). + +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 bootstrap( + app: Arc, + working_mode: WorkingMode, + info_flag: bool, + ) -> Result { + let model = Model::retrieve_model(&app.config, &app.config.model_id, ModelType::Chat)?; + + let mut functions = app.functions.clone(); + if working_mode.is_repl() { + functions.append_user_interaction_functions(); + } + + let mut mcp_runtime = McpRuntime::default(); + if let Some(registry) = &app.mcp_registry { + mcp_runtime.sync_from_registry(registry); + } + + Ok(Self { + app, + macro_flag: false, + info_flag, + working_mode, + model, + agent_variables: None, + role: None, + session: None, + rag: None, + agent: None, + last_message: None, + tool_scope: ToolScope { + functions, + mcp_runtime, + tool_tracker: ToolCallTracker::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; + 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 default_app_state() -> Arc { + Arc::new(AppState { + config: Arc::new(AppConfig::default()), + vault: Arc::new(crate::vault::Vault::default()), + mcp_factory: Arc::new(super::super::mcp_factory::McpFactory::default()), + rag_cache: Arc::new(RagCache::default()), + mcp_config: None, + mcp_log_path: None, + mcp_registry: None, + functions: Functions::default(), + }) + } + + fn create_test_ctx() -> RequestContext { + RequestContext::new(default_app_state(), WorkingMode::Cmd) + } + + #[test] + fn new_creates_clean_state() { + let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + + 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 mut ctx = RequestContext::new(default_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()); + } +} diff --git a/src/config/session.rs b/src/config/session.rs index c8e4204..42f19c8 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -67,12 +67,7 @@ pub struct Session { } impl Session { - #[allow(dead_code)] - pub fn new_from_ctx( - ctx: &request_context::RequestContext, - app: &AppConfig, - name: &str, - ) -> Self { + 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(), @@ -84,9 +79,8 @@ impl Session { session } - #[allow(dead_code)] pub fn load_from_ctx( - ctx: &request_context::RequestContext, + ctx: &RequestContext, app: &AppConfig, name: &str, path: &Path, @@ -680,7 +674,8 @@ impl AutoName { mod tests { use super::*; use crate::client::{Message, MessageContent, MessageRole, Model}; - use crate::config::{AppState, Config}; + use crate::config::{AppConfig, AppState, RequestContext, WorkingMode}; + use crate::function::Functions; use std::sync::Arc; #[test] @@ -694,17 +689,18 @@ mod tests { #[test] fn session_new_from_ctx_captures_save_session() { - let cfg = Config::default(); - let app_config = Arc::new(cfg.to_app_config()); + let app_config = Arc::new(AppConfig::default()); let app_state = Arc::new(AppState { config: app_config.clone(), - vault: cfg.vault.clone(), - mcp_factory: Arc::new(crate::config::mcp_factory::McpFactory::new()), - rag_cache: Arc::new(crate::config::rag_cache::RagCache::new()), + vault: Arc::new(crate::vault::Vault::default()), + mcp_factory: Arc::new(mcp_factory::McpFactory::default()), + rag_cache: Arc::new(rag_cache::RagCache::default()), mcp_config: None, mcp_log_path: None, + mcp_registry: None, + functions: Functions::default(), }); - let ctx = cfg.to_request_context(app_state); + let ctx = RequestContext::new(app_state, WorkingMode::Cmd); let session = Session::new_from_ctx(&ctx, &app_config, "test-session"); assert_eq!(session.name(), "test-session"); diff --git a/src/config/tool_scope.rs b/src/config/tool_scope.rs new file mode 100644 index 0000000..b499b23 --- /dev/null +++ b/src/config/tool_scope.rs @@ -0,0 +1,172 @@ +//! 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 +//! +//! `ToolScope` lives on [`RequestContext`](super::request_context::RequestContext) +//! and is built/replaced as the active scope changes (role swap, +//! session swap, agent enter/exit). The base `functions` are seeded +//! from [`AppState`](super::app_state::AppState) and per-scope filters +//! narrow the visible set. + +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>, +} + +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) { + self.servers.insert(name, handle); + } + + pub fn get(&self, name: &str) -> Option<&Arc> { + self.servers.get(name) + } + + pub fn server_names(&self) -> Vec { + 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> { + 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> { + 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::::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 { + 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 { + 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) + } +} diff --git a/src/function/mod.rs b/src/function/mod.rs index 0852f89..98e31b6 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -148,7 +148,7 @@ pub async fn eval_tool_calls( } if !output.is_empty() { - let (has_escalations, summary) = if ctx.current_depth() == 0 + let (has_escalations, summary) = if ctx.current_depth == 0 && let Some(queue) = ctx.root_escalation_queue() && queue.has_pending() { @@ -192,7 +192,7 @@ pub struct Functions { } impl Functions { - fn install_global_tools() -> Result<()> { + pub fn install_builtin_global_tools() -> Result<()> { info!( "Installing global built-in functions in {}", paths::functions_dir().display() @@ -241,7 +241,6 @@ impl Functions { } pub fn init(visible_tools: &[String]) -> Result { - Self::install_global_tools()?; Self::clear_global_functions_bin_dir()?; let declarations = Self { @@ -258,7 +257,6 @@ impl Functions { } pub fn init_agent(name: &str, global_tools: &[String]) -> Result { - Self::install_global_tools()?; Self::clear_agent_bin_dir(name)?; let global_tools_declarations = if !global_tools.is_empty() { @@ -943,7 +941,7 @@ impl ToolCall { pub async fn eval(&self, ctx: &mut RequestContext) -> Result { let agent = ctx.agent.clone(); let functions = ctx.tool_scope.functions.clone(); - let current_depth = ctx.current_depth(); + let current_depth = ctx.current_depth; let agent_name = agent.as_ref().map(|agent| agent.name().to_owned()); let (call_name, cmd_name, mut cmd_args, envs) = match agent.as_ref() { Some(agent) => self.extract_call_config_from_agent(&functions, agent)?, diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 8c19ba7..741257a 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -364,7 +364,7 @@ fn run_child_agent( input = input.merge_tool_results(output, tool_results); } - if let Some(supervisor) = child_ctx.supervisor().cloned() { + if let Some(supervisor) = child_ctx.supervisor.clone() { supervisor.read().cancel_all(); } @@ -441,7 +441,8 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { let (max_depth, current_depth) = { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active; Agent spawning not enabled"))?; let sup = supervisor.read(); @@ -455,7 +456,7 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { ), })); } - (sup.max_depth(), ctx.current_depth() + 1) + (sup.max_depth(), ctx.current_depth + 1) }; if current_depth > max_depth { @@ -481,10 +482,12 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { let child_app_state = Arc::new(AppState { config: Arc::new(app_config.as_ref().clone()), vault: ctx.app.vault.clone(), - mcp_factory: Default::default(), - rag_cache: Default::default(), + mcp_factory: ctx.app.mcp_factory.clone(), + rag_cache: ctx.app.rag_cache.clone(), mcp_config: ctx.app.mcp_config.clone(), mcp_log_path: ctx.app.mcp_log_path.clone(), + mcp_registry: ctx.app.mcp_registry.clone(), + functions: ctx.app.functions.clone(), }); let agent = Agent::init( app_config.as_ref(), @@ -509,18 +512,9 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { agent_id.clone(), ); child_ctx.rag = agent.rag(); - child_ctx - .agent_runtime - .as_mut() - .expect("child agent runtime should be initialized") - .rag = child_ctx.rag.clone(); child_ctx.agent = Some(agent); if should_init_supervisor { - child_ctx - .agent_runtime - .as_mut() - .expect("child agent runtime should be initialized") - .supervisor = Some(Arc::new(RwLock::new(Supervisor::new( + child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new( max_concurrent, max_depth, )))); @@ -574,7 +568,8 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { }; let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); @@ -596,7 +591,8 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result { let is_finished = { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let sup = supervisor.read(); @@ -625,7 +621,8 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result let handle = { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); @@ -659,7 +656,8 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result fn handle_list(ctx: &mut RequestContext) -> Result { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let sup = supervisor.read(); @@ -691,7 +689,8 @@ fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result { .ok_or_else(|| anyhow!("'id' is required"))?; let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); @@ -722,17 +721,19 @@ fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result .ok_or_else(|| anyhow!("'message' is required"))?; let sender = ctx - .self_agent_id() - .map(str::to_owned) + .self_agent_id + .clone() .or_else(|| ctx.agent.as_ref().map(|a| a.name().to_string())) .unwrap_or_else(|| "parent".to_string()); let inbox = ctx - .supervisor() + .supervisor + .as_ref() .and_then(|sup| sup.read().inbox(id).cloned()); let inbox = inbox.or_else(|| { - ctx.parent_supervisor() + ctx.parent_supervisor + .as_ref() .and_then(|sup| sup.read().inbox(id).cloned()) }); @@ -760,7 +761,7 @@ fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result } fn handle_check_inbox(ctx: &mut RequestContext) -> Result { - match ctx.inbox() { + match ctx.inbox.as_ref() { Some(inbox) => { let messages: Vec = inbox .drain() @@ -797,8 +798,8 @@ fn handle_reply_escalation(ctx: &mut RequestContext, args: &Value) -> Result Result { } let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); @@ -883,7 +885,8 @@ fn handle_task_create(ctx: &mut RequestContext, args: &Value) -> Result { fn handle_task_list(ctx: &mut RequestContext) -> Result { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let sup = supervisor.read(); @@ -917,7 +920,8 @@ async fn handle_task_complete(ctx: &mut RequestContext, args: &Value) -> Result< let (newly_runnable, dispatchable) = { let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); @@ -997,7 +1001,8 @@ fn handle_task_fail(ctx: &mut RequestContext, args: &Value) -> Result { .ok_or_else(|| anyhow!("'task_id' is required"))?; let supervisor = ctx - .supervisor() + .supervisor + .as_ref() .cloned() .ok_or_else(|| anyhow!("No supervisor active"))?; let mut sup = supervisor.write(); diff --git a/src/function/todo.rs b/src/function/todo.rs index 52fd5cd..ed8e76e 100644 --- a/src/function/todo.rs +++ b/src/function/todo.rs @@ -94,31 +94,23 @@ pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value) .strip_prefix(TODO_FUNCTION_PREFIX) .unwrap_or(cmd_name); + if ctx.agent.is_none() { + bail!("No active agent"); + } + match action { "init" => { let goal = args.get("goal").and_then(Value::as_str).unwrap_or_default(); - let agent = ctx.agent.as_mut(); - match agent { - Some(agent) => { - agent.init_todo_list(goal); - Ok(json!({"status": "ok", "message": "Initialized new todo list"})) - } - None => bail!("No active agent"), - } + ctx.init_todo_list(goal); + Ok(json!({"status": "ok", "message": "Initialized new todo list"})) } "add" => { let task = args.get("task").and_then(Value::as_str).unwrap_or_default(); if task.is_empty() { return Ok(json!({"error": "task description is required"})); } - let agent = ctx.agent.as_mut(); - match agent { - Some(agent) => { - let id = agent.add_todo(task); - Ok(json!({"status": "ok", "id": id})) - } - None => bail!("No active agent"), - } + let id = ctx.add_todo(task); + Ok(json!({"status": "ok", "id": id})) } "done" => { let id = args @@ -130,47 +122,26 @@ pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value) .map(|v| v as usize); match id { Some(id) => { - let agent = ctx.agent.as_mut(); - match agent { - Some(agent) => { - if agent.mark_todo_done(id) { - Ok( - json!({"status": "ok", "message": format!("Marked todo {id} as done")}), - ) - } else { - Ok(json!({"error": format!("Todo {id} not found")})) - } - } - None => bail!("No active agent"), + if ctx.mark_todo_done(id) { + Ok(json!({"status": "ok", "message": format!("Marked todo {id} as done")})) + } else { + Ok(json!({"error": format!("Todo {id} not found")})) } } None => Ok(json!({"error": "id is required and must be a number"})), } } "list" => { - let agent = ctx.agent.as_ref(); - match agent { - Some(agent) => { - let list = agent.todo_list(); - if list.is_empty() { - Ok(json!({"goal": "", "todos": []})) - } else { - Ok(serde_json::to_value(list) - .unwrap_or(json!({"error": "serialization failed"}))) - } - } - None => bail!("No active agent"), + let list = &ctx.todo_list; + if list.is_empty() { + Ok(json!({"goal": "", "todos": []})) + } else { + Ok(serde_json::to_value(list).unwrap_or(json!({"error": "serialization failed"}))) } } "clear" => { - let agent = ctx.agent.as_mut(); - match agent { - Some(agent) => { - agent.clear_todo_list(); - Ok(json!({"status": "ok", "message": "Todo list cleared"})) - } - None => bail!("No active agent"), - } + ctx.clear_todo_list(); + Ok(json!({"status": "ok", "message": "Todo list cleared"})) } _ => bail!("Unknown todo action: {action}"), } diff --git a/src/function/user_interaction.rs b/src/function/user_interaction.rs index 55e94bb..e747a76 100644 --- a/src/function/user_interaction.rs +++ b/src/function/user_interaction.rs @@ -128,7 +128,7 @@ pub async fn handle_user_tool( .strip_prefix(USER_FUNCTION_PREFIX) .unwrap_or(cmd_name); - let depth = ctx.current_depth(); + let depth = ctx.current_depth; if depth == 0 { handle_direct(action, args) @@ -213,8 +213,8 @@ async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> R }); let from_agent_id = ctx - .self_agent_id() - .map(ToOwned::to_owned) + .self_agent_id + .clone() .unwrap_or_else(|| "unknown".to_string()); let from_agent_name = ctx .agent diff --git a/src/main.rs b/src/main.rs index b2abf3f..b453566 100644 --- a/src/main.rs +++ b/src/main.rs @@ -22,8 +22,8 @@ use crate::client::{ use crate::config::paths; use crate::config::{ Agent, AppConfig, AppState, CODE_ROLE, Config, EXPLAIN_SHELL_ROLE, Input, RequestContext, - SHELL_ROLE, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, load_env_file, - macro_execute, + SHELL_ROLE, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, install_builtins, + list_agents, load_env_file, macro_execute, sync_models, }; use crate::render::{prompt_theme, render_error}; use crate::repl::Repl; @@ -82,45 +82,38 @@ async fn main() -> Result<()> { let log_path = setup_logger()?; + install_builtins()?; + if let Some(client_arg) = &cli.authenticate { - let config = Config::init_bare()?; - let (client_name, provider) = resolve_oauth_client(client_arg.as_deref(), &config.clients)?; + let cfg = Config::load_with_interpolation(true).await?; + let app_config = AppConfig::from_config(cfg)?; + let (client_name, provider) = + resolve_oauth_client(client_arg.as_deref(), &app_config.clients)?; oauth::run_oauth_flow(&*provider, &client_name).await?; return Ok(()); } if vault_flags { - return Vault::handle_vault_flags(cli, Config::init_bare()?); + let cfg = Config::load_with_interpolation(true).await?; + let app_config = AppConfig::from_config(cfg)?; + let vault = Vault::init(&app_config); + return Vault::handle_vault_flags(cli, &vault); } let abort_signal = create_abort_signal(); let start_mcp_servers = cli.agent.is_none() && cli.role.is_none(); - let cfg = Config::init( - working_mode, - info_flag, - start_mcp_servers, - log_path, - abort_signal.clone(), - ) - .await?; - let app_config: Arc = Arc::new(cfg.to_app_config()); - let (mcp_config, mcp_log_path) = match &cfg.mcp_registry { - Some(reg) => (reg.mcp_config().cloned(), reg.log_path().cloned()), - None => (None, None), - }; - let app_state: Arc = Arc::new(AppState { - config: app_config, - vault: cfg.vault.clone(), - mcp_factory: Default::default(), - rag_cache: Default::default(), - mcp_config, - mcp_log_path, - }); - let ctx = cfg.to_request_context(app_state); - log::debug!( - "ctx.tool_scope.mcp_runtime servers after sync: {:?}", - ctx.tool_scope.mcp_runtime.server_names() + let cfg = Config::load_with_interpolation(info_flag).await?; + let app_config: Arc = Arc::new(AppConfig::from_config(cfg)?); + let app_state: Arc = Arc::new( + AppState::init( + app_config, + log_path, + start_mcp_servers, + abort_signal.clone(), + ) + .await?, ); + let ctx = RequestContext::bootstrap(app_state, working_mode, info_flag)?; { let app = &*ctx.app.config; @@ -153,7 +146,7 @@ async fn run( ) -> Result<()> { if cli.sync_models { let url = ctx.app.config.sync_models_url(); - return Config::sync_models(&url, abort_signal.clone()).await; + return sync_models(&url, abort_signal.clone()).await; } if cli.list_models { diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index a03cac9..6dc22c7 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -1,21 +1,18 @@ -use crate::config::Config; +use crate::config::AppConfig; use crate::config::paths; use crate::utils::{AbortSignal, abortable_run_with_spinner}; +use crate::vault::Vault; use crate::vault::interpolate_secrets; use anyhow::{Context, Result, anyhow}; -use futures_util::future::BoxFuture; use futures_util::{StreamExt, TryStreamExt, stream}; use indoc::formatdoc; -use rmcp::model::{CallToolRequestParams, CallToolResult}; use rmcp::service::RunningService; use rmcp::transport::TokioChildProcess; use rmcp::{RoleClient, ServiceExt}; use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::fs::OpenOptions; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::Arc; use tokio::process::Command; @@ -82,7 +79,8 @@ impl McpRegistry { start_mcp_servers: bool, enabled_mcp_servers: Option, abort_signal: AbortSignal, - config: &Config, + app_config: &AppConfig, + vault: &Vault, ) -> Result { let mut registry = Self { log_path, @@ -115,7 +113,7 @@ impl McpRegistry { return Ok(registry); } - let (parsed_content, missing_secrets) = interpolate_secrets(&content, &config.vault); + let (parsed_content, missing_secrets) = interpolate_secrets(&content, vault); if !missing_secrets.is_empty() { return Err(anyhow!(formatdoc!( @@ -130,7 +128,7 @@ impl McpRegistry { serde_json::from_str(&parsed_content).with_context(err)?; registry.config = Some(mcp_servers_config); - if start_mcp_servers && config.mcp_server_support { + if start_mcp_servers && app_config.mcp_server_support { abortable_run_with_spinner( registry.start_select_mcp_servers(enabled_mcp_servers), "Loading MCP servers", @@ -249,65 +247,6 @@ impl McpRegistry { self.servers.keys().cloned().collect() } - #[allow(dead_code)] - pub async fn describe(&self, server_id: &str, tool: &str) -> Result { - let server = self - .servers - .iter() - .filter(|(id, _)| &server_id == id) - .map(|(_, s)| s.clone()) - .next() - .ok_or(anyhow!("{server_id} MCP server not found in config"))?; - - let tool_schema = server - .list_tools(None) - .await? - .tools - .into_iter() - .find(|it| it.name == tool) - .ok_or(anyhow!( - "{tool} not found in {server_id} MCP server catalog" - ))? - .input_schema; - Ok(json!({ - "type": "object", - "properties": { - "tool": { - "type": "string", - }, - "arguments": tool_schema - } - })) - } - - #[allow(dead_code)] - pub fn invoke( - &self, - server: &str, - tool: &str, - arguments: Value, - ) -> BoxFuture<'static, Result> { - let server = self - .servers - .get(server) - .cloned() - .with_context(|| format!("Invoked MCP server does not exist: {server}")); - - let tool = tool.to_owned(); - Box::pin(async move { - let server = server?; - let call_tool_request = CallToolRequestParams { - name: Cow::Owned(tool.to_owned()), - arguments: arguments.as_object().cloned(), - meta: None, - task: None, - }; - - let result = server.call_tool(call_tool_request).await?; - Ok(result) - }) - } - pub fn is_empty(&self) -> bool { self.servers.is_empty() } @@ -323,7 +262,7 @@ impl McpRegistry { pub(crate) async fn spawn_mcp_server( spec: &McpServer, - log_path: Option<&std::path::Path>, + log_path: Option<&Path>, ) -> Result> { let mut cmd = Command::new(&spec.command); if let Some(args) = &spec.args { diff --git a/src/rag/mod.rs b/src/rag/mod.rs index 5bfef46..1232452 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -20,6 +20,30 @@ use std::{ }; use tokio::time::sleep; +const RAG_TEMPLATE: &str = r#"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags) + + +__CONTEXT__ + + + +__SOURCES__ + + + +- If you don't know, just say so. +- If you are not sure, ask for clarification. +- Answer in the same language as the user query. +- If the context appears unreadable or of poor quality, tell the user then answer as best as you can. +- If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge. +- Answer directly and without using xml tags. +- When using information from the context, cite the relevant source from the section. + + + +__INPUT__ +"#; + pub struct Rag { app_config: Arc, name: String, @@ -64,7 +88,6 @@ impl Rag { pub async fn init( app: &AppConfig, - clients: &[ClientConfig], name: &str, save_path: &Path, doc_paths: &[String], @@ -85,7 +108,7 @@ impl Rag { top_k, embedding_model.max_batch_size(), ); - let mut rag = Self::create(app, clients, name, save_path, data)?; + let mut rag = Self::create(app, name, save_path, data)?; let mut paths = doc_paths.to_vec(); if paths.is_empty() { paths = add_documents()?; @@ -104,25 +127,14 @@ impl Rag { Ok(rag) } - pub fn load( - app: &AppConfig, - clients: &[ClientConfig], - name: &str, - path: &Path, - ) -> Result { + pub fn load(app: &AppConfig, name: &str, path: &Path) -> Result { let err = || format!("Failed to load rag '{name}' at '{}'", path.display()); let content = fs::read_to_string(path).with_context(err)?; let data: RagData = serde_yaml::from_str(&content).with_context(err)?; - Self::create(app, clients, name, path, data) + Self::create(app, name, path, data) } - pub fn create( - app: &AppConfig, - _clients: &[ClientConfig], - name: &str, - path: &Path, - data: RagData, - ) -> Result { + pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result { let hnsw = data.build_hnsw(); let bm25 = data.build_bm25(); let embedding_model = @@ -330,6 +342,29 @@ impl Rag { Ok((embeddings, sources, ids)) } + pub async fn search_with_template( + &self, + app: &AppConfig, + text: &str, + abort_signal: AbortSignal, + ) -> Result { + let (reranker_model, top_k) = self.get_config(); + let (embeddings, sources, ids) = self + .search(text, top_k, reranker_model.as_deref(), abort_signal) + .await?; + let rag_template = app.rag_template.as_deref().unwrap_or(RAG_TEMPLATE); + let text = if embeddings.is_empty() { + text.to_string() + } else { + rag_template + .replace("__CONTEXT__", &embeddings) + .replace("__SOURCES__", &sources) + .replace("__INPUT__", text) + }; + self.set_last_sources(&ids); + Ok(text) + } + fn resolve_source(&self, id: &DocumentId) -> String { let (file_index, _) = id.split(); self.data diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 01eba1d..a139dae 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -231,6 +231,7 @@ impl Repl { }) } + #[allow(clippy::await_holding_lock)] pub async fn run(&mut self) -> Result<()> { if AssertState::False(StateFlags::AGENT | StateFlags::RAG).assert(self.ctx.read().state()) { print!( @@ -770,11 +771,11 @@ pub async fn run_repl_command( "The todo system is not enabled for this agent. Set 'auto_continue: true' in the agent's config.yaml to enable it." ); } - if agent.todo_list().is_empty() { + if ctx.todo_list.is_empty() { println!("Todo list is already empty."); false } else { - agent.clear_todo_list(); + ctx.clear_todo_list(); println!("Todo list cleared."); true } @@ -824,7 +825,7 @@ pub async fn run_repl_command( _ => unknown_command()?, }, None => { - reset_agent_continuation(ctx); + reset_continuation(ctx); let input = Input::from_str(ctx, line, None); ask(ctx, abort_signal.clone(), input, true).await?; } @@ -884,14 +885,14 @@ async fn ask( if should_continue { let full_prompt = { + let todo_state = ctx.todo_list.render_for_model(); + let remaining = ctx.todo_list.incomplete_count(); + ctx.set_last_continuation_response(output.clone()); + ctx.increment_auto_continue_count(); let agent = ctx.agent.as_mut().expect("agent checked above"); - agent.set_last_continuation_response(output.clone()); - agent.increment_continuation(); - let count = agent.continuation_count(); + let count = ctx.auto_continue_count; let max = agent.max_auto_continues(); - let todo_state = agent.todo_list().render_for_model(); - let remaining = agent.todo_list().incomplete_count(); let prompt = agent.continuation_prompt(); let color = if app.light_theme() { @@ -911,7 +912,7 @@ async fn ask( let continuation_input = Input::from_str(ctx, &full_prompt, None); ask(ctx, abort_signal, continuation_input, false).await } else { - reset_agent_continuation(ctx); + reset_continuation(ctx); if ctx.maybe_autoname_session() { let color = if app.light_theme() { nu_ansi_term::Color::LightGray @@ -955,13 +956,13 @@ async fn ask( if agent_can_continue_after_compress { let full_prompt = { + let todo_state = ctx.todo_list.render_for_model(); + let remaining = ctx.todo_list.incomplete_count(); + ctx.increment_auto_continue_count(); let agent = ctx.agent.as_mut().expect("agent checked above"); - agent.increment_continuation(); - let count = agent.continuation_count(); + let count = ctx.auto_continue_count; let max = agent.max_auto_continues(); - let todo_state = agent.todo_list().render_for_model(); - let remaining = agent.todo_list().incomplete_count(); let prompt = agent.continuation_prompt(); let color = if app.light_theme() { @@ -990,20 +991,12 @@ async fn ask( fn agent_should_continue(ctx: &RequestContext) -> bool { ctx.agent.as_ref().is_some_and(|agent| { - agent.auto_continue_enabled() - && agent.continuation_count() < agent.max_auto_continues() - && agent.todo_list().has_incomplete() - }) + agent.auto_continue_enabled() && ctx.auto_continue_count < agent.max_auto_continues() + }) && ctx.todo_list.has_incomplete() } -fn reset_agent_continuation(ctx: &mut RequestContext) { - if ctx - .agent - .as_ref() - .is_some_and(|agent| agent.continuation_count() > 0) - { - ctx.agent.as_mut().unwrap().reset_continuation(); - } +fn reset_continuation(ctx: &mut RequestContext) { + ctx.reset_continuation_count(); } fn unknown_command() -> Result<()> { diff --git a/src/vault/mod.rs b/src/vault/mod.rs index e170075..2f691a8 100644 --- a/src/vault/mod.rs +++ b/src/vault/mod.rs @@ -5,7 +5,7 @@ pub use utils::create_vault_password_file; pub use utils::interpolate_secrets; use crate::cli::Cli; -use crate::config::Config; +use crate::config::AppConfig; use crate::vault::utils::ensure_password_file_initialized; use anyhow::{Context, Result}; use fancy_regex::Regex; @@ -26,7 +26,7 @@ pub type GlobalVault = Arc; impl Vault { pub fn init_bare() -> Self { - let vault_password_file = Config::default().vault_password_file(); + let vault_password_file = AppConfig::default().vault_password_file(); let local_provider = LocalProvider { password_file: Some(vault_password_file), git_branch: None, @@ -36,7 +36,7 @@ impl Vault { Self { local_provider } } - pub fn init(config: &Config) -> Self { + pub fn init(config: &AppConfig) -> Self { let vault_password_file = config.vault_password_file(); let mut local_provider = LocalProvider { password_file: Some(vault_password_file), @@ -130,25 +130,25 @@ impl Vault { Ok(secrets) } - pub fn handle_vault_flags(cli: Cli, config: Config) -> Result<()> { + pub fn handle_vault_flags(cli: Cli, vault: &Vault) -> Result<()> { if let Some(secret_name) = cli.add_secret { - config.vault.add_secret(&secret_name)?; + vault.add_secret(&secret_name)?; } if let Some(secret_name) = cli.get_secret { - config.vault.get_secret(&secret_name, true)?; + vault.get_secret(&secret_name, true)?; } if let Some(secret_name) = cli.update_secret { - config.vault.update_secret(&secret_name)?; + vault.update_secret(&secret_name)?; } if let Some(secret_name) = cli.delete_secret { - config.vault.delete_secret(&secret_name)?; + vault.delete_secret(&secret_name)?; } if cli.list_secrets { - config.vault.list_secrets(true)?; + vault.list_secrets(true)?; } Ok(())