diff --git a/config.example.toml b/config.example.toml index 9ef3b1c..e2ba4f2 100644 --- a/config.example.toml +++ b/config.example.toml @@ -101,6 +101,7 @@ Be concise and helpful.""" # [agent] # max_iterations = 25 # Agent loop cap (default 25) # empty_response_retry_limit = 3 # Recovery attempts for empty model responses (default 3; 0 = fail immediately) +# parse_retry_limit = 3 # Retries when API response is missing 'choices' field (default 3; 0 = fail immediately; exponential backoff: 1s, 2s, 4s...) # LangSmith observability (optional) # Traces every LLM call and tool execution for debugging in the LangSmith UI. diff --git a/docs/superpowers/plans/2026-07-06-compaction-context-management.md b/docs/superpowers/plans/2026-07-06-compaction-context-management.md new file mode 100644 index 0000000..f100ea4 --- /dev/null +++ b/docs/superpowers/plans/2026-07-06-compaction-context-management.md @@ -0,0 +1,1403 @@ +# Context-Aware Compaction & Memory Management Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) for syntax tracking. + +**Goal:** Add dynamic context window detection, system prompt preservation, structured YAML+state summaries, RAG-assisted compaction, and configurable hybrid search to RustFox's compaction system. + +**Architecture:** Six focused changes across the provider layer, memory layer, agent prompt builder, and agent loop. ProviderConfig gains a runtime context_window_cache with a background warmup task. The compact flow partitions system messages before summarization, injects retrieved RAG context and tool-group-aware truncation, then rebuilds system messages after. vec0 tables gain metadata columns for pre-filtering. + +**Tech Stack:** Rust 2021, tokio, rusqlite, sqlite-vec, reqwest + +--- + +### Files Touched + +| File | Change | +|------|--------| +| `src/config.rs` | Add `rrf_k`, `rrf_weight_fts`, `rrf_weight_vec` to `MemoryConfig` | +| `src/provider.rs` | Add `context_window_cache` to `ProviderConfig`, `fetch_context_window()` trait method, `effective_context_window()` on registry | +| `src/memory/mod.rs` | vec0 metadata column migration (DROP + recreate with `is_summarized`, `role`) | +| `src/memory/conversations.rs` | Configurable RRF weights in hybrid search queries | +| `src/memory/rag.rs` | New `retrieve_context_for_compaction()` function | +| `src/agent_prompt.rs` | Enhanced structured summary prompt, updated boundary marker, updated recovery nudge | +| `src/agent.rs` | System prompt preservation, RAG-aware compact flow, tool-group-aware truncation | +| `src/main.rs` | Background startup task to warm `context_window_cache` | + + +--- + +## Task 1: Config — Add RRF parameters to MemoryConfig + +**Files:** +- Modify: `src/config.rs:235-254` + +- [ ] **Step 1: Add rrf fields to MemoryConfig** + +```rust +#[derive(Debug, Deserialize, Clone)] +pub struct MemoryConfig { + // ... existing fields ... + #[serde(default = "default_rrf_k")] + pub rrf_k: f64, + #[serde(default = "default_rrf_weight_fts")] + pub rrf_weight_fts: f64, + #[serde(default = "default_rrf_weight_vec")] + pub rrf_weight_vec: f64, +} + +fn default_rrf_k() -> f64 { 60.0 } +fn default_rrf_weight_fts() -> f64 { 0.5 } +fn default_rrf_weight_vec() -> f64 { 0.5 } +``` + +- [ ] **Step 2: Run `cargo check` to verify compilation** + +Run: `cargo check 2>&1 | head -20` +Expected: compiles without errors (no warnings for dead_code on new fields is fine) + +- [ ] **Step 3: Commit** + +```bash +git add src/config.rs +git commit -m "feat(config): add rrf_k, rrf_weight_fts, rrf_weight_vec to MemoryConfig" +``` + +--- + +## Task 2: Provider — Add context_window_cache and fetch_context_window + +**Files:** +- Modify: `src/provider.rs` +- Add import: `use tokio::sync::RwLock;` at line 2 + +- [ ] **Step 1: Add context_window_cache to ProviderConfig** + +```rust +use std::sync::Arc; +use tokio::sync::RwLock; // add to imports at line 2 + +#[derive(Debug, Clone)] +pub struct ProviderConfig { + // ... existing fields ... + pub context_window: usize, + /// Runtime cache for the current model's context window, populated + /// asynchronously from the provider API. When None, falls back to + /// `context_window`. + pub context_window_cache: Arc>>, + // ... remaining fields ... +} +``` + +Initialize in `From<&ProviderSection>`: +```rust +impl From<&ProviderSection> for ProviderConfig { + fn from(s: &ProviderSection) -> Self { + Self { + // ... existing fields ... + context_window: s.context_window, + context_window_cache: Arc::new(RwLock::new(None)), + parse_retry_limit: 0, + } + } +} +``` + +- [ ] **Step 2: Add `fetch_context_window` to Provider trait** + +```rust +#[async_trait] +pub trait Provider: Send + Sync { + fn name(&self) -> &str; + fn default_model(&self) -> &str; + fn supports_vision(&self) -> bool; + fn config(&self) -> &ProviderConfig; + + // ... existing methods ... + + /// Fetch the context window for a given model from the provider API. + /// Returns None if the provider doesn't support runtime detection. + async fn fetch_context_window( + &self, + _client: &reqwest::Client, + _model: &str, + ) -> Option { + None // default: no API-based detection + } +} +``` + +- [ ] **Step 3: Implement for OpenRouterProvider** + +Add method after `list_models` (after line 383): + +```rust +async fn fetch_context_window( + &self, + client: &reqwest::Client, + model: &str, +) -> Option { + let url = format!("{}/models", self.config.base_url); + let mut req = client.get(&url); + if let Some(ref key) = self.config.api_key { + req = req.header("Authorization", format!("Bearer {}", key)); + } + let response = req.send().await.ok()?; + if !response.status().is_success() { + return None; + } + let list: serde_json::Value = response.json().await.ok()?; + let ctx = list["data"].as_array()? + .iter() + .find(|m| m["id"].as_str() == Some(model))? + .get("context_length")? + .as_u64()?; + Some(ctx as usize) +} +``` + +- [ ] **Step 4: Implement for OpenAICompatibleProvider** + +```rust +async fn fetch_context_window( + &self, + client: &reqwest::Client, + model: &str, +) -> Option { + // Same implementation as OpenRouterProvider — many OpenAI-compatible + // providers expose the same /v1/models endpoint with context_length. + let url = format!("{}/models", self.config.base_url); + let mut req = client.get(&url); + if let Some(ref key) = self.config.api_key { + req = req.header("Authorization", format!("Bearer {}", key)); + } + let response = req.send().await.ok()?; + if !response.status().is_success() { + return None; + } + let list: serde_json::Value = response.json().await.ok()?; + let ctx = list["data"].as_array()? + .iter() + .find(|m| m["id"].as_str() == Some(model))? + .get("context_length")? + .as_u64()?; + Some(ctx as usize) +} +``` + +- [ ] **Step 5: Add `effective_context_window` to ProviderRegistry** + +```rust +impl ProviderRegistry { + // ... existing methods ... + + /// Return the effective context window for a model: runtime cache + /// if populated, otherwise static config fallback. + pub fn effective_context_window(&self, model: &str) -> usize { + let (provider, _) = self.resolve_model(model); + let cached = provider.config().context_window_cache.try_read() + .ok() + .and_then(|c| *c); + cached.unwrap_or(provider.config().context_window) + } +} +``` + +Note: using `try_read()` instead of `read()` because this may be called from sync contexts in the agent loop where `.await` is not available in the hot path (the agent loop calls this without holding an async runtime lock for the hot path at line 638-643). + +- [ ] **Step 6: Add `Agent::refresh_context_window_cache` method** + +In `src/agent.rs`, add after `set_model` (after ~line 395): + +```rust +impl Agent { + /// Fetch the context window size for the current model from the + /// provider API and cache it. Non-fatal — uses static fallback on + /// failure. + pub async fn refresh_context_window_cache(&self) { + let model = self.current_model.read().await.clone(); + let (provider, actual_model) = self.registry.resolve_model(&model); + let client = reqwest::Client::new(); + if let Some(ctx) = provider.fetch_context_window(&client, actual_model).await { + let mut cache = provider.config().context_window_cache.write().await; + *cache = Some(ctx); + tracing::info!("Context window for {}: {} tokens", actual_model, ctx); + } + } +} +``` + +- [ ] **Step 7: Update agent loop context_window resolution** (agent.rs:638-643) + +Replace: +```rust +let context_window = { + let model = self.current_model.read().await; + let (provider, _) = self.registry.resolve_model(&model); + provider.config().context_window +}; +``` + +With: +```rust +let context_window = { + let model = self.current_model.read().await; + self.registry.effective_context_window(&model) +}; +``` + +- [ ] **Step 8: Run cargo check** + +Run: `cargo check 2>&1 | head -30` +Expected: compiles successfully + +- [ ] **Step 9: Commit** + +```bash +git add src/provider.rs src/agent.rs +git commit -m "feat(provider): runtime context_window_cache with fetch_context_window" +``` + +--- + +## Task 3: Memory — vec0 metadata columns migration + +**Files:** +- Modify: `src/memory/mod.rs` + +- [ ] **Step 1: Add schema detection + migration after dimension-check block** + +After the `}` closing the `need_migrate` else branch (after line 374), add: + +```rust +// Migration: add metadata columns (is_summarized, role) to vec0 tables +// for pre-filtering. ALTER TABLE is not supported for vec0, so we +// must DROP and recreate. +let has_meta = conn + .prepare("PRAGMA table_info(message_embeddings)") + .and_then(|mut stmt| { + let cols: Vec = stmt + .query_map([], |row| row.get(1))? + .collect::, _>>()?; + Ok(cols.contains(&"is_summarized".to_string())) + }) + .unwrap_or(false); + +if table_exists(conn, "message_embeddings") && !has_meta { + conn.execute_batch("DROP TABLE message_embeddings;")?; + conn.execute_batch(&format!( + "CREATE VIRTUAL TABLE message_embeddings USING vec0(\ + embedding float[{}], is_summarized integer, role text);", + dims + ))?; + info!("Migrated message_embeddings with metadata columns (is_summarized, role)"); +} +``` + +- [ ] **Step 2: Run cargo check** + +Run: `cargo check 2>&1 | head -20` +Expected: compiles successfully + +- [ ] **Step 3: Commit** + +```bash +git add src/memory/mod.rs +git commit -m "feat(memory): add vec0 metadata column migration for is_summarized and role" +``` + +--- + +## Task 4: Memory — Configurable RRF weights in hybrid search + +**Files:** +- Modify: `src/memory/conversations.rs` +- Modify: `src/config.rs` (already done in Task 1) + +- [ ] **Step 1: Update `search_messages_in_conversation` to accept RRF params** + +Change the signature to accept optional RRF overrides: + +```rust +pub async fn search_messages_in_conversation( + &self, + query: &str, + conversation_id: &str, + limit: usize, +) -> Result> { +``` + +No signature change needed — read RRF params from `self.config` (but MemoryStore doesn't hold config directly... let me check). + +Actually, looking at the MemoryConfig, it's part of Config which the Agent holds. MemoryStore doesn't hold a Config reference. The simplest approach: store the MemoryConfig in MemoryStore. + +Add to MemoryStore: +```rust +pub struct MemoryStore { + conn: Arc>, + pub embeddings: Arc, + config: MemoryConfig, // NEW +} +``` + +Update `open()` to accept config, update `open_in_memory()` to use defaults. + +Actually, that's a bigger refactor. Let me take a simpler approach: make `search_messages_in_conversation` and `search_messages` accept optional RRF parameters. + +```rust +pub async fn search_messages_in_conversation( + &self, + query: &str, + conversation_id: &str, + limit: usize, +) -> Result> { +``` + +No, let me check the actual approach. The RRF parameters are hardcoded in the SQL as `60` (the k constant) and `0.5` (the weight constants). A simpler approach: add parameters to the function and update the SQL to use parameters. + +Let me keep it simple — add default parameters: + +```rust +pub async fn search_messages_in_conversation( + &self, + query: &str, + conversation_id: &str, + limit: usize, +) -> Result> { +``` + +The SQL has: +```sql +coalesce(1.0 / (60 + fts.rank_number), 0.0) * 0.5 ++ coalesce(1.0 / (60 + vec.rank_number), 0.0) * 0.5 as combined_rank +``` + +The `60` is RRF k and the `0.5` values are FTS weight and vec weight. These need to be parameterized in the SQL. + +Wait, rusqlite doesn't support parameterizing within SQL expressions like `1.0 / (?1 + rank)`. Actually it does — you can use named or positional parameters in SQL expressions. The issue is that we already use `?1`, `?2`, `?3`, `?4` for other parameters. + +Let me look at the current params: +- `?1` = query_bytes (vec embedding blob) +- `?2` = search_limit +- `?3` = query text (FTS) +- `?4` = conversation_id + +I need to add `?5`, `?6`, `?7` for rrf_k, rrf_weight_fts, rrf_weight_vec. + +This is straightforward. Let me write the implementation. + +Actually, I realize the simplest approach that avoids touching MemoryStore's constructor: keep the SQL queries with the MemoryConfig defaults baked in as parameters. We'll add the parameters to the function signatures. + +Let me do: + +```rust +pub async fn search_messages_in_conversation( + &self, + query: &str, + conversation_id: &str, + limit: usize, +) -> Result> { + // Use default RRF parameters from config if available + // or hardcoded defaults (backward compatible) + let rrf_k = 60.0; + let rrf_weight_fts = 0.5; + let rrf_weight_vec = 0.5; + ... +``` + +Wait, but the spec says "configurable". The cleanest path without refactoring MemoryStore's constructor: pass the parameters through the function. But callers would need access to MemoryConfig... + +Let me check what callers exist: +1. `rag.rs:30`: `store.search_messages_in_conversation(&search_query, conversation_id, limit).await` +2. `conversations.rs:189`: definition +3. `conversations.rs:446`: test + +And `search_messages` callers (the global version): +1. No callers currently (it has `#[allow(dead_code)]`) + +Wait, I see `rag.rs` calls `search_messages_in_conversation`. And `auto_retrieve_context` in `rag.rs` gets `store` as a `&MemoryStore`. + +I think the best approach is to simply add MemoryConfig to MemoryStore. Let me do that. + +Add to `MemoryStore`: +```rust +pub struct MemoryStore { + conn: Arc>, + pub embeddings: Arc, + config: MemoryConfig, +} +``` + +Change `MemoryStore::open`: +```rust +pub fn open(path: &Path, embedding_config: Option, memory_config: MemoryConfig) -> Result { +``` + +But that requires updating all callers. Let me check where `open` is called: +1. `main.rs:118`: `MemoryStore::open(&config.memory.database_path, embedding_config)` +2. Tests: `open_in_memory()` + +And `open_in_memory`: +```rust +pub fn open_in_memory() -> Result { +``` + +I'll add a default parameter. But Rust doesn't have default function parameters... + +OK, let me take the cleaner approach: just pass RRF params through the search function signatures. The callers in `rag.rs` already have access to `store` — they can be updated to read from store. And `retrieve_context_for_compaction` is new anyway. + +Actually, let me add `rrf_k`, `rrf_weight_fts`, `rrf_weight_vec` as parameters with `impl Into>` or just take f64 and have the callers pass them. The callers in rag.rs can get them from... nothing yet. + +Let me think about this differently. The simplest approach: + +1. Add `MemoryConfig` to `MemoryStore` struct +2. Update `open()` to accept config +3. `open_in_memory()` uses default config +4. Search functions read from `self.config` +5. Update `main.rs` to pass `config.memory` to `MemoryStore::open` + +This is the correct approach and clean. Let me do it. + +```rust +pub struct MemoryStore { + conn: Arc>, + pub embeddings: Arc, + pub config: MemoryConfig, +} +``` + +- [ ] **Step 1: Add MemoryConfig field to MemoryStore** + +```rust +pub struct MemoryStore { + conn: Arc>, + pub embeddings: Arc, + pub config: MemoryConfig, +} +``` + +Update `open()`: +```rust +pub fn open(path: &Path, embedding_config: Option, memory_config: MemoryConfig) -> Result { + // ... existing code ... + let store = Self { + conn: Arc::new(Mutex::new(conn)), + embeddings: Arc::new(embeddings), + config: memory_config, + }; + // ... +} +``` + +Update `open_in_memory()`: +```rust +pub fn open_in_memory() -> Result { + // ... existing code ... + let store = Self { + conn: Arc::new(Mutex::new(conn)), + embeddings: Arc::new(embeddings), + config: MemoryConfig::default(), + }; + // ... +} +``` + +MemoryConfig needs a Default impl: + +```rust +impl Default for MemoryConfig { + fn default() -> Self { + Self { + database_path: PathBuf::new(), + rag_limit: 5, + max_raw_messages: 50, + summarize_threshold: 20, + summarize_cron: "0 3 * * * *".to_string(), + query_rewriter_enabled: false, + rrf_k: 60.0, + rrf_weight_fts: 0.5, + rrf_weight_vec: 0.5, + } + } +} + +// Note: The serde defaults (from `default_rag_limit`, `default_max_raw_messages`, etc.) +// are defined via `#[serde(default = "...")]` on each field and remain the source +// of truth for deserialization. The `Default` impl above is for `open_in_memory()` only. +``` + +- [ ] **Step 2: Update main.rs to pass memory config** + +In `src/main.rs`, change: +```rust +let memory = MemoryStore::open(&config.memory.database_path, embedding_config) + .context("Failed to initialize memory store")?; +``` + +To: +```rust +let memory = MemoryStore::open( + &config.memory.database_path, + embedding_config, + config.memory.clone(), +) + .context("Failed to initialize memory store")?; +``` + +- [ ] **Step 3: Update search_messages_in_conversation to use configurable RRF** + +In `src/memory/conversations.rs`, replace the hardcoded RRF values in the SQL with parameterized ones: + +Hybrid branch (line 202-230): +```sql +WITH vec_matches AS ( + SELECT rowid, distance, + row_number() OVER (ORDER BY distance) as rank_number + FROM message_embeddings + WHERE embedding MATCH ?1 + AND k = ?2 + AND is_summarized = 0 + AND role IN ('user', 'assistant') + ORDER BY distance + LIMIT ?2 +), +fts_matches AS ( + SELECT rowid, + row_number() OVER (ORDER BY rank) as rank_number + FROM messages_fts + WHERE messages_fts MATCH ?3 + LIMIT ?2 +) +SELECT m.role, m.content, m.tool_calls, m.tool_call_id, + coalesce(1.0 / (?5 + fts.rank_number), 0.0) * ?7 + + coalesce(1.0 / (?5 + vec.rank_number), 0.0) * ?6 as combined_rank +FROM messages m +LEFT JOIN vec_matches vec ON m.rowid = vec.rowid +LEFT JOIN fts_matches fts ON m.rowid = fts.rowid +WHERE (vec.rowid IS NOT NULL OR fts.rowid IS NOT NULL) + AND m.conversation_id = ?4 + AND m.role IN ('user', 'assistant') + AND (m.is_summarized IS NULL OR m.is_summarized = 0) +ORDER BY combined_rank DESC +LIMIT ?2 +``` + +And update the params: +```rust +let rrf_k = self.config.rrf_k; +let rrf_weight_fts = self.config.rrf_weight_fts; +let rrf_weight_vec = self.config.rrf_weight_vec; +let mut stmt = conn.prepare(sql)?; +let messages = stmt + .query_map( + rusqlite::params![query_bytes, search_limit, query, conversation_id, rrf_k, rrf_weight_vec, rrf_weight_fts], + parse_message_row, + )? + .collect::, _>>() + .context("Failed to hybrid-search messages in conversation")?; +``` + +Note: `?6 = rrf_weight_vec`, `?7 = rrf_weight_fts` — order matters since FTS weight applies to the `fts.rank_number` term. + +Wait, let me double check the order. `coalesce(1.0 / (?5 + fts.rank_number), 0.0) * ?7 + coalesce(1.0 / (?5 + vec.rank_number), 0.0) * ?6`: +- `?5` = rrf_k (for both FTS and vec) +- `?6` = rrf_weight_vec (applied to vec term) +- `?7` = rrf_weight_fts (applied to FTS term) + +That's clean. Actually wait, I realize there's an issue. The `k = ?2` in the vec_matches WHERE clause — that's the KNN k parameter (number of nearest neighbors), not RRF k. These are two different concepts! KNN k = how many nearest neighbors to retrieve, RRF k = the constant in the fusion formula. + +Let me rename to avoid confusion. The spec calls them `rrf_k` for the RRF constant. The KNN `LIMIT ?2` is `search_limit` which is `limit * 3`. These are different. The original code has `LIMIT ?2` in vec_matches which limits how many vectors to retrieve from the index (KNN search limit), and the outer `LIMIT ?2` limits final results after RRF fusion. + +So I need `?5` for rrf_k, not to be confused with the existing `?2` which is `search_limit`. + +OK, the param mapping is: +- `?1` = query_bytes +- `?2` = search_limit (KNN k, not RRF k!) +- `?3` = query text +- `?4` = conversation_id +- `?5` = rrf_k +- `?6` = rrf_weight_vec +- `?7` = rrf_weight_fts + +This is correct. Let me also update vec_matches to use metadata columns for pre-filtering (the spec says to do this): + +```sql +WHERE embedding MATCH ?1 + AND k = ?2 + AND is_summarized = 0 + AND role IN ('user', 'assistant') +``` + +This uses sqlite-vec's metadata column filtering. The `k` here is the KNN k, i.e., how many neighbors to return. + +Now for the global `search_messages` — same changes but without conversation_id: + +```sql +WITH vec_matches AS ( + SELECT rowid, distance, + row_number() OVER (ORDER BY distance) as rank_number + FROM message_embeddings + WHERE embedding MATCH ?1 + AND k = ?2 + AND is_summarized = 0 + AND role IN ('user', 'assistant') + ORDER BY distance + LIMIT ?2 +), +fts_matches AS ( + SELECT rowid, + row_number() OVER (ORDER BY rank) as rank_number + FROM messages_fts + WHERE messages_fts MATCH ?3 + LIMIT ?2 +) +SELECT m.role, m.content, m.tool_calls, m.tool_call_id, + coalesce(1.0 / (?4 + fts.rank_number), 0.0) * ?6 + + coalesce(1.0 / (?4 + vec.rank_number), 0.0) * ?5 as combined_rank +FROM messages m +LEFT JOIN vec_matches vec ON m.rowid = vec.rowid +LEFT JOIN fts_matches fts ON m.rowid = fts.rowid +WHERE vec.rowid IS NOT NULL OR fts.rowid IS NOT NULL +ORDER BY combined_rank DESC +LIMIT ?2 +``` + +With params: `rusqlite::params![query_bytes, search_limit, query, rrf_k, rrf_weight_vec, rrf_weight_fts]` + +- [ ] **Step 4: Run cargo check** + +Run: `cargo check 2>&1 | head -30` +Expected: compiles successfully + +Advisory: Remove the stale `#[allow(dead_code)]` annotation on `search_messages_in_conversation` (line 188 in conversations.rs) — it's already called by `auto_retrieve_context` and will have more callers after this task. + +- [ ] **Step 5: Commit** + +```bash +git add src/memory/conversations.rs src/memory/mod.rs src/main.rs +git commit -m "feat(memory): configurable RRF weights and vec0 metadata pre-filtering" +``` + +--- + +## Task 5: rag.rs — Add retrieve_context_for_compaction + +**Files:** +- Modify: `src/memory/rag.rs` + +- [ ] **Step 1: Write the failing test first** + +Add to the test module in `rag.rs`: + +```rust +#[tokio::test] +async fn test_retrieve_context_for_compaction_returns_none_for_short_query() { + let store = crate::memory::MemoryStore::open_in_memory().unwrap(); + let conv = store + .get_or_create_conversation("test", "compact_u1") + .await + .unwrap(); + let to_summarize = vec![crate::llm::ChatMessage { + role: "assistant".to_string(), + content: Some(crate::llm::MessageContent::from_text("some result")), + tool_calls: None, + tool_call_id: None, + }]; + let preserved = vec![]; + let result = retrieve_context_for_compaction(&store, &to_summarize, &preserved, &conv, 5) + .await + .unwrap(); + assert!(result.is_none(), "No user message should return None"); +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cargo test test_retrieve_context_for_compaction_returns_none_for_short_query -- --nocapture 2>&1 | tail -10` +Expected: FAIL with "function not found in module" + +- [ ] **Step 3: Add retrieve_context_for_compaction function** + +After `auto_retrieve_context` (after line 58), add: + +```rust +/// Retrieve context for compaction summarization. +/// +/// Uses the most recent user message (from both to_summarize and preserved +/// ranges) as a search query to find relevant historical context from the +/// conversation. Returns formatted snippets that help the summarizer write +/// a focused summary. Returns None when no suitable query is found or no +/// results are returned. +pub async fn retrieve_context_for_compaction( + store: &MemoryStore, + to_summarize: &[ChatMessage], + preserved: &[ChatMessage], + conversation_id: &str, + limit: usize, +) -> Result> { + // Find the most recent user message across both ranges. + // The most recent user message may be in preserved if compaction fires + // right after the user spoke before any tool calls happened. + let query = preserved + .iter() + .rev() + .chain(to_summarize.iter().rev()) + .find(|m| m.role == "user") + .map(|m| m.content.as_ref().map(|c| c.as_text()).unwrap_or_default()) + .unwrap_or_default(); + + if query.trim().len() < 5 { + return Ok(None); + } + + let results = store + .search_messages_in_conversation(&query, conversation_id, limit) + .await?; + + if results.is_empty() { + return Ok(None); + } + + let mut block = String::from( + "\n\ + Relevant context from conversation history for compaction:\n\n", + ); + + for msg in &results { + if let Some(content) = &msg.content { + let text = content.as_text(); + let snippet = crate::utils::strings::truncate_chars(&text, 300); + block.push_str(&format!("[{}] {}\n", msg.role, snippet)); + } + } + + block.push_str(""); + debug!( + "Compaction RAG: injected {} snippets for query: {:?}", + results.len(), + query + ); + Ok(Some(block)) +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cargo test test_retrieve_context_for_compaction_returns_none_for_short_query -- --nocapture 2>&1 | tail -10` +Expected: PASS (returns None because no user message in to_summarize or preserved) + +- [ ] **Step 5: Add second test for successful retrieval** + +```rust +#[tokio::test] +async fn test_retrieve_context_for_compaction_finds_user_message_in_preserved() { + let store = crate::memory::MemoryStore::open_in_memory().unwrap(); + let conv = store + .get_or_create_conversation("test", "compact_u2") + .await + .unwrap(); + + // Save a user message to the conversation + let msg = crate::llm::ChatMessage { + role: "user".to_string(), + content: Some(crate::llm::MessageContent::from_text("Tell me about Rust async")), + tool_calls: None, + tool_call_id: None, + }; + store.save_message(&conv, &msg).await.unwrap(); + + let to_summarize = vec![crate::llm::ChatMessage { + role: "assistant".to_string(), + content: Some(crate::llm::MessageContent::from_text("Here's how...")), + tool_calls: None, + tool_call_id: None, + }]; + let preserved = vec![crate::llm::ChatMessage { + role: "user".to_string(), + content: Some(crate::llm::MessageContent::from_text("Tell me about Rust async")), + tool_calls: None, + tool_call_id: None, + }]; + + let result = retrieve_context_for_compaction(&store, &to_summarize, &preserved, &conv, 5) + .await + .unwrap(); + // May return None if embeddings not available (FTS-only fallback), + // which is acceptable. The key is no panic. + let _ = result; +} +``` + +- [ ] **Step 6: Run tests** + +Run: `cargo test --test '*' -- --nocapture 2>&1 | tail -20` +Expected: both new tests pass + +- [ ] **Step 7: Commit** + +```bash +git add src/memory/rag.rs +git commit -m "feat(rag): add retrieve_context_for_compaction for RAG-aware compaction" +``` + +--- + +## Task 6: agent_prompt.rs — Enhanced structured summary prompt + +**Files:** +- Modify: `src/agent_prompt.rs` + +- [ ] **Step 1: Update `build_compact_summary_prompt`** + +Replace the function body (lines 440-465): + +```rust +pub fn build_compact_summary_prompt() -> ChatMessage { + let prompt_text = vec![ + "You are producing a compact state summary of the conversation below.", + "", + "OUTPUT FORMAT:", + "", + "## STATE", + "```yaml", + "stage: ", + "decisions:", + " - ", + "pending:", + " - ", + "last_action: ", + "last_action_result: ", + "conversation_phase: ", + "```", + "", + "## CONTEXT", + "- ", + "- ", + "- ", + "", + "CRITICAL RULES:", + "- State must be precise enough for the LLM to continue without re-reading history", + "- Include ALL pending items the user explicitly requested", + "- Include ALL error messages and their resolutions", + "- Be specific with file paths and tool names", + "- Do NOT call any tools. Respond with text only.", + ] + .join("\n"); + + ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(prompt_text)), + tool_calls: None, + tool_call_id: None, + } +} +``` + +- [ ] **Step 2: Update `build_compact_boundary_marker`** + +Change the format string (line 472): + +```rust +"★ COMPACT SUMMARY — previous {} messages → YAML state + narrative ({} messages) ★" +``` + +- [ ] **Step 3: Update `recovery_nudge_for`** + +Change the content strings (lines 109, 111): + +```rust +let content = if previous_is_tool { + "Continue from the tool result above. Read the ## STATE block in the compact summary for context. Either call the next required tool or provide a final answer.".to_string() +} else { + "Continue from the user's request above. Read the ## STATE block in the compact summary for context. Either call the next required tool or provide a final answer.".to_string() +}; +``` + +- [ ] **Step 4: Run cargo test to verify prompt tests still pass** + +Run: `cargo test --lib agent_prompt 2>&1 | tail -20` +Expected: all existing tests pass + +- [ ] **Step 5: Commit** + +```bash +git add src/agent_prompt.rs +git commit -m "feat(prompt): structured YAML state + narrative summary prompt" +``` + +--- + +## Task 7: agent.rs — System prompt preservation + RAG-aware compaction + +**Files:** +- Modify: `src/agent.rs` + +- [ ] **Step 1: Update `auto_compact_conversation` to partition system messages and pass memory** + +Replace the function (lines 1330-1362): + +```rust +/// Tier 3: Auto-compact via LLM summarization. +async fn auto_compact_conversation( + llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, + messages: &[ChatMessage], + _context_window: usize, +) -> Result> { + // 1. Separate system messages from the rest + let mut system_msgs: Vec = Vec::new(); + let non_system: Vec = messages + .iter() + .filter(|msg| { + if msg.role == "system" { + system_msgs.push(msg.clone()); + false + } else { + true + } + }) + .collect(); + + let tool_groups = crate::agent_prompt::find_tool_groups(&non_system); + + let preserve_count = PRESERVED_TOOL_GROUPS.min(tool_groups.len()); + let preserved_groups_start = tool_groups.len().saturating_sub(preserve_count); + + let summary_end = if preserved_groups_start > 0 { + let last_summary = &tool_groups[preserved_groups_start - 1]; + *last_summary + .tool_result_indices + .last() + .unwrap_or(&last_summary.assistant_idx) + + 1 + } else { + return Ok(messages.to_vec()); + }; + + let to_summarize = &non_system[..summary_end]; + let preserved = &non_system[summary_end..]; + + // 2. Summarize with RAG-aware compaction + let mut compacted = Self::summarize_and_replace( + llm, + memory, + conversation_id, + to_summarize, + preserved, + "Auto-compact", + "★ COMPACT SUMMARY ★", + ) + .await?; + + // 3. Prepend system messages back + let mut result = system_msgs; + result.append(&mut compacted); + Ok(result) +} +``` + +- [ ] **Step 2: Update `reactive_compact` to partition system messages and pass memory** + +Replace the function (lines 1369-1391): + +```rust +/// Tier 4: Reactive compact — emergency 413 recovery. +async fn reactive_compact( + llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, + messages: &[ChatMessage], + _context_window: usize, +) -> Result> { + const PRESERVE_COUNT: usize = 4; + + // 1. Separate system messages + let mut system_msgs: Vec = Vec::new(); + let non_system: Vec = messages + .iter() + .filter(|msg| { + if msg.role == "system" { + system_msgs.push(msg.clone()); + false + } else { + true + } + }) + .collect(); + + if non_system.len() <= PRESERVE_COUNT { + anyhow::bail!("Too few non-system messages for reactive compact"); + } + + let split = non_system.len().saturating_sub(PRESERVE_COUNT); + let to_summarize = &non_system[..split]; + let preserved = &non_system[split..]; + + let mut compacted = Self::summarize_and_replace( + llm, + memory, + conversation_id, + to_summarize, + preserved, + "Reactive compact", + "★ COMPACT SUMMARY (EMERGENCY) ★", + ) + .await?; + + let mut result = system_msgs; + result.append(&mut compacted); + Ok(result) +} +``` + +- [ ] **Step 3: Replace `summarize_and_replace` with RAG-aware + tool-group-aware truncation** + +Replace the function (lines 1395-1451): + +```rust +/// Shared helper for Tiers 3 and 4: send messages to LLM for +/// summarization, then assemble the compacted result. +async fn summarize_and_replace( + llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, + to_summarize: &[ChatMessage], + preserved: &[ChatMessage], + error_label: &str, + summary_label: &str, +) -> Result> { + // NEW: RAG retrieval for compaction (non-fatal — warn on error, continue) + let retrieved = match crate::memory::rag::retrieve_context_for_compaction( + memory, + to_summarize, + preserved, + conversation_id, + 5, + ) + .await + { + Ok(r) => r, + Err(e) => { + warn!("RAG retrieval for compaction failed: {}", e); + None + } + }; + + // Build compact messages: summary prompt + optional retrieved context + truncated input + let mut compact_msgs = Vec::new(); + compact_msgs.push(build_compact_summary_prompt()); + + if let Some(ref ctx) = retrieved { + compact_msgs.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(ctx.clone())), + tool_calls: None, + tool_call_id: None, + }); + } + + // Tool-group-aware truncation: keep the first N tool groups (origin) + // and the last M tool groups (recent), preserving tool-call→result pairs. + let groups = crate::agent_prompt::find_tool_groups(to_summarize); + let bookend_groups = 1usize; + let tail_groups = 3usize.min(groups.len().saturating_sub(bookend_groups)); + + let mut seen_indices = std::collections::HashSet::new(); + + // First bookend groups (conversation origin) + for group in groups.iter().take(bookend_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Last tail groups (recent flow) + for group in groups.iter().rev().take(tail_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Always include any non-assistant/non-tool messages (conversation opening, user messages, etc.) + for (idx, msg) in to_summarize.iter().enumerate() { + if msg.role != "assistant" && !msg.has_tool_calls() { + seen_indices.insert(idx); + } + } + + // Build sampled list in original order, inserting one truncation notice + let mut sampled: Vec = Vec::new(); + let mut inserted_notice = false; + for (idx, msg) in to_summarize.iter().enumerate() { + if seen_indices.contains(&idx) { + sampled.push(msg.clone()); + } else if !inserted_notice { + sampled.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(format!( + "[... {} messages omitted, see retrieved_context above ...]", + to_summarize.len() - seen_indices.len() + ))), + tool_calls: None, + tool_call_id: None, + }); + inserted_notice = true; + } + } + + // If no truncation happened, use the full to_summarize + if sampled.is_empty() { + sampled = to_summarize.to_vec(); + } + + let user_prompt = format!( + "Summarize the following conversation (sampled from {} messages):", + to_summarize.len() + ); + compact_msgs.push(ChatMessage { + role: "user".to_string(), + content: Some(MessageContent::Text(user_prompt)), + tool_calls: None, + tool_call_id: None, + }); + compact_msgs.extend(sampled); + + let summary_response = match llm.chat(&compact_msgs, &[]).await { + Ok(c) => c, + Err(e) => anyhow::bail!("{} LLM call failed: {}", error_label, e), + }; + + let summary_text = summary_response + .content + .as_ref() + .map(|c| c.as_text()) + .unwrap_or_default(); + + if summary_text.is_empty() { + anyhow::bail!("{} returned empty summary", error_label); + } + + let boundary = build_compact_boundary_marker(to_summarize.len(), 1); + let summary_msg = ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(format!( + "{}\n\n{}", + summary_label, summary_text + ))), + tool_calls: None, + tool_call_id: None, + }; + + let mut result: Vec = Vec::with_capacity(3 + preserved.len()); + result.push(boundary); + result.push(summary_msg); + result.extend(preserved.iter().cloned()); + + let nudge = recovery_nudge_for(&result); + result.push(nudge); + + Ok(result) +} +``` + +- [ ] **Step 4: Update callers of auto_compact_conversation and reactive_compact** + +In the agent loop at line 664, change: +```rust +Self::auto_compact_conversation(&self.llm, &messages, context_window).await +``` + +To: +```rust +Self::auto_compact_conversation( + &self.llm, + &self.memory, + &conversation_id, + &messages, + context_window, +) +.await +``` + +At line 817, change: +```rust +Self::reactive_compact(&self.llm, &messages, context_window).await +``` + +To: +```rust +Self::reactive_compact( + &self.llm, + &self.memory, + &conversation_id, + &messages, + context_window, +) +.await +``` + +- [ ] **Step 5: Run cargo check** + +Run: `cargo check 2>&1 | head -30` +Expected: compiles successfully + +- [ ] **Step 6: Commit** + +```bash +git add src/agent.rs +git commit -m "feat(agent): system prompt preservation and RAG-aware compaction" +``` + +--- + +## Task 8: main.rs — Background context_window cache warmup + +- [ ] **Step 1: Add startup background task** + +After `build_registry` (after line 76), before the info logging: + +```rust +// Spawn background task to warm context_window_cache for all providers +{ + let registry_clone = Arc::clone(®istry); + tokio::spawn(async move { + let client = reqwest::Client::new(); + for name in registry_clone.provider_names() { + if let Some(provider) = registry_clone.get_provider(&name) { + let model = provider.default_model(); + if let Some(ctx) = provider.fetch_context_window(&client, model).await { + let mut cache = provider.config().context_window_cache.write().await; + *cache = Some(ctx); + tracing::info!( + "Context window cache: {} / {} = {} tokens", + name, + model, + ctx + ); + } + } + } + }); +} +``` + +- [ ] **Step 2: Run cargo check** + +Run: `cargo check 2>&1 | head -10` +Expected: compiles successfully + +- [ ] **Step 3: Commit** + +```bash +git add src/main.rs +git commit -m "feat(main): background context_window cache warmup at startup" +``` + +--- + +## Task 9: Unit tests + +- [ ] **Step 1: Add provider tests for effective_context_window** + +In `src/provider.rs`, add to the test module: + +```rust +#[test] +fn effective_context_window_falls_back_to_static_when_cache_empty() { + let sections = vec![make_section( + "alpha", + ProviderType::OpenRouter, + "https://openrouter.ai/api/v1", + "anthropic/claude-sonnet-4-6", + )]; + let reg = build_registry(§ions, "alpha", 3).unwrap(); + let ctx = reg.effective_context_window("anthropic/claude-sonnet-4-6"); + assert_eq!(ctx, 512_000); // static fallback from section +} + +#[test] +fn effective_context_window_returns_cached_value_when_set() { + use tokio::sync::RwLock; + let sections = vec![make_section( + "alpha", + ProviderType::OpenRouter, + "https://openrouter.ai/api/v1", + "anthropic/claude-sonnet-4-6", + )]; + let reg = build_registry(§ions, "alpha", 3).unwrap(); + let provider = reg.get_provider("alpha").unwrap(); + // Set cache manually + *provider.config().context_window_cache.try_write().unwrap() = Some(200_000); + let ctx = reg.effective_context_window("anthropic/claude-sonnet-4-6"); + assert_eq!(ctx, 200_000); +} +``` + +- [ ] **Step 2: Run provider tests** + +Run: `cargo test --lib provider 2>&1 | tail -15` +Expected: all tests pass including new ones + +- [ ] **Step 3: Add agent_prompt test for new summary prompt keywords** + +In `src/agent_prompt.rs` test module: + +```rust +#[test] +fn compact_summary_prompt_contains_state_keywords() { + let msg = build_compact_summary_prompt(); + let text = msg.content.as_ref().unwrap().as_text(); + assert!(text.contains("## STATE"), "Should have STATE section"); + assert!(text.contains("## CONTEXT"), "Should have CONTEXT section"); + assert!(text.contains("decisions:"), "Should have decisions field"); + assert!(text.contains("pending:"), "Should have pending field"); + assert!(text.contains("last_action:"), "Should have last_action field"); +} + +#[test] +fn compact_boundary_marker_contains_state_hint() { + let msg = build_compact_boundary_marker(10, 1); + let text = msg.content.as_ref().unwrap().as_text(); + assert!(text.contains("STATE"), "Should hint at state format"); +} +``` + +- [ ] **Step 4: Run agent_prompt tests** + +Run: `cargo test --lib agent_prompt 2>&1 | tail -15` +Expected: all tests pass including new ones + +- [ ] **Step 5: Run full test suite** + +Run: `cargo test 2>&1 | tail -20` +Expected: all tests pass + +- [ ] **Step 6: Run clippy** + +Run: `cargo clippy -- -D warnings 2>&1 | tail -20` +Expected: no warnings + +- [ ] **Step 7: Commit tests** + +```bash +git add src/provider.rs src/agent_prompt.rs +git commit -m "test: add tests for context_window cache and structured summary prompt" +``` + +--- + +## Self-Review + +**1. Spec coverage:** +- Section 1 (Dynamic context window): Task 2 (fetch_context_window, cache, effective_context_window) + Task 8 (startup warmup) ✓ +- Section 2 (System prompt preservation): Task 7 (partition system msgs in auto_compact + reactive_compact) ✓ +- Section 3 (Structured summary prompt): Task 6 (build_compact_summary_prompt, boundary_marker, recovery_nudge_for) ✓ +- Section 4 (RAG-aware compaction): Task 5 (retrieve_context_for_compaction) + Task 7 (summarize_and_replace with RAG + truncation) ✓ +- Section 5 (sqlite-vec + FTS5 optimizations): Task 1 (configurable RRF) + Task 3 (vec0 metadata columns) + Task 4 (parameterized RRF in SQL) ✓ +- Testing strategy: Task 9 ✓ + +**2. Placeholder scan:** No TBD, TODO, or incomplete code patterns found. All steps have complete code blocks. + +**3. Type consistency:** `fetch_context_window` returns `Option` in trait definition and all implementations ✓. `summarize_and_replace` signature consistently has `memory: &MemoryStore, conversation_id: &str` before `to_summarize` ✓. `effective_context_window` takes `&str` and returns `usize` ✓. diff --git a/docs/superpowers/specs/2026-07-06-compaction-context-management-design.md b/docs/superpowers/specs/2026-07-06-compaction-context-management-design.md new file mode 100644 index 0000000..d1c2553 --- /dev/null +++ b/docs/superpowers/specs/2026-07-06-compaction-context-management-design.md @@ -0,0 +1,543 @@ +# Context-Aware Compaction & Memory Management for RustFox + +## Overview + +RustFox's compaction system (Tiers 1-4) currently drops the system prompt, sends all old messages verbatim to the summarizer LLM, produces free-form summaries, and relies on a static `context_window` from config. This spec addresses all four gaps with targeted changes (Approach C). + +## Files Touched + +| File | Changes | +|------|---------| +| `src/provider.rs` | Add `context_window_cache`, `fetch_context_window()` trait method, `effective_context_window()` | +| `src/agent.rs` | System prompt preservation in auto_compact/reactive_compact, RAG-aware compact flow | +| `src/agent_prompt.rs` | Enhanced structured summary prompt | +| `src/memory/rag.rs` | New `retrieve_context_for_compaction()` function | +| `src/memory/conversations.rs` | Configurable RRF weights | +| `src/memory/mod.rs` | Optional metadata columns on vec0 (migration: DROP + recreate with `is_summarized`, `role`) | +| `src/config.rs` | Add `rrf_k`, `rrf_weight_fts`, `rrf_weight_vec` to memory config | +| `src/main.rs` | Background startup task to warm `context_window_cache` | + +--- + +## Section 1: Dynamic Context Window Detection + +### Problem + +`context_window` is a static field in `config.toml`. When the user switches models via `/model`, the compaction thresholds (20%/60%/85%) don't adjust — they use the original model's window. The current model is already tracked in `current_model: RwLock`. + +### Solution + +**Runtime context window cache** that overrides the static config value when available, falls back to static. + +**`ProviderConfig` changes** (provider.rs:17-30): +```rust +pub struct ProviderConfig { + // … existing fields … + pub context_window: usize, // static fallback + pub context_window_cache: Arc>>, // runtime override (NEW) +} +``` + +**New trait method** on `Provider`: +```rust +async fn fetch_context_window( + &self, + client: &Client, + model: &str, +) -> Option; +``` + +**OpenRouterProvider implementation** (provider.rs:358-383): The existing `list_models()` already calls `GET /api/v1/models` which returns `context_length` per model. Add a new method that parses `context_length`: + +``` +GET /api/v1/models → JSON array of { id, context_length, ... } +``` + +Find the entry where `id` matches the current model, extract `context_length`. Return `None` on error or no match. + +**OpenAICompatibleProvider**: Tries the same pattern (many OpenAI-compatible providers return model metadata). Returns `None` on failure. + +**OllamaProvider**: Returns `None` — Ollama doesn't report context length. + +**`ProviderRegistry` new method**: +```rust +pub fn effective_context_window(&self, model: &str) -> usize { + let (provider, _) = self.resolve_model(model); + let cached = provider.config().context_window_cache.read(); + cached.unwrap_or(provider.config().context_window) +} +``` + +**Trigger points for cache update:** + +1. **Startup** (`main.rs` after `build_registry`): Spawn a background task that iterates all registered providers, calling `fetch_context_window` for each provider's default model and writing results to `context_window_cache`. No startup delay — the static config value is used for the first LLM call. + + ```rust + // In main.rs, after registry is built: + let registry_clone = Arc::clone(®istry); + tokio::spawn(async move { + let client = reqwest::Client::new(); + for name in registry_clone.provider_names() { + if let Some(provider) = registry_clone.get_provider(&name) { + let model = provider.default_model(); + if let Some(ctx) = provider.fetch_context_window(&client, model).await { + // write to provider's context_window_cache (via ProviderConfig) + let mut cache = provider.config().context_window_cache.write(); + *cache = Some(ctx); + } + } + } + }); + ``` + +2. **Model switch** (`set_model()` in agent.rs:328): `Agent` already holds `self.registry: Arc`. After persisting the new model to config, call `self.refresh_context_window_cache()` which resolves the new model through the registry and calls `fetch_context_window`: + + ```rust + pub async fn refresh_context_window_cache(&self) { + let model = self.current_model.read().await.clone(); + let (provider, actual_model) = self.registry.resolve_model(&model); + let client = reqwest::Client::new(); + if let Some(ctx) = provider.fetch_context_window(&client, actual_model).await { + let mut cache = provider.config().context_window_cache.write(); + *cache = Some(ctx); + } + } + ``` + +3. **`/model` command handler**: Calls `agent.set_model(new_model).await` which internally calls `refresh_context_window_cache` (integrated into `set_model`). + +**Agent loop** (agent.rs:638-643): Replace: +```rust +// BEFORE: +let context_window = provider.config().context_window; + +// AFTER: +let context_window = registry.effective_context_window(&model); +``` + +**Edge cases:** +- API call fails → silently keep static value (log at debug) +- Model not found in API response → keep static value +- Model switches between static and dynamic providers → cache handles each independently +- Race: context_window_cache uses `Arc>>` so writes never block reads + +--- + +## Section 2: System Prompt Preservation + +### Problem + +`auto_compact_conversation` (agent.rs:1330) includes `messages[0]` (the system message) in `to_summarize`. The summarizer LLM receives it and replaces everything with a summary. After compaction, the system prompt is gone until the next `process_message()` refreshes it in-memory. If the LLM tries a tool call after compaction (within the same agentic loop), it runs without system instructions. + +### Solution + +**In `auto_compact_conversation`** (agent.rs:1330-1362): + +```rust +async fn auto_compact_conversation(...) -> Result> { + // 1. Separate system messages from the rest + let mut system_msgs: Vec = Vec::new(); + let mut non_system: Vec = Vec::new(); + for msg in messages { + if msg.role == "system" { + system_msgs.push(msg.clone()); + } else { + non_system.push(msg.clone()); + } + } + + // 2. Compute tool groups and summary split on non-system only + let tool_groups = find_tool_groups(&non_system); + let preserve_count = PRESERVED_TOOL_GROUPS.min(tool_groups.len()); + let preserved_groups_start = tool_groups.len().saturating_sub(preserve_count); + + if preserved_groups_start == 0 { + // Not enough groups to compact — return original (including system) + return Ok(messages.to_vec()); + } + + // summary_end relative to non_system + let last_summary = &tool_groups[preserved_groups_start - 1]; + let summary_end = *last_summary.tool_result_indices.last() + .unwrap_or(&last_summary.assistant_idx) + 1; + + let to_summarize = &non_system[..summary_end]; + let preserved = &non_system[summary_end..]; + + // 3. Summarize (existing logic) + let mut compacted = Self::summarize_and_replace( + llm, to_summarize, preserved, ... + ).await?; + + // 4. Prepend system messages back + let mut result = system_msgs; + result.append(&mut compacted); + Ok(result) +} +``` + +**Same pattern in `reactive_compact`** (agent.rs:1369-1391): Currently uses a simple `PRESERVE_COUNT = 4` from the end. Apply the same partition logic. + +**Why this is safe:** +- System messages are read-only context (never tool calls) +- The LLM never modifies them, they don't participate in tool groups +- Partitioning them out before compaction preserves them exactly +- On the next `process_message()` iteration, the system prompt is refreshed anyway by `build_system_prompt()` + +--- + +## Section 3: Enhanced Structured Summary Prompt + +### Problem + +`build_compact_summary_prompt()` (agent_prompt.rs:440-465) produces a free-form narrative. The subsequent LLM must re-read the entire narrative to extract state — defeating some of the benefit of compaction. There's no machine-readable state block. + +### Solution + +**Replace `build_compact_summary_prompt()`** with a version that requests both a structured YAML state block and brief narrative: + +```rust +pub fn build_compact_summary_prompt() -> ChatMessage { + let prompt_text = vec![ + "You are producing a compact state summary of the conversation below.", + "", + "OUTPUT FORMAT:", + "", + "## STATE", + "```yaml", + "stage: ", + "decisions:", + " - ", + "pending:", + " - ", + "last_action: ", + "last_action_result: ", + "conversation_phase: ", + "```", + "", + "## CONTEXT", + "- ", + "- ", + "- ", + "", + "CRITICAL RULES:", + "- State must be precise enough for the LLM to continue without re-reading history", + "- Include ALL pending items the user explicitly requested", + "- Include ALL error messages and their resolutions", + "- Be specific with file paths and tool names", + "- Do NOT call any tools. Respond with text only.", + ].join("\n"); + + ChatMessage { role: "system", content: Some(MessageContent::Text(prompt_text)), ... } +} +``` + +**Update `build_compact_boundary_marker()`** (agent_prompt.rs:468): +``` +"★ COMPACT SUMMARY — previous {} msgs → YAML state + narrative ({} msgs) ★" +``` + +**Update `recovery_nudge_for()`** (agent_prompt.rs:105): Add a hint to read the STATE block: +``` +"Continue from the [user's request|tool result] above. Read the ## STATE block in the compact summary for context. Either call the next required tool or provide a final answer." +``` + +--- + +## Section 4: RAG-Aware Tier 3/4 Compaction + +### Problem + +`summarize_and_replace` (agent.rs:1395) sends ALL messages in `to_summarize` verbatim to the summarizer LLM. For a conversation with 50+ tool-call groups, this means 15K-30K tokens just to generate a summary — burning tokens and losing precision in the "middle" of the input. + +### Solution + +**Before calling the summarizer LLM, retrieve relevant context from the database** using the existing hybrid search. + +**New function in `src/memory/rag.rs`:** + +```rust +/// Retrieve context for compaction summarization. +/// +/// Uses the most recent user message (from both to_summarize and preserved +/// ranges) as a search query to find relevant historical context from the +/// conversation. Returns formatted snippets that help the summarizer write +/// a focused summary. +pub async fn retrieve_context_for_compaction( + store: &MemoryStore, + to_summarize: &[ChatMessage], + preserved: &[ChatMessage], + conversation_id: &str, + limit: usize, +) -> Result> { + // Find the most recent user message across both ranges. + // The most recent user message may be in `preserved` (the last few tool + // groups) if compaction fires right after the user spoke before any + // tool calls happened. + let query = preserved.iter().rev() + .chain(to_summarize.iter().rev()) + .find(|m| m.role == "user") + .map(|m| m.content.as_ref().map(|c| c.as_text()).unwrap_or_default()) + .unwrap_or_default(); + + if query.trim().len() < 5 { + return Ok(None); + } + + let results = store + .search_messages_in_conversation(&query, conversation_id, limit) + .await?; + + if results.is_empty() { + return Ok(None); + } + + let mut block = String::from("\nRelevant context from conversation history:\n\n"); + for msg in &results { + if let Some(content) = &msg.content { + let text = content.as_text(); + let snippet = truncate_chars(&text, 300); + block.push_str(&format!("[{}] {}\n", msg.role, snippet)); + } + } + block.push_str(""); + Ok(Some(block)) +} +``` + +**Modified `summarize_and_replace`** (agent.rs:1395): + +```rust +async fn summarize_and_replace( + llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, + to_summarize: &[ChatMessage], + preserved: &[ChatMessage], + error_label: &str, + summary_label: &str, +) -> Result> { + + // NEW: RAG retrieval for compaction (non-fatal — warn on error, continue) + let retrieved = match memory::rag::retrieve_context_for_compaction( + memory, to_summarize, preserved, conversation_id, 5 + ).await { + Ok(r) => r, + Err(e) => { + warn!("RAG retrieval for compaction failed: {}", e); + None + } + }; + + // Build compact messages: summary prompt + retrieved context + truncated input + let mut compact_msgs = Vec::new(); + compact_msgs.push(build_compact_summary_prompt()); + + if let Some(ref ctx) = retrieved { + compact_msgs.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(ctx.clone())), + tool_calls: None, + tool_call_id: None, + }); + } + + // Tool-group-aware truncation: keep the first N tool groups (conversation + // origin) and the last M tool groups (recent flow). This preserves the + // structural integrity of tool-call→result pairs — unlike raw-index + // truncation which could split a tool group and leave a dangling call + // with no result. + let groups = find_tool_groups(to_summarize); + let bookend_groups = 1usize; + let tail_groups = 3usize.min(groups.len().saturating_sub(bookend_groups)); + + let mut sampled: Vec = Vec::new(); + let mut seen_indices = std::collections::HashSet::new(); + + // First bookend groups (conversation origin) + for group in groups.iter().take(bookend_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Last tail groups (recent flow) + for group in groups.iter().rev().take(tail_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Always include any non-assistant/non-tool messages at the beginning + // (conversation opening, user's first message, etc.) + for (idx, msg) in to_summarize.iter().enumerate() { + if msg.role != "assistant" && !msg.has_tool_calls() { + seen_indices.insert(idx); + } + } + + // Build sampled list in original order, inserting truncation notice + let mut inserted_notice = false; + for (idx, msg) in to_summarize.iter().enumerate() { + if seen_indices.contains(&idx) { + sampled.push(msg.clone()); + } else if !inserted_notice { + sampled.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text( + format!("[... {} messages omitted, see retrieved_context above ...]", + to_summarize.len() - seen_indices.len()) + )), + tool_calls: None, + tool_call_id: None, + }); + inserted_notice = true; + } + } + + let user_prompt = format!("Summarize the following conversation (sampled from {} messages):", to_summarize.len()); + compact_msgs.push(ChatMessage { + role: "user".to_string(), + content: Some(MessageContent::Text(user_prompt)), + tool_calls: None, + tool_call_id: None, + }); + compact_msgs.extend(sampled); + + // Existing LLM call + let summary_response = match llm.chat(&compact_msgs, &[]).await { ... }; + // ... rest of existing logic ... +} +``` + +**New parameter** to `summarize_and_replace`: `memory: &MemoryStore, conversation_id: &str`. + +**Updated callers:** +- `auto_compact_conversation` — passes `memory + conversation_id` +- `reactive_compact` — passes `memory + conversation_id` + +**Existing callers of `summarize_and_replace`** (none other than the two above). + +**Token impact:** + +| Scenario | Before | After | Savings | +|----------|--------|-------|---------| +| 50 msgs, ~15K tokens | 15K sent to LLM | ~3K (retrieved + bookends) | 80% | +| 100 msgs, ~30K tokens | 30K sent | ~5K | 83% | +| 5 msgs, ~2K tokens | N/A — won't trigger | N/A | N/A | + +**Edge cases:** +- `to_summarize` has no user messages → `retrieve_context_for_compaction` returns `None`, compact uses full `to_summarize` (fallback) +- Embeddings disabled (FTS5-only) → search falls back to FTS5-only, still works +- Search returns duplicates of what's already in bookends → RRF naturally ranks them, deduplication by `rowid` +- `to_summarize` and `preserved` both lack a user message → `retrieve_context_for_compaction` returns `None`, fallback to no RAG context +- Very small `to_summarize` (< 3 tool groups) → no truncation needed, send all + +--- + +## Section 5: sqlite-vec + FTS5 Optimizations + +### Problem + +The current hybrid search (conversations.rs:189-267) works correctly but the RRF parameters are hardcoded and the vec0 table doesn't use metadata columns for pre-filtering. + +### Solution + +**Configurable RRF parameters** in `config.rs`: + +```rust +pub struct MemoryConfig { + // ... existing fields ... + pub rrf_k: f64, // default: 60.0 + pub rrf_weight_fts: f64, // default: 0.5 + pub rrf_weight_vec: f64, // default: 0.5 +} +``` + +Update `search_messages_in_conversation` and `search_messages` to accept these parameters instead of hardcoded values. + +**Optional vec0 metadata columns** (migration in `mod.rs:332-337`): + +```sql +-- Current: +CREATE VIRTUAL TABLE message_embeddings USING vec0(embedding float[N]); + +-- New (metadata columns for pre-filtering): +CREATE VIRTUAL TABLE message_embeddings USING vec0( + embedding float[N], + is_summarized integer, + role text +); +``` + +Metadata columns allow KNN to filter at the vector level: +```sql +WHERE embedding MATCH ?1 + AND k = ?2 + AND is_summarized = 0 + AND role IN ('user', 'assistant') +``` + +This avoids post-filtering where valid KNN results get discarded by the WHERE clause after distance calculation. Migration re-creates the table on dimension change (already handled). + +**Explicit migration** in `run_migrations` (mod.rs:350-374): Since `ALTER TABLE` is not supported for vec0 virtual tables, add a schema check on every startup: + +```rust +// After the existing dimension-check migration block, add: +let has_metadata_columns = conn + .prepare("PRAGMA table_info(message_embeddings)") + .and_then(|mut stmt| { + let cols: Vec = stmt.query_map([], |row| row.get(1))? + .collect::, _>>()?; + Ok(cols.contains(&"is_summarized".to_string())) + }) + .unwrap_or(false); + +if table_exists(conn, "message_embeddings") && !has_metadata_columns { + conn.execute_batch("DROP TABLE message_embeddings;")?; + conn.execute_batch(&format!( + "CREATE VIRTUAL TABLE message_embeddings USING vec0(\ + embedding float[{}], is_summarized integer, role text);", + dims + ))?; + info!("Migrated message_embeddings with metadata columns"); +} +``` + +This is safe because `message_embeddings` is a derived index — dropping it loses no message data (messages are in the `messages` table). Embeddings will be regenerated on the next `save_message` call. + +**Why this is safe:** +- `message_embeddings` is a derived index, not source-of-truth +- Messages are stored in the `messages` table (persistent) +- Vectors are regenerated on next `save_message` (mod.rs:67-75) +- The migration runs once; subsequent startups skip it +- Backward compatible with FTS5-only fallback + +--- + +## Testing Strategy + +**Context window detection:** +- Unit test: `fetch_context_window` returns `None` on API failure +- Unit test: `effective_context_window` falls back to static when cache is empty +- Unit test: `effective_context_window` returns cached value when set + +**System prompt preservation:** +- Unit test: `auto_compact_conversation` preserves system messages in output +- Unit test: system messages are at indices 0..N in result (before any summary) +- Unit test: `reactive_compact` same behavior + +**Structured summary:** +- Unit test: `build_compact_summary_prompt` produces prompt with STATE and CONTEXT keywords +- Unit test: `build_compact_boundary_marker` contains "STATE" hint + +**RAG-aware compaction:** +- Integration test: `retrieve_context_for_compaction` returns snippets for real conversation +- Integration test: compact with retrieval returns valid state block +- Edge case test: empty retrieval falls back to full to_summarize + +**sqlite-vec metadata:** +- Integration test: metadata columns are queryable in vec0 +- Integration test: hybrid search with metadata pre-filter returns same results as post-filter diff --git a/src/agent.rs b/src/agent.rs index fe6d65f..df3e6f1 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -28,6 +28,10 @@ use crate::scheduler::Scheduler; use crate::skills::{format_listed_section, SkillRegistry}; use crate::tools; +/// Number of context snippets to retrieve from conversation history for +/// compaction summarization. +const COMPACTION_RAG_LIMIT: usize = 5; + /// A request dispatched from a fire closure to the background job runner. pub struct ScheduledJobRequest { pub incoming: IncomingMessage, @@ -385,6 +389,20 @@ impl Agent { Ok(()) } + /// Fetch the context window size for the current model from the + /// provider API and cache it. Non-fatal — uses static fallback on + /// failure. + pub async fn refresh_context_window_cache(&self) { + let model = self.current_model.read().await.clone(); + let (provider, actual_model) = self.registry.resolve_model(&model); + let client = reqwest::Client::new(); + if let Some(ctx) = provider.fetch_context_window(&client, actual_model).await { + let mut cache = provider.config().context_window_cache.write().await; + *cache = Some(ctx); + tracing::info!("Context window for {}: {} tokens", actual_model, ctx); + } + } + /// Process an incoming message and return the response text pub(crate) fn now_iso8601_static() -> String { chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true) @@ -638,8 +656,7 @@ impl Agent { // Resolve context_window once before the loop (can't .await inside the loop) let context_window = { let model = self.current_model.read().await; - let (provider, _) = self.registry.resolve_model(&model); - provider.config().context_window + self.registry.effective_context_window(&model) }; for iteration in 0..max_iterations { @@ -660,8 +677,14 @@ impl Agent { let messages_before = messages.clone(); let messages_before_bytes = estimate_prompt_bytes(&messages_before); - if let Ok(compacted) = - Self::auto_compact_conversation(&self.llm, &messages, context_window).await + if let Ok(compacted) = Self::auto_compact_conversation( + &self.llm, + &self.memory, + &conversation_id, + &messages, + context_window, + ) + .await { conv_meta.last_compact_turn = iteration as usize; messages = compacted; @@ -814,7 +837,14 @@ impl Agent { if is_413 && !conv_meta.has_attempted_reactive_compact { conv_meta.has_attempted_reactive_compact = true; let messages_before_compact = messages.len(); - match Self::reactive_compact(&self.llm, &messages, context_window).await + match Self::reactive_compact( + &self.llm, + &self.memory, + &conversation_id, + &messages, + context_window, + ) + .await { Ok(compacted) => { let compacted_len = compacted.len(); @@ -1329,10 +1359,27 @@ impl Agent { /// split decisions internally. async fn auto_compact_conversation( llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, messages: &[ChatMessage], _context_window: usize, ) -> Result> { - let tool_groups = crate::agent_prompt::find_tool_groups(messages); + // 1. Separate system messages from the rest + let mut system_msgs: Vec = Vec::new(); + let non_system: Vec = messages + .iter() + .filter(|&msg| { + if msg.role == "system" { + system_msgs.push(msg.clone()); + false + } else { + true + } + }) + .cloned() + .collect(); + + let tool_groups = crate::agent_prompt::find_tool_groups(&non_system); let preserve_count = PRESERVED_TOOL_GROUPS.min(tool_groups.len()); let preserved_groups_start = tool_groups.len().saturating_sub(preserve_count); @@ -1348,17 +1395,25 @@ impl Agent { return Ok(messages.to_vec()); }; - let to_summarize = &messages[..summary_end]; - let preserved = &messages[summary_end..]; + let to_summarize = &non_system[..summary_end]; + let preserved = &non_system[summary_end..]; - Self::summarize_and_replace( + // 2. Summarize with RAG-aware compaction + let mut compacted = Self::summarize_and_replace( llm, + memory, + conversation_id, to_summarize, preserved, "Auto-compact", "★ COMPACT SUMMARY ★", ) - .await + .await?; + + // 3. Prepend system messages back + let mut result = system_msgs; + result.append(&mut compacted); + Ok(result) } /// Tier 4: Reactive compact — emergency 413 recovery. @@ -1368,50 +1423,159 @@ impl Agent { /// split decisions internally. async fn reactive_compact( llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, messages: &[ChatMessage], _context_window: usize, ) -> Result> { const PRESERVE_COUNT: usize = 4; - if messages.len() <= PRESERVE_COUNT { - anyhow::bail!("Too few messages for reactive compact"); + + // 1. Separate system messages + let mut system_msgs: Vec = Vec::new(); + let non_system: Vec = messages + .iter() + .filter(|&msg| { + if msg.role == "system" { + system_msgs.push(msg.clone()); + false + } else { + true + } + }) + .cloned() + .collect(); + + if non_system.len() <= PRESERVE_COUNT { + anyhow::bail!("Too few non-system messages for reactive compact"); } - let split = messages.len().saturating_sub(PRESERVE_COUNT); - let to_summarize = &messages[..split]; - let preserved = &messages[split..]; + let split = non_system.len().saturating_sub(PRESERVE_COUNT); + let to_summarize = &non_system[..split]; + let preserved = &non_system[split..]; - Self::summarize_and_replace( + let mut compacted = Self::summarize_and_replace( llm, + memory, + conversation_id, to_summarize, preserved, "Reactive compact", "★ COMPACT SUMMARY (EMERGENCY) ★", ) - .await + .await?; + + let mut result = system_msgs; + result.append(&mut compacted); + Ok(result) } /// Shared helper for Tiers 3 and 4: send messages to LLM for /// summarization, then assemble the compacted result. async fn summarize_and_replace( llm: &LlmClient, + memory: &MemoryStore, + conversation_id: &str, to_summarize: &[ChatMessage], preserved: &[ChatMessage], error_label: &str, summary_label: &str, ) -> Result> { - let summary_prompt = build_compact_summary_prompt(); - let mut compact_msgs = Vec::with_capacity(to_summarize.len() + 2); - compact_msgs.push(summary_prompt); + // NEW: RAG retrieval for compaction (non-fatal — warn on error, continue) + let retrieved = match crate::memory::rag::retrieve_context_for_compaction( + memory, + to_summarize, + preserved, + conversation_id, + COMPACTION_RAG_LIMIT, + ) + .await + { + Ok(r) => r, + Err(e) => { + warn!("RAG retrieval for compaction failed: {}", e); + None + } + }; + + // Build compact messages: summary prompt + optional retrieved context + truncated input + let mut compact_msgs = Vec::new(); + compact_msgs.push(build_compact_summary_prompt()); + + if let Some(ref ctx) = retrieved { + compact_msgs.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(ctx.clone())), + tool_calls: None, + tool_call_id: None, + }); + } + + // Tool-group-aware truncation + let groups = crate::agent_prompt::find_tool_groups(to_summarize); + let bookend_groups = 1usize; + let tail_groups = 3usize.min(groups.len().saturating_sub(bookend_groups)); + + let mut seen_indices = std::collections::HashSet::new(); + + // First bookend groups (conversation origin) + for group in groups.iter().take(bookend_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Last tail groups (recent flow) + for group in groups.iter().rev().take(tail_groups) { + seen_indices.insert(group.assistant_idx); + for &ti in &group.tool_result_indices { + seen_indices.insert(ti); + } + } + + // Always include non-assistant/non-tool messages + for (idx, msg) in to_summarize.iter().enumerate() { + if msg.role != "assistant" && !msg.has_tool_calls() { + seen_indices.insert(idx); + } + } + + // Build sampled list in original order with truncation notice + let mut sampled: Vec = Vec::new(); + let mut inserted_notice = false; + for (idx, msg) in to_summarize.iter().enumerate() { + if seen_indices.contains(&idx) { + sampled.push(msg.clone()); + } else if !inserted_notice { + sampled.push(ChatMessage { + role: "system".to_string(), + content: Some(MessageContent::Text(format!( + "[... {} messages omitted, see retrieved_context above ...]", + to_summarize.len() - seen_indices.len() + ))), + tool_calls: None, + tool_call_id: None, + }); + inserted_notice = true; + } + } + + // If no truncation happened, use the full to_summarize + if sampled.is_empty() { + sampled = to_summarize.to_vec(); + } + + let user_prompt = format!( + "Summarize the following conversation (sampled from {} messages):", + to_summarize.len() + ); compact_msgs.push(ChatMessage { role: "user".to_string(), - content: Some(MessageContent::Text(format!( - "Summarize the following conversation portion ({} messages):", - to_summarize.len(), - ))), + content: Some(MessageContent::Text(user_prompt)), tool_calls: None, tool_call_id: None, }); - compact_msgs.extend(to_summarize.iter().cloned()); + compact_msgs.extend(sampled); let summary_response = match llm.chat(&compact_msgs, &[]).await { Ok(c) => c, diff --git a/src/agent_prompt.rs b/src/agent_prompt.rs index 57267ec..1bc6209 100644 --- a/src/agent_prompt.rs +++ b/src/agent_prompt.rs @@ -106,9 +106,9 @@ pub fn recovery_nudge_for(messages: &[ChatMessage]) -> ChatMessage { let previous_is_tool = messages.last().is_some_and(|msg| msg.role == "tool"); let content = if previous_is_tool { - "Continue from the tool result above. Either call the next required tool or provide a concise user-visible final answer.".to_string() + "Continue from the tool result above. Read the ## STATE block in the compact summary for context. Either call the next required tool or provide a final answer.".to_string() } else { - "Continue from the user's request above. Either call the next required tool or provide a concise user-visible final answer.".to_string() + "Continue from the user's request above. Read the ## STATE block in the compact summary for context. Either call the next required tool or provide a final answer.".to_string() }; ChatMessage { @@ -439,20 +439,33 @@ pub fn should_auto_compact( /// Returns a system-role message instructing the LLM to summarize. pub fn build_compact_summary_prompt() -> ChatMessage { let prompt_text = vec![ - "Your task is to create a detailed summary of the conversation so far.", + "You are producing a compact state summary of the conversation below.", "", - "Your summary must include these sections:", + "OUTPUT FORMAT:", "", - "1. Primary Request and Intent: What was the user's original request?", - "2. Key Technical Concepts: Technologies, frameworks, approaches discussed", - "3. Files and Code Sections: Files read, created, or modified", - "4. Errors and Fixes: Errors encountered and how they were fixed", - "5. All User Messages: List ALL user messages verbatim (not tool results)", - "6. Pending Tasks: What tasks were explicitly requested but not completed", - "7. Current Work: Exactly what was being worked on before this summary", - "8. Next Step: The next logical action based on the most recent user request", + "## STATE", + "```yaml", + "stage: ", + "decisions:", + " - ", + "pending:", + " - ", + "last_action: ", + "last_action_result: ", + "conversation_phase: ", + "```", "", - "IMPORTANT: Do NOT call any tools. Respond with text only.", + "## CONTEXT", + "- ", + "- ", + "- ", + "", + "CRITICAL RULES:", + "- State must be precise enough for the LLM to continue without re-reading history", + "- Include ALL pending items the user explicitly requested", + "- Include ALL error messages and their resolutions", + "- Be specific with file paths and tool names", + "- Do NOT call any tools. Respond with text only.", ] .join("\n"); @@ -469,7 +482,7 @@ pub fn build_compact_boundary_marker(original_count: usize, compacted_count: usi ChatMessage { role: "system".to_string(), content: Some(MessageContent::Text(format!( - "★ COMPACT SUMMARY — previous {} messages compressed into this summary (now {} messages) ★", + "★ COMPACT SUMMARY — previous {} messages → YAML state + narrative ({} messages) ★", original_count, compacted_count, ))), tool_calls: None, @@ -942,4 +955,25 @@ mod tests { }); assert!(has_boundary); } + + #[test] + fn compact_summary_prompt_contains_state_keywords() { + let msg = build_compact_summary_prompt(); + let text = msg.content.as_ref().unwrap().as_text(); + assert!(text.contains("## STATE"), "Should have STATE section"); + assert!(text.contains("## CONTEXT"), "Should have CONTEXT section"); + assert!(text.contains("decisions:"), "Should have decisions field"); + assert!(text.contains("pending:"), "Should have pending field"); + assert!( + text.contains("last_action:"), + "Should have last_action field" + ); + } + + #[test] + fn compact_boundary_marker_contains_state_hint() { + let msg = build_compact_boundary_marker(10, 1); + let text = msg.content.as_ref().unwrap().as_text(); + assert!(text.contains("state"), "Should hint at state format"); + } } diff --git a/src/config.rs b/src/config.rs index 8fc02e1..19d6eee 100644 --- a/src/config.rs +++ b/src/config.rs @@ -251,6 +251,48 @@ pub struct MemoryConfig { /// Can be toggled per-user at runtime via the `/query-rewrite` command. #[serde(default)] pub query_rewriter_enabled: bool, + /// RRF (Reciprocal Rank Fusion) parameter `k` — the rank offset used to + /// smooth the combined score from full-text and vector search results. + #[serde(default = "default_rrf_k")] + pub rrf_k: f64, + /// Weight assigned to the full-text search (FTS) rank in RRF scoring. + /// Must be between 0.0 and 1.0. The vector weight is derived from the + /// configured `rrf_weight_vec`. + #[serde(default = "default_rrf_weight_fts")] + pub rrf_weight_fts: f64, + /// Weight assigned to the vector-search rank in RRF scoring. + /// Must be between 0.0 and 1.0. The FTS weight is derived from the + /// configured `rrf_weight_fts`. + #[serde(default = "default_rrf_weight_vec")] + pub rrf_weight_vec: f64, +} + +impl Default for MemoryConfig { + fn default() -> Self { + Self { + database_path: PathBuf::new(), + rag_limit: default_rag_limit(), + max_raw_messages: default_max_raw_messages(), + summarize_threshold: default_summarize_threshold(), + summarize_cron: default_summarize_cron(), + query_rewriter_enabled: false, + rrf_k: default_rrf_k(), + rrf_weight_fts: default_rrf_weight_fts(), + rrf_weight_vec: default_rrf_weight_vec(), + } + } +} + +fn default_rrf_k() -> f64 { + 60.0 +} + +fn default_rrf_weight_fts() -> f64 { + 0.5 +} + +fn default_rrf_weight_vec() -> f64 { + 0.5 } #[derive(Debug, Deserialize, Clone)] @@ -281,6 +323,8 @@ pub struct AgentConfig { pub max_iterations: u32, #[serde(default = "default_empty_response_retry_limit")] pub empty_response_retry_limit: u32, + #[serde(default = "default_parse_retry_limit")] + pub parse_retry_limit: u32, } #[derive(Debug, Deserialize, Clone)] @@ -390,6 +434,9 @@ fn default_memory_config() -> MemoryConfig { summarize_threshold: default_summarize_threshold(), summarize_cron: default_summarize_cron(), query_rewriter_enabled: false, + rrf_k: default_rrf_k(), + rrf_weight_fts: default_rrf_weight_fts(), + rrf_weight_vec: default_rrf_weight_vec(), } } @@ -413,10 +460,15 @@ fn default_empty_response_retry_limit() -> u32 { 3 } +fn default_parse_retry_limit() -> u32 { + 3 +} + fn default_agent_config() -> AgentConfig { AgentConfig { max_iterations: default_max_iterations(), empty_response_retry_limit: default_empty_response_retry_limit(), + parse_retry_limit: default_parse_retry_limit(), } } @@ -494,6 +546,11 @@ impl Config { self.agent.empty_response_retry_limit } + /// Parse retry limit for missing 'choices' field (from [agent] parse_retry_limit, default 3). + pub fn parse_retry_limit(&self) -> u32 { + self.agent.parse_retry_limit + } + /// Resolve the home root and every data path, create directories, and write /// the resolved paths back into the config fields. Unset paths are /// materialized to absolute paths under the home root; absolute overrides @@ -1002,6 +1059,40 @@ mod tests { assert_eq!(cfg.empty_response_retry_limit(), 0); } + #[test] + fn test_agent_parse_retry_limit_defaults_to_three() { + let toml = r#" + [telegram] + bot_token = "tok" + allowed_user_ids = [1] + [openrouter] + api_key = "key" + [sandbox] + allowed_directory = "/tmp" + "#; + let cfg: Config = toml::from_str(toml).unwrap(); + assert_eq!(cfg.agent.parse_retry_limit, 3); + assert_eq!(cfg.parse_retry_limit(), 3); + } + + #[test] + fn test_agent_parse_retry_limit_can_be_configured_to_zero() { + let toml = r#" + [telegram] + bot_token = "tok" + allowed_user_ids = [1] + [openrouter] + api_key = "key" + [sandbox] + allowed_directory = "/tmp" + [agent] + parse_retry_limit = 0 + "#; + let cfg: Config = toml::from_str(toml).unwrap(); + assert_eq!(cfg.agent.parse_retry_limit, 0); + assert_eq!(cfg.parse_retry_limit(), 0); + } + #[test] fn test_provider_section_parses_ollama() { let toml = r#" diff --git a/src/main.rs b/src/main.rs index 94a4dcf..233e077 100644 --- a/src/main.rs +++ b/src/main.rs @@ -65,8 +65,12 @@ async fn main() -> Result<()> { // Build provider registry from config let (provider_sections, default_provider, fallback_chain) = config.build_providers(); let registry = Arc::new( - provider::build_registry(&provider_sections, &default_provider) - .context("Failed to build LLM provider registry")?, + provider::build_registry( + &provider_sections, + &default_provider, + config.parse_retry_limit(), + ) + .context("Failed to build LLM provider registry")?, ); info!( " Providers: {} (default: {}, fallback: {} model(s))", @@ -75,6 +79,29 @@ async fn main() -> Result<()> { fallback_chain.len() ); + // Spawn background task to warm context_window_cache for all providers + { + let registry_clone = Arc::clone(®istry); + tokio::spawn(async move { + let client = reqwest::Client::new(); + for name in registry_clone.provider_names() { + if let Some(provider) = registry_clone.get_provider(&name) { + let model = provider.default_model(); + if let Some(ctx) = provider.fetch_context_window(&client, model).await { + let mut cache = provider.config().context_window_cache.write().await; + *cache = Some(ctx); + tracing::info!( + "Context window cache: {} / {} = {} tokens", + name, + model, + ctx + ); + } + } + } + }); + } + info!("Configuration loaded successfully"); let default_provider_obj = registry .get_provider(registry.default_provider_name()) @@ -115,8 +142,12 @@ async fn main() -> Result<()> { }); // Initialize memory store (SQLite + vector embeddings) - let memory = MemoryStore::open(&config.memory.database_path, embedding_config) - .context("Failed to initialize memory store")?; + let memory = MemoryStore::open( + &config.memory.database_path, + embedding_config, + config.memory.clone(), + ) + .context("Failed to initialize memory store")?; info!(" Database: {}", config.memory.database_path.display()); // Refresh any expiring OAuth tokens before connecting to MCP servers diff --git a/src/memory/conversations.rs b/src/memory/conversations.rs index bf0dc70..d9608d0 100644 --- a/src/memory/conversations.rs +++ b/src/memory/conversations.rs @@ -102,8 +102,8 @@ impl MemoryStore { if let Some(ref emb) = embedding { let embedding_bytes = f32_slice_to_bytes(emb); conn.execute( - "INSERT INTO message_embeddings (rowid, embedding) VALUES (?1, ?2)", - rusqlite::params![rowid, embedding_bytes], + "INSERT INTO message_embeddings (rowid, embedding, is_summarized, role) VALUES (?1, ?2, 0, ?3)", + rusqlite::params![rowid, embedding_bytes, &message.role], )?; } @@ -205,6 +205,8 @@ impl MemoryStore { row_number() OVER (ORDER BY distance) as rank_number FROM message_embeddings WHERE embedding MATCH ?1 + AND is_summarized = 0 + AND role IN ('user', 'assistant') ORDER BY distance LIMIT ?2 ), @@ -216,8 +218,8 @@ impl MemoryStore { LIMIT ?2 ) SELECT m.role, m.content, m.tool_calls, m.tool_call_id, - coalesce(1.0 / (60 + fts.rank_number), 0.0) * 0.5 - + coalesce(1.0 / (60 + vec.rank_number), 0.0) * 0.5 as combined_rank + coalesce(1.0 / (?5 + fts.rank_number), 0.0) * ?7 + + coalesce(1.0 / (?5 + vec.rank_number), 0.0) * ?6 as combined_rank FROM messages m LEFT JOIN vec_matches vec ON m.rowid = vec.rowid LEFT JOIN fts_matches fts ON m.rowid = fts.rowid @@ -230,10 +232,21 @@ impl MemoryStore { "; let search_limit = (limit * 3) as i64; + let rrf_k = self.config.rrf_k; + let rrf_weight_fts = self.config.rrf_weight_fts; + let rrf_weight_vec = self.config.rrf_weight_vec; let mut stmt = conn.prepare(sql)?; let messages = stmt .query_map( - rusqlite::params![query_bytes, search_limit, query, conversation_id], + rusqlite::params![ + query_bytes, + search_limit, + query, + conversation_id, + rrf_k, + rrf_weight_vec, + rrf_weight_fts + ], parse_message_row, )? .collect::, _>>() @@ -283,6 +296,8 @@ impl MemoryStore { row_number() OVER (ORDER BY distance) as rank_number FROM message_embeddings WHERE embedding MATCH ?1 + AND is_summarized = 0 + AND role IN ('user', 'assistant') ORDER BY distance LIMIT ?2 ), @@ -294,22 +309,34 @@ impl MemoryStore { LIMIT ?2 ) SELECT m.role, m.content, m.tool_calls, m.tool_call_id, - coalesce(1.0 / (60 + fts.rank_number), 0.0) * 0.5 - + coalesce(1.0 / (60 + vec.rank_number), 0.0) * 0.5 as combined_rank + coalesce(1.0 / (?4 + fts.rank_number), 0.0) * ?6 + + coalesce(1.0 / (?4 + vec.rank_number), 0.0) * ?5 as combined_rank FROM messages m LEFT JOIN vec_matches vec ON m.rowid = vec.rowid LEFT JOIN fts_matches fts ON m.rowid = fts.rowid - WHERE vec.rowid IS NOT NULL OR fts.rowid IS NOT NULL + WHERE (vec.rowid IS NOT NULL OR fts.rowid IS NOT NULL) + AND m.role IN ('user', 'assistant') + AND (m.is_summarized IS NULL OR m.is_summarized = 0) ORDER BY combined_rank DESC LIMIT ?2 "; - let search_limit = (limit * 3) as i64; + let rrf_k = self.config.rrf_k; + let rrf_weight_fts = self.config.rrf_weight_fts; + let rrf_weight_vec = self.config.rrf_weight_vec; let mut stmt = conn.prepare(sql)?; let messages = stmt - .query_map(rusqlite::params![query_bytes, search_limit, query], |row| { - parse_message_row(row) - })? + .query_map( + rusqlite::params![ + query_bytes, + search_limit, + query, + rrf_k, + rrf_weight_vec, + rrf_weight_fts + ], + parse_message_row, + )? .collect::, _>>() .context("Failed to hybrid-search messages")?; @@ -321,6 +348,8 @@ impl MemoryStore { FROM messages m JOIN messages_fts fts ON m.rowid = fts.rowid WHERE messages_fts MATCH ?1 + AND m.role IN ('user', 'assistant') + AND (m.is_summarized IS NULL OR m.is_summarized = 0) ORDER BY fts.rank LIMIT ?2 "; diff --git a/src/memory/mod.rs b/src/memory/mod.rs index fbc54fa..00c1aa3 100644 --- a/src/memory/mod.rs +++ b/src/memory/mod.rs @@ -12,6 +12,7 @@ use std::sync::Arc; use tokio::sync::Mutex; use tracing::info; +use crate::config::MemoryConfig; use crate::memory::embeddings::{EmbeddingConfig, EmbeddingEngine}; /// Thread-safe SQLite memory store with hybrid vector+FTS5 search @@ -19,13 +20,18 @@ use crate::memory::embeddings::{EmbeddingConfig, EmbeddingEngine}; pub struct MemoryStore { conn: Arc>, pub embeddings: Arc, + pub config: MemoryConfig, } impl MemoryStore { /// Open or create the SQLite database at the given path. /// If `embedding_config` is provided, vector search is enabled alongside FTS5. /// If None, falls back to FTS5-only search. - pub fn open(path: &Path, embedding_config: Option) -> Result { + pub fn open( + path: &Path, + embedding_config: Option, + memory_config: MemoryConfig, + ) -> Result { // Register sqlite-vec extension before opening any connection unsafe { type VecInitFn = unsafe extern "C" fn( @@ -57,6 +63,7 @@ impl MemoryStore { let store = Self { conn: Arc::new(Mutex::new(conn)), embeddings: Arc::new(embeddings), + config: memory_config, }; info!("Memory store initialized at: {}", path.display()); @@ -89,6 +96,7 @@ impl MemoryStore { let store = Self { conn: Arc::new(Mutex::new(conn)), embeddings: Arc::new(embeddings), + config: MemoryConfig::default(), }; Ok(store) } @@ -373,6 +381,52 @@ impl MemoryStore { } } + // Migration: ensure message_embeddings has metadata columns (is_summarized, role) + // for pre-filtering, and no existing rows with NULL metadata. + // ALTER TABLE is not supported for vec0, so we must DROP and recreate. + if table_exists(conn, "message_embeddings") { + let cols: Vec = conn + .prepare("PRAGMA table_info(message_embeddings)") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get(1))? + .collect::, _>>() + }) + .unwrap_or_default(); + + let has_meta = cols.contains(&"is_summarized".to_string()); + + let needs_recreate = if has_meta { + // Columns exist but old rows may have NULL metadata because the + // original INSERT didn't write metadata columns. Check for NULLs. + conn.query_row( + "SELECT COUNT(*) FROM message_embeddings WHERE is_summarized IS NULL", + [], + |row| row.get::<_, i64>(0), + ) + .map(|count| count > 0) + .unwrap_or(false) + } else { + true + }; + + if needs_recreate { + conn.execute_batch("DROP TABLE message_embeddings;")?; + conn.execute_batch(&format!( + "CREATE VIRTUAL TABLE message_embeddings USING vec0(\ + embedding float[{}], is_summarized integer, role text);", + dims + ))?; + info!( + "Migrated message_embeddings with metadata columns (is_summarized, role){}", + if has_meta { + " and rebuilt existing data" + } else { + "" + } + ); + } + } + Ok(()) } } diff --git a/src/memory/rag.rs b/src/memory/rag.rs index ef8f48a..cea1bc1 100644 --- a/src/memory/rag.rs +++ b/src/memory/rag.rs @@ -2,6 +2,7 @@ use anyhow::Result; use tracing::debug; use super::MemoryStore; +use crate::llm::ChatMessage; /// Auto-retrieve semantically relevant past messages from a conversation /// and format them as a `` block for the system prompt. @@ -57,6 +58,62 @@ pub async fn auto_retrieve_context( Ok(Some(block)) } +/// Retrieve context for compaction summarization. +/// +/// Uses the most recent user message (from both to_summarize and preserved +/// ranges) as a search query to find relevant historical context from the +/// conversation. Returns formatted snippets that help the summarizer write +/// a focused summary. Returns None when no suitable query is found or no +/// results are returned. +pub async fn retrieve_context_for_compaction( + store: &MemoryStore, + to_summarize: &[ChatMessage], + preserved: &[ChatMessage], + conversation_id: &str, + limit: usize, +) -> Result> { + let query = preserved + .iter() + .rev() + .chain(to_summarize.iter().rev()) + .find(|m| m.role == "user") + .and_then(|m| m.content.as_ref().map(|c| c.as_text())) + .unwrap_or_default(); + + if query.trim().len() < 5 { + return Ok(None); + } + + let results = store + .search_messages_in_conversation(&query, conversation_id, limit) + .await?; + + if results.is_empty() { + return Ok(None); + } + + let mut block = String::from( + "\n\ + Relevant context from conversation history for compaction:\n\n", + ); + + for msg in &results { + if let Some(content) = &msg.content { + let text = content.as_text(); + let snippet = crate::utils::strings::truncate_chars(&text, 300); + block.push_str(&format!("[{}] {}\n", msg.role, snippet)); + } + } + + block.push_str(""); + tracing::debug!( + "Compaction RAG: injected {} snippets for query: {:?}", + results.len(), + query + ); + Ok(Some(block)) +} + #[cfg(test)] mod tests { use super::*; @@ -180,4 +237,61 @@ mod tests { .unwrap(); let _ = result; // Just verify no panic } + + #[tokio::test] + async fn test_retrieve_context_for_compaction_returns_none_for_short_query() { + let store = MemoryStore::open_in_memory().unwrap(); + let conv = store + .get_or_create_conversation("test", "compact_u1") + .await + .unwrap(); + let to_summarize = vec![ChatMessage { + role: "assistant".to_string(), + content: Some(MessageContent::from_text("some result")), + tool_calls: None, + tool_call_id: None, + }]; + let preserved = vec![]; + let result = retrieve_context_for_compaction(&store, &to_summarize, &preserved, &conv, 5) + .await + .unwrap(); + assert!(result.is_none(), "No user message should return None"); + } + + #[tokio::test] + async fn test_retrieve_context_for_compaction_finds_user_message_in_preserved() { + let store = MemoryStore::open_in_memory().unwrap(); + let conv = store + .get_or_create_conversation("test", "compact_u2") + .await + .unwrap(); + + let msg = ChatMessage { + role: "user".to_string(), + content: Some(MessageContent::from_text("Tell me about Rust async")), + tool_calls: None, + tool_call_id: None, + }; + store.save_message(&conv, &msg).await.unwrap(); + + let to_summarize = vec![ChatMessage { + role: "assistant".to_string(), + content: Some(MessageContent::from_text("Here's how...")), + tool_calls: None, + tool_call_id: None, + }]; + let preserved = vec![ChatMessage { + role: "user".to_string(), + content: Some(MessageContent::from_text("Tell me about Rust async")), + tool_calls: None, + tool_call_id: None, + }]; + + let result = retrieve_context_for_compaction(&store, &to_summarize, &preserved, &conv, 5) + .await + .unwrap(); + // May return None if embeddings not available (FTS-only fallback), + // which is acceptable. The key is no panic. + let _ = result; + } } diff --git a/src/provider.rs b/src/provider.rs index 193baf0..46d5386 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -1,10 +1,12 @@ use std::collections::HashMap; use std::sync::Arc; +use tokio::sync::RwLock; use anyhow::{Context, Result}; use async_trait::async_trait; use crate::config::{ProviderSection, ProviderType}; +use crate::llm::internal::{ChatResponse, Choice}; use crate::llm::{ChatCompletion, ChatMessage, ToolDefinition}; /// Runtime configuration for a single LLM provider. @@ -23,6 +25,13 @@ pub struct ProviderConfig { pub max_tokens: u32, pub discover_models: bool, pub context_window: usize, + /// Runtime cache for the current model's context window, populated + /// asynchronously from the provider API. When None, falls back to + /// `context_window`. + pub context_window_cache: Arc>>, + /// Number of retries when the response is missing the `choices` field. + /// Defaults to 3; set to 0 to disable. + pub parse_retry_limit: u32, } impl From<&ProviderSection> for ProviderConfig { @@ -37,10 +46,101 @@ impl From<&ProviderSection> for ProviderConfig { max_tokens: s.max_tokens, discover_models: s.discover_models, context_window: s.context_window, + context_window_cache: Arc::new(RwLock::new(None)), + parse_retry_limit: 0, // overwritten by build_registry } } } +/// Internal helper: send a chat completion request with retry logic for +/// missing/empty `choices` field. Returns the first valid `Choice` on success. +/// +/// The request is re-sent on each retry (up to `parse_retry_limit` times) with +/// exponential backoff (1s, 2s, 4s...). Only retries when the JSON response +/// is valid but `choices` is missing or empty; other errors (network, HTTP +/// status, JSON parse) are returned immediately. +async fn chat_completion_with_retry( + client: &reqwest::Client, + url: &str, + request: &crate::llm::internal::ChatRequest, + api_key: Option<&str>, + provider_name: &str, + parse_retry_limit: u32, +) -> Result { + let mut last_error: Option = None; + + for attempt in 0..=parse_retry_limit { + let mut req = client.post(url).json(request); + if let Some(key) = api_key { + req = req.header("Authorization", format!("Bearer {}", key)); + } + + let response = match req.send().await { + Ok(r) => r, + Err(e) => return Err(e).context("Failed to send request"), + }; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow::anyhow!( + "{} API error ({}): {}", + provider_name, + status, + body + )); + } + + let body = match response.bytes().await { + Ok(b) => b, + Err(e) => return Err(e).context("Failed to read response body"), + }; + + let value: serde_json::Value = match serde_json::from_slice(&body) { + Ok(v) => v, + Err(e) => return Err(e).context("Failed to parse JSON response"), + }; + + let has_choices = value + .get("choices") + .and_then(|c| c.as_array()) + .is_some_and(|c| !c.is_empty()); + + if has_choices { + let chat_response: ChatResponse = serde_json::from_value(value)?; + return chat_response + .choices + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("No response from {}", provider_name)); + } + + // Missing or empty choices + let backoff_s = 1u64 << attempt; + let err = anyhow::anyhow!( + "Response from {} missing or empty 'choices' field – attempt {}/{}, retry limit {}", + provider_name, + attempt + 1, + parse_retry_limit + 1, + parse_retry_limit + ); + tracing::warn!( + "Retry {}/{}: {} - backing off {}s", + attempt + 1, + parse_retry_limit, + err, + backoff_s + ); + last_error = Some(err); + + if attempt < parse_retry_limit { + tokio::time::sleep(std::time::Duration::from_secs(backoff_s)).await; + } + } + + Err(last_error.unwrap()) +} + /// Unified LLM provider abstraction. /// /// Each implementation owns its own HTTP shaping (URL, headers, model list @@ -63,6 +163,12 @@ pub trait Provider: Send + Sync { ) -> Result; async fn list_models(&self, client: &reqwest::Client) -> Result>; + + /// Fetch the context window for a given model from the provider API. + /// Returns None if the provider doesn't support runtime detection. + async fn fetch_context_window(&self, _client: &reqwest::Client, _model: &str) -> Option { + None // default: no API-based detection + } } /// Holds the set of providers available to the agent and the default provider @@ -133,17 +239,32 @@ impl ProviderRegistry { let provider = &self.providers[&self.default_provider]; provider.config().context_window } + + /// Return the effective context window for a model: runtime cache + /// if populated, otherwise static config fallback. + pub fn effective_context_window(&self, model: &str) -> usize { + let (provider, _) = self.resolve_model(model); + let cached = provider + .config() + .context_window_cache + .try_read() + .ok() + .and_then(|c| *c); + cached.unwrap_or(provider.config().context_window) + } } /// Build a ProviderRegistry from config sections. pub fn build_registry( sections: &[ProviderSection], default_name: &str, + parse_retry_limit: u32, ) -> Result { let mut providers: HashMap> = HashMap::new(); for section in sections { - let cfg: ProviderConfig = ProviderConfig::from(section); + let mut cfg: ProviderConfig = ProviderConfig::from(section); + cfg.parse_retry_limit = parse_retry_limit; let provider: Arc = match section.provider_type { ProviderType::OpenRouter => Arc::new(OpenRouterProvider::new(cfg)), ProviderType::OpenAICompatible => Arc::new(OpenAICompatibleProvider::new(cfg)), @@ -223,32 +344,19 @@ impl Provider for OpenRouterProvider { }; let url = format!("{}/chat/completions", self.config.base_url); - let mut req = client.post(&url).json(&request); - if let Some(ref key) = self.config.api_key { - req = req.header("Authorization", format!("Bearer {}", key)); - } - - let response = req - .send() - .await - .context("Failed to send request to OpenRouter")?; - - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("OpenRouter API error ({}): {}", status, body); - } - - let chat_response: crate::llm::internal::ChatResponse = response - .json() - .await - .context("Failed to parse OpenRouter response")?; - - let mut choice = chat_response - .choices - .into_iter() - .next() - .context("No response from OpenRouter")?; + let api_key = self.config.api_key.as_deref(); + let provider_name = self.config.name.as_str(); + let parse_retry_limit = self.config.parse_retry_limit; + + let mut choice = chat_completion_with_retry( + client, + &url, + &request, + api_key, + provider_name, + parse_retry_limit, + ) + .await?; // Kimi tool-call fallback let has_tool_calls = choice @@ -298,6 +406,26 @@ impl Provider for OpenRouterProvider { .unwrap_or_default(); Ok(models) } + + async fn fetch_context_window(&self, client: &reqwest::Client, model: &str) -> Option { + let url = format!("{}/models", self.config.base_url); + let mut req = client.get(&url); + if let Some(ref key) = self.config.api_key { + req = req.header("Authorization", format!("Bearer {}", key)); + } + let response = req.send().await.ok()?; + if !response.status().is_success() { + return None; + } + let list: serde_json::Value = response.json().await.ok()?; + let ctx = list["data"] + .as_array()? + .iter() + .find(|m| m["id"].as_str() == Some(model))? + .get("context_length")? + .as_u64()?; + Some(ctx as usize) + } } // === OpenAICompatibleProvider === @@ -349,34 +477,19 @@ impl Provider for OpenAICompatibleProvider { }; let url = format!("{}/chat/completions", self.config.base_url); - let mut req = client.post(&url).json(&request); - if let Some(ref key) = self.config.api_key { - req = req.header("Authorization", format!("Bearer {}", key)); - } - - let response = req - .send() - .await - .context("Failed to send request to OpenAI-compatible provider")?; - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "Provider '{}' error ({}): {}", - self.config.name, - status, - body - ); - } - - let chat_response: crate::llm::internal::ChatResponse = response - .json() - .await - .context("Failed to parse OpenAI-compatible provider response")?; - let choice = - chat_response.choices.into_iter().next().ok_or_else(|| { - anyhow::anyhow!("No response from provider '{}'", self.config.name) - })?; + let api_key = self.config.api_key.as_deref(); + let provider_name = self.config.name.as_str(); + let parse_retry_limit = self.config.parse_retry_limit; + + let choice = chat_completion_with_retry( + client, + &url, + &request, + api_key, + provider_name, + parse_retry_limit, + ) + .await?; Ok(ChatCompletion { message: choice.message, @@ -416,6 +529,26 @@ impl Provider for OpenAICompatibleProvider { .unwrap_or_default(); Ok(models) } + + async fn fetch_context_window(&self, client: &reqwest::Client, model: &str) -> Option { + let url = format!("{}/models", self.config.base_url); + let mut req = client.get(&url); + if let Some(ref key) = self.config.api_key { + req = req.header("Authorization", format!("Bearer {}", key)); + } + let response = req.send().await.ok()?; + if !response.status().is_success() { + return None; + } + let list: serde_json::Value = response.json().await.ok()?; + let ctx = list["data"] + .as_array()? + .iter() + .find(|m| m["id"].as_str() == Some(model))? + .get("context_length")? + .as_u64()?; + Some(ctx as usize) + } } // === OllamaProvider === @@ -544,6 +677,8 @@ mod tests { assert_eq!(cfg.default_model, "anthropic/claude-sonnet-4-6"); assert_eq!(cfg.max_tokens, 1024); assert_eq!(cfg.context_window, 512_000); + // build_registry overwrites parse_retry_limit; From impl sets 0 as sentinel + assert_eq!(cfg.parse_retry_limit, 0); } #[test] @@ -562,7 +697,7 @@ mod tests { "llama3.1", ), ]; - let reg = build_registry(§ions, "alpha").unwrap(); + let reg = build_registry(§ions, "alpha", 3).unwrap(); assert_eq!(reg.provider_count(), 2); assert!(reg.get_provider("alpha").is_some()); assert!(reg.get_provider("beta").is_some()); @@ -571,7 +706,7 @@ mod tests { #[test] fn build_registry_fails_with_no_providers() { - let result = build_registry(&[], "alpha"); + let result = build_registry(&[], "alpha", 3); assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("No LLM providers")); } @@ -584,7 +719,7 @@ mod tests { "https://openrouter.ai/api/v1", "anthropic/claude-sonnet-4-6", )]; - let result = build_registry(§ions, "missing"); + let result = build_registry(§ions, "missing", 3); assert!(result.is_err()); assert!(result .unwrap_err() @@ -608,7 +743,7 @@ mod tests { "llama3.1", ), ]; - let reg = build_registry(§ions, "alpha").unwrap(); + let reg = build_registry(§ions, "alpha", 3).unwrap(); let (provider, model) = reg.resolve_model("beta/llama3.1:8b"); assert_eq!(provider.name(), "beta"); assert_eq!(model, "llama3.1:8b"); @@ -622,7 +757,7 @@ mod tests { "https://openrouter.ai/api/v1", "anthropic/claude-sonnet-4-6", )]; - let reg = build_registry(§ions, "alpha").unwrap(); + let reg = build_registry(§ions, "alpha", 3).unwrap(); let (provider, model) = reg.resolve_model("moonshotai/kimi-k2.5"); assert_eq!(provider.name(), "alpha"); assert_eq!(model, "moonshotai/kimi-k2.5"); @@ -636,7 +771,7 @@ mod tests { "https://openrouter.ai/api/v1", "anthropic/claude-sonnet-4-6", )]; - let reg = build_registry(§ions, "alpha").unwrap(); + let reg = build_registry(§ions, "alpha", 3).unwrap(); let (provider, model) = reg.resolve_model("claude-sonnet-4-6"); assert_eq!(provider.name(), "alpha"); assert_eq!(model, "claude-sonnet-4-6"); @@ -658,12 +793,41 @@ mod tests { "llama3.1", ), ]; - let reg = build_registry(§ions, "alpha").unwrap(); + let reg = build_registry(§ions, "alpha", 3).unwrap(); let mut names = reg.provider_names(); names.sort(); assert_eq!(names, vec!["alpha".to_string(), "beta".to_string()]); } + #[test] + fn effective_context_window_falls_back_to_static_when_cache_empty() { + let sections = vec![make_section( + "alpha", + ProviderType::OpenRouter, + "https://openrouter.ai/api/v1", + "anthropic/claude-sonnet-4-6", + )]; + let reg = build_registry(§ions, "alpha", 3).unwrap(); + let ctx = reg.effective_context_window("anthropic/claude-sonnet-4-6"); + assert_eq!(ctx, 512_000); // static fallback from section + } + + #[test] + fn effective_context_window_returns_cached_value_when_set() { + let sections = vec![make_section( + "alpha", + ProviderType::OpenRouter, + "https://openrouter.ai/api/v1", + "anthropic/claude-sonnet-4-6", + )]; + let reg = build_registry(§ions, "alpha", 3).unwrap(); + let provider = reg.get_provider("alpha").unwrap(); + // Set cache manually + *provider.config().context_window_cache.try_write().unwrap() = Some(200_000); + let ctx = reg.effective_context_window("anthropic/claude-sonnet-4-6"); + assert_eq!(ctx, 200_000); + } + #[test] fn ollama_discovery_url_strips_v1_suffix() { let cfg = ProviderConfig { @@ -676,6 +840,8 @@ mod tests { max_tokens: 1024, discover_models: true, context_window: 512_000, + context_window_cache: Arc::new(RwLock::new(None)), + parse_retry_limit: 3, }; let p = OllamaProvider::new(cfg); assert_eq!(p.discovery_url(), "http://localhost:11434/api/tags"); @@ -693,6 +859,8 @@ mod tests { max_tokens: 1024, discover_models: true, context_window: 512_000, + context_window_cache: Arc::new(RwLock::new(None)), + parse_retry_limit: 3, }; let p = OllamaProvider::new(cfg); assert_eq!(p.discovery_url(), "http://localhost:11434/api/tags"); @@ -710,6 +878,8 @@ mod tests { max_tokens: 1024, discover_models: true, context_window: 512_000, + context_window_cache: Arc::new(RwLock::new(None)), + parse_retry_limit: 3, }; let p = OllamaProvider::new(cfg); assert_eq!(p.discovery_url(), "http://localhost:11434/api/tags");