trending/dbx
1
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-29 07:23:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai): context window management and token usage tracking

Browse Source 
* @
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
This commit is contained in:
Abeautifulsnow
2026-06-15 18:55:22 +08:00
committed by GitHub
parent 6672acf42c
commit c300336281
12 changed files with 665 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -61,4 +61,4 @@ test.pdb
portable/
DBX_*_x64-portable.zip
.agents/skills
.pnpm-store/
.pnpm-store/*
apps/desktop/src/components/editor/AiAssistant.vue
+97 -33
View File
@@ -48,6 +48,8 @@ interface ChatMessage {
reasoning?: string;
isThinking?: boolean;
agentSteps?: AiAgentStepItem[];
/** Hidden system-generated context summary; not rendered in chat UI but included in LLM history. */
kind?: "contextSummary";
}
const props = defineProps<{
@@ -76,6 +78,10 @@ const showConversationList = ref(false);
const promptTextareaRef = ref<HTMLTextAreaElement | null>(null);
const promptCompositionActive = ref(false);
const shikiCodeHighlighter = ref<AiCodeHighlighter>();
const agentTokens = ref<{ input: number; output: number } | null>(null);
/** Deferred context compaction info; applied after stream ends to avoid shifting assistantIdx. */
const pendingCompaction = ref<{ summary: string; compactedMessages: number } | null>(null);
// 新增：输入框拖拽调整相关常量
const AI_TEXTAREA_MIN_ROWS = 3;
@@ -126,8 +132,30 @@ function selectAction(action: AiAction) {
}
}
/** Messages visible in the chat UI (excludes hidden context summaries). */
const visibleMessages = computed(() => messages.value.filter((m) => m.kind !== "contextSummary"));
function messagesForAgentHistory(historyMessages: ChatMessage[]): AiMessage[] {
let latestSummaryIndex = -1;
for (let i = historyMessages.length - 1; i >= 0; i--) {
if (historyMessages[i].kind === "contextSummary") {
latestSummaryIndex = i;
break;
}
}
if (latestSummaryIndex < 0) {
return historyMessages.map((m) => ({ role: m.role, content: m.content }));
}
const compactedHistory = historyMessages.slice(latestSummaryIndex);
const firstMsg = historyMessages[0];
if (firstMsg && firstMsg.role === "user" && firstMsg.kind !== "contextSummary") {
return [{ role: "user" as const, content: firstMsg.content }, ...compactedHistory.map((m) => ({ role: m.role, content: m.content }))];
}
return compactedHistory.map((m) => ({ role: m.role, content: m.content }));
}
const chatTitle = computed(() => {
const first = messages.value.find((m) => m.role === "user");
const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary");
return first ? first.content.slice(0, 30) : t("ai.newChat");
});
@@ -305,6 +333,32 @@ function parseExplainFromData(explainData: unknown, dbType: string): ParsedExpla
}
}
function agentEventToStep(event: AgentEvent, index: number): AiAgentStepItem | undefined {
if (event.type === "context_compacted") {
return {
key: `compact-${index}`,
labelKey: "ai.agentSteps.contextCompacted",
tone: "active",
toolResult: `Compacted ${event.compacted_messages} messages. Estimated prompt tokens: ${event.estimated_before.toLocaleString()} -> ${event.estimated_after.toLocaleString()}. Summary: ${event.summary_tokens.toLocaleString()} tokens.`,
isError: false,
};
}
if (event.type !== "tool_call_start" && event.type !== "tool_call_end") return undefined;
return {
key: `${event.tool_call_id || ""}-${event.type}`,
labelKey: event.type === "tool_call_start" ? "ai.agentSteps.callingTool" : event.is_error ? "ai.agentSteps.toolError" : "ai.agentSteps.toolDone",
tone: (event.type === "tool_call_start" ? "active" : event.is_error ? "danger" : "success") as AiAgentStepTone,
titleKey: undefined,
titleParams: { tool: event.tool_name || "" },
toolName: event.tool_name,
toolArgs: event.type === "tool_call_start" ? (event.args as Record<string, unknown>) : undefined,
toolResult: event.type === "tool_call_end" && !event.is_error ? extractToolResultContent(event.result) : undefined,
isError: event.type === "tool_call_end" ? event.is_error : undefined,
};
}
function toggleReasoning(index: number) {
const next = new Set(expandedReasoning.value);
if (next.has(index)) {
@@ -536,14 +590,12 @@ async function send() {
const sessionId = uuid();
currentSessionId.value = sessionId;
const agentEvents: AgentEvent[] = [];
agentTokens.value = null;
try {
const context = await buildAiContext(props.tab, props.connection, {
mentionedTables,
});
const history: AiMessage[] = messages.value.slice(0, -2).map((m) => ({
role: m.role,
content: m.content,
}));
const history: AiMessage[] = messagesForAgentHistory(messages.value.slice(0, -2));
await runAgentStream(
{
config: settings.aiConfig,
@@ -561,23 +613,27 @@ async function send() {
if (event.type === "reasoning_delta" && event.delta) {
appendAssistantReasoning(assistantIdx, event.delta);
}
if (event.type === "agent_end") {
if (event.input_tokens || event.output_tokens) {
agentTokens.value = { input: event.input_tokens ?? 0, output: event.output_tokens ?? 0 };
}
}
if (event.type === "context_compacted") {
const msg = messages.value[assistantIdx];
if (msg) {
if (!msg.agentSteps) msg.agentSteps = [];
const step = agentEventToStep(event, agentEvents.length - 1);
if (step) msg.agentSteps.push(step);
}
pendingCompaction.value = { summary: event.summary, compactedMessages: event.compacted_messages };
}
// Real-time agent step rendering
if (event.type === "tool_call_start" || event.type === "tool_call_end") {
const msg = messages.value[assistantIdx];
if (msg) {
if (!msg.agentSteps) msg.agentSteps = [];
msg.agentSteps.push({
key: `${event.tool_call_id || ""}-${event.type}`,
labelKey: event.type === "tool_call_start" ? "ai.agentSteps.callingTool" : event.is_error ? "ai.agentSteps.toolError" : "ai.agentSteps.toolDone",
tone: (event.type === "tool_call_start" ? "active" : event.is_error ? "danger" : "success") as AiAgentStepTone,
titleKey: undefined,
titleParams: { tool: event.tool_name || "" },
toolName: event.tool_name,
toolArgs: event.type === "tool_call_start" ? (event.args as Record<string, unknown>) : undefined,
toolResult: event.type === "tool_call_end" && !event.is_error ? extractToolResultContent(event.result) : undefined,
isError: event.type === "tool_call_end" ? event.is_error : undefined,
explainData: event.type === "tool_call_end" ? extractExplainData(event.result) : undefined,
});
const step = agentEventToStep(event, agentEvents.length - 1);
if (step) msg.agentSteps.push(step);
}
}
scrollToBottom();
@@ -593,20 +649,7 @@ async function send() {
isGenerating.value = false;
// Render agent tool call steps from agent events
if (msg && agentEvents.length > 0 && !msg.agentSteps?.length) {
msg.agentSteps = agentEvents
.filter((e) => e.type === "tool_call_start" || e.type === "tool_call_end")
.map((e) => ({
key: `${e.tool_call_id || ""}-${e.type}`,
labelKey: e.type === "tool_call_start" ? "ai.agentSteps.callingTool" : e.is_error ? "ai.agentSteps.toolError" : "ai.agentSteps.toolDone",
tone: e.type === "tool_call_start" ? "active" : e.is_error ? "danger" : "success",
titleKey: undefined,
titleParams: { tool: e.tool_name || "" },
toolName: e.tool_name,
toolArgs: e.type === "tool_call_start" ? (e.args as Record<string, unknown>) : undefined,
toolResult: e.type === "tool_call_end" && !e.is_error ? extractToolResultContent(e.result) : undefined,
isError: e.type === "tool_call_end" ? e.is_error : undefined,
explainData: e.type === "tool_call_end" ? extractExplainData(e.result) : undefined,
}));
msg.agentSteps = agentEvents.map((e, index) => agentEventToStep(e, index)).filter((step): step is AiAgentStepItem => Boolean(step));
}
// Fallback: use aiAgentPlan for backward compatibility
if (msg && !msg.agentSteps?.length) {
@@ -622,6 +665,20 @@ async function send() {
}
activeAction.value = "generate";
currentSessionId.value = "";
// Apply deferred context compaction after streaming so assistantIdx stays stable.
// Visible chat history is kept for the user; future LLM history starts from this hidden summary.
if (pendingCompaction.value) {
const { summary, compactedMessages } = pendingCompaction.value;
pendingCompaction.value = null;
const insertAt = Math.min(1 + compactedMessages, messages.value.length - 1);
if (summary) {
messages.value.splice(insertAt, 0, {
role: "user",
content: summary,
kind: "contextSummary",
});
}
}
persistConversation();
scrollToBottom();
}
@@ -665,7 +722,7 @@ function clearMessages() {
async function persistConversation() {
if (!messages.value.length || !props.connection) return;
if (!conversationId.value) conversationId.value = uuid();
const first = messages.value.find((m) => m.role === "user");
const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary");
await saveAiConversation({
id: conversationId.value,
title: first ? first.content.slice(0, 50) : "Untitled",
@@ -675,6 +732,7 @@ async function persistConversation() {
role: m.role,
content: m.content,
...(m.reasoning ? { reasoning: m.reasoning } : {}),
...(m.kind ? { kind: m.kind } : {}),
})),
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
@@ -692,7 +750,10 @@ function selectConversation(conv: AiConversation) {
role: m.role as "user" | "assistant",
content: m.content,
reasoning: m.reasoning,
kind: m.kind,
}));
agentTokens.value = null;
pendingCompaction.value = null;
showConversationList.value = false;
scrollToBottom();
}
@@ -854,7 +915,7 @@ const messageRenderer = computed(() => {
</div>
<ScrollArea v-else ref="scrollRef" class="min-h-0 flex-1 overflow-hidden">
<div class="flex flex-col gap-3 p-3">
<template v-for="(msg, i) in messages" :key="i">
<template v-for="(msg, i) in visibleMessages" :key="i">
<div v-if="msg.role === 'user'" class="flex justify-end">
<div class="max-w-[85%] rounded-lg bg-primary px-3 py-2 text-xs text-primary-foreground">
{{ msg.content }}
@@ -940,6 +1001,9 @@ const messageRenderer = computed(() => {
<Loader2 class="h-3.5 w-3.5 animate-spin" />
<span>{{ t("ai.thinking") }}</span>
</div>
<div v-if="agentTokens && !isGenerating" class="flex items-center gap-1 text-[10px] text-muted-foreground px-2 pb-1">
<span>&#8593;{{ agentTokens.input.toLocaleString() }} &#8595;{{ agentTokens.output.toLocaleString() }} tokens</span>
</div>
</div>
</ScrollArea>
apps/desktop/src/components/editor/EditorSettingsDialog.vue
+13 -1
View File
@@ -1065,6 +1065,7 @@ const aiEditApiStyle = ref<AiApiStyle>(settingsStore.aiConfig.apiStyle || "compl
const aiEditProxyEnabled = ref(!!settingsStore.aiConfig.proxyEnabled);
const aiEditProxyUrl = ref(settingsStore.aiConfig.proxyUrl || "");
const aiEditEnableThinking = ref(settingsStore.aiConfig.enableThinking ?? true);
const aiEditContextWindow = ref<number | undefined>(settingsStore.aiConfig.contextWindow);
const aiModelOptions = ref<AiModelInfo[]>([]);
const aiModelLoading = ref(false);
@@ -1138,6 +1139,7 @@ function currentAiEditConfig() {
proxyEnabled: aiEditProxyEnabled.value,
proxyUrl: aiEditProxyUrl.value,
enableThinking: aiEditEnableThinking.value,
contextWindow: aiEditContextWindow.value || undefined,
};
}
@@ -1211,6 +1213,7 @@ function syncAiEditState() {
aiEditProxyEnabled.value = !!settingsStore.aiConfig.proxyEnabled;
aiEditProxyUrl.value = settingsStore.aiConfig.proxyUrl || "";
aiEditEnableThinking.value = settingsStore.aiConfig.enableThinking ?? true;
aiEditContextWindow.value = settingsStore.aiConfig.contextWindow;
aiTestResult.value = "";
aiTestError.value = "";
aiTestLatency.value = null;
@@ -1240,7 +1243,8 @@ function aiHasChanges(): boolean {
aiEditApiStyle.value !== (settingsStore.aiConfig.apiStyle || "completions") ||
aiEditProxyEnabled.value !== !!settingsStore.aiConfig.proxyEnabled ||
aiEditProxyUrl.value !== (settingsStore.aiConfig.proxyUrl || "") ||
aiEditEnableThinking.value !== (settingsStore.aiConfig.enableThinking ?? true)
aiEditEnableThinking.value !== (settingsStore.aiConfig.enableThinking ?? true) ||
aiEditContextWindow.value !== settingsStore.aiConfig.contextWindow
);
}
@@ -2363,6 +2367,14 @@ watch(
</div>
</div>
<div class="grid grid-cols-3 items-start gap-3">
<Label class="text-right text-xs">{{ t("ai.contextWindow") }}</Label>
<div class="col-span-2">
<Input v-model.number="aiEditContextWindow" type="number" min="1000" step="1000" class="h-8 text-xs" :placeholder="t('ai.contextWindowAuto')" />
<p class="mt-1 text-xs text-muted-foreground">{{ t("ai.contextWindowHint") }}</p>
</div>
</div>
<div class="grid grid-cols-3 items-center gap-3">
<Label class="text-right text-xs">{{ t("ai.proxy") }}</Label>
<label class="col-span-2 flex items-center gap-2 text-xs text-muted-foreground">
apps/desktop/src/i18n/locales/en.ts
+4
View File
@@ -982,6 +982,7 @@
callingTool: "Calling tool...",
toolDone: "Tool completed",
toolError: "Tool error",
contextCompacted: "Context compressed",
},
agentStepTitles: {
riskCheck: "Risk check: {action} · {category} · {environment} · {reasons}",
@@ -1000,6 +1001,9 @@
enableThinkingOn: "Enabled",
enableThinkingOff: "Disabled",
enableThinkingHint: "This option only takes effect on /chat/completions APIs and supported models. When disabled, it can significantly reduce token usage, but the quality of generated results may decrease slightly.",
contextWindow: "Context Window",
contextWindowAuto: "Auto (detect from model name)",
contextWindowHint: "Tokens. Leave empty to auto-detect. Set manually for local/custom models.",
actions: {
generate: "Generate SQL",
explain: "Explain SQL",
apps/desktop/src/i18n/locales/zh-CN.ts
+4
View File
@@ -981,6 +981,7 @@
callingTool: "正在调用工具...",
toolDone: "工具调用完成",
toolError: "工具调用出错",
contextCompacted: "上下文已压缩",
},
agentStepTitles: {
riskCheck: "风险检查：{action} · {category} · {environment} · {reasons}",
@@ -999,6 +1000,9 @@
enableThinkingOn: "已启用",
enableThinkingOff: "已禁用",
enableThinkingHint: "此选项仅对 /chat/completions API 且部分支持的模型生效。设为禁用后可大幅节省 token，但生成结果质量可能会略微下降。",
contextWindow: "上下文窗口",
contextWindowAuto: "自动（根据模型名称推断）",
contextWindowHint: "单位：token。留空自动推断，本地/自定义模型建议手动设置。",
actions: {
generate: "生成 SQL",
explain: "解释 SQL",
apps/desktop/src/lib/tauri.ts
+3 -1
View File
@@ -288,7 +288,8 @@ export type AgentEvent =
| { type: "tool_call_start"; tool_call_id: string; tool_name: string; args: Record<string, unknown> }
| { type: "tool_call_end"; tool_call_id: string; tool_name: string; result: unknown; is_error: boolean }
| { type: "turn_end"; turn: number }
| { type: "agent_end"; total_tokens?: number }
| { type: "agent_end"; input_tokens?: number; output_tokens?: number }
| { type: "context_compacted"; summary: string; summary_tokens: number; compacted_messages: number; estimated_before: number; estimated_after: number }
| { type: "error"; message: string };
export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: AgentEvent) => void, mode?: string, _signal?: AbortSignal): Promise<string> {
@@ -424,6 +425,7 @@ export interface AiChatMessage {
role: string;
content: string;
reasoning?: string;
kind?: "contextSummary";
}
export interface AiConversation {
apps/desktop/src/stores/settingsStore.ts
+2
View File
@@ -25,6 +25,7 @@ export interface AiConfig {
proxyEnabled?: boolean;
proxyUrl?: string;
enableThinking?: boolean;
contextWindow?: number;
}
export interface AiTestConnectionResult {
@@ -179,6 +180,7 @@ export function normalizeAiConfig(config: Partial<AiConfig> | null | undefined):
proxyEnabled: !!config?.proxyEnabled,
proxyUrl: config?.proxyUrl ?? "",
enableThinking: config?.enableThinking ?? true,
contextWindow: config?.contextWindow ?? undefined,
};
}
crates/dbx-core/src/agent_events.rs
+9 -1
View File
@@ -17,7 +17,15 @@ pub enum AgentEvent {
/// A turn in the agent loop has ended.
TurnEnd { turn: u32 },
/// The agent loop has finished.
AgentEnd { total_tokens: Option<u64> },
AgentEnd { input_tokens: Option<u32>, output_tokens: Option<u32> },
/// Context was compacted to stay within context window limits.
ContextCompacted {
summary: String,
summary_tokens: u32,
compacted_messages: usize,
estimated_before: u32,
estimated_after: u32,
},
/// An error occurred during the agent loop.
Error { message: String },
}
crates/dbx-core/src/agent_loop.rs
+464 -25
View File
@@ -1,3 +1,4 @@
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use futures::future::join_all;
@@ -10,10 +11,15 @@ 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 {
@@ -69,6 +75,7 @@ pub async fn run_agent_loop(
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
@@ -77,38 +84,84 @@ pub async fn run_agent_loop(
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 });
// Build the LLM request with tools
let request =
build_tool_request(config, system_prompt, &conversation_messages, &tools, max_tokens, temperature);
let mut stream_result: Option<(Vec<ToolCall>, Option<TokenUsage>, String)> = None;
let mut last_stream_error: Option<String> = None;
// Stream the LLM response, collecting text and tool_calls
let accumulated_text = Arc::new(Mutex::new(String::new()));
let session_id = format!("agent-turn-{turn}");
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);
let acc = accumulated_text.clone();
let on_event2 = on_event.clone();
let on_chunk = move |chunk: AiStreamChunk| {
if !chunk.delta.is_empty() {
if let Ok(mut text) = acc.try_lock() {
text.push_str(&chunk.delta);
// 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() });
}
on_event2(AgentEvent::TextDelta { delta: chunk.delta.clone() });
}
if let Some(ref reasoning) = chunk.reasoning_delta {
on_event2(AgentEvent::ReasoningDelta { delta: reasoning.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())
);
};
// Call the LLM with tool support
let collected_tool_calls =
stream_with_tools(config, &request, &session_id, &tools, cancelled, on_chunk).await?;
if let Some(usage) = turn_usage {
total_usage.add(&usage);
}
on_event(AgentEvent::TurnEnd { turn });
let accumulated_text = accumulated_text.lock().await.clone();
// Add assistant message to conversation (including tool_use blocks)
conversation_messages.push(AiMessage {
role: "assistant".to_string(),
@@ -194,7 +247,7 @@ pub async fn run_agent_loop(
});
conversation_messages.push(AiMessage {
role: "tool".to_string(),
content: result.content.clone(),
content: compact_tool_result_for_context(&tc.name, &result.content),
tool_call_id: Some(tc.id.clone()),
tool_calls: Vec::new(),
});
@@ -203,7 +256,10 @@ pub async fn run_agent_loop(
final_text = accumulated_text;
}
on_event(AgentEvent::AgentEnd { total_tokens: None });
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)
}
@@ -238,7 +294,7 @@ async fn stream_with_tools(
tools: &[ToolDefinition],
cancelled: &Notify,
on_chunk: impl Fn(AiStreamChunk) + Send + Sync + 'static,
) -> Result<Vec<ToolCall>, String> {
) -> 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());
@@ -247,6 +303,23 @@ async fn stream_with_tools(
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
@@ -277,7 +350,7 @@ async fn run_agent_loop_text_only(
let result = ai::complete(&request).await?;
on_event(AgentEvent::TextDelta { delta: result.clone() });
on_event(AgentEvent::AgentEnd { total_tokens: None });
on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
Ok(result)
}
@@ -320,3 +393,369 @@ async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str)
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}")
}
crates/dbx-core/src/ai.rs
+55 -12
View File
@@ -1,3 +1,4 @@
use crate::token_usage::TokenUsage;
use futures::StreamExt;
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
@@ -88,6 +89,8 @@ pub struct AiConfig {
pub proxy_url: String,
#[serde(default = "default_enable_thinking")]
pub enable_thinking: bool,
#[serde(default)]
pub context_window: Option<u32>,
}
fn default_enable_thinking() -> bool {
@@ -147,6 +150,10 @@ pub struct AiChatMessage {
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub covered_messages: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -1394,7 +1401,7 @@ async fn stream_claude_with_tools(
tools: &[crate::agent_events::ToolDefinition],
cancelled: &Notify,
on_event: &impl Fn(StreamToolEvent),
) -> Result<(), String> {
) -> Result<Option<TokenUsage>, String> {
let mut messages: Vec<serde_json::Value> = Vec::new();
let mut pending_tool_results: Vec<serde_json::Value> = Vec::new();
for m in &request.messages {
@@ -1466,6 +1473,7 @@ async fn stream_claude_with_tools(
// Track the current content block index and type for tool_use blocks
let mut current_block_index: Option<u32> = None;
let mut current_block_type: Option<String> = None;
let mut token_usage: Option<TokenUsage> = None;
loop {
tokio::select! {
@@ -1489,6 +1497,20 @@ async fn stream_claude_with_tools(
let event_type = event["type"].as_str().unwrap_or("");
match event_type {
// message_start carries input_tokens (prompt cost)
"message_start" => {
if let Some(i) = event["message"]["usage"]["input_tokens"].as_u64() {
let existing_output = token_usage.as_ref().map(|u| u.output_tokens).unwrap_or(0);
token_usage = Some(TokenUsage { input_tokens: i as u32, output_tokens: existing_output });
}
}
// message_delta carries output_tokens (generation cost)
"message_delta" => {
if let Some(o) = event["usage"]["output_tokens"].as_u64() {
let existing_input = token_usage.as_ref().map(|u| u.input_tokens).unwrap_or(0);
token_usage = Some(TokenUsage { input_tokens: existing_input, output_tokens: o as u32 });
}
}
"content_block_start" => {
let idx = event["index"].as_u64().unwrap_or(0) as u32;
let block_type = event["content_block"]["type"].as_str().unwrap_or("");
@@ -1556,7 +1578,7 @@ async fn stream_claude_with_tools(
}
}
Ok(())
Ok(token_usage)
}
/// Streaming OpenAI-compatible call with tool support.
@@ -1567,7 +1589,7 @@ async fn stream_openai_with_tools(
tools: &[crate::agent_events::ToolDefinition],
cancelled: &Notify,
on_event: &impl Fn(StreamToolEvent),
) -> Result<(), String> {
) -> Result<Option<TokenUsage>, String> {
let headers = maybe_bearer_headers(&request.config)?;
let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })];
@@ -1606,6 +1628,7 @@ async fn stream_openai_with_tools(
"tools": tool_json,
"tool_choice": "auto",
"stream": true,
"stream_options": { "include_usage": true },
});
add_temperature_if_supported(&mut body, request);
@@ -1624,6 +1647,7 @@ async fn stream_openai_with_tools(
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut token_usage: Option<TokenUsage> = None;
loop {
tokio::select! {
@@ -1644,6 +1668,15 @@ async fn stream_openai_with_tools(
}
if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
// Token usage (final chunk from OpenAI with include_usage)
if let Some(usage) = event.get("usage") {
if let (Some(p), Some(c)) = (
usage.get("prompt_tokens").and_then(|v| v.as_u64()),
usage.get("completion_tokens").and_then(|v| v.as_u64()),
) {
token_usage = Some(TokenUsage { input_tokens: p as u32, output_tokens: c as u32 });
}
}
// Reasoning
if let Some(reasoning) = openai_stream_reasoning(&event) {
on_event(StreamToolEvent::Chunk(AiStreamChunk {
@@ -1686,7 +1719,7 @@ async fn stream_openai_with_tools(
}
}
Ok(())
Ok(token_usage)
}
/// Streaming Gemini call with tool support.
@@ -1697,7 +1730,7 @@ async fn stream_gemini_with_tools(
tools: &[crate::agent_events::ToolDefinition],
cancelled: &Notify,
on_event: &impl Fn(StreamToolEvent),
) -> Result<(), String> {
) -> Result<Option<TokenUsage>, String> {
let mut contents: Vec<serde_json::Value> = Vec::new();
let mut pending_function_responses: Vec<serde_json::Value> = Vec::new();
for m in &request.messages {
@@ -1773,6 +1806,7 @@ async fn stream_gemini_with_tools(
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut tool_call_idx: u32 = 0;
let mut token_usage: Option<TokenUsage> = None;
loop {
tokio::select! {
@@ -1787,6 +1821,13 @@ async fn stream_gemini_with_tools(
let Some(data) = stream_data_payload(&line) else { continue };
if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
// Token usage (overwrite each chunk, keep last value)
if let (Some(p), Some(c)) = (
event["usageMetadata"]["promptTokenCount"].as_u64(),
event["usageMetadata"]["candidatesTokenCount"].as_u64(),
) {
token_usage = Some(TokenUsage { input_tokens: p as u32, output_tokens: c as u32 });
}
if let Some(candidates) = event["candidates"].as_array() {
if let Some(parts) = candidates[0]["content"]["parts"].as_array() {
for part in parts {
@@ -1827,11 +1868,11 @@ async fn stream_gemini_with_tools(
}
}
Ok(())
Ok(token_usage)
}
/// Public entry point: stream an LLM call with tool support, accumulating tool calls.
/// Returns completed tool calls when the stream finishes.
/// Returns completed tool calls and token usage when the stream finishes.
pub async fn stream_with_tools(
config: &AiConfig,
request: &AiCompletionRequest,
@@ -1839,7 +1880,7 @@ pub async fn stream_with_tools(
tools: &[crate::agent_events::ToolDefinition],
cancelled: &Notify,
on_chunk: impl Fn(AiStreamChunk),
) -> Result<Vec<crate::agent_events::ToolCall>, String> {
) -> Result<(Vec<crate::agent_events::ToolCall>, Option<TokenUsage>), String> {
validate_config(config)?;
let stream_timeout = if config.enable_thinking { 600 } else { 120 };
@@ -1847,7 +1888,7 @@ pub async fn stream_with_tools(
let accumulator = Arc::new(std::sync::Mutex::new(StreamingToolCallAccumulator::new()));
match config.provider {
let token_usage = match config.provider {
AiProvider::Claude => {
stream_claude_with_tools(&client, session_id, request, tools, cancelled, &|event| {
accumulator.lock().unwrap().process(event, &on_chunk);
@@ -1866,13 +1907,15 @@ pub async fn stream_with_tools(
})
.await?
}
}
};
Ok(Arc::try_unwrap(accumulator)
let tool_calls = Arc::try_unwrap(accumulator)
.expect("stream_with_tools: accumulator Arc should have single owner")
.into_inner()
.expect("stream_with_tools: accumulator Mutex should not be poisoned")
.finalize())
.finalize();
Ok((tool_calls, token_usage))
}
// ---------------------------------------------------------------------------
crates/dbx-core/src/lib.rs
+1
View File
@@ -48,6 +48,7 @@ pub mod table_export;
pub mod table_import;
pub mod table_structure_sql;
pub mod text_export;
pub mod token_usage;
pub mod transfer;
pub mod types;
pub mod update;
crates/dbx-core/src/token_usage.rs
+12
View File
@@ -0,0 +1,12 @@
#[derive(Debug, Clone, Default)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl TokenUsage {
pub fn add(&mut self, other: &TokenUsage) {
self.input_tokens += other.input_tokens;
self.output_tokens += other.output_tokens;
}
}