ooodc/PicoBot
ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

Compare commits

base: ooodc:2a02adc7c3016fe4e649bcc3ec872585d2a437cf
Branches Tags
ooodc:main
..
compare: ooodc:4071dcc6eeb1a9f3fedb1837d4e1780efde5bc9e
Branches Tags
ooodc:main

No commits in common. "2a02adc7c3016fe4e649bcc3ec872585d2a437cf" and "4071dcc6eeb1a9f3fedb1837d4e1780efde5bc9e" have entirely different histories.
2a02adc7c3 ... 4071dcc6ee

16 changed files with 167 additions and 1164 deletions
Show Stats Expand all files Collapse all files

741
src/agent/context_compressor.rs
View File

@ -5,6 +5,7 @@ use crate::bus::{
use crate::config::LLMProviderConfig;
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
use crate::text::{char_count, take_prefix_chars};
use crate::agent::{AgentError, AgentRuntimeConfig};
const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
@ -12,145 +13,6 @@ const CJK_CHARS_PER_TOKEN: f64 = 2.0;
const OTHER_CHARS_PER_TOKEN: f64 = 4.0;
const JSON_OVERHEAD_PER_MESSAGE: usize = 50;
/// System context marker for the heavy compression (older segment) summary.
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER: &str = "history_compaction_older";
/// System context marker for the light compression (newer segment) summary.
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_newer";
/// Default threshold ratio: compress when estimated tokens exceed 70% of context window.
const DEFAULT_THRESHOLD_RATIO: f64 = 0.7;
/// Default budget split: older segment gets 30% of summary budget, newer gets 70%.
const OLDER_BUDGET_RATIO: f64 = 0.3;
const NEWER_BUDGET_RATIO: f64 = 0.7;
// ============================================================================
// HistoryUnit — atomic message units for compression
// ============================================================================
/// An indivisible unit of conversation history. Compression operates at the
/// unit level, so tool call sequences can never be split across boundaries.
#[derive(Debug, Clone)]
enum HistoryUnit {
/// Always-preserved system messages (agent_prompt, scheduled_prompt).
SystemGuard(ChatMessage),
/// A user message.
UserMessage(ChatMessage),
/// An atomic tool call sequence: assistant(tool_calls) + all following
/// tool-result messages. Must be preserved or compressed as a whole.
ToolRound {
assistant: ChatMessage,
results: Vec<ChatMessage>,
},
/// A plain assistant text response (no tool_calls).
AssistantText(ChatMessage),
}
impl HistoryUnit {
/// Estimate tokens for this unit alone.
fn estimate_tokens(&self) -> usize {
match self {
HistoryUnit::SystemGuard(msg) | HistoryUnit::UserMessage(msg) | HistoryUnit::AssistantText(msg) => {
estimate_tokens(std::slice::from_ref(msg))
}
HistoryUnit::ToolRound { assistant, results } => {
let mut all = vec![assistant.clone()];
all.extend(results.clone());
estimate_tokens(&all)
}
}
}
/// Expand this unit back into a flat list of ChatMessages.
#[allow(dead_code)]
fn into_messages(self) -> Vec<ChatMessage> {
match self {
HistoryUnit::SystemGuard(msg)
| HistoryUnit::UserMessage(msg)
| HistoryUnit::AssistantText(msg) => vec![msg],
HistoryUnit::ToolRound { assistant, results } => {
let mut msgs = vec![assistant];
msgs.extend(results);
msgs
}
}
}
/// Returns true if this unit is a ToolRound.
#[allow(dead_code)]
fn is_tool_round(&self) -> bool {
matches!(self, HistoryUnit::ToolRound { .. })
}
}
// ============================================================================
// Unit parser — one forward pass, O(n)
// ============================================================================
fn is_system_guard(msg: &ChatMessage) -> bool {
msg.role == "system"
&& (msg.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
|| msg.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
}
fn is_assistant_with_tool_calls(msg: &ChatMessage) -> bool {
msg.role == "assistant"
&& msg
.tool_calls
.as_ref()
.map_or(false, |calls| !calls.is_empty())
}
/// Parse a flat message list into atomic units. Orphaned tool results
/// (without a preceding assistant with tool_calls) are silently dropped.
fn parse_to_units(messages: &[ChatMessage]) -> Vec<HistoryUnit> {
let mut units = Vec::new();
let mut i = 0;
while i < messages.len() {
let msg = &messages[i];
if is_system_guard(msg) {
units.push(HistoryUnit::SystemGuard(msg.clone()));
i += 1;
} else if msg.role == "user" {
units.push(HistoryUnit::UserMessage(msg.clone()));
i += 1;
} else if is_assistant_with_tool_calls(msg) {
let assistant = msg.clone();
let mut results = Vec::new();
i += 1;
while i < messages.len() && messages[i].role == "tool" {
results.push(messages[i].clone());
i += 1;
}
units.push(HistoryUnit::ToolRound { assistant, results });
} else if msg.role == "assistant" {
units.push(HistoryUnit::AssistantText(msg.clone()));
i += 1;
} else if msg.role == "tool" {
// Orphaned tool result — drop with a warning
tracing::warn!(
tool_call_id = ?msg.tool_call_id,
message_index = i,
"Dropping orphaned tool result during unit parsing"
);
i += 1;
} else {
// system messages that aren't guards — keep as plain messages
units.push(HistoryUnit::AssistantText(msg.clone()));
i += 1;
}
}
units
}
/// Estimate total tokens from a slice of units.
fn estimate_tokens_from_units(units: &[HistoryUnit]) -> usize {
units.iter().map(|u| u.estimate_tokens()).sum()
}
/// Check if a character is CJK (Chinese, Japanese, Korean)
fn is_cjk_char(c: char) -> bool {
matches!(c,
@ -235,7 +97,7 @@ impl Default for ContextCompressionConfig {
pub struct ContextCompressor {
config: ContextCompressionConfig,
context_window: usize,
/// Threshold ratio to trigger compression (70% of context window).
/// Threshold ratio to trigger compression (50% of context window)
threshold_ratio: f64,
}
@ -251,10 +113,6 @@ impl ContextCompressor {
.clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS)
}
// =========================================================================
// Transcript building (shared helpers, unchanged from original)
// =========================================================================
fn format_transcript_entry(message: &ChatMessage) -> String {
let role = match message.role.as_str() {
"assistant" => "Assistant",
@ -356,86 +214,6 @@ Be concise, aim for {} characters or less.
{}
"#,
target_chars, transcript
)
}
// =========================================================================
// Two-segment compression: heavy (older) + light (newer) prompts
// =========================================================================
/// Prompt for the older segment — aggressive compression, keep only essentials.
fn build_heavy_summary_prompt(transcript: &str, target_chars: usize) -> String {
format!(
r#"You are a conversation compaction engine. Aggressively summarize the following OLDER conversation segment. These events happened earlier in the session and will be marked as "较早的操作记录（已压缩）".
A newer portion (not shown here) follows this segment. The model can use tools to re-read files for exact data, so you do NOT need to preserve full tool outputs.
=== MUST PRESERVE (keep verbatim or near-verbatim) ===
- All file paths (absolute and relative)
- All URLs
- Key numeric results and conclusions (one sentence per result)
- Error messages (exact text if important)
- User preferences and decisions
- The original user task or goal
=== SHOULD CONDENSE ===
- Each tool call → one line summary: "tool_name: 关键参数 → 结果概要"
- Multiple similar operations → group: "对 ./data/ 目录下的 10 个 CSV 文件执行了统计分析"
- Verbose tool output → drop entirely unless it contains a critical result or error
=== SHOULD DROP ===
- Full tool output (hundreds of lines of raw data)
- Repeated attempts of the same operation
- Greetings, filler text, markdown formatting
Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.
---
OLDER SEGMENT (events from earlier in the session):
{}
"#,
target_chars, transcript
)
}
/// Prompt for the newer segment — lighter compression, keep more detail.
fn build_light_summary_prompt(transcript: &str, target_chars: usize) -> String {
format!(
r#"You are a conversation compaction engine. Lightly summarize the following RECENT conversation segment. These events happened just before the current moment and will be marked as "近期操作记录（轻度压缩，保留了更多细节）".
An older segment (already summarized separately) precedes this. Keep enough detail so the model can continue the task seamlessly without re-reading files for data that appears in these recent results.
=== MUST PRESERVE (keep with more detail than a normal summary) ===
- All file paths, URLs, and identifiers
- The exact sequence of recent operations (step by step)
- Tool parameters (especially file paths, search queries, command strings)
- Key outputs and results (shortened but keep the substance)
- Error messages (complete, not truncated)
- Current task status and what should happen next
=== SHOULD CONDENSE ===
- Tool outputs → truncate to the most meaningful parts (key data, conclusions, not raw output)
- Long text → keep the essence but not every word
- Repeated similar outputs → note the pattern but keep a representative example
=== SHOULD DROP ===
- Completely irrelevant debug output
- Boilerplate text with no informational value
- Trivial operations that have no bearing on the task
Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.
---
RECENT SEGMENT (events from just before the current moment):
{}
"#,
target_chars, transcript
)
@ -517,339 +295,12 @@ RECENT SEGMENT (events from just before the current moment):
Ok(take_prefix_chars(transcript, target_chars))
}
// =========================================================================
// Two-segment compression
// =========================================================================
/// Find a safe split point in the unit list. The split ensures:
/// - Accumulated tokens up to `ratio` of total are placed in the older segment.
/// - The split never lands on a SystemGuard (they're always preserved).
/// - The split lands on a unit boundary (ToolRound, AssistantText, or UserMessage).
fn find_safe_split_point(&self, units: &[HistoryUnit], ratio: f64) -> usize {
let total_tokens = estimate_tokens_from_units(units);
let target = (total_tokens as f64 * ratio) as usize;
let mut accumulated = 0;
for (i, unit) in units.iter().enumerate() {
// Never split on a SystemGuard — they must be preserved
if matches!(unit, HistoryUnit::SystemGuard(_)) {
continue;
}
accumulated += unit.estimate_tokens();
if accumulated >= target {
return i + 1;
}
}
// Fallback: put everything in the older segment
units.len()
}
/// Summarize a transcript with a custom prompt builder function.
async fn summarize_with_prompt(
&self,
provider: &dyn LLMProvider,
transcript: &str,
target_chars: usize,
build_prompt: fn(&str, usize) -> String,
) -> Result<String, AgentError> {
let request = ChatCompletionRequest {
messages: vec![
Message::system("You are a helpful assistant."),
Message::user(build_prompt(transcript, target_chars)),
],
temperature: Some(0.3),
max_tokens: Some(1000),
tools: None,
};
let response = provider
.chat(request)
.await
.map_err(|e| AgentError::LlmError(e.to_string()))?;
Ok(response.content)
}
/// Summarize a segment of units with chunked fallback.
async fn summarize_units_segment(
&self,
provider: &dyn LLMProvider,
messages: &[ChatMessage],
transcript: &str,
target_chars: usize,
build_prompt: fn(&str, usize) -> String,
) -> Result<String, AgentError> {
if messages.is_empty() {
return Ok(String::new());
}
if char_count(transcript) <= target_chars {
return self
.summarize_with_prompt(provider, transcript, target_chars, build_prompt)
.await;
}
// Use chunked summarization for oversized transcripts
let target = target_chars.max(1);
let mut layer = Self::chunk_messages_for_summary(messages, target);
if layer.is_empty() {
layer.push(transcript.to_string());
}
for _ in 0..6 {
if layer.len() == 1 && char_count(&layer[0]) <= target {
return self
.summarize_with_prompt(provider, &layer[0], target, build_prompt)
.await;
}
let per_chunk_target = (target / layer.len().max(1)).max(500).min(target);
let mut summaries = Vec::with_capacity(layer.len());
for chunk in &layer {
summaries.push(
self.summarize_with_prompt(provider, chunk, per_chunk_target, build_prompt)
.await?,
);
}
if summaries.len() == 1 {
let summary = summaries.pop().unwrap_or_default();
if char_count(&summary) <= target {
return Ok(summary);
}
layer = Self::split_text_chunks(&summary, target);
continue;
}
let merged = summaries.join("\n\n");
if char_count(&merged) <= target {
return self
.summarize_with_prompt(provider, &merged, target, build_prompt)
.await;
}
layer = Self::split_text_chunks(&merged, target);
}
Ok(take_prefix_chars(transcript, target))
}
/// Main entry point for two-segment compression.
///
/// Splits history into older and newer segments, then summarizes each
/// with LLM using different prompts and budget allocations. The result
/// contains no tool_calls or tool-result messages — only system summaries
/// and the original user message — eliminating any risk of API 400 errors
/// from orphaned tool call sequences.
pub async fn compress_two_segment(
&self,
history: &[ChatMessage],
provider_config: &LLMProviderConfig,
) -> Result<Vec<ChatMessage>, AgentError> {
let tokens = estimate_tokens(history);
if tokens <= self.threshold() {
tracing::info!(
tokens = tokens,
threshold = self.threshold(),
msg_count = history.len(),
"Two-segment compression not needed (under threshold)"
);
return Ok(history.to_vec());
}
tracing::info!(
tokens = tokens,
threshold = self.threshold(),
msg_count = history.len(),
"Starting two-segment compression"
);
// Step 1: Parse into atomic units
let units = parse_to_units(history);
// Step 2: Separate system guards + user messages from compressible units
let mut system_guards: Vec<ChatMessage> = Vec::new();
let mut user_messages: Vec<ChatMessage> = Vec::new();
let mut compressible: Vec<HistoryUnit> = Vec::new();
for unit in units {
match unit {
HistoryUnit::SystemGuard(msg) => system_guards.push(msg),
HistoryUnit::UserMessage(msg) => user_messages.push(msg),
other => compressible.push(other),
}
}
// Keep only the latest user message in full; older ones go into compression
let latest_user_msg = user_messages.pop();
// Step 3: Find split point in compressible units
let total_compressible = estimate_tokens_from_units(&compressible);
let split_point = if total_compressible > 0 {
self.find_safe_split_point(&compressible, 0.5)
} else {
0
};
let older_units: Vec<&HistoryUnit> = compressible[..split_point].iter().collect();
let newer_units: Vec<&HistoryUnit> = compressible[split_point..].iter().collect();
// Step 4: Build transcripts
let older_messages: Vec<ChatMessage> = older_units
.iter()
.flat_map(|u| match u {
HistoryUnit::ToolRound { assistant, results } => {
let mut msgs = vec![assistant.clone()];
msgs.extend(results.clone());
msgs
}
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
HistoryUnit::UserMessage(msg) => {
// Older user messages go into the compressible transcript
vec![msg.clone()]
}
_ => vec![],
})
.collect();
let newer_messages: Vec<ChatMessage> = newer_units
.iter()
.flat_map(|u| match u {
HistoryUnit::ToolRound { assistant, results } => {
let mut msgs = vec![assistant.clone()];
msgs.extend(results.clone());
msgs
}
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
_ => vec![],
})
.collect();
// Include older user messages in the older transcript
let older_user_msgs: Vec<ChatMessage> = user_messages
.iter()
.map(|m| {
let mut msg = m.clone();
msg.role = "user".to_string();
msg
})
.collect();
let all_older_messages: Vec<ChatMessage> = older_user_msgs
.iter()
.chain(older_messages.iter())
.cloned()
.collect();
let older_transcript = Self::build_transcript(&all_older_messages);
let newer_transcript = Self::build_transcript(&newer_messages);
// Step 5: Budget allocation
let total_budget = self.config.summary_max_chars;
let older_budget = (total_budget as f64 * OLDER_BUDGET_RATIO) as usize;
let newer_budget = (total_budget as f64 * NEWER_BUDGET_RATIO) as usize;
// Step 6: Create provider and run both summaries (can be parallel)
let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
let provider = create_provider(runtime_config.provider)
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
let (older_result, newer_result) = if older_units.is_empty() && newer_units.is_empty() {
(Ok(String::new()), Ok(String::new()))
} else if older_units.is_empty() {
let result = self
.summarize_units_segment(
provider.as_ref(),
&newer_messages,
&newer_transcript,
newer_budget,
Self::build_light_summary_prompt,
)
.await;
(Ok(String::new()), result)
} else if newer_units.is_empty() {
let result = self
.summarize_units_segment(
provider.as_ref(),
&all_older_messages,
&older_transcript,
older_budget,
Self::build_heavy_summary_prompt,
)
.await;
(result, Ok(String::new()))
} else {
let older_fut = self.summarize_units_segment(
provider.as_ref(),
&all_older_messages,
&older_transcript,
older_budget,
Self::build_heavy_summary_prompt,
);
let newer_fut = self.summarize_units_segment(
provider.as_ref(),
&newer_messages,
&newer_transcript,
newer_budget,
Self::build_light_summary_prompt,
);
tokio::join!(older_fut, newer_fut)
};
let older_summary = older_result?;
let newer_summary = newer_result?;
// Step 7: Assemble compressed history
let mut compressed: Vec<ChatMessage> = Vec::with_capacity(4);
// System guards first
compressed.extend(system_guards);
// Latest user message
if let Some(user_msg) = latest_user_msg {
compressed.push(user_msg);
}
// Heavy compression summary (older)
if !older_summary.is_empty() {
compressed.push(ChatMessage::system_with_context(
format!("## 较早的操作记录（已压缩）\n\n{}", older_summary),
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER.to_string()),
));
}
// Light compression summary (newer)
if !newer_summary.is_empty() {
compressed.push(ChatMessage::system_with_context(
format!(
"## 近期操作记录（轻度压缩，保留了更多细节）\n\n{}\n\n---\n以上为近期操作记录。如需精确数据，可使用工具重新读取相关文件。",
newer_summary
),
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER.to_string()),
));
}
tracing::info!(
original_tokens = tokens,
final_tokens = estimate_tokens(&compressed),
final_msg_count = compressed.len(),
older_units = older_units.len(),
newer_units = newer_units.len(),
"Two-segment compression completed"
);
Ok(compressed)
}
// =========================================================================
// Legacy methods (kept for backward compatibility during migration)
// =========================================================================
/// Create a new compressor with the given context window size.
pub fn new(context_window: usize) -> Self {
Self {
config: ContextCompressionConfig::default(),
context_window,
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
threshold_ratio: 0.5,
}
}
@ -872,11 +323,11 @@ RECENT SEGMENT (events from just before the current moment):
Self {
config,
context_window,
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
threshold_ratio: 0.5,
}
}
/// Get the compression threshold in tokens (70% of context window).
/// Get the compression threshold in tokens.
fn threshold(&self) -> usize {
(self.context_window as f64 * self.threshold_ratio) as usize
}
@ -1256,7 +707,7 @@ mod tests {
#[test]
fn test_threshold() {
let compressor = ContextCompressor::new(128_000);
assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000
assert_eq!(compressor.threshold(), 64_000);
}
#[test]
@ -1305,184 +756,4 @@ mod tests {
assert!(chunks.iter().all(|chunk| char_count(chunk) <= 10));
assert_eq!(chunks.concat(), "user: xxxxxxxxxxxxxxxxxxxxxxxxx");
}
// =========================================================================
// HistoryUnit / parse_to_units tests
// =========================================================================
#[test]
fn test_parse_to_units_system_guards_preserved() {
let messages = vec![
ChatMessage::system_with_context(
"agent prompt",
Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
),
ChatMessage::user("hello"),
ChatMessage::assistant("hi"),
];
let units = parse_to_units(&messages);
// SystemGuard + UserMessage + AssistantText = 3
assert_eq!(units.len(), 3);
assert!(matches!(units[0], HistoryUnit::SystemGuard(_)));
assert!(matches!(units[1], HistoryUnit::UserMessage(_)));
assert!(matches!(units[2], HistoryUnit::AssistantText(_)));
}
#[test]
fn test_parse_to_units_tool_round_atomic() {
let messages = vec![
ChatMessage::user("read a file"),
ChatMessage::assistant_with_tool_calls(
"let me read",
vec![crate::domain::messages::ToolCall {
id: "call_1".to_string(),
name: "file_read".to_string(),
arguments: serde_json::json!({"path": "/data/test.txt"}),
}],
),
ChatMessage::tool("call_1", "file_read", "file contents here"),
ChatMessage::assistant("the file says hello"),
];
let units = parse_to_units(&messages);
assert_eq!(units.len(), 3); // UserMessage + ToolRound + AssistantText
assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
assert!(matches!(units[1], HistoryUnit::ToolRound { .. }));
assert!(matches!(units[2], HistoryUnit::AssistantText(_)));
// Verify ToolRound contains both assistant and results
if let HistoryUnit::ToolRound { assistant, results } = &units[1] {
assert_eq!(assistant.role, "assistant");
assert!(assistant.tool_calls.is_some());
assert_eq!(results.len(), 1);
assert_eq!(results[0].role, "tool");
assert_eq!(results[0].tool_call_id, Some("call_1".to_string()));
}
}
#[test]
fn test_parse_to_units_orphaned_tool_result_dropped() {
let messages = vec![
ChatMessage::user("hello"),
// Orphaned tool result — no preceding assistant with tool_calls
ChatMessage::tool("orphan_1", "bash", "some output"),
ChatMessage::assistant("done"),
];
let units = parse_to_units(&messages);
// Orphaned tool result should be dropped, leaving UserMessage + AssistantText
assert_eq!(units.len(), 2);
assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
assert!(matches!(units[1], HistoryUnit::AssistantText(_)));
}
#[test]
fn test_parse_to_units_multiple_tool_results_same_round() {
let messages = vec![
ChatMessage::user("do multiple things"),
ChatMessage::assistant_with_tool_calls(
"doing things",
vec![
crate::domain::messages::ToolCall {
id: "call_a".to_string(),
name: "bash".to_string(),
arguments: serde_json::json!({"command": "ls"}),
},
crate::domain::messages::ToolCall {
id: "call_b".to_string(),
name: "file_read".to_string(),
arguments: serde_json::json!({"path": "/data/test.txt"}),
},
],
),
ChatMessage::tool("call_a", "bash", "file1.txt\nfile2.txt"),
ChatMessage::tool("call_b", "file_read", "hello world"),
];
let units = parse_to_units(&messages);
assert_eq!(units.len(), 2); // UserMessage + ToolRound
if let HistoryUnit::ToolRound { results, .. } = &units[1] {
assert_eq!(results.len(), 2);
assert_eq!(results[0].tool_call_id, Some("call_a".to_string()));
assert_eq!(results[1].tool_call_id, Some("call_b".to_string()));
} else {
panic!("Expected ToolRound");
}
}
#[test]
fn test_parse_to_units_empty_input() {
let units = parse_to_units(&[]);
assert!(units.is_empty());
}
#[test]
fn test_find_safe_split_point() {
let compressor = ContextCompressor::new(100_000);
// Build a list of units with known token sizes
let messages: Vec<ChatMessage> = (0..10)
.map(|i| ChatMessage::assistant(&format!("message content number {}", i)))
.collect();
let units = parse_to_units(&messages);
// All units are AssistantText, split at 50% token ratio
let split = compressor.find_safe_split_point(&units, 0.5);
// Should split somewhere in the middle (not 0, not len())
assert!(split > 0 && split < units.len(),
"split {} should be between 0 and {}", split, units.len());
}
#[test]
fn test_compress_two_segment_no_tool_calls_in_output() {
// This test verifies the critical invariant:
// compress_two_segment output MUST NOT contain any tool_calls or tool messages.
let compressor = ContextCompressor::new(128_000);
// Build a simple history that's well under the threshold
// so compress_two_segment returns it unchanged (no LLM call needed).
let history = vec![
ChatMessage::system_with_context(
"You are a helpful assistant.",
Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
),
ChatMessage::user("hello"),
ChatMessage::assistant("hi there"),
];
// Dummy config — won't be used because history is under threshold.
let config = LLMProviderConfig {
provider_type: "openai".to_string(),
name: "test".to_string(),
base_url: "http://localhost".to_string(),
api_key: "sk-test".to_string(),
extra_headers: std::collections::HashMap::new(),
llm_timeout_secs: 120,
memory_maintenance_timeout_secs: 300,
model_id: "test-model".to_string(),
temperature: None,
max_tokens: None,
context_window_tokens: Some(128_000),
model_extra: std::collections::HashMap::new(),
max_tool_iterations: 100,
tool_result_max_chars: 100_000,
context_tool_result_trim_chars: 2_000,
max_images_in_context: 10,
max_image_age_rounds: 50,
};
let runtime = tokio::runtime::Runtime::new().unwrap();
let result = runtime.block_on(compressor.compress_two_segment(&history, &config));
assert!(result.is_ok());
let compressed = result.unwrap();
// Under threshold: should return unchanged (3 messages)
assert_eq!(compressed.len(), 3);
// Critical invariant: NO tool_calls or tool_call_id anywhere
for msg in &compressed {
assert!(msg.tool_calls.is_none(),
"compress_two_segment output should never contain tool_calls");
assert!(msg.tool_call_id.is_none(),
"compress_two_segment output should never contain tool_call_id");
}
}
}

2
src/channels/cli.rs
View File

@ -65,7 +65,7 @@ impl Default for CliChannel {
#[async_trait]
impl Channel for CliChannel {
fn name(&self) -> &str {
"websocket"
"cli"
}
fn is_running(&self) -> bool {

28
src/channels/manager.rs
View File

@ -15,19 +15,19 @@ use crate::protocol::Channel as ProtocolChannel;
pub struct ChannelManager {
channels: Arc<RwLock<HashMap<String, Arc<dyn Channel + Send + Sync>>>>,
bus: Arc<MessageBus>,
websocket_channel: Arc<CliChannel>,
cli_channel: Arc<CliChannel>,
}
impl ChannelManager {
pub fn new() -> Self {
let websocket_channel = Arc::new(CliChannel::new());
let cli_channel = Arc::new(CliChannel::new());
let mut channels: HashMap<String, Arc<dyn Channel + Send + Sync>> = HashMap::new();
channels.insert("websocket".to_string(), websocket_channel.clone());
channels.insert("cli".to_string(), cli_channel.clone());
Self {
channels: Arc::new(RwLock::new(channels)),
bus: MessageBus::new(100),
websocket_channel,
cli_channel,
}
}
@ -36,8 +36,8 @@ impl ChannelManager {
self.bus.clone()
}
pub fn websocket_channel(&self) -> Arc<CliChannel> {
self.websocket_channel.clone()
pub fn cli_channel(&self) -> Arc<CliChannel> {
self.cli_channel.clone()
}
/// Initialize all Channel instances from config
@ -142,7 +142,16 @@ impl ChannelManager {
let mut seen = HashSet::new();
let mut channels: Vec<ProtocolChannel> = Vec::new();
// 所有注册的通道（websocket, feishu, wechat 等）
// 1. WebSocket 通道 — Web 前端自己的连接，始终存在
seen.insert("websocket".to_string());
channels.push(ProtocolChannel {
id: "websocket".to_string(),
name: "WebSocket".to_string(),
description: Some("Web 前端通道".to_string()),
is_writable: true,
});
// 2. 所有动态注册的通道（cli, feishu, wechat 等）
for (name, _channel) in self.channels().await {
if seen.contains(&name) {
continue;
@ -163,6 +172,7 @@ impl ChannelManager {
fn channel_display_name(name: &str) -> String {
match name {
"websocket" => "WebSocket".to_string(),
"cli" => "命令行".to_string(),
"feishu" => "飞书".to_string(),
"wechat" => "微信".to_string(),
other => other.to_string(),
@ -252,7 +262,7 @@ mod tests {
.collect::<Vec<_>>();
names.sort();
assert_eq!(names, vec!["backup", "primary", "websocket"]);
assert_eq!(names, vec!["backup", "cli", "primary"]);
assert_eq!(manager.get_channel("primary").await.unwrap().name(), "primary");
assert_eq!(manager.get_channel("backup").await.unwrap().name(), "backup");
}
@ -313,7 +323,7 @@ mod tests {
.collect::<Vec<_>>();
names.sort();
assert_eq!(names, vec!["websocket", "wechat_main"]);
assert_eq!(names, vec!["cli", "wechat_main"]);
assert_eq!(manager.get_channel("wechat_main").await.unwrap().name(), "wechat_main");
}
}

12
src/command/handlers/mod.rs
View File

@ -51,13 +51,9 @@ pub async fn get_messages_from_session(
.map(|m| m.clone())
.unwrap_or_default())
}
None => {
tracing::warn!(
channel = %channel_name,
chat_id = %chat_id,
"No in-memory session, returning empty message list"
);
Ok(Vec::new())
}
None => Err(CommandError::new(
"SESSION_NOT_FOUND",
format!("Session not found for channel: {}", channel_name),
)),
}
}

117
src/gateway/compaction.rs
View File

@ -6,60 +6,97 @@ use crate::agent::AgentError;
use super::session::Session;
/// Run two-segment history compression synchronously.
///
/// Unlike the previous background approach (tokio::spawn), this holds the
/// session lock during the LLM calls (2–5 seconds). Since the agent loop
/// has already finished by this point there is no response-time impact, and
/// the synchronous guarantee means the next execution always starts with
/// freshly compacted history.
pub(crate) async fn schedule_background_history_compaction(
session: Arc<Mutex<Session>>,
chat_id: impl Into<String>,
) -> Result<(), AgentError> {
let chat_id = chat_id.into();
let mut session_guard = session.lock().await;
session_guard.ensure_persistent_session(&chat_id)?;
session_guard.ensure_chat_loaded(&chat_id)?;
let snapshot = {
let mut session_guard = session.lock().await;
let session_record = session_guard.ensure_persistent_session(&chat_id)?;
session_guard.ensure_chat_loaded(&chat_id)?;
let history = session_guard.get_or_create_history(&chat_id).clone();
let compressor = session_guard.compressor().clone();
let history = session_guard.get_or_create_history(&chat_id).clone();
let compressor = session_guard.compressor().clone();
if !compressor.should_compress(&history) {
return Ok(());
}
if !compressor.should_compress(&history) {
return Ok(());
}
if !session_guard.try_start_background_compaction(&chat_id) {
return Ok(());
}
let store = session_guard.store();
let session_id = session_guard.persistent_session_id(&chat_id);
let provider_config = session_guard.provider_config().clone();
(
session_guard.store(),
session_guard.persistent_session_id(&chat_id),
session_record.message_count,
history,
compressor,
session_guard.provider_config().clone(),
)
};
tracing::info!(
chat_id = %chat_id,
msg_count = history.len(),
"Starting synchronous two-segment compression"
);
let (
store,
session_id,
snapshot_end_seq,
history,
compressor,
provider_config,
) = snapshot;
let session_for_task = session.clone();
let chat_id_for_task = chat_id.clone();
// Synchronous compression — holds lock during LLM calls.
// compress_two_segment guarantees the result contains no tool_calls,
// so there is no risk of orphaned tool call sequences.
let compressed = compressor
.compress_two_segment(&history, &provider_config)
.await?;
tokio::spawn(async move {
tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Starting background history compaction");
// Replace the entire history with the compressed result.
// Since we hold the lock, no concurrent modifications can occur.
store
.replace_active_history(&session_id, &compressed)
.map_err(|e| AgentError::Other(format!("replace_active_history error: {}", e)))?;
let compaction_result = compressor
.build_compaction_plan(&history, &provider_config)
.await;
let mut committed = false;
tracing::info!(
chat_id = %chat_id,
compressed_msg_count = compressed.len(),
"Two-segment compression committed"
);
match compaction_result {
Ok(Some(plan)) => match store.compact_active_history(
&session_id,
snapshot_end_seq,
&plan.preserved_system_messages,
&plan.summary_message,
&plan.preserved_messages,
) {
Ok(true) => {
committed = true;
tracing::info!(
chat_id = %chat_id_for_task,
snapshot_end_seq,
compressed_turns = plan.compressed_turns,
preserved_turns = plan.preserved_turns,
"Background history compaction committed"
);
}
Ok(false) => {
tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Background history compaction skipped due to stale snapshot");
}
Err(error) => {
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction commit failed");
}
},
Ok(None) => {
tracing::debug!(chat_id = %chat_id_for_task, "Background history compaction not needed after snapshot analysis");
}
Err(error) => {
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction build failed");
}
}
session_guard.reload_chat_history(&chat_id)?;
let mut session_guard = session_for_task.lock().await;
if committed {
if let Err(error) = session_guard.reload_chat_history(&chat_id_for_task) {
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Failed to reload history after background compaction");
}
}
session_guard.finish_background_compaction(&chat_id_for_task);
});
Ok(())
}

65
src/gateway/http.rs
View File

@ -62,7 +62,7 @@ pub struct SaveConfigResponse {
pub async fn get_config(
State(state): State<Arc<GatewayState>>,
) -> Json<Config> {
Json(mask_config(&*state.config.read().await))
Json(mask_config(&state.config))
}
/// PUT /api/config — Save config to file, preserving original api_keys if masked
@ -72,29 +72,26 @@ pub async fn save_config(
) -> Result<Json<SaveConfigResponse>, (StatusCode, String)> {
// Merge: preserve original api_keys if the submitted ones are masked
let mut new_config = req.config;
// Read old values under read lock (not held across disk I/O)
{
let cfg = state.config.read().await;
for (name, provider) in new_config.providers.iter_mut() {
if is_masked_key(&provider.api_key) {
if let Some(original) = cfg.providers.get(name) {
provider.api_key = original.api_key.clone();
}
for (name, provider) in new_config.providers.iter_mut() {
if is_masked_key(&provider.api_key) {
// Restore original api_key
if let Some(original) = state.config.providers.get(name) {
provider.api_key = original.api_key.clone();
}
}
for (name, channel) in new_config.channels.iter_mut() {
if let Some(feishu) = channel.as_feishu_mut() {
if is_masked_key(&feishu.app_secret) {
if let Some(original_channel) = cfg.channels.get(name) {
if let Some(original_feishu) = original_channel.as_feishu() {
feishu.app_secret = original_feishu.app_secret.clone();
}
}
// Merge: preserve original app_secrets if the submitted ones are masked
for (name, channel) in new_config.channels.iter_mut() {
if let Some(feishu) = channel.as_feishu_mut() {
if is_masked_key(&feishu.app_secret) {
if let Some(original_channel) = state.config.channels.get(name) {
if let Some(original_feishu) = original_channel.as_feishu() {
feishu.app_secret = original_feishu.app_secret.clone();
}
}
}
}
} // read lock released here
}
// Validate timezone
if let Err(e) = new_config.time.parse_timezone() {
@ -106,19 +103,13 @@ pub async fn save_config(
.map(std::path::PathBuf::from)
.unwrap_or_else(|_| get_default_config_path());
// Serialize and write to disk (no lock held)
// Serialize and write
let json = serde_json::to_string_pretty(&new_config)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Serialize error: {}", e)))?;
std::fs::write(&config_path, &json)
std::fs::write(&config_path, json)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Write error: {}", e)))?;
// Update in-memory config (write lock, held only for assignment)
{
let mut cfg = state.config.write().await;
*cfg = new_config.clone();
}
tracing::info!(path = %config_path.display(), "Config saved via API");
Ok(Json(SaveConfigResponse {
@ -160,25 +151,3 @@ pub async fn restart(
message: "服务正在重启...".to_string(),
}))
}
/// GET /api/mcp/status — Return MCP server connection status
pub async fn mcp_status(
State(state): State<Arc<GatewayState>>,
) -> Json<crate::mcp::client::McpStatusResponse> {
let status = match &state.mcp_manager {
Some(manager) => {
// Clone mcp_servers before await to avoid holding read lock across it
let mcp_servers = state.config.read().await.mcp_servers.clone();
manager.get_status(&mcp_servers).await
}
None => crate::mcp::client::McpStatusResponse {
enabled: false,
total_servers: 0,
connected_servers: 0,
failed_servers: 0,
total_tools: 0,
servers: vec![],
},
};
Json(status)
}

52
src/gateway/mod.rs
View File

@ -52,10 +52,10 @@ use session_message_sender::BusSessionMessageSender;
use session::SessionManager;
use static_files::static_handler;
use tokio::sync::{watch, RwLock};
use tokio::sync::watch;
pub struct GatewayState {
pub config: Arc<RwLock<Config>>,
pub config: Config,
pub session_manager: SessionManager,
pub channel_manager: ChannelManager,
pub bus: Arc<MessageBus>,
@ -107,7 +107,7 @@ impl GatewayState {
let cancel_manager = CancelManager::new();
Ok(Self {
config: Arc::new(RwLock::new(config)),
config,
session_manager,
channel_manager,
bus,
@ -122,19 +122,18 @@ impl GatewayState {
pub async fn start_message_processing(&self) {
let bus_for_outbound = self.bus.clone();
// Read config under read lock
let cfg = self.config.read().await;
let max_concurrent = cfg.gateway.max_concurrent_requests;
let provider_config = match cfg.get_provider_config("default") {
// Create semaphore for controlling concurrent requests
let max_concurrent = self.config.gateway.max_concurrent_requests;
let semaphore = Arc::new(Semaphore::new(max_concurrent));
// Spawn inbound processor with semaphore-controlled concurrency
let provider_config = match self.config.get_provider_config("default") {
Ok(config) => config,
Err(e) => {
tracing::error!(error = %e, "Failed to get provider config");
return;
}
};
drop(cfg); // release read lock before spawning long-running tasks
let semaphore = Arc::new(Semaphore::new(max_concurrent));
let inbound_processor =
InboundProcessor::new(self.bus.clone(), self.session_manager.clone(), semaphore, provider_config, self.cancel_manager.clone());
tokio::spawn(inbound_processor.run());
@ -171,30 +170,23 @@ pub async fn run(
let state = Arc::new(GatewayState::from_config(config, restart_tx)?);
// Get provider config for channels
let cfg = state.config.read().await;
let provider_config = cfg.get_provider_config("default")?;
let provider_config = state.config.get_provider_config("default")?;
// Initialize and start channels
state
.channel_manager
.init(&*cfg, provider_config.clone())
.init(&state.config, provider_config.clone())
.await?;
drop(cfg);
state.channel_manager.start_all().await?;
// Start message processing (inbound processor + outbound dispatcher)
state.start_message_processing().await;
let (scheduler_shutdown_tx, scheduler_shutdown_rx) = tokio::sync::watch::channel(false);
let scheduler_enabled = {
let cfg = state.config.read().await;
cfg.scheduler.enabled
};
if scheduler_enabled {
let scheduler_cfg = state.config.read().await.scheduler.clone();
if state.config.scheduler.enabled {
let scheduler = Scheduler::new(
state.bus.clone(),
scheduler_cfg,
state.config.scheduler.clone(),
timezone,
state.session_manager.store(),
AgentTaskExecutor::new(state.session_manager.clone()),
@ -207,12 +199,8 @@ pub async fn run(
}
// CLI args override config file values
let (bind_host, bind_port) = {
let cfg = state.config.read().await;
let h = host.unwrap_or_else(|| cfg.gateway.host.clone());
let p = port.unwrap_or(cfg.gateway.port);
(h, p)
};
let bind_host = host.unwrap_or_else(|| state.config.gateway.host.clone());
let bind_port = port.unwrap_or(state.config.gateway.port);
// 使用嵌入的静态文件（编译时打包进二进制）
// 开发模式下可通过 STATIC_DIR 环境变量使用磁盘文件
@ -223,7 +211,6 @@ pub async fn run(
.route("/health", routing::get(http::health))
.route("/api/config", routing::get(http::get_config).put(http::save_config))
.route("/api/restart", routing::post(http::restart))
.route("/api/mcp/status", routing::get(http::mcp_status))
.route("/ws", routing::get(ws::ws_handler))
.fallback(static_handler)
.with_state(state.clone())
@ -233,7 +220,6 @@ pub async fn run(
.route("/health", routing::get(http::health))
.route("/api/config", routing::get(http::get_config).put(http::save_config))
.route("/api/restart", routing::post(http::restart))
.route("/api/mcp/status", routing::get(http::mcp_status))
.route("/ws", routing::get(ws::ws_handler))
.fallback_service(ServeDir::new(&static_dir))
.with_state(state.clone())
@ -267,8 +253,8 @@ pub async fn run(
cancel_manager.cancel_all().await;
let _ = scheduler_shutdown_tx.send(true);
if let Some(ref mgr) = mcp_manager {
tracing::info!("Shutting down MCP servers before shutdown");
let _ = mgr.shutdown_all().await;
tracing::info!("Disconnecting MCP servers before shutdown");
let _ = mgr.disconnect_all().await;
}
let _ = channel_manager.stop_all().await;
let _ = result_tx.send(false);
@ -280,8 +266,8 @@ pub async fn run(
cancel_manager.cancel_all().await;
let _ = scheduler_shutdown_tx.send(true);
if let Some(ref mgr) = mcp_manager {
tracing::info!("Shutting down MCP servers before restart");
let _ = mgr.shutdown_all().await;
tracing::info!("Disconnecting MCP servers before restart");
let _ = mgr.disconnect_all().await;
}
let _ = channel_manager.stop_all().await;
let _ = result_tx.send(true);

2
src/gateway/session.rs
View File

@ -526,12 +526,10 @@ impl Session {
&self.compressor
}
#[allow(dead_code)]
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
self.history.try_start_background_compaction(chat_id)
}
#[allow(dead_code)]
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
self.history.finish_background_compaction(chat_id);
}

2
src/gateway/session_history.rs
View File

@ -235,12 +235,10 @@ impl SessionHistory {
Ok(())
}
#[allow(dead_code)]
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
self.compression_in_flight.insert(chat_id.to_string())
}
#[allow(dead_code)]
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
self.compression_in_flight.remove(chat_id);
}

40
src/gateway/ws.rs
View File

@ -40,7 +40,7 @@ use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::mpsc;
const WS_CHANNEL_NAME: &str = "websocket";
const CLI_CHANNEL_NAME: &str = "cli";
/// Default media directory for WebSocket uploads
fn default_ws_media_dir() -> PathBuf {
@ -134,13 +134,9 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
let cli_sessions = state.session_manager.cli_sessions();
let store = state.session_manager.store();
// 1. 查询 websocket 和 cli 两个通道的 Sessions（兼容旧版本 cli 通道创建的会话）
let mut websocket_sessions = store.list_sessions("websocket", false)
// 1. 先查询 websocket 通道的 Sessions
let websocket_sessions = store.list_sessions("websocket", false)
.unwrap_or_default();
let cli_channel_sessions = store.list_sessions("cli", false)
.unwrap_or_default();
websocket_sessions.extend(cli_channel_sessions);
websocket_sessions.sort_by_key(|s| -(s.last_active_at));
// 2. 如果没有，自动创建一个默认 Session
let initial_record = if websocket_sessions.is_empty() {
@ -161,7 +157,7 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
let mut current_topic_id: Option<String> = None;
state
.channel_manager
.websocket_channel()
.cli_channel()
.register_connection(
current_session_id.clone(),
runtime_session_id.clone(),
@ -182,20 +178,16 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
.send(WsOutbound::ChannelList { channels })
.await;
// 3. 发送合并后的 Session 列表（已在上面合并了 websocket + cli 通道）
// 如果刚创建了新会话，确保它也在列表中
let has_initial = websocket_sessions.iter().any(|s| s.id == initial_record.id);
if !has_initial {
websocket_sessions.push(initial_record);
websocket_sessions.sort_by_key(|s| -(s.last_active_at));
}
// 3. 重新查询 websocket 通道的 Session 列表（包含刚创建的）
let final_sessions = store.list_sessions("websocket", false)
.unwrap_or_default();
tracing::info!("Sending {} sessions to client", websocket_sessions.len());
for s in &websocket_sessions {
tracing::info!("Sending {} websocket sessions to client", final_sessions.len());
for s in &final_sessions {
tracing::info!(" - {}: {} (channel: {})", s.id, s.title, s.channel_name);
}
let session_summaries: Vec<crate::protocol::SessionSummary> = websocket_sessions
let session_summaries: Vec<crate::protocol::SessionSummary> = final_sessions
.into_iter()
.map(|s| crate::protocol::SessionSummary {
session_id: s.id,
@ -281,7 +273,7 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
state
.channel_manager
.websocket_channel()
.cli_channel()
.unregister_connection(&runtime_session_id)
.await;
tracing::info!(session_id = %runtime_session_id, current_session_id = %current_session_id, "CLI session ended");
@ -309,7 +301,7 @@ async fn handle_inbound(
state
.channel_manager
.websocket_channel()
.cli_channel()
.register_connection(
chat_id.clone(),
runtime_session_id.to_string(),
@ -323,7 +315,7 @@ async fn handle_inbound(
state
.bus
.publish_inbound(InboundMessage {
channel: WS_CHANNEL_NAME.to_string(),
channel: CLI_CHANNEL_NAME.to_string(),
sender_id,
chat_id,
content,
@ -376,7 +368,7 @@ async fn handle_inbound(
let store = state.session_manager.store();
let skills = state.session_manager.skills();
let skills_for_handler = skills.clone();
let provider_config = state.config.read().await.get_provider_config("default")
let provider_config = state.config.get_provider_config("default")
.map_err(|e| AgentError::Other(e.to_string()))?;
let prompt_repository = state.session_manager.store().clone();
@ -458,7 +450,7 @@ async fn handle_inbound(
current_topic_id = ?current_topic_id,
"Building CommandContext for WebSocket command"
);
let mut cmd_ctx = CommandContext::new("websocket", "websocket")
let mut cmd_ctx = CommandContext::new("websocket", "cli")
.with_session_id(current_session_id.as_str())
.with_chat_id(current_session_id.as_str());
// 只在有 topic_id 时才设置
@ -481,7 +473,7 @@ async fn handle_inbound(
*current_session_id = session_id.clone();
state
.channel_manager
.websocket_channel()
.cli_channel()
.register_connection(
session_id.clone(),
runtime_session_id.to_string(),

120
src/mcp/client.rs
View File

@ -61,10 +61,6 @@ pub struct McpClientManager {
clients: RwLock<HashMap<String, Arc<McpClient>>>,
/// Server information cache keyed by server key
server_info: RwLock<HashMap<String, McpServerInfo>>,
/// Count of active stdio (child process) connections
stdio_client_count: std::sync::atomic::AtomicUsize,
/// Connection errors per server key (last error message)
connection_errors: RwLock<HashMap<String, String>>,
}
impl McpClientManager {
@ -73,8 +69,6 @@ impl McpClientManager {
Self {
clients: RwLock::new(HashMap::new()),
server_info: RwLock::new(HashMap::new()),
stdio_client_count: std::sync::atomic::AtomicUsize::new(0),
connection_errors: RwLock::new(HashMap::new()),
}
}
@ -145,12 +139,7 @@ impl McpClientManager {
attempts = MAX_RETRIES,
"Failed to connect to MCP server after all retries"
);
// Record error for status reporting
self.connection_errors.write().await.insert(key.clone(), e.to_string());
failed += 1;
} else {
// Clear any previous error on successful connection
self.connection_errors.write().await.remove(&key);
}
}
@ -236,9 +225,6 @@ impl McpClientManager {
// Use default client handler (empty tuple)
let client = ().serve(transport).await?;
// Track that we have a stdio (child process) connection
self.stdio_client_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(client)
}
@ -375,116 +361,10 @@ impl McpClientManager {
Ok(())
}
/// Shut down all MCP connections with proper cleanup
///
/// Drops all client connections and waits for child processes to terminate
/// if stdio transports were in use. This prevents race conditions during
/// gateway restart where old MCP processes may still be running when
/// new ones start.
pub async fn shutdown_all(&self) -> anyhow::Result<()> {
let stdio_count = self.stdio_client_count.load(std::sync::atomic::Ordering::SeqCst);
// Drop all clients (triggers cancellation + graceful shutdown in rmcp)
self.disconnect_all().await?;
// If stdio connections were active, wait for child processes to be killed.
// rmcp's RunningService::drop() triggers async cancellation with:
// - 2 second graceful drain period
// - 3 second process kill timeout
// Total: ~5 seconds. We add 1 second buffer.
if stdio_count > 0 {
tracing::info!(
stdio_count,
"Waiting for MCP child processes to terminate (up to 6s)..."
);
tokio::time::sleep(std::time::Duration::from_secs(6)).await;
tracing::info!("MCP child process cleanup wait complete");
self.stdio_client_count.store(0, std::sync::atomic::Ordering::SeqCst);
}
Ok(())
}
/// Check if any servers are connected
pub async fn has_connections(&self) -> bool {
!self.clients.read().await.is_empty()
}
/// Get current MCP connection status for all configured servers
pub async fn get_status(
&self,
configured_servers: &HashMap<String, crate::mcp::McpServerConfig>,
) -> McpStatusResponse {
let info_map = self.server_info.read().await;
let errors = self.connection_errors.read().await;
let clients = self.clients.read().await;
let mut total_servers = 0usize;
let mut connected_servers = 0usize;
let mut failed_servers = 0usize;
let mut total_tools = 0usize;
let mut servers = Vec::new();
for (key, config) in configured_servers {
total_servers += 1;
let name = config.effective_name(key);
let transport_type = config.transport_type.clone();
let is_active = config.is_active;
let connected = clients.contains_key(key);
let tool_count = info_map.get(key).map(|info| info.tools.len()).unwrap_or(0);
let error = errors.get(key).cloned();
if connected {
connected_servers += 1;
} else if is_active && error.is_some() {
failed_servers += 1;
}
total_tools += tool_count;
servers.push(McpServerStatus {
key: key.clone(),
name,
transport_type,
is_active,
connected,
tool_count,
error,
});
}
McpStatusResponse {
enabled: !configured_servers.is_empty(),
total_servers,
connected_servers,
failed_servers,
total_tools,
servers,
}
}
}
/// Status of a single MCP server connection
#[derive(Debug, Clone, serde::Serialize)]
pub struct McpServerStatus {
pub key: String,
pub name: String,
pub transport_type: String,
pub is_active: bool,
pub connected: bool,
pub tool_count: usize,
pub error: Option<String>,
}
/// Overall MCP status response
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct McpStatusResponse {
pub enabled: bool,
pub total_servers: usize,
pub connected_servers: usize,
pub failed_servers: usize,
pub total_tools: usize,
pub servers: Vec<McpServerStatus>,
}
impl Default for McpClientManager {

2
src/mcp/mod.rs
View File

@ -16,5 +16,5 @@ pub mod client;
pub mod tool_adapter;
pub use config::{McpConfig, McpServerConfig, McpTransportConfig};
pub use client::{McpClientManager, McpClient, McpServerInfo, McpInitializer, McpServerStatus, McpStatusResponse};
pub use client::{McpClientManager, McpClient, McpServerInfo, McpInitializer};
pub use tool_adapter::{McpToolWrapper, register_mcp_tools};

56
src/storage/mod.rs
View File

@ -814,59 +814,6 @@ impl SessionStore {
Ok(true)
}
/// Replace the entire active history for a session.
///
/// This is a simpler alternative to `compact_active_history` for when the
/// compressor has already produced a complete, validated message list
/// (e.g. two-segment compression). It replaces all existing messages
/// with the new list in a single transaction.
pub fn replace_active_history(
&self,
session_id: &str,
messages: &[ChatMessage],
) -> Result<(), StorageError> {
let conn = self.pool.get()?;
let tx = conn.unchecked_transaction()?;
let now = current_timestamp();
// Delete all existing messages for this session
tx.execute(
"DELETE FROM messages WHERE session_id = ?1",
params![session_id],
)?;
// Insert new messages with sequential seq numbers
let mut active_user_turn_count = 0_i64;
for (i, message) in messages.iter().enumerate() {
let seq = (i + 1) as i64;
if message.role == "user" {
active_user_turn_count += 1;
}
insert_message_with_seq(&tx, session_id, seq, message)?;
}
tx.execute(
"
UPDATE sessions
SET message_count = ?2,
user_turn_count = ?3,
updated_at = ?4,
last_active_at = ?4,
archived_at = NULL
WHERE id = ?1 AND deleted_at IS NULL
",
params![
session_id,
messages.len() as i64,
active_user_turn_count,
now,
],
)?;
tx.commit()?;
Ok(())
}
pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> {
let now = current_timestamp();
let conn = self.pool.get()?;
@ -1601,7 +1548,7 @@ impl SessionStore {
}
pub fn persistent_session_id(channel_name: &str, chat_id: &str) -> String {
if channel_name == "cli" || channel_name == "websocket" {
if channel_name == "cli" {
chat_id.to_string()
} else {
format!("{}:{}", channel_name, chat_id)
@ -2288,7 +2235,6 @@ mod tests {
#[test]
fn test_persistent_session_id_for_cli_and_channel() {
assert_eq!(persistent_session_id("cli", "abc"), "abc");
assert_eq!(persistent_session_id("websocket", "websocket:abc"), "websocket:abc");
assert_eq!(persistent_session_id(TEST_CHANNEL, "abc"), "test-channel:abc");
}

17
src/storage/ports.rs
View File

@ -54,15 +54,6 @@ pub trait ConversationRepository: Send + Sync + 'static {
summary_message: &ChatMessage,
preserved_messages: &[ChatMessage],
) -> Result<bool, StorageError>;
/// Replace the entire active history for a session with new messages.
/// Used when the compressor has already produced a complete, validated
/// message list (e.g. two-segment compression).
fn replace_active_history(
&self,
session_id: &str,
messages: &[ChatMessage],
) -> Result<(), StorageError>;
}
pub trait PromptInjectionRepository: Send + Sync + 'static {
@ -244,14 +235,6 @@ impl ConversationRepository for super::SessionStore {
preserved_messages,
)
}
fn replace_active_history(
&self,
session_id: &str,
messages: &[ChatMessage],
) -> Result<(), StorageError> {
super::SessionStore::replace_active_history(self, session_id, messages)
}
}
impl PromptInjectionRepository for super::SessionStore {

1
tests/test_request_format.rs
View File

@ -126,7 +126,6 @@ fn test_tool_call_outbound_serialization() {
topic_id: None,
timestamp: None,
reasoning_content: None,
user_message_id: None,
};
let json = serde_json::to_string(&msg).unwrap();

74
web/src/components/Settings/ConfigPage.tsx
View File

@ -20,7 +20,6 @@ interface ImageContextConfig { max_images_in_context: number; max_image_age_roun
interface SubagentsConfig { enabled: boolean; sources: string[] }
interface ClientConfig { gateway_url: string }
interface McpServerConfig {
name?: string
type: 'stdio' | 'streamableHttp' | 'http'
is_active: boolean
command?: string
@ -30,25 +29,6 @@ interface McpServerConfig {
headers?: Record<string, string>
description?: string
}
interface McpServerStatus {
key: string
name: string
transport_type: string
is_active: boolean
connected: boolean
tool_count: number
error?: string
}
interface McpStatusResponse {
enabled: boolean
total_servers: number
connected_servers: number
failed_servers: number
total_tools: number
servers: McpServerStatus[]
}
interface AppConfig {
providers: Record<string, ProviderConfig>
models: Record<string, ModelConfig>
@ -304,14 +284,6 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
const [dirty, setDirty] = useState(false)
const [showRestartDialog, setShowRestartDialog] = useState(false)
const [restarting, setRestarting] = useState(false)
const [mcpStatus, setMcpStatus] = useState<McpStatusResponse | null>(null)
const fetchMcpStatus = useCallback(async () => {
try {
const resp = await fetch('/api/mcp/status')
if (resp.ok) setMcpStatus(await resp.json())
} catch { /* ignore fetch errors */ }
}, [])
const handleClose = useCallback(() => {
if (dirty && !confirm('有未保存的更改，确定要关闭吗？')) return
@ -326,11 +298,6 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
}).catch(e => { setError('加载配置失败: ' + e.message); setLoading(false) })
}, [])
// Fetch MCP status when MCP tab is selected
useEffect(() => {
if (activeTab === 'mcp') fetchMcpStatus()
}, [activeTab, fetchMcpStatus])
// ESC to close
useEffect(() => {
const h = (e: KeyboardEvent) => { if (e.key === 'Escape') handleClose() }
@ -354,8 +321,9 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
})
const data = await resp.json()
if (!resp.ok) throw new Error(data.message || data.error || '保存失败')
// Config is now synced to both disk and in-memory state,
// so the local state is already correct. No need to re-fetch.
// Reload config from server to get masked values
const refreshed = await fetch('/api/config').then(r => r.json())
setConfig(refreshed)
setDirty(false)
// Show restart confirmation dialog
setShowRestartDialog(true)
@ -693,7 +661,6 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
const renderMcp = () => {
const entries = Object.entries(config.mcpServers)
const statusFor = (key: string) => mcpStatus?.servers?.find(s => s.key === key)
const addMcp = () => {
const name = prompt('MCP 服务器名称:')?.trim()
if (name && !config.mcpServers[name]) {
@ -704,37 +671,9 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
const updMcp = (name: string, patch: Partial<McpServerConfig>) => update('mcpServers', { ...config.mcpServers, [name]: { ...config.mcpServers[name], ...patch } })
return (
<div className="space-y-4">
{/* MCP Status Summary */}
{mcpStatus && mcpStatus.enabled && (
<div className="flex items-center gap-3 p-3 rounded-lg bg-[var(--bg-tertiary)] text-xs">
<div className="flex items-center gap-1.5">
<span className={`inline-block w-2 h-2 rounded-full ${mcpStatus.connected_servers > 0 ? 'bg-green-400' : 'bg-gray-400'}`} />
<span className="text-[var(--text-secondary)]">{mcpStatus.connected_servers}/{mcpStatus.total_servers} 已连接</span>
</div>
{mcpStatus.failed_servers > 0 && (
<span className="text-red-400">{mcpStatus.failed_servers} 失败</span>
)}
<span className="text-[var(--text-muted)]">{mcpStatus.total_tools} 个工具</span>
<button onClick={fetchMcpStatus} className="ml-auto px-2 py-1 rounded text-[var(--text-muted)] hover:text-[var(--accent-cyan)] transition-colors" title="刷新状态">
<RefreshCw className="h-3 w-3" />
</button>
</div>
)}
{entries.map(([name, s]) => {
const st = statusFor(name)
return (
{entries.map(([name, s]) => (
<div key={name} className="rounded-xl border border-[var(--border-color)] bg-[var(--bg-secondary)]/60 overflow-hidden">
<div className="flex items-center gap-2 p-3 border-b border-[var(--border-color)]">
{st ? (
st.connected
? <span className="inline-flex items-center gap-1 text-xs text-green-400"><span className="w-2 h-2 rounded-full bg-green-400" /> {st.tool_count} 工具</span>
: st.error
? <span className="inline-flex items-center gap-1 text-xs text-red-400" title={st.error}><span className="w-2 h-2 rounded-full bg-red-400" /> 错误</span>
: <span className="inline-flex items-center gap-1 text-xs text-gray-400"><span className="w-2 h-2 rounded-full bg-gray-400" /> 未连接</span>
) : null}
<span className="flex-1 text-sm font-medium text-[var(--text-primary)]">{name}</span>
<button onClick={() => delMcp(name)} className="p-1 rounded text-[var(--text-muted)] hover:text-red-400 transition-colors"><Trash2 className="h-3.5 w-3.5" /></button>
</div>
<MapEntryHeader name={name} onDelete={() => delMcp(name)} />
<div className="p-4 space-y-3">
<Field label="传输类型">
<select value={s.type} onChange={e => updMcp(name, { type: e.target.value as McpServerConfig['type'] })} className={selectCls}>
@ -783,8 +722,7 @@ export function ConfigPage({ onClose, onSaveConnection }: ConfigPageProps) {
)}
</div>
</div>
)
})}
))}
<button onClick={addMcp} className="flex items-center gap-2 px-4 py-2.5 rounded-xl border border-dashed border-[var(--border-color)] text-[var(--text-muted)] hover:text-[var(--accent-cyan)] hover:border-[var(--accent-cyan)]/30 transition-colors text-sm w-full justify-center">
<Plus className="h-4 w-4" /> 添加 MCP 服务器
</button>