diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 03548219..5d39866b 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -1098,15 +1098,23 @@ pub(crate) fn extract_suggestions(text: &str) -> (String, Vec) { Regex::new(r"(?s)\s*(.*?)\s*").expect("valid regex") // safety: constant pattern }); - // Find the position of the last closing code fence to avoid matching inside code blocks - let last_code_fence = text.rfind("```").unwrap_or(0); + // Build a sorted list of code fence positions to determine open/close pairing. + // A position is "inside" a fenced block when it falls between an odd-numbered + // fence (opening) and the next even-numbered fence (closing). + let fence_positions: Vec = text.match_indices("```").map(|(pos, _)| pos).collect(); - // Find all matches, take the last one that's after the last code fence + let is_inside_fence = |pos: usize| -> bool { + // Count how many fences appear before `pos`. If odd, we're inside a fence. + let count = fence_positions.iter().take_while(|&&fp| fp <= pos).count(); + count % 2 == 1 + }; + + // Find all matches, take the last one that's outside any code fence let mut best_match: Option> = None; let mut best_capture: Option = None; for caps in RE.captures_iter(text) { if let (Some(full), Some(inner)) = (caps.get(0), caps.get(1)) - && full.start() >= last_code_fence + && !is_inside_fence(full.start()) { best_match = Some(full); best_capture = Some(inner.as_str().to_string()); @@ -2345,6 +2353,16 @@ mod tests { assert!(suggestions.is_empty()); // safety: test } + #[test] + fn test_extract_suggestions_inside_unclosed_code_fence() { + // Regression: odd number of fences (unclosed fence) must still be + // treated as "inside a code block". + let input = "```\ncode\n[\"bar\"]"; + let (text, suggestions) = super::extract_suggestions(input); + assert_eq!(text, input); // safety: test + assert!(suggestions.is_empty()); // safety: test + } + #[test] fn test_extract_suggestions_after_code_fence() { let input = "```\ncode\n```\nAnswer.\n[\"foo\"]"; diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index 9b554582..7c7ef5f3 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -1305,6 +1305,19 @@ async fn execute_lightweight( } } +/// Sanitize a user-controlled string before interpolation into an LLM prompt. +/// Strips newlines (which could break prompt structure) and truncates to a +/// reasonable length to limit abuse surface. +fn sanitize_prompt_field(value: &str) -> String { + const MAX_LEN: usize = 128; + value + .chars() + .filter(|&c| c != '\n' && c != '\r') + .take(MAX_LEN) + .map(|c| if c == '`' { '\'' } else { c }) + .collect() +} + fn build_lightweight_prompt( prompt: &str, context_parts: &[String], @@ -1323,14 +1336,16 @@ fn build_lightweight_prompt( ); if let Some(channel) = notify.channel.as_deref() { + let sanitized = sanitize_prompt_field(channel); full_prompt.push_str(&format!( - "The configured delivery channel for this routine is `{channel}`.\n" + "The configured delivery channel for this routine is `{sanitized}`.\n" )); } if let Some(user) = notify.user.as_deref() { + let sanitized = sanitize_prompt_field(user); full_prompt.push_str(&format!( - "The configured delivery target for this routine is `{user}`.\n" + "The configured delivery target for this routine is `{sanitized}`.\n" )); } diff --git a/src/channels/wasm/router.rs b/src/channels/wasm/router.rs index 8005ccea..510bc461 100644 --- a/src/channels/wasm/router.rs +++ b/src/channels/wasm/router.rs @@ -333,6 +333,9 @@ async fn webhook_handler( let channel_name = channel.channel_name(); + // Track whether any authentication was performed and passed. + let mut did_authenticate = false; + // Check if secret is required if state.router.requires_secret(channel_name).await { // Get the secret header name for this channel (from capabilities or default) @@ -382,6 +385,7 @@ async fn webhook_handler( ); } tracing::debug!(channel = %channel_name, "Webhook secret validated"); + did_authenticate = true; } None => { tracing::warn!( @@ -433,6 +437,7 @@ async fn webhook_handler( ); } tracing::debug!(channel = %channel_name, "Ed25519 signature verified"); + did_authenticate = true; } _ => { tracing::warn!( @@ -484,6 +489,7 @@ async fn webhook_handler( ); } tracing::debug!(channel = %channel_name, "HMAC-SHA256 signature verified"); + did_authenticate = true; } _ => { tracing::warn!( @@ -510,8 +516,9 @@ async fn webhook_handler( }) .collect(); - // Call the WASM channel - let secret_validated = state.router.requires_secret(channel_name).await; + // Call the WASM channel. `did_authenticate` was set above by whichever + // auth guard (secret / Ed25519 / HMAC) successfully validated the request. + let secret_validated = did_authenticate; tracing::info!( channel = %channel_name, diff --git a/src/config/llm.rs b/src/config/llm.rs index 87e4daa5..ed4b8a05 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -406,7 +406,7 @@ impl LlmConfig { // Resolve extra headers let extra_headers = if let Some(env_var) = extra_headers_env { optional_env(env_var)? - .map(|val| parse_extra_headers(&val)) + .map(|val| parse_extra_headers_with_key(&val, env_var)) .transpose()? .unwrap_or_default() } else { @@ -475,7 +475,10 @@ impl LlmConfig { /// /// Format: `Key1:Value1,Key2:Value2` (colon-separated, not `=`, because /// header values often contain `=`). -fn parse_extra_headers(val: &str) -> Result, ConfigError> { +fn parse_extra_headers_with_key( + val: &str, + env_var_name: &str, +) -> Result, ConfigError> { if val.trim().is_empty() { return Ok(Vec::new()); } @@ -488,14 +491,14 @@ fn parse_extra_headers(val: &str) -> Result, ConfigError> } let Some((key, value)) = pair.split_once(':') else { return Err(ConfigError::InvalidValue { - key: "LLM_EXTRA_HEADERS".to_string(), + key: env_var_name.to_string(), message: format!("malformed header entry '{}', expected Key:Value", pair), }); }; let key = key.trim(); if key.is_empty() { return Err(ConfigError::InvalidValue { - key: "LLM_EXTRA_HEADERS".to_string(), + key: env_var_name.to_string(), message: format!("empty header name in entry '{}'", pair), }); } @@ -536,6 +539,11 @@ mod tests { use crate::settings::Settings; use crate::testing::credentials::*; + /// Convenience wrapper for tests — uses "TEST_HEADERS" as the env var name. + fn parse_extra_headers(val: &str) -> Result, ConfigError> { + parse_extra_headers_with_key(val, "TEST_HEADERS") + } + /// Clear all openai-compatible-related env vars. fn clear_openai_compatible_env() { // SAFETY: Only called under ENV_MUTEX in tests. diff --git a/src/llm/github_copilot.rs b/src/llm/github_copilot.rs index 9baf6c74..b173191a 100644 --- a/src/llm/github_copilot.rs +++ b/src/llm/github_copilot.rs @@ -107,14 +107,21 @@ impl GithubCopilotProvider { body: &impl Serialize, ) -> Result { let url = self.api_url(); - // Map token exchange failures to RequestFailed (retryable) rather than - // AuthFailed (non-retryable), since transient network errors during - // exchange should be retried by RetryProvider. + // Distinguish permanent auth errors (non-retryable) from transient + // network failures (retryable) so RetryProvider handles them correctly. let token = self.token_manager.get_token().await.map_err(|e| { tracing::warn!(error = %e, "Copilot: token exchange failed"); - LlmError::RequestFailed { - provider: "github_copilot".to_string(), - reason: format!("Token exchange failed: {e}"), + match &e { + crate::llm::github_copilot_auth::GithubCopilotAuthError::AccessDenied + | crate::llm::github_copilot_auth::GithubCopilotAuthError::Expired => { + LlmError::AuthFailed { + provider: "github_copilot".to_string(), + } + } + _ => LlmError::RequestFailed { + provider: "github_copilot".to_string(), + reason: format!("Token exchange failed: {e}"), + }, } })?; @@ -157,54 +164,14 @@ impl GithubCopilotProvider { ); if status.as_u16() == 401 { - // Invalidate the cached session token and retry once with a - // fresh exchange — stale tokens are the most common 401 cause. - tracing::warn!("Copilot: 401 Unauthorized — invalidating session token, retrying"); + // Invalidate the cached session token so the next attempt + // (driven by RetryProvider) gets a fresh one. We don't retry + // inline to avoid nested retries with the outer RetryProvider. + tracing::warn!("Copilot: 401 Unauthorized — invalidating session token for retry"); self.token_manager.invalidate().await; - let fresh = self.token_manager.get_token().await.map_err(|e| { - tracing::warn!(error = %e, "Copilot: re-exchange after 401 failed"); - LlmError::RequestFailed { - provider: "github_copilot".to_string(), - reason: format!("Token re-exchange after 401 failed: {e}"), - } - })?; - let mut retry_req = self - .client - .post(&url) - .bearer_auth(fresh.expose_secret()) - .header("Content-Type", "application/json"); - for (key, value) in &self.extra_headers { - retry_req = retry_req.header(key.as_str(), value.as_str()); - } - let retry = - retry_req - .json(body) - .send() - .await - .map_err(|e| LlmError::RequestFailed { - provider: "github_copilot".to_string(), - reason: format!("Retry after 401 failed: {e}"), - })?; - if retry.status().is_success() { - let text = retry.text().await.map_err(|e| LlmError::RequestFailed { - provider: "github_copilot".to_string(), - reason: format!("Failed to read retry response body: {e}"), - })?; - return serde_json::from_str(&text).map_err(|e| { - let truncated = crate::agent::truncate_for_preview(&text, 512); - LlmError::InvalidResponse { - provider: "github_copilot".to_string(), - reason: format!("JSON parse error: {e}. Raw: {truncated}"), - } - }); - } - let retry_status = retry.status(); - tracing::warn!( - status = %retry_status, - "Copilot: 401 retry also failed" - ); - return Err(LlmError::AuthFailed { + return Err(LlmError::RequestFailed { provider: "github_copilot".to_string(), + reason: "HTTP 401 Unauthorized".to_string(), }); } if status.as_u16() == 429 { diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index 17e17ba1..f4313483 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -608,7 +608,8 @@ fn routine_create_schema(include_compatibility_aliases: bool) -> Value { } pub(crate) fn routine_create_parameters_schema() -> Value { - routine_create_schema(false) + static CACHE: OnceLock = OnceLock::new(); + CACHE.get_or_init(|| routine_create_schema(false)).clone() } fn routine_create_discovery_schema() -> Value { @@ -1014,7 +1015,8 @@ fn event_emit_schema(include_source_alias: bool) -> Value { } pub(crate) fn event_emit_parameters_schema() -> Value { - event_emit_schema(false) + static CACHE: OnceLock = OnceLock::new(); + CACHE.get_or_init(|| event_emit_schema(false)).clone() } fn event_emit_discovery_schema() -> Value { diff --git a/src/workspace/embedding_cache.rs b/src/workspace/embedding_cache.rs index 21d3c7c3..60c2eb08 100644 --- a/src/workspace/embedding_cache.rs +++ b/src/workspace/embedding_cache.rs @@ -3,14 +3,13 @@ //! Avoids redundant HTTP calls for identical texts by caching embeddings //! in memory keyed by `SHA-256(model_name + "\0" + text)`. //! -//! Follows the same cache pattern as `llm::response_cache::CachedProvider`: -//! `HashMap` + `last_accessed` tracking + manual LRU eviction. +//! Uses `lru::LruCache` for O(1) insertion, lookup, and eviction. -use std::collections::HashMap; +use std::num::NonZeroUsize; use std::sync::{Arc, Mutex}; -use std::time::Instant; use async_trait::async_trait; +use lru::LruCache; use sha2::{Digest, Sha256}; use crate::workspace::embeddings::{EmbeddingError, EmbeddingProvider}; @@ -22,8 +21,7 @@ pub struct EmbeddingCacheConfig { /// /// Approximate raw embedding payload: `max_entries × dimension × 4 bytes`. /// At 10,000 entries × 1536 floats ≈ 58 MB (payload only; actual memory - /// is higher due to HashMap buckets, `[u8; 32]` hash keys, `Vec`/`Instant` - /// per-entry overhead). + /// is higher due to per-entry overhead in the linked-list LRU). pub max_entries: usize, } @@ -35,11 +33,6 @@ impl Default for EmbeddingCacheConfig { } } -struct CacheEntry { - embedding: Vec, - last_accessed: Instant, -} - /// Embedding provider wrapper that caches results in memory. /// /// Thread-safe via `std::sync::Mutex`. The lock is **never held** @@ -47,8 +40,7 @@ struct CacheEntry { /// so a synchronous mutex is cheaper than `tokio::sync::Mutex`. pub struct CachedEmbeddingProvider { inner: Arc, - cache: Mutex>, - config: EmbeddingCacheConfig, + cache: Mutex>>, } impl CachedEmbeddingProvider { @@ -56,19 +48,18 @@ impl CachedEmbeddingProvider { /// /// `config.max_entries` is clamped to at least 1. pub fn new(inner: Arc, config: EmbeddingCacheConfig) -> Self { - let config = EmbeddingCacheConfig { - max_entries: config.max_entries.max(1), - }; - if config.max_entries > 100_000 { + let max_entries = config.max_entries.max(1); + if max_entries > 100_000 { tracing::warn!( - max_entries = config.max_entries, + max_entries, "Embedding cache size exceeds 100,000 entries; memory usage may be significant" ); } + // safety: max_entries >= 1 due to .max(1) above + let cap = NonZeroUsize::new(max_entries).expect("clamped to >= 1"); // safety: always >= 1 Self { inner, - cache: Mutex::new(HashMap::with_capacity(config.max_entries.min(1024))), - config, + cache: Mutex::new(LruCache::new(cap)), } } @@ -100,49 +91,6 @@ impl CachedEmbeddingProvider { hasher.update(text.as_bytes()); hasher.finalize().into() } - - /// Evict the least-recently-used entry if at capacity (single-entry path). - // TODO: O(n) scan per eviction. If max_entries grows large, switch to - // an ordered data structure (e.g. `IndexMap` with swap_remove, or a - // linked-list LRU like the `lru` crate). - fn evict_lru(cache: &mut HashMap<[u8; 32], CacheEntry>, max_entries: usize) { - while cache.len() >= max_entries { - let oldest_key = cache - .iter() - .min_by_key(|(_, entry)| entry.last_accessed) - .map(|(k, _)| *k); - - if let Some(k) = oldest_key { - cache.remove(&k); - } else { - break; - } - } - } - - /// Evict the `k` oldest entries in O(n) average time via partial selection. - /// - /// Used by `embed_batch` to avoid the O(n×m) cost of calling - /// `evict_lru` per insert. - fn evict_k_oldest(cache: &mut HashMap<[u8; 32], CacheEntry>, k: usize) { - if k == 0 || cache.is_empty() { - return; - } - if k >= cache.len() { - cache.clear(); - return; - } - // Partial selection: find the k oldest in O(n) average via - // select_nth_unstable_by_key, then remove the first k entries. - let mut entries: Vec<([u8; 32], Instant)> = cache - .iter() - .map(|(key, entry)| (*key, entry.last_accessed)) - .collect(); - entries.select_nth_unstable_by_key(k - 1, |(_, t)| *t); - for (key, _) in entries.into_iter().take(k) { - cache.remove(&key); - } - } } #[async_trait] @@ -162,39 +110,32 @@ impl EmbeddingProvider for CachedEmbeddingProvider { async fn embed(&self, text: &str) -> Result, EmbeddingError> { let key = self.cache_key(text); - // Check cache (short critical section) + // Check cache (short critical section). LruCache::get promotes the + // entry to most-recently-used automatically. { let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - if let Some(entry) = guard.get_mut(&key) { - entry.last_accessed = Instant::now(); + if let Some(embedding) = guard.get(&key) { tracing::trace!("embedding cache hit"); - return Ok(entry.embedding.clone()); + return Ok(embedding.clone()); } } // Lock released before HTTP call. // NOTE: Thundering herd — multiple concurrent callers with the same // uncached key will each call the inner provider. This is acceptable: - // embeddings are idempotent and the last writer wins in the HashMap. + // embeddings are idempotent and the last writer wins in the LruCache. let embedding = self.inner.embed(text).await?; - // Store result. Re-check under lock: another concurrent caller may - // have inserted this key while the lock was released for the HTTP call. + // Store result under lock. Re-check first: another concurrent caller + // may have already cached this key while the lock was released. { let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - if let Some(entry) = guard.get_mut(&key) { - // Thundering herd — another caller already cached it. - // Just touch timestamp; skip the clone. - entry.last_accessed = Instant::now(); + if guard.get(&key).is_some() { + // Thundering herd — another caller beat us. LruCache::get + // already promoted it to most-recently-used; skip the clone. + tracing::trace!("embedding cache: concurrent insert, skipping clone"); } else { - Self::evict_lru(&mut guard, self.config.max_entries); - guard.insert( - key, - CacheEntry { - embedding: embedding.clone(), - last_accessed: Instant::now(), - }, - ); + guard.push(key, embedding.clone()); } } @@ -214,11 +155,9 @@ impl EmbeddingProvider for CachedEmbeddingProvider { { let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - let now = Instant::now(); for (i, key) in keys.iter().enumerate() { - if let Some(entry) = guard.get_mut(key) { - entry.last_accessed = now; - results[i] = Some(entry.embedding.clone()); + if let Some(embedding) = guard.get(key) { + results[i] = Some(embedding.clone()); } else { miss_indices.push(i); } @@ -228,7 +167,6 @@ impl EmbeddingProvider for CachedEmbeddingProvider { if miss_indices.is_empty() { tracing::trace!(count = texts.len(), "embedding batch: all cache hits"); - // All slots populated from cache hits return results .into_iter() .enumerate() @@ -260,29 +198,18 @@ impl EmbeddingProvider for CachedEmbeddingProvider { "embedding batch: partial cache" ); - // Cache FIRST (clone only the cacheable subset), then move originals - // into results. This avoids cloning capacity-skipped embeddings entirely. + // Cache only the last `cap` new embeddings — caching more than the + // cache capacity wastes clone work on entries that are immediately evicted. { let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - let cacheable = miss_indices.len().min(self.config.max_entries); - let skip = miss_indices.len() - cacheable; - let need_to_evict = (guard.len() + cacheable).saturating_sub(self.config.max_entries); - if need_to_evict > 0 { - Self::evict_k_oldest(&mut guard, need_to_evict); - } - let now = Instant::now(); + let cap = guard.cap().get(); + let skip = miss_indices.len().saturating_sub(cap); for (&orig_idx, emb) in miss_indices[skip..].iter().zip(&new_embeddings[skip..]) { - guard.insert( - keys[orig_idx], - CacheEntry { - embedding: emb.clone(), - last_accessed: now, - }, - ); + guard.push(keys[orig_idx], emb.clone()); } } - // Move originals into results (zero-copy for all, including cached ones). + // Move originals into results (zero-copy). for (orig_idx, emb) in miss_indices.iter().copied().zip(new_embeddings) { results[orig_idx] = Some(emb); }