diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 50d2363..7c4712f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -227,6 +227,47 @@ fn lock_workspace_root(state: &tauri::State<'_, WorkspaceState>) -> Result e, + Err(_) => return, + }; + + let cutoff = std::time::SystemTime::now() - Duration::from_secs(7 * 24 * 3600); + let mut removed = 0u32; + + while let Ok(Some(entry)) = entries.next_entry().await { + let meta = match entry.metadata().await { + Ok(m) => m, + Err(_) => continue, + }; + if !meta.is_dir() { + continue; + } + let modified = meta.modified().unwrap_or(UNIX_EPOCH); + if modified < cutoff { + if let Err(e) = tokio::fs::remove_dir_all(entry.path()).await { + eprintln!( + "[cleanup] failed to remove expired tool-results dir {}: {}", + entry.path().display(), + e + ); + } else { + removed += 1; + } + } + } + + if removed > 0 { + eprintln!( + "[cleanup] removed {} expired tool-results session dir(s) from {}", + removed, + results_dir.display() + ); + } +} + #[tauri::command] async fn open_workspace( root: String, @@ -327,6 +368,12 @@ async fn open_workspace( eprintln!("[thought_reconcile] skipped: {e}"); } }); + + let cleanup_root = canonical_root.clone(); + tokio::spawn(async move { + cleanup_expired_tool_results(&cleanup_root).await; + }); + Ok(nodes) } diff --git a/src-tauri/src/llm/agent_loop.rs b/src-tauri/src/llm/agent_loop.rs index 4b70dc2..9f5809f 100644 --- a/src-tauri/src/llm/agent_loop.rs +++ b/src-tauri/src/llm/agent_loop.rs @@ -18,6 +18,7 @@ use super::approval::ToolApprovalState; use super::context_guard::{ContextGuard, PrecomputedSummary}; use super::memory; use super::provider::{LlmProvider, NormalizedToolCall}; +use super::tool_result_processor::{self, ToolResultProcessor}; use super::{LlmChatMessage, LlmToolCall, LlmToolCallFunction}; use crate::tools::context::ToolContextFactory; use crate::tools::registry::ToolRegistry; @@ -40,7 +41,9 @@ WEB: When the user provides a specific URL (http/https link), always use `web.re Only use `web.search` when no URL is given and you need to find relevant pages by keyword. \ PDF: When `web.read_page` results mention a PDF link or the page is an academic paper with a PDF download, \ immediately call `web.read_pdf` with the PDF URL to extract the full text — do NOT tell the user to download it themselves. \ -RESULT MATCHING: Each tool result is prefixed with [call:ID] to help you match results to calls when the same tool is invoked multiple times."; +RESULT MATCHING: Each tool result is prefixed with [call:ID] to help you match results to calls when the same tool is invoked multiple times. \ +RECALL: When a tool result shows [summarized from N chars | ref:XXX], the full raw content is stored on disk. \ +If the summary lacks detail you need, call `tool.recall` with that ref ID to retrieve the original content."; /// Agent Loop 上限配置；任一项达到上限即终止循环并 emit `llm:agent-done`。 #[allow(dead_code)] @@ -49,13 +52,17 @@ pub struct AgentLoopConfig { pub max_tool_calls: u16, /// 每轮模型流式请求的默认超时（毫秒）；现代码从调用点的 ai.request.timeout_ms 传入。 pub timeout_ms: u64, - /// 整轮中累计追加给模型的 tool 结果总字符数上限（防止上下文爆炸）。 - pub max_tool_result_chars: usize, + /// Per-result truncation threshold (chars). Results exceeding this are + /// truncated with a marker. Not used as a loop termination condition. + pub max_single_result_chars: usize, /// Iter 5 #4: 本轮 agent loop 内 ToolContext.nesting_depth 的赋值。 /// 主对话默认 0；skill 子轮次为 1（由 [`crate::skills::runtime::run_skill_with_depth`] 设置）。 pub nesting_depth: u8, /// Provider context window size (tokens). Used by ContextGuard to trim history. pub max_context_tokens: Option, + /// Tool results longer than this (chars) are summarized before entering + /// the message array. Set to 0 to disable front-load summarization. + pub summarize_threshold: usize, } impl Default for AgentLoopConfig { @@ -63,9 +70,10 @@ impl Default for AgentLoopConfig { Self { max_tool_calls: 25, timeout_ms: 60_000, - max_tool_result_chars: 24_000, + max_single_result_chars: 12_000, nesting_depth: 0, max_context_tokens: None, + summarize_threshold: tool_result_processor::DEFAULT_SUMMARIZE_THRESHOLD, } } } @@ -129,31 +137,46 @@ pub async fn run_agent_stream( memory_manager: SharedMemoryManager, ) -> String { let mut messages = initial_messages; - let mut total_tool_result_chars: usize = 0; let mut tool_call_count: u16 = 0; + let effective_context_tokens = config.max_context_tokens + .or_else(|| provider.model_context_window().map(|w| w as u64)); let context_guard = if config.nesting_depth > 0 { - ContextGuard::new(config.max_context_tokens) + ContextGuard::new(effective_context_tokens) } else { - ContextGuard::with_provider(config.max_context_tokens, provider.clone()) + ContextGuard::with_provider(effective_context_tokens, provider.clone()) }; let mut loop_detector = LoopDetector::new(); let mut pending_summary: Option>> = None; + let results_dir = if config.nesting_depth == 0 { + Some(workspace_root.join(".knowforge").join("tool-results")) + } else { + None + }; + let result_processor: Option = if config.summarize_threshold > 0 { + Some(ToolResultProcessor::new( + provider.clone(), + config.summarize_threshold, + results_dir, + session_id.clone(), + )) + } else { + None + }; + let mut iteration: u32 = 0; loop { iteration += 1; let est_tokens: usize = messages.iter().map(|m| m.content.len() / 3).sum(); eprintln!( - "[agent_loop] session={} iter={} msgs={} est_tokens={} tool_calls_so_far={}/{} result_chars={}/{}", + "[agent_loop] session={} iter={} msgs={} est_tokens={} tool_calls_so_far={}/{}", &session_id[..8.min(session_id.len())], iteration, messages.len(), est_tokens, tool_call_count, config.max_tool_calls, - total_tool_result_chars, - config.max_tool_result_chars, ); if cancel.is_cancelled() { @@ -261,7 +284,7 @@ pub async fn run_agent_stream( } // 5. 并行执行工具（跳过循环调用；每个工具有独立超时，支持取消） - let tool_timeout = Duration::from_millis(config.timeout_ms); + let default_tool_timeout = Duration::from_millis(config.timeout_ms); let results = join_all(normalized_calls.iter().enumerate().map(|(idx, tc)| { let skip = looped.get(idx).copied().unwrap_or(false); let cancel = cancel.clone(); @@ -280,6 +303,11 @@ pub async fn run_agent_stream( if skip { return (Err(format!("loop detected: '{}' called too many times with same arguments", tc.name)), 0u64); } + let tool_timeout = registry + .get(&tc.name) + .and_then(|t| t.timeout_ms()) + .map(Duration::from_millis) + .unwrap_or(default_tool_timeout); let nesting_depth = config.nesting_depth; let exec_start = std::time::Instant::now(); let result = tokio::select! { @@ -354,27 +382,70 @@ pub async fn run_agent_stream( }); // 6. 把每个 tool 结果以 role=tool 的消息追加到历史 - for (i, tc) in normalized_calls.iter().enumerate() { - let raw_content = match results.get(i) { + // When a result processor is available, long results are summarized + // in parallel before being appended (front-load compression). + let user_goal = result_processor + .as_ref() + .and_then(|_| tool_result_processor::extract_user_goal(&messages)); + + let raw_contents: Vec = normalized_calls + .iter() + .enumerate() + .map(|(i, _tc)| match results.get(i) { Some((Ok(val), _)) => val.to_string(), Some((Err(e), _)) => format!("error: {}", e), None => "error: no result".to_string(), + }) + .collect(); + + let processed: Vec> = + if let Some(ref proc) = result_processor { + let futs: Vec<_> = normalized_calls + .iter() + .enumerate() + .map(|(i, tc)| { + let proc = proc.clone(); + let name = tc.name.clone(); + let id = tc.id.clone(); + let raw = raw_contents[i].clone(); + let goal = user_goal.clone(); + async move { + Some(proc.process(&name, &id, &raw, goal.as_deref()).await) + } + }) + .collect(); + join_all(futs).await + } else { + vec![None; normalized_calls.len()] }; - let fenced = fence_if_external(&tc.name, &raw_content); - let prefixed = format!("[call:{}] {}", tc.id, fenced); - let remaining = config - .max_tool_result_chars - .saturating_sub(total_tool_result_chars); - let content = if prefixed.len() > remaining { - let mut end = remaining; - while end > 0 && !prefixed.is_char_boundary(end) { - end -= 1; + + for (i, tc) in normalized_calls.iter().enumerate() { + let effective_content = if let Some(Some(pr)) = processed.get(i) { + if pr.was_summarized { + eprintln!( + "[agent_loop] session={} tool={} summarized {}->{} chars", + &session_id[..8.min(session_id.len())], + tc.name, + pr.original_len, + pr.content.len(), + ); } - prefixed[..end].to_string() + pr.content.clone() + } else { + raw_contents[i].clone() + }; + + let fenced = fence_if_external(&tc.name, &effective_content); + let prefixed = format!("[call:{}] {}", tc.id, fenced); + let content = if prefixed.len() > config.max_single_result_chars { + let end = find_char_boundary(&prefixed, config.max_single_result_chars); + format!( + "{}\n[… truncated, showing first {} of {} chars]", + &prefixed[..end], end, prefixed.len() + ) } else { prefixed }; - total_tool_result_chars = total_tool_result_chars.saturating_add(content.len()); let mut tool_msg = provider.build_tool_result_message(&tc.id, &tc.name, &content); @@ -420,14 +491,8 @@ pub async fn run_agent_stream( } // 7b. Budget exhausted → graceful summary instead of silent truncation - if tool_call_count >= config.max_tool_calls - || total_tool_result_chars >= config.max_tool_result_chars - { - let reason = if tool_call_count >= config.max_tool_calls { - format!("tool_calls {}/{}", tool_call_count, config.max_tool_calls) - } else { - format!("result_chars {}/{}", total_tool_result_chars, config.max_tool_result_chars) - }; + if tool_call_count >= config.max_tool_calls { + let reason = format!("tool_calls {}/{}", tool_call_count, config.max_tool_calls); eprintln!( "[agent_loop] session={} iter={} BUDGET EXHAUSTED ({}), requesting final summary...", &session_id[..8.min(session_id.len())], iteration, reason, @@ -469,7 +534,7 @@ pub async fn run_agent_stream( } let pressure = context_guard.budget_pressure(&messages); - if pressure > 0.7 { + if pressure > 0.5 { eprintln!( "[agent_loop] session={} iter={} context pressure={:.2}, pre-summarizing", &session_id[..8.min(session_id.len())], iteration, pressure, @@ -577,6 +642,7 @@ pub(crate) async fn execute_tool( app_bundle_resource_dir, nesting_depth, ); + ctx.session_id = session_id.to_string(); ctx.call_id = Some(tc.id.clone()); ctx.provider = provider; if let Some(ec) = app.try_state::>() { @@ -661,6 +727,14 @@ fn fence_if_external(tool_name: &str, content: &str) -> String { } } +fn find_char_boundary(s: &str, target: usize) -> usize { + let mut end = target.min(s.len()); + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + end +} + #[allow(dead_code)] pub fn manifest_to_tool(manifest: &ToolManifest) -> Value { json!({ @@ -885,6 +959,69 @@ mod fence_tests { } } +#[cfg(test)] +mod truncation_tests { + use super::*; + + #[test] + fn single_result_truncation_adds_marker() { + let long = "x".repeat(15_000); + let config = AgentLoopConfig { + max_single_result_chars: 12_000, + ..Default::default() + }; + let prefixed = format!("[call:abc] {}", long); + let content = if prefixed.len() > config.max_single_result_chars { + let end = find_char_boundary(&prefixed, config.max_single_result_chars); + format!( + "{}\n[… truncated, showing first {} of {} chars]", + &prefixed[..end], end, prefixed.len() + ) + } else { + prefixed.clone() + }; + assert!(content.contains("[… truncated, showing first 12000 of")); + assert!(content.len() < prefixed.len()); + } + + #[test] + fn short_result_passes_through() { + let short = "hello world"; + let config = AgentLoopConfig { + max_single_result_chars: 12_000, + ..Default::default() + }; + let prefixed = format!("[call:abc] {}", short); + let content = if prefixed.len() > config.max_single_result_chars { + let end = find_char_boundary(&prefixed, config.max_single_result_chars); + format!( + "{}\n[… truncated, showing first {} of {} chars]", + &prefixed[..end], end, prefixed.len() + ) + } else { + prefixed.clone() + }; + assert_eq!(content, prefixed); + assert!(!content.contains("truncated")); + } + + #[test] + fn find_char_boundary_respects_utf8() { + let s = "你好世界测试数据"; // 8 CJK chars, 24 bytes + // target=5 falls in the middle of a 3-byte char + let boundary = find_char_boundary(s, 5); + assert!(s.is_char_boundary(boundary)); + assert!(boundary <= 5); + assert_eq!(boundary, 3); // first char is 3 bytes + + // target=0 + assert_eq!(find_char_boundary(s, 0), 0); + + // target beyond length + assert_eq!(find_char_boundary(s, 100), s.len()); + } +} + #[cfg(test)] mod loop_detector_tests { use super::*; diff --git a/src-tauri/src/llm/context_guard.rs b/src-tauri/src/llm/context_guard.rs index 65ed013..454e9a7 100644 --- a/src-tauri/src/llm/context_guard.rs +++ b/src-tauri/src/llm/context_guard.rs @@ -12,9 +12,21 @@ const MAX_SUMMARY_INPUT_CHARS: usize = 6000; const MAX_CONTENT_PER_MESSAGE: usize = 500; const SUMMARY_SYSTEM: &str = "\ -Summarize the following conversation excerpt in 2-3 concise sentences. \ -Focus on: what the user asked, what tools were called, and key findings. \ -Output only the summary, nothing else."; +Summarize the following conversation into a structured summary.\n\ +You MUST preserve:\n\ +- What the user wants to accomplish (goal)\n\ +- Any constraints the user specified\n\ +- What tools were called and their key findings\n\ +- Decisions made so far\n\ +- Open questions or unresolved issues\n\ +- What should happen next\n\n\ +Format:\n\ +[Goal] ...\n\ +[Findings] ...\n\ +[Decisions] ...\n\ +[Open] ...\n\ +[Next] ...\n\n\ +Be concise. Each section 1-2 sentences max. Output only the summary."; #[derive(Clone)] pub struct ContextGuard { @@ -25,7 +37,7 @@ pub struct ContextGuard { pub struct PrecomputedSummary { pub summary_text: String, - pub original_msg_count: usize, + pub summarized_up_to: usize, } impl ContextGuard { @@ -120,7 +132,7 @@ impl ContextGuard { Some(PrecomputedSummary { summary_text, - original_msg_count: messages.len(), + summarized_up_to: tail_boundary, }) } @@ -129,12 +141,11 @@ impl ContextGuard { messages: &mut Vec, cached: &PrecomputedSummary, ) -> bool { - if messages.len() != cached.original_msg_count { + if messages.len() < cached.summarized_up_to { return false; } - let tail_boundary = Self::find_tail_boundary(messages); - let removable_indices: Vec = (0..tail_boundary.min(messages.len())) + let removable_indices: Vec = (0..cached.summarized_up_to.min(messages.len())) .filter(|&i| messages[i].role != "system") .collect(); @@ -205,18 +216,54 @@ impl ContextGuard { fn phase1_remove_tool_results(&self, messages: &mut Vec, budget: usize) { let tail_boundary = Self::find_tail_boundary(messages); + + // Pass 1: trim raw (non-summarized) tool results let mut i = 0; while i < tail_boundary.min(messages.len()) && Self::estimate_total(messages) > budget { - if messages[i].role == "system" { - i += 1; - continue; - } - if messages[i].role == "tool" { - messages.remove(i); - continue; + if messages[i].role == "tool" && messages[i].content.len() > 40 { + if !messages[i].content.starts_with(super::tool_result_processor::SUMMARIZED_MARKER) { + let orig_len = messages[i].content.len(); + messages[i].content = format!( + "[tool result trimmed, was {} chars]", + orig_len + ); + } } i += 1; } + + // Pass 2: if still over budget, degrade summarized results that have + // a stored ref — the model can still recall them via tool.recall. + if Self::estimate_total(messages) > budget { + let mut i = 0; + while i < tail_boundary.min(messages.len()) && Self::estimate_total(messages) > budget { + if messages[i].role == "tool" + && messages[i].content.starts_with(super::tool_result_processor::SUMMARIZED_MARKER) + { + if let Some(stored_marker) = Self::extract_stored_ref_marker(&messages[i].content) { + messages[i].content = stored_marker; + } + } + i += 1; + } + } + } + + fn extract_stored_ref_marker(content: &str) -> Option { + let ref_start = content.find("| ref:")?; + let ref_value_start = ref_start + "| ref:".len(); + let ref_end = content[ref_value_start..].find(']')?; + let ref_id = &content[ref_value_start..ref_value_start + ref_end]; + + let chars_start = super::tool_result_processor::SUMMARIZED_MARKER.len(); + let chars_end = content[chars_start..].find(' ').unwrap_or(0); + let orig_chars = &content[chars_start..chars_start + chars_end]; + + Some(format!( + "{}{}]", + super::tool_result_processor::STORED_REF_MARKER, + format!("{}, was {} chars", ref_id, orig_chars) + )) } async fn phase1_5_summarize(&self, messages: &mut Vec, budget: usize) { @@ -397,8 +444,12 @@ mod tests { } #[test] - fn removes_old_tool_results_first() { - let guard = ContextGuard::new(Some(50)); + fn degrades_old_tool_results_first() { + // budget = max_tokens(692) - RESERVE(512) = 180. + // Total before trim ≈ 206 > 180 → triggers Phase 1. + // After degrading 300-char tool results to ~36-char placeholders, total ≈ 140 < 180. + // Phase 2 does not kick in. + let guard = ContextGuard::new(Some(692)); let mut msgs = vec![ sys("system prompt"), user("q1"), @@ -411,9 +462,160 @@ mod tests { assistant("a3"), tool(&"z".repeat(30)), ]; + let original_len = msgs.len(); guard.trim_if_needed(&mut msgs); assert!(msgs.iter().any(|m| m.role == "system")); assert!(msgs.iter().any(|m| m.content == "a3")); + // tool messages are degraded, not removed + let tool_count = msgs.iter().filter(|m| m.role == "tool").count(); + assert!(tool_count > 0, "tool messages should be degraded, not deleted"); + // at least one tool result should be trimmed to placeholder + let has_placeholder = msgs.iter().any(|m| m.role == "tool" && m.content.starts_with("[tool result trimmed")); + assert!(has_placeholder); + // message count stays the same (degrade, not delete) + assert_eq!(msgs.len(), original_len); + } + + #[test] + fn degraded_tool_result_preserves_structure() { + // budget = 620 - 512 = 108. Total with 500-char tool ≈ 164 > 108 → triggers trim. + // After degrading to placeholder (~36 chars), total ≈ 48 < 108. Phase 2 skipped. + let guard = ContextGuard::new(Some(620)); + let mut msgs = vec![ + sys("sys"), + user("q1"), + assistant("a1"), + tool(&"r".repeat(500)), + user("q2"), + assistant("a2"), + user("q3"), + assistant("a3"), + ]; + guard.trim_if_needed(&mut msgs); + // tool message still present + assert!(msgs.iter().any(|m| m.role == "tool")); + // its content is a placeholder + let tool_msg = msgs.iter().find(|m| m.role == "tool").unwrap(); + assert!(tool_msg.content.starts_with("[tool result trimmed, was 500 chars]")); + } + + #[test] + fn small_tool_result_not_degraded() { + // Budget tight but tool result is tiny (2 chars < 40 threshold). + // Phase 1 skips it, Phase 2 degrades user/assistant content instead. + let guard = ContextGuard::new(Some(30)); + let small_content = "ok"; + let mut msgs = vec![ + sys("sys"), + user(&"long ".repeat(50)), + assistant("a1"), + tool(small_content), + user("q2"), + assistant("a2"), + ]; + guard.trim_if_needed(&mut msgs); + let tool_msg = msgs.iter().find(|m| m.role == "tool").unwrap(); + assert_eq!(tool_msg.content, small_content); + } + + #[test] + fn phase1_skips_already_summarized_tool_result() { + // Budget very tight: 580 - 512 = 68 token budget. + // Two tool results compete for space. The summarized one (with + // SUMMARIZED_MARKER) should survive; the raw one should be trimmed. + let guard = ContextGuard::new(Some(580)); + let summarized = format!( + "{}500 chars]\nKey finding: X is important.", + super::super::tool_result_processor::SUMMARIZED_MARKER, + ); + let mut msgs = vec![ + sys("sys"), + user("q1"), + assistant("a1"), + tool(&summarized), + tool(&"x".repeat(300)), + user("q2"), + assistant("a2"), + user("q3"), + assistant("a3"), + ]; + guard.trim_if_needed(&mut msgs); + let tool_msgs: Vec<&LlmChatMessage> = msgs.iter().filter(|m| m.role == "tool").collect(); + // The summarized tool result should be preserved as-is + assert!(tool_msgs.iter().any(|m| m.content.contains("Key finding"))); + // The raw tool result should be trimmed + assert!(tool_msgs.iter().any(|m| m.content.starts_with("[tool result trimmed"))); + } + + #[test] + fn phase1_pass2_degrades_summarized_with_ref_to_stored_marker() { + let guard = ContextGuard::new(Some(10000)); + let summarized_with_ref = format!( + "{}12474 chars | ref:019717ab]\n{}", + super::super::tool_result_processor::SUMMARIZED_MARKER, + "x".repeat(600), + ); + let summarized_no_ref = format!( + "{}500 chars]\nAnother finding.", + super::super::tool_result_processor::SUMMARIZED_MARKER, + ); + let mut msgs = vec![ + sys("sys"), + user("q1"), + assistant("a1"), + tool(&summarized_with_ref), + tool(&summarized_no_ref), + tool(&"r".repeat(500)), + user("q2"), + assistant("a2"), + user("q3"), + assistant("a3"), + ]; + + let total_before = ContextGuard::estimate_total(&msgs); + // Budget that pass 1 alone can't satisfy (need pass 2 too). + // Pass 1 saves ~120 tokens by trimming the raw 500-char result. + // Pass 2 saves ~150 tokens by degrading summarized_with_ref. + let budget = total_before - 200; + guard.phase1_remove_tool_results(&mut msgs, budget); + + let tool_msgs: Vec<&LlmChatMessage> = msgs.iter().filter(|m| m.role == "tool").collect(); + assert!( + tool_msgs.iter().any(|m| m.content.starts_with("[tool result trimmed")), + "pass 1 should have trimmed the raw result" + ); + assert!( + tool_msgs.iter().any(|m| m.content.starts_with( + super::super::tool_result_processor::STORED_REF_MARKER + )), + "pass 2 should have degraded summarized-with-ref, got: {:?}", + tool_msgs.iter().map(|m| &m.content).collect::>() + ); + let stored = tool_msgs.iter().find(|m| m.content.contains("019717ab")).unwrap(); + assert!(stored.content.contains("was 12474 chars")); + assert!( + tool_msgs.iter().any(|m| m.content.contains("Another finding")), + "summarized-without-ref should be preserved" + ); + } + + #[test] + fn extract_stored_ref_marker_works() { + let content = format!( + "{}12474 chars | ref:019717ab]\nKey finding.", + super::super::tool_result_processor::SUMMARIZED_MARKER, + ); + let marker = ContextGuard::extract_stored_ref_marker(&content).unwrap(); + assert_eq!(marker, "[tool result stored | ref:019717ab, was 12474 chars]"); + } + + #[test] + fn extract_stored_ref_marker_returns_none_without_ref() { + let content = format!( + "{}500 chars]\nSome summary.", + super::super::tool_result_processor::SUMMARIZED_MARKER, + ); + assert!(ContextGuard::extract_stored_ref_marker(&content).is_none()); } #[test] @@ -549,4 +751,79 @@ mod tests { let degraded = msgs.iter().any(|m| m.content == "[content trimmed]"); assert!(degraded); } + + #[test] + fn cached_summary_applies_when_messages_grew() { + let guard = ContextGuard::new(Some(4096)); + // Simulate a cached summary that covered messages 0..5 + // (tail_boundary was 5 when pre_summarize ran). + let cached = PrecomputedSummary { + summary_text: "User asked about X, tool returned Y.".to_string(), + summarized_up_to: 5, + }; + let mut msgs = vec![ + sys("system prompt"), // 0: system, skipped + user("q1"), // 1: removed + assistant("a1"), // 2: removed + tool("result1"), // 3: removed + user("q2"), // 4: removed + // --- summarized_up_to = 5 --- + assistant("a2"), // 5: kept (after boundary) + user("q3"), // 6: kept (new since snapshot) + assistant("a3"), // 7: kept (new since snapshot) + ]; + let applied = guard.apply_cached_summary(&mut msgs, &cached); + assert!(applied); + // system prompt preserved, 4 non-system removed, summary inserted, 3 tail kept + assert!(msgs.iter().any(|m| m.content.contains("User asked about X"))); + assert!(msgs.iter().any(|m| m.content == "system prompt")); + assert!(msgs.iter().any(|m| m.content == "a3")); + assert!(!msgs.iter().any(|m| m.content == "q1")); + } + + #[test] + fn cached_summary_rejects_when_messages_shrunk() { + let guard = ContextGuard::new(Some(4096)); + let cached = PrecomputedSummary { + summary_text: "summary".to_string(), + summarized_up_to: 10, + }; + // Only 5 messages — fewer than summarized_up_to + let mut msgs = vec![ + sys("sys"), + user("q1"), + assistant("a1"), + user("q2"), + assistant("a2"), + ]; + let applied = guard.apply_cached_summary(&mut msgs, &cached); + assert!(!applied); + assert_eq!(msgs.len(), 5); + } + + #[test] + fn cached_summary_preserves_system_messages() { + let guard = ContextGuard::new(Some(4096)); + let cached = PrecomputedSummary { + summary_text: "conversation summary".to_string(), + summarized_up_to: 4, + }; + let mut msgs = vec![ + sys("core system prompt"), // 0: system, preserved + sys("extra system"), // 1: system, preserved + user("q1"), // 2: removed + assistant("a1"), // 3: removed + // --- summarized_up_to = 4 --- + user("q2"), // 4: kept + assistant("a2"), // 5: kept + ]; + let applied = guard.apply_cached_summary(&mut msgs, &cached); + assert!(applied); + let system_count = msgs.iter().filter(|m| m.role == "system").count(); + // 2 original system msgs + 1 summary system msg = 3 + assert_eq!(system_count, 3); + assert!(msgs.iter().any(|m| m.content == "core system prompt")); + assert!(msgs.iter().any(|m| m.content == "extra system")); + assert!(msgs.iter().any(|m| m.content.contains("conversation summary"))); + } } diff --git a/src-tauri/src/llm/mod.rs b/src-tauri/src/llm/mod.rs index f1d8faf..40d124f 100644 --- a/src-tauri/src/llm/mod.rs +++ b/src-tauri/src/llm/mod.rs @@ -6,6 +6,7 @@ pub(crate) mod context_guard; pub(crate) mod planning; pub(crate) mod provider; pub(crate) mod provider_impl; +pub(crate) mod tool_result_processor; pub mod memory; pub use provider::{ @@ -997,6 +998,7 @@ pub async fn start_chat_stream( } let tools_json = provider.convert_tools(&manifests); let loop_config = agent_loop::AgentLoopConfig { + timeout_ms: ai.request.timeout_ms, max_context_tokens: ai.request.max_context_tokens, ..Default::default() }; diff --git a/src-tauri/src/llm/provider.rs b/src-tauri/src/llm/provider.rs index 854b8ef..c306422 100644 --- a/src-tauri/src/llm/provider.rs +++ b/src-tauri/src/llm/provider.rs @@ -61,6 +61,12 @@ pub trait LlmProvider: Send + Sync { fn provider_name(&self) -> &'static str; fn is_remote(&self) -> bool; + + /// Return the model's context window size in tokens, if known. + /// Used by ContextGuard as fallback when user doesn't configure max_context_tokens. + fn model_context_window(&self) -> Option { + None + } } pub fn resolve_model_name(last_used: Option<&str>, default_model: &str) -> Option { diff --git a/src-tauri/src/llm/provider_impl.rs b/src-tauri/src/llm/provider_impl.rs index 21a647c..ce35d87 100644 --- a/src-tauri/src/llm/provider_impl.rs +++ b/src-tauri/src/llm/provider_impl.rs @@ -239,8 +239,9 @@ impl LlmProvider for UnifiedProvider { let msg_count = messages.len(); let est_tokens: usize = messages_json.iter().map(|v| v.to_string().len() / 3).sum(); + let idle_timeout = std::time::Duration::from_millis(self.timeout_ms); eprintln!( - "[provider] session={} chat_stream start: model={} msgs={} est_tokens={} timeout={}ms tools={}", + "[provider] session={} chat_stream start: model={} msgs={} est_tokens={} idle_timeout={}ms tools={}", &session_id[..8.min(session_id.len())], self.model, msg_count, @@ -250,12 +251,10 @@ impl LlmProvider for UnifiedProvider { ); let request_start = std::time::Instant::now(); - let req = self.build_auth_headers( - self.client - .post(&url) - .timeout(std::time::Duration::from_millis(self.timeout_ms)), - ) - .json(&body); + // No request-level timeout — use idle detection in the SSE loop instead. + // The connect timeout on the shared client (15s) still protects against + // unreachable hosts. + let req = self.build_auth_headers(self.client.post(&url)).json(&body); let resp = match req.send().await { Ok(r) => r, Err(e) => { @@ -294,6 +293,7 @@ impl LlmProvider for UnifiedProvider { let mut line_buf = String::new(); let mut accumulated_content = String::new(); let mut pending_tool_calls: Vec = Vec::new(); + let tools_provided = tools.as_ref().map_or(false, |t| !t.is_empty()); let mut last_chunk_at = std::time::Instant::now(); let mut chunk_count: u64 = 0; @@ -311,6 +311,23 @@ impl LlmProvider for UnifiedProvider { emit_error(app, session_id, Some("cancelled"), "Request aborted"); return Err("cancelled".to_string()); } + _ = tokio::time::sleep_until(tokio::time::Instant::from_std(last_chunk_at + idle_timeout)) => { + let idle_secs = last_chunk_at.elapsed().as_secs_f64(); + let msg = format!( + "SSE idle timeout: no data received for {:.0}s", + idle_secs, + ); + eprintln!( + "[provider] session={} {} after {:.1}s total, {} chunks, content_len={}", + &session_id[..8.min(session_id.len())], + msg, + request_start.elapsed().as_secs_f64(), + chunk_count, + accumulated_content.len(), + ); + emit_error(app, session_id, Some("idle_timeout"), &msg); + return Err(msg); + } item = stream.next() => { match item { None => { @@ -378,25 +395,36 @@ impl LlmProvider for UnifiedProvider { emit_chunk(app, session_id, &text); } if let Some(tc_deltas) = delta.tool_calls { - for tc_delta in tc_deltas { - let idx = tc_delta.index; - while pending_tool_calls.len() <= idx { - pending_tool_calls.push(PendingToolCall { - id: String::new(), - name: String::new(), - arguments_buf: String::new(), - }); - } - let pending = &mut pending_tool_calls[idx]; - if let Some(id) = tc_delta.id { - pending.id = id; - } - if let Some(func) = tc_delta.function { - if let Some(name) = func.name { - pending.name = name; + if tools_provided { + for tc_delta in tc_deltas { + let idx = tc_delta.index; + while pending_tool_calls.len() <= idx { + pending_tool_calls.push(PendingToolCall { + id: String::new(), + name: String::new(), + arguments_buf: String::new(), + }); + } + let pending = &mut pending_tool_calls[idx]; + if let Some(id) = tc_delta.id { + pending.id = id; + } + if let Some(func) = tc_delta.function { + if let Some(name) = func.name { + pending.name = name; + } + if let Some(args) = func.arguments { + pending.arguments_buf.push_str(&args); + } } - if let Some(args) = func.arguments { - pending.arguments_buf.push_str(&args); + } + } else { + for tc_delta in tc_deltas { + if let Some(func) = tc_delta.function { + if let Some(args) = func.arguments { + accumulated_content.push_str(&args); + emit_chunk(app, session_id, &args); + } } } } @@ -584,6 +612,60 @@ impl LlmProvider for UnifiedProvider { fn is_remote(&self) -> bool { self.is_remote } + + fn model_context_window(&self) -> Option { + infer_context_window(&self.model) + } +} + +const MAX_INFERRED_CONTEXT: usize = 524_288; + +fn infer_context_window(model: &str) -> Option { + let m = model.to_lowercase(); + + let raw = if m.contains("qwen-long") { + 10_000_000 + } else if m.contains("qwen-max") || m.contains("qwen-plus") { + 131_072 + } else if m.contains("qwen-turbo") { + 131_072 + } else if m.contains("qwen2") || m.contains("qwen3") { + 131_072 + } else if m.contains("qwen") { + 32_768 + } else if m.contains("gpt-4o") || m.contains("gpt-4-turbo") { + 128_000 + } else if m.contains("gpt-4") { + 8_192 + } else if m.contains("gpt-3.5") { + 16_385 + } else if m.contains("llama-4") || m.contains("llama4") { + if m.contains("scout") { + 10_000_000 + } else { + 1_000_000 + } + } else if m.contains("llama-3") || m.contains("llama3") { + if m.contains("3.1") || m.contains("3.2") || m.contains("3.3") { + 128_000 + } else { + 8_192 + } + } else if m.contains("llama2") || m.contains("llama-2") { + 4_096 + } else if m.contains("deepseek") { + 128_000 + } else if m.contains("gemma4") || m.contains("gemma-4") { + 262_144 + } else if m.contains("gemma3") || m.contains("gemma-3") { + 128_000 + } else if m.contains("gemma") { + 8_192 + } else { + return None; + }; + + Some(raw.min(MAX_INFERRED_CONTEXT)) } #[cfg(test)] @@ -708,6 +790,34 @@ mod tests { assert_eq!(tools[0]["function"]["name"], "web-search"); } + #[test] + fn infer_context_window_known_models() { + assert_eq!(infer_context_window("qwen-plus"), Some(131_072)); + assert_eq!(infer_context_window("qwen-turbo-latest"), Some(131_072)); + assert_eq!(infer_context_window("qwen3-235b-a22b"), Some(131_072)); + assert_eq!(infer_context_window("gpt-4o"), Some(128_000)); + assert_eq!(infer_context_window("gpt-4"), Some(8_192)); + assert_eq!(infer_context_window("gpt-3.5-turbo"), Some(16_385)); + assert_eq!(infer_context_window("deepseek-chat"), Some(128_000)); + assert_eq!(infer_context_window("llama3.3-70b"), Some(128_000)); + assert_eq!(infer_context_window("llama3-8b"), Some(8_192)); + assert_eq!(infer_context_window("gemma3-27b"), Some(128_000)); + } + + #[test] + fn infer_context_window_clamped_to_max() { + // qwen-long has 10M native window but should be clamped to 512K + assert_eq!(infer_context_window("qwen-long"), Some(MAX_INFERRED_CONTEXT)); + assert_eq!(infer_context_window("llama4-scout"), Some(MAX_INFERRED_CONTEXT)); + assert_eq!(MAX_INFERRED_CONTEXT, 524_288); + } + + #[test] + fn infer_context_window_unknown_returns_none() { + assert_eq!(infer_context_window("my-custom-model"), None); + assert_eq!(infer_context_window("some-random-name"), None); + } + #[test] fn serialize_messages_empty_content_no_tool_calls() { let messages = vec![LlmChatMessage { diff --git a/src-tauri/src/llm/tool_result_processor.rs b/src-tauri/src/llm/tool_result_processor.rs new file mode 100644 index 0000000..9d0972d --- /dev/null +++ b/src-tauri/src/llm/tool_result_processor.rs @@ -0,0 +1,559 @@ +//! Front-load summarization of tool results before they enter the LLM context. +//! +//! Long tool results (web pages, search hits) are compressed here so that +//! ContextGuard only has to deal with already-compact messages downstream. +//! +//! When a `results_dir` is configured, raw results exceeding the summarize +//! threshold are persisted to disk before summarization. The summarized marker +//! then contains a `ref:` that the `tool_result.recall` tool can use to +//! retrieve the original content on demand. + +use std::path::PathBuf; +use std::sync::Arc; + +use serde_json::Value; + +use super::provider::{CompletionOverrides, LlmProvider}; +use super::LlmChatMessage; + +pub const DEFAULT_SUMMARIZE_THRESHOLD: usize = 3000; +const MAX_RAW_INPUT_FOR_SUMMARY: usize = 4000; +const MAX_GOAL_CHARS: usize = 200; + +/// Marker prefix injected into summarized results so that ContextGuard can +/// recognise them and skip aggressive trimming. +pub const SUMMARIZED_MARKER: &str = "[summarized from "; + +/// Marker used by ContextGuard to detect stored-ref results that can be +/// recalled via `tool_result.recall`. +pub const STORED_REF_MARKER: &str = "[tool result stored | ref:"; + +const SUMMARIZE_SYSTEM: &str = "\ +You are a tool result summarizer. Extract only the information relevant to the user's task.\n\ +Output a concise summary (under 800 chars) containing:\n\ +1. Key facts and findings relevant to the task\n\ +2. Important data points, names, URLs worth keeping\n\ +3. Whether more reading/searching is needed\n\ +Output ONLY the summary, nothing else."; + +/// Maximum number of search/semantic hits to keep in rule-based extraction. +const RULE_BASED_TOP_K: usize = 5; +/// Maximum snippet length per hit in rule-based extraction. +const RULE_BASED_SNIPPET_LEN: usize = 200; + +// ───────────────────────────────────────────────────────────────────────────── + +#[derive(Clone)] +pub struct ToolResultProcessor { + provider: Arc, + summarize_threshold: usize, + results_dir: Option, + session_id: String, +} + +#[derive(Clone)] +pub struct ProcessedResult { + pub content: String, + pub was_summarized: bool, + pub original_len: usize, +} + +/// Name of the recall tool — used to skip re-summarization of recalled results. +pub const RECALL_TOOL_NAME: &str = "tool.recall"; + +impl ToolResultProcessor { + pub fn new( + provider: Arc, + summarize_threshold: usize, + results_dir: Option, + session_id: String, + ) -> Self { + Self { + provider, + summarize_threshold, + results_dir, + session_id, + } + } + + /// Process a single tool result: summarize if long, pass through if short. + /// + /// When `results_dir` is configured and the result exceeds the threshold, + /// the raw content is persisted to disk before summarization. The resulting + /// marker includes a `ref:` that `tool_result.recall` can use to retrieve + /// the original. + pub async fn process( + &self, + tool_name: &str, + call_id: &str, + raw_content: &str, + user_goal: Option<&str>, + ) -> ProcessedResult { + let original_len = raw_content.len(); + + if tool_name == RECALL_TOOL_NAME { + return ProcessedResult { + content: raw_content.to_string(), + was_summarized: false, + original_len, + }; + } + + if original_len <= self.summarize_threshold { + return ProcessedResult { + content: raw_content.to_string(), + was_summarized: false, + original_len, + }; + } + + let summary = if tool_name.starts_with("web.") { + self.llm_summarize(tool_name, raw_content, user_goal).await + } else { + rule_based_extract(tool_name, raw_content) + }; + + match summary { + Some(text) => { + let ref_tag = match self.persist_raw(call_id, tool_name, raw_content).await { + Some(short_id) => format!(" | ref:{}", short_id), + None => String::new(), + }; + ProcessedResult { + content: format!( + "{}{} chars{}]\n{}", + SUMMARIZED_MARKER, original_len, ref_tag, text + ), + was_summarized: true, + original_len, + } + } + None => ProcessedResult { + content: raw_content.to_string(), + was_summarized: false, + original_len, + }, + } + } + + /// Persist raw tool result to `{results_dir}/{session_id}/{call_id}.json`. + /// Returns the short ref ID (first 8 chars of call_id) on success, None on + /// failure or if no results_dir is configured. + async fn persist_raw( + &self, + call_id: &str, + tool_name: &str, + raw_content: &str, + ) -> Option { + let results_dir = self.results_dir.as_ref()?; + let session_dir = results_dir.join(&self.session_id); + + if let Err(e) = tokio::fs::create_dir_all(&session_dir).await { + eprintln!( + "[tool_result_processor] failed to create dir {}: {}", + session_dir.display(), + e + ); + return None; + } + + let file_path = session_dir.join(format!("{}.json", call_id)); + let ts = chrono::Utc::now().to_rfc3339(); + let record = serde_json::json!({ + "call_id": call_id, + "tool_name": tool_name, + "ts": ts, + "content": raw_content, + "len": raw_content.len(), + }); + + match serde_json::to_string(&record) { + Ok(json_str) => { + if let Err(e) = tokio::fs::write(&file_path, json_str).await { + eprintln!( + "[tool_result_processor] failed to write {}: {}", + file_path.display(), + e + ); + return None; + } + } + Err(e) => { + eprintln!("[tool_result_processor] failed to serialize: {}", e); + return None; + } + } + + let short_id = &call_id[..call_id.len().min(8)]; + eprintln!( + "[tool_result_processor] persisted {} ({} chars) → {}", + tool_name, + raw_content.len(), + file_path.display() + ); + Some(short_id.to_string()) + } + + async fn llm_summarize( + &self, + tool_name: &str, + raw_content: &str, + user_goal: Option<&str>, + ) -> Option { + let truncated = truncate_at_boundary(raw_content, MAX_RAW_INPUT_FOR_SUMMARY); + let goal_line = user_goal + .map(|g| format!("User's task context: {}\n\n", truncate_at_boundary(g, MAX_GOAL_CHARS))) + .unwrap_or_default(); + + let user_prompt = format!( + "{}Tool: {}\nRaw result ({} chars):\n{}", + goal_line, + tool_name, + raw_content.len(), + truncated, + ); + + let messages = vec![ + LlmChatMessage { + role: "system".to_string(), + content: SUMMARIZE_SYSTEM.to_string(), + ..Default::default() + }, + LlmChatMessage { + role: "user".to_string(), + content: user_prompt, + ..Default::default() + }, + ]; + + let overrides = CompletionOverrides { + temperature: Some(0.0), + ..Default::default() + }; + + self.provider + .chat_completion(&messages, Some(&overrides)) + .await + .ok() + .filter(|t| !t.trim().is_empty()) + } +} + +// ─── Rule-based extraction ────────────────────────────────────────────────── + +fn rule_based_extract(tool_name: &str, raw_content: &str) -> Option { + match tool_name { + "vault.search_keyword" => extract_vault_keyword_hits(raw_content), + "vault.semantic_search" => extract_semantic_hits(raw_content), + _ => None, // fallback: no rule-based extraction, keep raw + } +} + +/// Extract top-K snippets from vault.search_keyword JSON result. +/// +/// Expected shape: `{"snippets": [{"rel_path": "...", "snippet": "...", ...}, ...]}` +fn extract_vault_keyword_hits(raw: &str) -> Option { + let val: Value = serde_json::from_str(raw).ok()?; + let snippets = val.get("snippets")?.as_array()?; + + let mut lines = Vec::new(); + for item in snippets.iter().take(RULE_BASED_TOP_K) { + let path = item.get("rel_path").and_then(|v| v.as_str()).unwrap_or("?"); + let snippet = item + .get("snippet") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let short = truncate_at_boundary(snippet, RULE_BASED_SNIPPET_LEN); + lines.push(format!("- {}: {}", path, short)); + } + + if lines.is_empty() { + return None; + } + + let total = snippets.len(); + let mut out = lines.join("\n"); + if total > RULE_BASED_TOP_K { + out.push_str(&format!("\n({} more results omitted)", total - RULE_BASED_TOP_K)); + } + Some(out) +} + +/// Extract top-K hits from vault.semantic_search JSON result. +/// +/// Expected shape: `{"hits": [{"rel_path": "...", "snippet": "...", "score": 0.8, ...}, ...]}` +fn extract_semantic_hits(raw: &str) -> Option { + let val: Value = serde_json::from_str(raw).ok()?; + let hits = val.get("hits")?.as_array()?; + + let mut lines = Vec::new(); + for item in hits.iter().take(RULE_BASED_TOP_K) { + let path = item.get("rel_path").and_then(|v| v.as_str()).unwrap_or("?"); + let snippet = item + .get("snippet") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let score = item + .get("score") + .and_then(|v| v.as_f64()) + .map(|s| format!(" ({:.2})", s)) + .unwrap_or_default(); + let short = truncate_at_boundary(snippet, RULE_BASED_SNIPPET_LEN); + lines.push(format!("- {}{}: {}", path, score, short)); + } + + if lines.is_empty() { + return None; + } + + let total = hits.len(); + let mut out = lines.join("\n"); + if total > RULE_BASED_TOP_K { + out.push_str(&format!("\n({} more results omitted)", total - RULE_BASED_TOP_K)); + } + Some(out) +} + +// ─── Helpers ──────────────────────────────────────────────────────────────── + +/// Extract the latest user message content (truncated) as a task-context hint. +pub fn extract_user_goal(messages: &[LlmChatMessage]) -> Option { + messages + .iter() + .rev() + .find(|m| m.role == "user") + .map(|m| truncate_at_boundary(&m.content, MAX_GOAL_CHARS).to_string()) +} + +fn truncate_at_boundary(s: &str, max: usize) -> &str { + if s.len() <= max { + return s; + } + let mut end = max; + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + &s[..end] +} + +// ─── Tests ────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + + fn make_processor(threshold: usize) -> ToolResultProcessor { + ToolResultProcessor::new( + Arc::new(FakeProvider), + threshold, + None, + "test-session".to_string(), + ) + } + + fn make_processor_with_dir(threshold: usize, dir: PathBuf) -> ToolResultProcessor { + ToolResultProcessor::new( + Arc::new(FakeProvider), + threshold, + Some(dir), + "test-session".to_string(), + ) + } + + #[test] + fn short_result_passes_through() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let proc = make_processor(3000); + let result = proc.process("note.read", "call-1", "short content", None).await; + assert_eq!(result.content, "short content"); + assert!(!result.was_summarized); + assert_eq!(result.original_len, 13); + }); + } + + #[test] + fn long_non_web_non_vault_passes_through() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let proc = make_processor(100); + let long = "x".repeat(200); + let result = proc.process("note.read", "call-2", &long, None).await; + assert_eq!(result.content, long); + assert!(!result.was_summarized); + }); + } + + #[test] + fn recall_tool_result_passes_through() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let proc = make_processor(10); + let long = "x".repeat(200); + let result = proc + .process(RECALL_TOOL_NAME, "call-3", &long, None) + .await; + assert_eq!(result.content, long); + assert!(!result.was_summarized); + }); + } + + #[test] + fn persist_raw_writes_file() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let tmp = tempfile::tempdir().unwrap(); + let proc = make_processor_with_dir(10, tmp.path().to_path_buf()); + let result = proc + .process("web.search", "abcdef12-3456", &"y".repeat(100), None) + .await; + assert!(result.was_summarized); + assert!(result.content.contains("ref:abcdef12")); + + let file = tmp + .path() + .join("test-session") + .join("abcdef12-3456.json"); + assert!(file.exists()); + let stored: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(&file).unwrap()).unwrap(); + assert_eq!(stored["call_id"], "abcdef12-3456"); + assert_eq!(stored["tool_name"], "web.search"); + assert_eq!(stored["len"], 100); + }); + } + + #[test] + fn vault_keyword_rule_extraction() { + let raw = serde_json::json!({ + "snippets": [ + {"rel_path": "a.md", "snippet": "x".repeat(300), "kind": "Match"}, + {"rel_path": "b.md", "snippet": "y".repeat(300), "kind": "Match"}, + {"rel_path": "c.md", "snippet": "z".repeat(50), "kind": "Match"}, + ] + }) + .to_string(); + + let result = extract_vault_keyword_hits(&raw).unwrap(); + assert!(result.contains("a.md")); + assert!(result.contains("b.md")); + assert!(result.contains("c.md")); + // Long snippets should be truncated + assert!(!result.contains(&"x".repeat(300))); + } + + #[test] + fn vault_keyword_more_than_top_k() { + let snippets: Vec = (0..8) + .map(|i| { + serde_json::json!({ + "rel_path": format!("{}.md", i), + "snippet": format!("content {}", i), + "kind": "Match" + }) + }) + .collect(); + let raw = serde_json::json!({ "snippets": snippets }).to_string(); + let result = extract_vault_keyword_hits(&raw).unwrap(); + assert!(result.contains("3 more results omitted")); + } + + #[test] + fn semantic_hits_extraction() { + let raw = serde_json::json!({ + "hits": [ + {"rel_path": "note1.md", "snippet": "some content", "score": 0.92}, + {"rel_path": "note2.md", "snippet": "other content", "score": 0.85}, + ] + }) + .to_string(); + + let result = extract_semantic_hits(&raw).unwrap(); + assert!(result.contains("note1.md")); + assert!(result.contains("(0.92)")); + assert!(result.contains("note2.md")); + } + + #[test] + fn extract_user_goal_finds_last_user_msg() { + let messages = vec![ + LlmChatMessage { + role: "user".to_string(), + content: "first question".to_string(), + ..Default::default() + }, + LlmChatMessage { + role: "assistant".to_string(), + content: "answer".to_string(), + ..Default::default() + }, + LlmChatMessage { + role: "user".to_string(), + content: "second question".to_string(), + ..Default::default() + }, + ]; + assert_eq!( + extract_user_goal(&messages), + Some("second question".to_string()) + ); + } + + #[test] + fn truncate_at_boundary_cjk() { + let s = "你好世界测试"; + let t = truncate_at_boundary(s, 7); + assert!(t.len() <= 7); + assert_eq!(t, "你好"); + } + + // Minimal fake provider for unit tests (LLM summarization not tested here) + struct FakeProvider; + + #[async_trait::async_trait] + impl LlmProvider for FakeProvider { + async fn chat_stream( + &self, + _app: &tauri::AppHandle, + _session_id: &str, + _messages: Vec, + _tools: Option>, + _cancel: tokio_util::sync::CancellationToken, + ) -> Result { + unimplemented!() + } + + async fn chat_completion( + &self, + _messages: &[LlmChatMessage], + _overrides: Option<&CompletionOverrides>, + ) -> Result { + Ok("fake summary".to_string()) + } + + async fn list_models(&self) -> Result, String> { + Ok(vec![]) + } + + fn convert_tools(&self, _manifests: &[Value]) -> Vec { + vec![] + } + + fn build_tool_result_message( + &self, + _call_id: &str, + _tool_name: &str, + _content: &str, + ) -> LlmChatMessage { + LlmChatMessage::default() + } + + fn provider_name(&self) -> &'static str { + "fake" + } + + fn is_remote(&self) -> bool { + false + } + } +} diff --git a/src-tauri/src/skills/mod.rs b/src-tauri/src/skills/mod.rs index 69bad91..9ab347a 100644 --- a/src-tauri/src/skills/mod.rs +++ b/src-tauri/src/skills/mod.rs @@ -174,8 +174,8 @@ fn web_research_manifest() -> SkillManifest { "note.append".to_string(), "vault.search_keyword".to_string(), ], - max_tool_calls: 25, - timeout_secs: 180, + max_tool_calls: 50, + timeout_secs: 300, ui_entry: SkillUiEntry::ConversationMode, tags: vec!["research".to_string(), "web".to_string()], auto_invocable: true, @@ -321,8 +321,8 @@ mod mod_tests { assert!(m.allowed_tools.contains(&"web.read_page".to_string())); assert!(m.allowed_tools.contains(&"web.read_pdf".to_string())); assert!(m.allowed_tools.contains(&"note.create".to_string())); - assert_eq!(m.max_tool_calls, 25); - assert_eq!(m.timeout_secs, 180); + assert_eq!(m.max_tool_calls, 50); + assert_eq!(m.timeout_secs, 300); assert_eq!(m.max_tool_result_chars, 20000); assert!(m.auto_invocable); } diff --git a/src-tauri/src/skills/runtime.rs b/src-tauri/src/skills/runtime.rs index d6b3527..9c19bbd 100644 --- a/src-tauri/src/skills/runtime.rs +++ b/src-tauri/src/skills/runtime.rs @@ -163,9 +163,10 @@ pub async fn run_skill_with_depth( let config = AgentLoopConfig { max_tool_calls: manifest.max_tool_calls, timeout_ms: manifest.timeout_secs.saturating_mul(1000), - max_tool_result_chars: manifest.max_tool_result_chars as usize, + max_single_result_chars: manifest.max_tool_result_chars as usize, nesting_depth, max_context_tokens, + summarize_threshold: crate::llm::tool_result_processor::DEFAULT_SUMMARIZE_THRESHOLD, }; let conv_id = skill_conversation_id(&manifest.id, &parent_conversation_id); diff --git a/src-tauri/src/skills/skill_tool.rs b/src-tauri/src/skills/skill_tool.rs index 92f4655..01f397e 100644 --- a/src-tauri/src/skills/skill_tool.rs +++ b/src-tauri/src/skills/skill_tool.rs @@ -47,6 +47,7 @@ pub struct SkillAsTool { manifest: ToolManifest, skill_id: String, skill_name: String, + overall_timeout_ms: u64, app: AppHandle, semaphore: Arc, } @@ -104,10 +105,12 @@ impl SkillAsTool { tags: vec!["skill".to_string()], deprecated: None, }; + let overall_timeout_ms = (skill.timeout_secs as u64).saturating_mul(2000); Arc::new(Self { manifest, skill_id: skill.id.clone(), skill_name: skill.name.clone(), + overall_timeout_ms, app, semaphore, }) @@ -124,6 +127,10 @@ impl Tool for SkillAsTool { ToolCategory::Skill } + fn timeout_ms(&self) -> Option { + Some(self.overall_timeout_ms) + } + async fn invoke(&self, ctx: &ToolContext, input: Value) -> ToolResult { if ctx.nesting_depth >= MAX_NESTING_DEPTH { return tool_err( @@ -243,7 +250,7 @@ impl Tool for SkillAsTool { (manifest.timeout_secs as u64).saturating_mul(2), ); eprintln!( - "[skill_tool] skill={} session={} starting, timeout={}s max_tool_calls={} max_result_chars={}", + "[skill_tool] skill={} session={} starting, timeout={}s max_tool_calls={} max_single_result_chars={}", self.skill_id, &session_id[..8.min(session_id.len())], overall_timeout.as_secs(), manifest.max_tool_calls, manifest.max_tool_result_chars, ); diff --git a/src-tauri/src/tools/built_in/mod.rs b/src-tauri/src/tools/built_in/mod.rs index 46f0505..93c158c 100644 --- a/src-tauri/src/tools/built_in/mod.rs +++ b/src-tauri/src/tools/built_in/mod.rs @@ -9,3 +9,4 @@ pub mod web_ops; pub mod memory_ops; pub mod web_search; pub mod webview_renderer; +pub mod tool_result_recall; diff --git a/src-tauri/src/tools/built_in/note_ops.rs b/src-tauri/src/tools/built_in/note_ops.rs index dc16f21..23eefac 100644 --- a/src-tauri/src/tools/built_in/note_ops.rs +++ b/src-tauri/src/tools/built_in/note_ops.rs @@ -571,7 +571,11 @@ impl NoteCreateTool { name: "note.create".to_string(), version: "1.0.0".to_string(), protocol_version: "1.0".to_string(), - description: "在工作区创建一篇新的 Markdown 笔记文件（含 frontmatter）".to_string(), + description: "Create a new Markdown note file in the workspace. Use for \ + structured, long-form content: research reports, meeting notes, \ + technical analyses, tutorials, or any document that deserves its \ + own file and path. NOT for short fleeting ideas (use thought.create \ + for those).".to_string(), input_schema: serde_json::json!({ "type": "object", "required": ["rel_path", "title"], diff --git a/src-tauri/src/tools/built_in/thought_ops.rs b/src-tauri/src/tools/built_in/thought_ops.rs index 9affb1a..2133a41 100644 --- a/src-tauri/src/tools/built_in/thought_ops.rs +++ b/src-tauri/src/tools/built_in/thought_ops.rs @@ -177,7 +177,12 @@ impl ThoughtCreateTool { name: "thought.create".to_string(), version: "1.0.0".to_string(), protocol_version: "1.0".to_string(), - description: "Create a new independent Thought entry. A thought captures the user's thinking, insight, inspiration, or hypothesis about a topic. Do NOT use for behavioral instructions, personal preferences, or memory directives (e.g. \"remember...\", \"always...\", \"never...\") — those belong in memory.save.".to_string(), + description: "Create a short, independent Thought entry (a few sentences). \ + A thought captures a fleeting idea, spark of inspiration, or \ + quick hypothesis — content too small to warrant its own note \ + file. Do NOT use for structured documents, research reports, \ + or long-form content (use note.create for those). Do NOT use \ + for behavioral instructions or preferences (use memory.save).".to_string(), input_schema: serde_json::json!({ "type": "object", "required": ["content"], diff --git a/src-tauri/src/tools/built_in/tool_result_recall.rs b/src-tauri/src/tools/built_in/tool_result_recall.rs new file mode 100644 index 0000000..b38bd2c --- /dev/null +++ b/src-tauri/src/tools/built_in/tool_result_recall.rs @@ -0,0 +1,207 @@ +use std::path::PathBuf; + +use async_trait::async_trait; +use serde_json::Value; + +use crate::tools::context::ToolContext; +use crate::tools::types::{ + ApprovalPolicy, Effect, Risk, Tool, ToolCategory, ToolError, ToolErrorCode, ToolManifest, + ToolMetrics, ToolResult, +}; + +pub struct ToolResultRecallTool { + manifest: ToolManifest, +} + +impl ToolResultRecallTool { + pub fn new() -> Self { + Self { + manifest: ToolManifest { + name: "tool.recall".to_string(), + version: "1.0.0".to_string(), + protocol_version: "1.0".to_string(), + description: "Retrieve the full raw content of a previously summarized tool \ + result. Use when the summary (marked with [summarized from ... | \ + ref:XXX]) is insufficient and you need the original details." + .to_string(), + input_schema: serde_json::json!({ + "type": "object", + "required": ["ref"], + "properties": { + "ref": { + "type": "string", + "description": "The ref ID from the [summarized from ... | ref:XXX] marker" + } + }, + "additionalProperties": false + }), + output_schema: serde_json::json!({ + "type": "object", + "properties": { + "content": { "type": "string" }, + "tool_name": { "type": "string" }, + "len": { "type": "integer" } + } + }), + effects: vec![Effect::Read], + risk: Risk::Safe, + privacy_aware: true, + requires_workspace: true, + default_approval: ApprovalPolicy::Auto, + examples: vec![], + tags: vec!["utility".to_string(), "recall".to_string()], + deprecated: None, + }, + } + } +} + +#[async_trait] +impl Tool for ToolResultRecallTool { + fn manifest(&self) -> &ToolManifest { + &self.manifest + } + + fn category(&self) -> ToolCategory { + ToolCategory::Utility + } + + async fn invoke(&self, ctx: &ToolContext, input: Value) -> ToolResult { + let ref_id = match input.get("ref").and_then(|v| v.as_str()) { + Some(r) => r, + None => { + return ToolResult::Err { error: ToolError { + code: ToolErrorCode::InvalidInput, + message: "missing required field: ref".to_string(), + retryable: false, + cause: None, + } }; + } + }; + + let session_dir = build_session_dir(ctx); + + let file_path = match find_by_prefix(&session_dir, ref_id).await { + Some(p) => p, + None => { + return ToolResult::Err { error: ToolError { + code: ToolErrorCode::NotFound, + message: format!( + "no stored tool result matching ref '{}' in {}", + ref_id, + session_dir.display() + ), + retryable: false, + cause: None, + } }; + } + }; + + let raw = match tokio::fs::read_to_string(&file_path).await { + Ok(s) => s, + Err(e) => { + return ToolResult::Err { error: ToolError { + code: ToolErrorCode::Internal, + message: format!("failed to read {}: {}", file_path.display(), e), + retryable: false, + cause: None, + } }; + } + }; + + let record: Value = match serde_json::from_str(&raw) { + Ok(v) => v, + Err(e) => { + return ToolResult::Err { error: ToolError { + code: ToolErrorCode::Internal, + message: format!("corrupted stored result: {}", e), + retryable: false, + cause: None, + } }; + } + }; + + let content = record + .get("content") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let tool_name = record + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + let len = record + .get("len") + .and_then(|v| v.as_u64()) + .unwrap_or(content.len() as u64); + + ToolResult::Ok { + data: serde_json::json!({ + "content": content, + "tool_name": tool_name, + "len": len, + }), + redacted_count: 0, + warnings: vec![], + metrics: ToolMetrics { + duration_ms: 0, + ..Default::default() + }, + } + } +} + +fn build_session_dir(ctx: &ToolContext) -> PathBuf { + let sid = if ctx.session_id.is_empty() { + &ctx.conversation_id + } else { + &ctx.session_id + }; + ctx.workspace_root + .join(".knowforge") + .join("tool-results") + .join(sid) +} + +async fn find_by_prefix(session_dir: &PathBuf, ref_id: &str) -> Option { + let mut entries = match tokio::fs::read_dir(session_dir).await { + Ok(e) => e, + Err(_) => return None, + }; + + while let Ok(Some(entry)) = entries.next_entry().await { + let name = entry.file_name(); + let name_str = name.to_string_lossy(); + if name_str.starts_with(ref_id) && name_str.ends_with(".json") { + return Some(entry.path()); + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn find_by_prefix_matches() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let tmp = tempfile::tempdir().unwrap(); + let file = tmp.path().join("abcdef12-3456-7890.json"); + tokio::fs::write( + &file, + r#"{"call_id":"abcdef12-3456-7890","tool_name":"web.search","content":"hello","len":5}"#, + ) + .await + .unwrap(); + + let found = find_by_prefix(&tmp.path().to_path_buf(), "abcdef12").await; + assert!(found.is_some()); + assert_eq!(found.unwrap(), file); + + let not_found = find_by_prefix(&tmp.path().to_path_buf(), "zzz").await; + assert!(not_found.is_none()); + }); + } +} diff --git a/src-tauri/src/tools/context.rs b/src-tauri/src/tools/context.rs index e3ee2cf..b5020ae 100644 --- a/src-tauri/src/tools/context.rs +++ b/src-tauri/src/tools/context.rs @@ -12,6 +12,9 @@ use crate::semantic_index::EmbeddingCache; pub struct ToolContext { pub workspace_root: PathBuf, pub conversation_id: String, + /// The agent_loop's session ID — unique per stream call. Used for + /// tool-result externalization: `{workspace}/.knowforge/tool-results/{session_id}/`. + pub session_id: String, /// agent_loop 分配的工具调用追踪 ID；默认 None，由 execute_tool / invoke_tool 设置。 pub call_id: Option, pub audit_sink: Arc, @@ -102,6 +105,7 @@ impl ToolContextFactory { ToolContext { workspace_root, conversation_id: conversation_id.to_string(), + session_id: String::new(), call_id: None, audit_sink: self.audit_sink.clone(), privacy_filter: self.privacy_filter.clone(), diff --git a/src-tauri/src/tools/mod.rs b/src-tauri/src/tools/mod.rs index 493c30f..eca5160 100644 --- a/src-tauri/src/tools/mod.rs +++ b/src-tauri/src/tools/mod.rs @@ -117,6 +117,9 @@ pub fn register_builtin_tools( registry.register(Arc::new(built_in::web_download::WebDownloadTool::new()))?; registry.register(Arc::new(built_in::web_download::WebReadPdfTool::new()))?; + // Tool result recall + registry.register(Arc::new(built_in::tool_result_recall::ToolResultRecallTool::new()))?; + Ok(()) } @@ -136,9 +139,9 @@ mod mod_tests { "register_builtin_tools failed: {:?}", result.err() ); - // 确认工具总数：1(time.now) + 8(P1) + 4(P3 写操作) + 2(memory) + 4(P4 网络) = 19 + // 确认工具总数：1(time.now) + 8(P1) + 4(P3 写操作) + 2(memory) + 4(P4 网络) + 1(recall) = 20 let tools = registry.list_for_llm(crate::tools::registry::ToolScope::Global); - assert_eq!(tools.len(), 19, "expected 19 registered tools, got {}", tools.len()); + assert_eq!(tools.len(), 20, "expected 20 registered tools, got {}", tools.len()); } #[test] @@ -148,8 +151,8 @@ mod mod_tests { let all = registry.list_for_llm(crate::tools::registry::ToolScope::Global); let core = registry.list_for_llm_filtered(&crate::tools::registry::ToolFilter::core()); assert!(core.len() < all.len(), "core ({}) should be less than all ({})", core.len(), all.len()); - // NoteRead(5) + Utility(1 time.now + 2 memory) = 8 - assert_eq!(core.len(), 8, "core should have 8 tools (5 NoteRead + 3 Utility)"); + // NoteRead(5) + Utility(1 time.now + 2 memory + 1 recall) = 9 + assert_eq!(core.len(), 9, "core should have 9 tools (5 NoteRead + 4 Utility)"); } #[test] @@ -185,5 +188,6 @@ mod mod_tests { check("time.now", ToolCategory::Utility); check("memory.save", ToolCategory::Utility); check("memory.forget", ToolCategory::Utility); + check("tool.recall", ToolCategory::Utility); } } diff --git a/src-tauri/src/tools/types.rs b/src-tauri/src/tools/types.rs index 0464466..d1dd174 100644 --- a/src-tauri/src/tools/types.rs +++ b/src-tauri/src/tools/types.rs @@ -184,4 +184,8 @@ pub trait Tool: Send + Sync { fn category(&self) -> ToolCategory { ToolCategory::NoteRead } + + fn timeout_ms(&self) -> Option { + None + } }