Abeautifulsnow
c300336281
* @
feat(ai): add conversation context management and token usage tracking (Phase 4)
## Background
Phase 4 of AI agent loop upgrade: token usage extraction from LLM responses,
cost estimation display, and automatic context window management (summarization
when approaching limits).
## Changes
### Backend (Rust)
- New token_usage::TokenUsage — cross-turn token accumulation
- New agent_loop::maybe_compact — proactive context compaction: token estimation
→ cut-point selection → LLM summarization → fallback summary
→ tool-pair integrity protection → drain/insert
- New agent_loop::context_window_for_model — auto-detect context window by model
- New agent_loop::is_context_length_error — detect upstream context-length errors
- New agent_loop::compact_tool_result_for_context — JSON/text sampling for large results
- Modified agent_loop::run_agent_loop — pre-turn compaction check, retry on context error
- Modified agent_events::AgentEnd — split total_tokens into input_tokens/output_tokens
- New agent_events::ContextCompacted — compaction event with metrics
- New AiConfig.context_window — user-configurable context window size
- New AiChatMessage.kind — context summary message marker
- Modified stream_with_tools & 3 provider stream functions — extract TokenUsage
- Claude: message_start(input) + message_delta(output)
- OpenAI: stream_options { include_usage } → final usage chunk
- Gemini: usageMetadata per chunk, keep last
### Frontend (TypeScript/Vue)
- AiAssistant.vue: token + cost display row, estimateCost() with 7 model families,
contextCompacted UI step, pendingCompaction deferred apply,
messagesForAgentHistory post-compaction history reconstruction
- EditorSettingsDialog.vue: contextWindow input (min=1000, auto-detect)
- settingsStore.ts: AiConfig.contextWindow field
- tauri.ts: AgentEvent type updates, AiChatMessage.kind
- i18n: 6 new keys in en.ts and zh-CN.ts
## Fixed Issues
- messages[0] dropped from LLM context after compaction
- estimate_text_tokens hard threshold causing 2.8x estimate jump
- context_compacted event field semantics mismatch with frontend insert position
- validate_summary English-only keywords causing fallback on non-English LLMs
- Stale agentTokens after conversation switch
@
* remove cost computation logic
762 lines
28 KiB
Rust
762 lines
28 KiB
Rust
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::Arc;
|
|
|
|
use futures::future::join_all;
|
|
use futures::FutureExt;
|
|
use serde_json::json;
|
|
use tokio::sync::Notify;
|
|
|
|
use crate::agent_events::{AgentEvent, ToolCall, ToolDefinition, ToolResult};
|
|
use crate::agent_tools;
|
|
use crate::ai::{self, AiCompletionRequest, AiConfig, AiMessage, AiProvider, AiStreamChunk};
|
|
use crate::connection::AppState;
|
|
use crate::models::connection::DatabaseType;
|
|
use crate::token_usage::TokenUsage;
|
|
use tokio::sync::Mutex;
|
|
|
|
/// Maximum number of agent loop turns to prevent infinite loops.
|
|
const MAX_AGENT_TURNS: u32 = 10;
|
|
const MAX_TOOL_RESULT_CONTEXT_CHARS: usize = 12_000;
|
|
const TOOL_RESULT_HEAD_CHARS: usize = 4_000;
|
|
const TOOL_RESULT_TAIL_CHARS: usize = 4_000;
|
|
const TOOL_RESULT_SAMPLE_ITEMS: usize = 5;
|
|
|
|
/// Context for an agent loop run.
|
|
pub struct AgentLoopContext {
|
|
pub state: Arc<AppState>,
|
|
pub connection_id: String,
|
|
pub database: String,
|
|
pub db_type: DatabaseType,
|
|
}
|
|
|
|
/// Check if the provider supports function calling / tool use.
|
|
/// Returns false for providers that are known to lack reliable tool support.
|
|
fn provider_supports_function_calling(config: &AiConfig) -> bool {
|
|
match config.provider {
|
|
// Ollama function calling support varies by model/version; conservative default is false.
|
|
// Users with capable models can override via openai-compatible with an Ollama endpoint.
|
|
AiProvider::Ollama => false,
|
|
_ => true,
|
|
}
|
|
}
|
|
|
|
/// Run the agent loop: call LLM with tools, execute tool calls, feed results back, repeat.
|
|
///
|
|
/// The `on_event` callback receives streaming events for the frontend.
|
|
/// Returns the final accumulated assistant text.
|
|
///
|
|
/// If the provider does not support function calling (e.g., Ollama), automatically
|
|
/// degrades to a text-only completion with schema context injected into the system prompt.
|
|
pub async fn run_agent_loop(
|
|
config: &AiConfig,
|
|
system_prompt: &str,
|
|
messages: &[AiMessage],
|
|
agent_ctx: &AgentLoopContext,
|
|
on_event: impl Fn(AgentEvent) + Send + Sync + Clone + 'static,
|
|
cancelled: &Notify,
|
|
max_tokens: Option<u32>,
|
|
temperature: Option<f32>,
|
|
is_agent_mode: bool,
|
|
) -> Result<String, String> {
|
|
// Auto-degrade: providers without function calling fall back to text-only completion.
|
|
if !provider_supports_function_calling(config) {
|
|
return run_agent_loop_text_only(
|
|
config,
|
|
system_prompt,
|
|
messages,
|
|
agent_ctx,
|
|
on_event,
|
|
cancelled,
|
|
max_tokens,
|
|
temperature,
|
|
)
|
|
.await;
|
|
}
|
|
let tools = if is_agent_mode { agent_tools::all_tools(agent_ctx.db_type) } else { agent_tools::read_only_tools() };
|
|
let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
|
|
let mut final_text = String::new();
|
|
let mut total_usage = TokenUsage::default();
|
|
|
|
for turn in 0..MAX_AGENT_TURNS {
|
|
// Check for cancellation before each turn
|
|
if cancelled.notified().now_or_never().is_some() {
|
|
on_event(AgentEvent::Error { message: "Agent loop cancelled".to_string() });
|
|
break;
|
|
}
|
|
|
|
// Check and maybe compact context
|
|
maybe_compact(config, system_prompt, &tools, &mut conversation_messages, max_tokens, &on_event, false).await;
|
|
|
|
on_event(AgentEvent::TurnStart { turn });
|
|
|
|
let mut stream_result: Option<(Vec<ToolCall>, Option<TokenUsage>, String)> = None;
|
|
let mut last_stream_error: Option<String> = None;
|
|
|
|
for attempt in 0..2 {
|
|
// Build the LLM request with tools. Rebuild after retry compaction so the request
|
|
// reflects the latest conversation_messages.
|
|
let request =
|
|
build_tool_request(config, system_prompt, &conversation_messages, &tools, max_tokens, temperature);
|
|
|
|
// Stream the LLM response, collecting text and tool_calls.
|
|
let accumulated_text = Arc::new(Mutex::new(String::new()));
|
|
let emitted_any_chunk = Arc::new(AtomicBool::new(false));
|
|
let session_id =
|
|
if attempt == 0 { format!("agent-turn-{turn}") } else { format!("agent-turn-{turn}-retry") };
|
|
|
|
let acc = accumulated_text.clone();
|
|
let emitted = emitted_any_chunk.clone();
|
|
let on_event2 = on_event.clone();
|
|
let on_chunk = move |chunk: AiStreamChunk| {
|
|
if !chunk.delta.is_empty() {
|
|
emitted.store(true, Ordering::Relaxed);
|
|
if let Ok(mut text) = acc.try_lock() {
|
|
text.push_str(&chunk.delta);
|
|
}
|
|
on_event2(AgentEvent::TextDelta { delta: chunk.delta.clone() });
|
|
}
|
|
if let Some(ref reasoning) = chunk.reasoning_delta {
|
|
emitted.store(true, Ordering::Relaxed);
|
|
on_event2(AgentEvent::ReasoningDelta { delta: reasoning.clone() });
|
|
}
|
|
};
|
|
|
|
match stream_with_tools(config, &request, &session_id, &tools, cancelled, on_chunk).await {
|
|
Ok((tool_calls, usage)) => {
|
|
let accumulated_text = accumulated_text.lock().await.clone();
|
|
stream_result = Some((tool_calls, usage, accumulated_text));
|
|
break;
|
|
}
|
|
Err(err)
|
|
if attempt == 0 && is_context_length_error(&err) && !emitted_any_chunk.load(Ordering::Relaxed) =>
|
|
{
|
|
last_stream_error = Some(err);
|
|
let compacted = maybe_compact(
|
|
config,
|
|
system_prompt,
|
|
&tools,
|
|
&mut conversation_messages,
|
|
max_tokens,
|
|
&on_event,
|
|
true,
|
|
)
|
|
.await;
|
|
if compacted.is_some() {
|
|
continue;
|
|
}
|
|
break;
|
|
}
|
|
Err(err) => return Err(err),
|
|
}
|
|
}
|
|
|
|
let Some((collected_tool_calls, turn_usage, accumulated_text)) = stream_result else {
|
|
return Err(
|
|
last_stream_error.unwrap_or_else(|| "LLM request failed after context compaction retry".to_string())
|
|
);
|
|
};
|
|
|
|
if let Some(usage) = turn_usage {
|
|
total_usage.add(&usage);
|
|
}
|
|
|
|
on_event(AgentEvent::TurnEnd { turn });
|
|
|
|
// Add assistant message to conversation (including tool_use blocks)
|
|
conversation_messages.push(AiMessage {
|
|
role: "assistant".to_string(),
|
|
content: accumulated_text.clone(),
|
|
tool_call_id: None,
|
|
tool_calls: collected_tool_calls
|
|
.iter()
|
|
.map(|tc| ai::ToolCallRef { id: tc.id.clone(), name: tc.name.clone(), arguments: tc.arguments.clone() })
|
|
.collect(),
|
|
});
|
|
|
|
if collected_tool_calls.is_empty() {
|
|
// No tool calls -- we're done
|
|
final_text = accumulated_text;
|
|
break;
|
|
}
|
|
|
|
// Execute each tool call
|
|
// Emit all ToolCallStart events first
|
|
for tc in &collected_tool_calls {
|
|
on_event(AgentEvent::ToolCallStart {
|
|
tool_call_id: tc.id.clone(),
|
|
tool_name: tc.name.clone(),
|
|
args: tc.arguments.clone(),
|
|
});
|
|
}
|
|
|
|
// Execute tool calls: parallel for read tools, sequential for execute_query
|
|
let state2 = Arc::clone(&agent_ctx.state);
|
|
let conn2 = agent_ctx.connection_id.clone();
|
|
let db2 = agent_ctx.database.clone();
|
|
let db_type = agent_ctx.db_type;
|
|
|
|
// Split by index into parallel and sequential groups using tool metadata
|
|
let tool_parallel_map: std::collections::HashMap<&str, bool> =
|
|
tools.iter().map(|t| (t.name, t.parallel_ok)).collect();
|
|
let (parallel_indices, sequential_indices): (Vec<usize>, Vec<usize>) = (0..collected_tool_calls.len())
|
|
.partition(|&i| *tool_parallel_map.get(collected_tool_calls[i].name.as_str()).unwrap_or(&false));
|
|
|
|
let make_tc =
|
|
|tc: &ToolCall| ToolCall { id: tc.id.clone(), name: tc.name.clone(), arguments: tc.arguments.clone() };
|
|
|
|
// Run parallel group
|
|
let parallel_futures: Vec<_> = parallel_indices
|
|
.iter()
|
|
.map(|&i| {
|
|
let tc = make_tc(&collected_tool_calls[i]);
|
|
let state = Arc::clone(&state2);
|
|
let conn = conn2.clone();
|
|
let db = db2.clone();
|
|
async move { agent_tools::execute_tool(&tc, &state, &conn, &db, &db_type).await }
|
|
})
|
|
.collect();
|
|
let parallel_results = join_all(parallel_futures).await;
|
|
|
|
// Run sequential group one-by-one
|
|
let mut sequential_results = Vec::with_capacity(sequential_indices.len());
|
|
for &i in &sequential_indices {
|
|
let tc = make_tc(&collected_tool_calls[i]);
|
|
sequential_results.push(agent_tools::execute_tool(&tc, &state2, &conn2, &db2, &db_type).await);
|
|
}
|
|
|
|
// Merge results back into original order
|
|
let mut results: Vec<Option<ToolResult>> = vec![None; collected_tool_calls.len()];
|
|
for (pos, &i) in parallel_indices.iter().enumerate() {
|
|
results[i] = Some(parallel_results[pos].clone());
|
|
}
|
|
for (pos, &i) in sequential_indices.iter().enumerate() {
|
|
results[i] = Some(sequential_results[pos].clone());
|
|
}
|
|
let results: Vec<ToolResult> = results.into_iter().map(|r| r.unwrap()).collect();
|
|
|
|
// Process results in order, emitting ToolCallEnd events
|
|
for (tc, result) in collected_tool_calls.iter().zip(results) {
|
|
on_event(AgentEvent::ToolCallEnd {
|
|
tool_call_id: tc.id.clone(),
|
|
tool_name: tc.name.clone(),
|
|
result: match &result.explain_data {
|
|
Some(ed) => json!({ "content": result.content, "explain_data": ed }),
|
|
None => json!({ "content": result.content }),
|
|
},
|
|
is_error: result.is_error,
|
|
});
|
|
conversation_messages.push(AiMessage {
|
|
role: "tool".to_string(),
|
|
content: compact_tool_result_for_context(&tc.name, &result.content),
|
|
tool_call_id: Some(tc.id.clone()),
|
|
tool_calls: Vec::new(),
|
|
});
|
|
}
|
|
|
|
final_text = accumulated_text;
|
|
}
|
|
|
|
on_event(AgentEvent::AgentEnd {
|
|
input_tokens: if total_usage.input_tokens > 0 { Some(total_usage.input_tokens) } else { None },
|
|
output_tokens: if total_usage.output_tokens > 0 { Some(total_usage.output_tokens) } else { None },
|
|
});
|
|
Ok(final_text)
|
|
}
|
|
|
|
/// Build an LLM request that includes tool definitions.
|
|
fn build_tool_request(
|
|
config: &AiConfig,
|
|
system_prompt: &str,
|
|
messages: &[AiMessage],
|
|
_tools: &[ToolDefinition], // Tools are injected in ai::stream_with_tools, not via AiCompletionRequest.
|
|
max_tokens: Option<u32>,
|
|
temperature: Option<f32>,
|
|
) -> AiCompletionRequest {
|
|
// Note: tools are passed via the body, not via AiCompletionRequest.
|
|
// The actual injection happens in stream_with_tools.
|
|
AiCompletionRequest {
|
|
config: config.clone(),
|
|
system_prompt: system_prompt.to_string(),
|
|
messages: messages.to_vec(),
|
|
max_tokens: max_tokens.or(Some(4096)),
|
|
temperature: temperature.or(Some(0.2)),
|
|
}
|
|
}
|
|
|
|
/// Stream an LLM response with tool support, parsing tool_calls from SSE deltas.
|
|
///
|
|
/// True streaming: text, reasoning, and tool call arguments are all emitted
|
|
/// incrementally as they arrive from the provider.
|
|
async fn stream_with_tools(
|
|
config: &AiConfig,
|
|
request: &AiCompletionRequest,
|
|
session_id: &str,
|
|
tools: &[ToolDefinition],
|
|
cancelled: &Notify,
|
|
on_chunk: impl Fn(AiStreamChunk) + Send + Sync + 'static,
|
|
) -> Result<(Vec<ToolCall>, Option<TokenUsage>), String> {
|
|
// Return early if the user cancelled before the LLM call started.
|
|
if cancelled.notified().now_or_never().is_some() {
|
|
return Err("Agent loop cancelled".to_string());
|
|
}
|
|
|
|
ai::stream_with_tools(config, request, session_id, &tools, cancelled, on_chunk).await
|
|
}
|
|
|
|
fn is_context_length_error(error: &str) -> bool {
|
|
let lower = error.to_lowercase();
|
|
[
|
|
"context length",
|
|
"context_length",
|
|
"maximum context",
|
|
"max context",
|
|
"token limit",
|
|
"too many tokens",
|
|
"prompt is too long",
|
|
"input is too long",
|
|
"reduce the length",
|
|
]
|
|
.iter()
|
|
.any(|marker| lower.contains(marker))
|
|
}
|
|
|
|
/// Text-only fallback for providers that don't support function calling.
|
|
///
|
|
/// Injects database schema context into the system prompt so the LLM can still
|
|
/// give informed answers, then performs a single non-streaming completion.
|
|
async fn run_agent_loop_text_only(
|
|
config: &AiConfig,
|
|
system_prompt: &str,
|
|
messages: &[AiMessage],
|
|
agent_ctx: &AgentLoopContext,
|
|
on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
|
|
_cancelled: &Notify,
|
|
max_tokens: Option<u32>,
|
|
temperature: Option<f32>,
|
|
) -> Result<String, String> {
|
|
// Build a schema-enriched system prompt so the LLM can answer schema questions
|
|
// even without tool access.
|
|
let enriched_prompt = build_schema_prompt(agent_ctx, system_prompt).await;
|
|
|
|
let request = AiCompletionRequest {
|
|
config: config.clone(),
|
|
system_prompt: enriched_prompt,
|
|
messages: messages.to_vec(),
|
|
max_tokens: max_tokens.or(Some(4096)),
|
|
temperature: temperature.or(Some(0.2)),
|
|
};
|
|
|
|
// Use a non-streaming completion as the simplest fallback.
|
|
let result = ai::complete(&request).await?;
|
|
|
|
on_event(AgentEvent::TextDelta { delta: result.clone() });
|
|
on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
|
|
Ok(result)
|
|
}
|
|
|
|
/// Build a system prompt enriched with database schema information
|
|
/// for text-only mode where the LLM cannot use tools.
|
|
async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str) -> String {
|
|
let mut enriched = system_prompt.to_string();
|
|
|
|
// Fetch real schema data using the same core functions the tools would use
|
|
let tables_result = crate::schema::list_tables_core(
|
|
&agent_ctx.state,
|
|
&agent_ctx.connection_id,
|
|
&agent_ctx.database,
|
|
"",
|
|
None,
|
|
Some(50), // smaller limit for prompt injection
|
|
None,
|
|
)
|
|
.await;
|
|
|
|
match tables_result {
|
|
Ok(tables) if !tables.is_empty() => {
|
|
enriched.push_str("\n\n## Database Schema (for context — no tools available)\n");
|
|
enriched.push_str(&format!("Database: {}\n", agent_ctx.database));
|
|
enriched.push_str("Tables:\n");
|
|
for t in &tables {
|
|
enriched.push_str(&format!(" - {} ({})", t.name, t.table_type));
|
|
if let Some(ref comment) = t.comment {
|
|
if !comment.trim().is_empty() {
|
|
enriched.push_str(&format!(" — {}", comment.trim()));
|
|
}
|
|
}
|
|
enriched.push('\n');
|
|
}
|
|
}
|
|
_ => {
|
|
enriched.push_str("\n\n(Note: Unable to load database schema for this request.)\n");
|
|
}
|
|
}
|
|
|
|
enriched
|
|
}
|
|
|
|
/// Estimate text tokens conservatively for mixed English, Chinese, SQL, and JSON content.
|
|
fn estimate_text_tokens(text: &str) -> u32 {
|
|
let chars = text.chars().count() as u32;
|
|
if chars == 0 {
|
|
return 0;
|
|
}
|
|
|
|
let non_ascii = text.chars().filter(|c| !c.is_ascii()).count() as u32;
|
|
let alpha = non_ascii as f32 / chars as f32;
|
|
let ascii_est = (text.len() as f32 / 3.5).ceil();
|
|
let nonascii_est = (chars as f32 * 1.2).ceil();
|
|
let estimated = (ascii_est * (1.0 - alpha) + nonascii_est * alpha).ceil() as u32;
|
|
estimated.max(1)
|
|
}
|
|
|
|
fn estimate_message_tokens(message: &AiMessage) -> u32 {
|
|
let mut tokens = estimate_text_tokens(&message.content) + 4;
|
|
|
|
if let Some(tool_call_id) = &message.tool_call_id {
|
|
tokens += estimate_text_tokens(tool_call_id) + 2;
|
|
}
|
|
|
|
for tool_call in &message.tool_calls {
|
|
tokens += estimate_text_tokens(&tool_call.id) + estimate_text_tokens(&tool_call.name) + 4;
|
|
if let Ok(args) = serde_json::to_string(&tool_call.arguments) {
|
|
tokens += estimate_text_tokens(&args);
|
|
}
|
|
}
|
|
|
|
tokens
|
|
}
|
|
|
|
/// Estimate tokens for a slice of messages.
|
|
fn estimate_tokens(messages: &[AiMessage]) -> u32 {
|
|
messages.iter().map(estimate_message_tokens).sum()
|
|
}
|
|
|
|
fn estimate_tool_schema_tokens(tools: &[ToolDefinition]) -> u32 {
|
|
tools
|
|
.iter()
|
|
.map(|tool| {
|
|
let schema_tokens =
|
|
serde_json::to_string(&tool.parameters).map(|schema| estimate_text_tokens(&schema)).unwrap_or_default();
|
|
estimate_text_tokens(tool.name) + estimate_text_tokens(tool.description) + schema_tokens + 16
|
|
})
|
|
.sum()
|
|
}
|
|
|
|
fn estimate_current_prompt_tokens(system_prompt: &str, tools: &[ToolDefinition], messages: &[AiMessage]) -> u32 {
|
|
estimate_text_tokens(system_prompt) + estimate_tool_schema_tokens(tools) + estimate_tokens(messages) + 16
|
|
}
|
|
|
|
/// Returns the context window size for a given model name.
|
|
fn context_window_for_model(model: &str) -> u32 {
|
|
let m = model.to_lowercase();
|
|
// GPT-4.1 family: 1M context
|
|
if m.contains("gpt-4.1") {
|
|
return 1_000_000;
|
|
}
|
|
if m.contains("claude") {
|
|
200_000
|
|
} else if m.contains("o1") || m.starts_with("o3") || m.starts_with("o4") {
|
|
200_000
|
|
} else if m.contains("gpt-4") {
|
|
128_000
|
|
} else if m.contains("gemini") {
|
|
1_000_000
|
|
} else if m.contains("deepseek") || m.contains("qwen") {
|
|
128_000
|
|
} else {
|
|
128_000
|
|
}
|
|
}
|
|
|
|
fn prompt_budget(window: u32, max_tokens: Option<u32>) -> u32 {
|
|
let output_reserve = max_tokens.unwrap_or(4096).min(window / 2);
|
|
let safety_reserve = (window / 10).clamp(2048, 16_384).min(window / 2);
|
|
window.saturating_sub(output_reserve).saturating_sub(safety_reserve)
|
|
}
|
|
|
|
fn keep_recent_budget(prompt_budget: u32) -> u32 {
|
|
if prompt_budget <= 4096 {
|
|
prompt_budget / 2
|
|
} else {
|
|
(prompt_budget * 6 / 10).clamp(4096, 50_000)
|
|
}
|
|
}
|
|
|
|
const COMPACT_SYSTEM_PROMPT: &str = "\
|
|
You are a conversation summarizer. Produce a concise structured summary of the conversation \
|
|
provided. Format:\n\
|
|
## Progress\n## Key Decisions\n## Critical Context\n## Next Steps\n\
|
|
Be factual. No commentary.";
|
|
|
|
async fn maybe_compact(
|
|
config: &AiConfig,
|
|
system_prompt: &str,
|
|
tools: &[ToolDefinition],
|
|
messages: &mut Vec<AiMessage>,
|
|
max_tokens: Option<u32>,
|
|
on_event: &(impl Fn(AgentEvent) + Send + Sync),
|
|
force: bool,
|
|
) -> Option<u32> {
|
|
let window = config.context_window.unwrap_or_else(|| context_window_for_model(&config.model));
|
|
let budget = prompt_budget(window, max_tokens);
|
|
let estimated_before = estimate_current_prompt_tokens(system_prompt, tools, messages);
|
|
|
|
if !force && estimated_before <= budget {
|
|
return None;
|
|
}
|
|
|
|
if messages.len() <= 2 {
|
|
return None;
|
|
}
|
|
|
|
// Find cut point: keep a dynamic budget of recent messages and summarize older context.
|
|
let keep_recent_tokens = if force { keep_recent_budget(budget) / 2 } else { keep_recent_budget(budget) };
|
|
let mut recent_tokens = 0u32;
|
|
let mut cut = messages.len();
|
|
for i in (0..messages.len()).rev() {
|
|
let t = estimate_message_tokens(&messages[i]);
|
|
if recent_tokens + t > keep_recent_tokens && i > 0 {
|
|
cut = i + 1;
|
|
break;
|
|
}
|
|
recent_tokens += t;
|
|
cut = i;
|
|
}
|
|
|
|
if cut >= messages.len() {
|
|
cut = messages.len().saturating_sub(1);
|
|
}
|
|
|
|
let cut = adjust_cut_for_tool_pair_integrity(messages, cut);
|
|
|
|
// Always keep messages[0] (the original user question) verbatim outside the summary.
|
|
// Only summarize messages[1..cut].
|
|
if cut <= 1 {
|
|
return None;
|
|
}
|
|
let summary_start = 1usize;
|
|
|
|
let compacted_messages = cut - summary_start;
|
|
|
|
let convo_text: String = messages[summary_start..cut].iter().map(format_message_for_summary).collect();
|
|
|
|
let summary_request = AiCompletionRequest {
|
|
config: config.clone(),
|
|
system_prompt: COMPACT_SYSTEM_PROMPT.to_string(),
|
|
messages: vec![AiMessage {
|
|
role: "user".to_string(),
|
|
content: format!("<conversation>\n{convo_text}</conversation>\n\nSummarize the above."),
|
|
tool_call_id: None,
|
|
tool_calls: Vec::new(),
|
|
}],
|
|
max_tokens: Some(1024),
|
|
temperature: Some(0.1),
|
|
};
|
|
|
|
let summary = match ai::complete(&summary_request).await {
|
|
Ok(s) => s,
|
|
Err(_) => fallback_summary(messages, cut),
|
|
};
|
|
|
|
let summary = if validate_summary(&summary) { summary } else { fallback_summary(messages, cut) };
|
|
|
|
let summary_tokens = estimate_text_tokens(&summary) + 4;
|
|
|
|
let summary_content = format!(
|
|
"[SYSTEM-GENERATED CONTEXT SUMMARY - earlier conversation compressed; background only, not a new user request]\n\n{summary}"
|
|
);
|
|
|
|
messages.drain(summary_start..cut);
|
|
messages.insert(
|
|
summary_start,
|
|
AiMessage {
|
|
role: "user".to_string(),
|
|
content: summary_content.clone(),
|
|
tool_call_id: None,
|
|
tool_calls: Vec::new(),
|
|
},
|
|
);
|
|
|
|
let estimated_after = estimate_current_prompt_tokens(system_prompt, tools, messages);
|
|
|
|
on_event(AgentEvent::ContextCompacted {
|
|
summary: summary_content,
|
|
summary_tokens,
|
|
compacted_messages,
|
|
estimated_before,
|
|
estimated_after,
|
|
});
|
|
Some(summary_tokens)
|
|
}
|
|
|
|
fn compact_tool_result_for_context(tool_name: &str, content: &str) -> String {
|
|
if content.chars().count() <= MAX_TOOL_RESULT_CONTEXT_CHARS {
|
|
return content.to_string();
|
|
}
|
|
|
|
if let Ok(value) = serde_json::from_str::<serde_json::Value>(content) {
|
|
return compact_json_tool_result(tool_name, content, &value);
|
|
}
|
|
|
|
compact_text_tool_result(tool_name, content)
|
|
}
|
|
|
|
fn compact_json_tool_result(tool_name: &str, original: &str, value: &serde_json::Value) -> String {
|
|
let compacted = match value {
|
|
serde_json::Value::Array(items) => json!({
|
|
"type": "array",
|
|
"totalItems": items.len(),
|
|
"head": items.iter().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>(),
|
|
"tail": items.iter().rev().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>().into_iter().rev().collect::<Vec<_>>(),
|
|
}),
|
|
serde_json::Value::Object(map) => {
|
|
let mut object = serde_json::Map::new();
|
|
object.insert("type".to_string(), json!("object"));
|
|
object.insert("keys".to_string(), json!(map.keys().cloned().collect::<Vec<_>>()));
|
|
for (key, field_value) in map {
|
|
match field_value {
|
|
serde_json::Value::Array(items) if items.len() > TOOL_RESULT_SAMPLE_ITEMS * 2 => {
|
|
object.insert(
|
|
key.clone(),
|
|
json!({
|
|
"totalItems": items.len(),
|
|
"head": items.iter().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>(),
|
|
"tail": items.iter().rev().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>().into_iter().rev().collect::<Vec<_>>(),
|
|
}),
|
|
);
|
|
}
|
|
_ => {
|
|
object.insert(key.clone(), field_value.clone());
|
|
}
|
|
}
|
|
}
|
|
serde_json::Value::Object(object)
|
|
}
|
|
_ => value.clone(),
|
|
};
|
|
|
|
let compacted_text = serde_json::to_string_pretty(&compacted).unwrap_or_else(|_| compacted.to_string());
|
|
if compacted_text.chars().count() <= MAX_TOOL_RESULT_CONTEXT_CHARS {
|
|
format!(
|
|
"[TOOL RESULT COMPACTED FOR CONTEXT]\nTool: {tool_name}\nOriginal chars: {}\nCompaction: parsed JSON with sampled arrays/fields. UI events preserve the full result.\n\n{}",
|
|
original.chars().count(),
|
|
compacted_text
|
|
)
|
|
} else {
|
|
compact_text_tool_result(tool_name, original)
|
|
}
|
|
}
|
|
|
|
fn compact_text_tool_result(tool_name: &str, content: &str) -> String {
|
|
let original_chars = content.chars().count();
|
|
let head = content.chars().take(TOOL_RESULT_HEAD_CHARS).collect::<String>();
|
|
let tail_chars = content.chars().rev().take(TOOL_RESULT_TAIL_CHARS).collect::<Vec<_>>();
|
|
let tail = tail_chars.into_iter().rev().collect::<String>();
|
|
format!(
|
|
"[TOOL RESULT COMPACTED FOR CONTEXT]\nTool: {tool_name}\nOriginal chars: {original_chars}\nCompaction: kept the head and tail; middle omitted. UI events preserve the full result.\n\n{head}\n\n...[middle omitted from tool result context]...\n\n{tail}"
|
|
)
|
|
}
|
|
|
|
fn validate_summary(summary: &str) -> bool {
|
|
let trimmed = summary.trim();
|
|
trimmed.len() >= 50 && trimmed.len() <= 6000
|
|
}
|
|
|
|
fn fallback_summary(messages: &[AiMessage], cut: usize) -> String {
|
|
// summary_start=1: messages[0] is kept verbatim, only messages[1..cut] are compacted.
|
|
let compacted_messages = &messages[1..cut];
|
|
let tool_calls = compacted_messages.iter().filter(|m| !m.tool_calls.is_empty()).count();
|
|
let tool_results = compacted_messages.iter().filter(|m| m.role == "tool").count();
|
|
let user_messages = compacted_messages.iter().filter(|m| m.role == "user").count();
|
|
let assistant_messages = compacted_messages.iter().filter(|m| m.role == "assistant").count();
|
|
|
|
let recent_roles = compacted_messages
|
|
.iter()
|
|
.rev()
|
|
.take(8)
|
|
.map(|message| format!("{}{}", message.role, if message.content.is_empty() { "" } else { ": content" }))
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
|
|
[
|
|
"## Progress".to_string(),
|
|
"- Context summarized by fallback generator because the LLM summary was unavailable or low quality.".to_string(),
|
|
"## Key Decisions".to_string(),
|
|
format!(
|
|
"- Compacted {} messages: {} user, {} assistant, {} tool results, {} assistant tool-call messages.",
|
|
compacted_messages.len(), user_messages, assistant_messages, tool_results, tool_calls
|
|
),
|
|
"## Critical Context".to_string(),
|
|
format!("- Recent compacted roles: {recent_roles}"),
|
|
"## Next Steps".to_string(),
|
|
"- Continue from the remaining recent conversation and recover any missing detail from tool handles or source paths if needed.".to_string(),
|
|
]
|
|
.join("\n")
|
|
}
|
|
|
|
fn adjust_cut_for_tool_pair_integrity(messages: &[AiMessage], mut cut: usize) -> usize {
|
|
if cut >= messages.len() {
|
|
return cut;
|
|
}
|
|
|
|
while cut < messages.len() && messages[cut].role == "tool" {
|
|
let Some(origin) = find_originating_assistant(messages, cut) else {
|
|
break;
|
|
};
|
|
if origin >= cut {
|
|
break;
|
|
}
|
|
cut = origin;
|
|
}
|
|
|
|
cut
|
|
}
|
|
|
|
fn find_originating_assistant(messages: &[AiMessage], tool_index: usize) -> Option<usize> {
|
|
let tool_call_id = messages.get(tool_index)?.tool_call_id.as_deref();
|
|
|
|
for i in (0..tool_index).rev() {
|
|
let message = &messages[i];
|
|
if message.role != "assistant" {
|
|
continue;
|
|
}
|
|
|
|
if let Some(tool_call_id) = tool_call_id {
|
|
if message.tool_calls.iter().any(|tool_call| tool_call.id == tool_call_id) {
|
|
return Some(i);
|
|
}
|
|
} else if !message.tool_calls.is_empty() {
|
|
return Some(i);
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn format_message_for_summary(message: &AiMessage) -> String {
|
|
let mut header = format!("[{}", message.role);
|
|
if let Some(tool_call_id) = &message.tool_call_id {
|
|
header.push_str(&format!(" tool_call_id={tool_call_id}"));
|
|
}
|
|
if !message.tool_calls.is_empty() {
|
|
let tool_names =
|
|
message.tool_calls.iter().map(|tool_call| tool_call.name.as_str()).collect::<Vec<_>>().join(", ");
|
|
header.push_str(&format!(" tool_calls={tool_names}"));
|
|
}
|
|
header.push(']');
|
|
|
|
format!("{header}: {}\n", summarize_message_content(&message.content))
|
|
}
|
|
|
|
fn summarize_message_content(content: &str) -> String {
|
|
let char_count = content.chars().count();
|
|
if char_count <= 4000 {
|
|
return content.to_string();
|
|
}
|
|
|
|
let head = content.chars().take(1500).collect::<String>();
|
|
let tail_chars = content.chars().rev().take(1500).collect::<Vec<_>>();
|
|
let tail = tail_chars.into_iter().rev().collect::<String>();
|
|
format!("{head}\n\n...[middle omitted for summary input]...\n\n{tail}")
|
|
}
|