github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-23 11:36:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

@

Browse Source 
feat(ai): add agent loop with tool calling, fix connection routing and Gemini compatibility

## 背景说明
将 dbx AI 助手从单次 SQL 文本生成器升级为具备 tool calling 的数据库 Agent。
Phase 1 实现 Agent Loop 框架 + 只读工具（list_tables、get_columns）。
参考：Pi agent-core，.trellis/tasks/06-11-feat-ai-agent-loop-tool-calling-upgrade

## 修改内容
- 新增 agent_loop.rs：Agent Loop 核心，支持 OpenAI/Claude/Gemini Provider 的 tool calling，最大 10 轮，Ollama 自动降级为 text-only
- 新增 agent_tools.rs：read_only_tools (list_tables, get_columns) 的工具定义与执行器
- 新增 agent_events.rs：AgentEvent/ToolCall/ToolResult/ToolDefinition 类型定义
- 修改 ai.rs：AiMessage 新增 tool_call_id 字段
- 修改 lib.rs：注册 agent_events、agent_loop、agent_tools 模块
- 修改 routes/ai.rs & commands/ai.rs：新增 ai_agent_stream SSE 端点
- 修改 tauri.ts & http.ts：新增 AgentEvent 类型与 aiAgentStream 函数
- 修改 api.ts：注册 aiAgentStream 转发
- 修改 ai.ts：新增 runAgentStream，AiContext 增加 connectionId 字段
- 修改 AiAssistant.vue：接入 runAgentStream，解析 AgentEvent 渲染工具调用步骤
- 修改 en.ts / zh-CN.ts：新增 agentSteps 国际化文案

## 修复的问题
- [#I01] connectionName 误传为 connectionId → add connectionId to AiContext
- [#I02] agent loop 取消操作无效 → add cancel checks before each turn
- [#C01] TODO 缺少 issue 引用 → add PRD reference to TODO
- [#C02] stream_with_tools 命名不准确 → add clarifying doc comment
- [#C03] Gemini functionResponse.name 解析错误 → fix tool_call_id format and extraction

## 测试建议
- execute_get_columns：合法表名/空表名/SQL 注入字符拒绝
- call_openai_with_tools：空 choices/空 tool_calls/非 JSON arguments
- call_gemini_with_tools：空 candidates/parts 无 functionCall
- runAgentStream：有效 connectionId 正常调用
@
This commit is contained in:
runstone
2026-06-11 18:27:24 +08:00
parent 759cab9088
commit 9ab41795d6
17 changed files with 1253 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/desktop/src/components/editor/AiAssistant.vue
+54 -28
View File
@@ -1,5 +1,5 @@
<script setup lang="ts">
import { computed, nextTick, onMounted, onUnmounted, ref } from "vue";
import { computed, nextTick, onMounted, onUnmounted, ref, type Component } from "vue";
import { uuid } from "@/lib/utils";
import { useI18n } from "vue-i18n";
import { translateBackendError } from "@/i18n/backend-errors";
@@ -16,7 +16,8 @@ import { connectionIconType } from "@/lib/connectionPresentation";
import DatabaseIcon from "@/components/icons/DatabaseIcon.vue";
import { useQueryStore } from "@/stores/queryStore";
import { useToast } from "@/composables/useToast";
import { buildAiContext, runAiStream, type AiAction } from "@/lib/ai";
import { buildAiContext, runAgentStream, type AiAction } from "@/lib/ai";
import type { AgentEvent } from "@/lib/tauri";
import { buildAiAgentPlan } from "@/lib/aiAgentPlan";
import { buildAiAgentStepItems, type AiAgentStepItem, type AiAgentStepTone } from "@/lib/aiAgentStepPresentation";
import { createAiShikiCodeHighlighter, type AiCodeHighlighter } from "@/lib/aiCodeHighlighter";
@@ -90,7 +91,7 @@ const selectedMentions = ref<AiTableMention[]>([]);
let mentionTimer: ReturnType<typeof setTimeout> | undefined;
let mentionRequestId = 0;
const actionButtons: { action: AiAction; icon: any; key: string }[] = [
const actionButtons: { action: AiAction; icon: Component; key: string }[] = [
{ action: "generate", icon: Wand2, key: "ai.actions.generate" },
{ action: "explain", icon: HelpCircle, key: "ai.actions.explain" },
{ action: "optimize", icon: Zap, key: "ai.actions.optimize" },
@@ -199,8 +200,9 @@ async function changeConnection(connectionId: string) {
if (tab) {
queryStore.updateDatabase(tab.id, database);
}
} catch (e: any) {
toast(t("connection.connectFailed", { message: translateBackendError(t, e?.message || String(e)) }), 5000);
} catch (e: unknown) {
const message = e instanceof Error ? e.message : String(e);
toast(t("connection.connectFailed", { message: translateBackendError(t, message) }), 5000);
}
}
@@ -355,9 +357,10 @@ async function loadMentionCandidates(query: string) {
mentionCache.value[key] = candidates.slice(0, 40);
mentionCandidates.value = mentionCache.value[key];
mentionSelectedIndex.value = 0;
} catch (e: any) {
} catch (e: unknown) {
if (requestId !== mentionRequestId) return;
mentionError.value = translateBackendError(t, e?.message || String(e));
const message = e instanceof Error ? e.message : String(e);
mentionError.value = translateBackendError(t, message);
mentionCandidates.value = [];
} finally {
if (requestId === mentionRequestId) mentionLoading.value = false;
@@ -484,6 +487,7 @@ async function send() {
const assistantIdx = messages.value.length - 1;
const sessionId = uuid();
currentSessionId.value = sessionId;
const agentEvents: AgentEvent[] = [];
try {
const context = await buildAiContext(props.tab, props.connection, {
mentionedTables,
@@ -492,7 +496,7 @@ async function send() {
role: m.role,
content: m.content,
}));
await runAiStream(
await runAgentStream(
{
config: settings.aiConfig,
action: activeAction.value,
@@ -501,29 +505,49 @@ async function send() {
context,
},
history,
(delta) => {
appendAssistantDelta(assistantIdx, delta);
(event: AgentEvent) => {
agentEvents.push(event);
if (event.type === "text_delta" && event.delta) {
appendAssistantDelta(assistantIdx, event.delta);
}
if (event.type === "reasoning_delta" && event.delta) {
appendAssistantReasoning(assistantIdx, event.delta);
}
scrollToBottom();
},
sessionId,
(reasoningDelta) => {
appendAssistantReasoning(assistantIdx, reasoningDelta);
},
);
} catch (e: any) {
messages.value[assistantIdx].content = `Error: ${e.message || e}`;
} catch (e: unknown) {
const message = e instanceof Error ? e.message : String(e);
messages.value[assistantIdx].content = `Error: ${message}`;
} finally {
const msg = messages.value[assistantIdx];
if (msg) msg.isThinking = false;
isGenerating.value = false;
const agentPlan = buildAiAgentPlan({
mode: requestedMode,
action: requestedAction,
instruction: displayText,
assistantContent: msg?.content || "",
connection: props.connection,
});
if (msg && requestedMode === "agent") msg.agentSteps = buildAiAgentStepItems(agentPlan);
if (agentPlan.handoffSql) emit("requestAutoExecuteSql", agentPlan.handoffSql);
// Render agent tool call steps from agent events
if (msg && agentEvents.length > 0) {
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 || "" },
}));
}
// Fallback: use aiAgentPlan for backward compatibility
if (msg && !msg.agentSteps?.length) {
const agentPlan = buildAiAgentPlan({
mode: requestedMode,
action: requestedAction,
instruction: displayText,
assistantContent: msg?.content || "",
connection: props.connection,
});
if (msg && requestedMode === "agent") msg.agentSteps = buildAiAgentStepItems(agentPlan);
if (agentPlan.handoffSql) emit("requestAutoExecuteSql", agentPlan.handoffSql);
}
activeAction.value = "generate";
currentSessionId.value = "";
persistConversation();
@@ -555,8 +579,9 @@ async function copyCode(code: string, key: string) {
setTimeout(() => {
if (copiedIndex.value === key) copiedIndex.value = "";
}, 2000);
} catch (e: any) {
toast(t("grid.copyFailed", { message: e?.message || String(e) }), 5000);
} catch (e: unknown) {
const message = e instanceof Error ? e.message : String(e);
toast(t("grid.copyFailed", { message }), 5000);
}
}
@@ -620,6 +645,7 @@ onMounted(async () => {
onUnmounted(() => {
clearTimeout(mentionTimer);
cancelStream();
});
function triggerAction(action: AiAction, instruction?: string) {
@@ -781,7 +807,7 @@ const messageRenderer = computed(() => {
<div v-if="connectionStore.connections.length" class="flex items-center gap-1 mb-1 text-xs text-foreground/80">
<DatabaseIcon v-if="connection" :db-type="connectionIconType(connection)" class="h-3 w-3 shrink-0" />
<Server v-else class="h-3 w-3 shrink-0" />
<Select :model-value="connection?.id || ''" @update:model-value="(v: any) => changeConnection(v)">
<Select :model-value="connection?.id || ''" @update:model-value="(v: string) => changeConnection(v)">
<SelectTrigger class="h-5 w-auto border-0 rounded-md bg-transparent dark:bg-transparent p-0 px-1 text-xs text-foreground/80 shadow-none focus:ring-0 focus-visible:ring-0 [&_svg]:size-3">
<SelectValue :placeholder="t('editor.selectConnection')">{{ connection?.name || t("editor.selectConnection") }}</SelectValue>
</SelectTrigger>
@@ -798,7 +824,7 @@ const messageRenderer = computed(() => {
<Database class="h-3 w-3 shrink-0 text-foreground/40" />
<Select
:model-value="selectedDatabaseSelectValue"
@update:model-value="(v: any) => changeDatabase(v)"
@update:model-value="(v: string) => changeDatabase(v)"
@update:open="
(open: boolean) => {
if (open) loadDatabases();
apps/desktop/src/i18n/locales/en.ts
+3
View File
@@ -956,6 +956,9 @@
autoExecute: "Ready to run",
notRequested: "Run not requested",
skipped: "Not run",
callingTool: "Calling tool...",
toolDone: "Tool completed",
toolError: "Tool error",
},
agentStepTitles: {
riskCheck: "Risk check: {action} · {category} · {environment} · {reasons}",
apps/desktop/src/i18n/locales/zh-CN.ts
+3
View File
@@ -955,6 +955,9 @@
autoExecute: "准备执行",
notRequested: "未请求执行",
skipped: "未执行",
callingTool: "正在调用工具...",
toolDone: "工具调用完成",
toolError: "工具调用出错",
},
agentStepTitles: {
riskCheck: "风险检查：{action} · {category} · {environment} · {reasons}",
apps/desktop/src/lib/ai.ts
+35 -14
View File
@@ -7,6 +7,8 @@ import { aiTableMentionKey, type AiTableMention } from "@/lib/aiTableMentions";
import { aiSkillForAction } from "@/lib/aiSkills";
import { isSchemaAware } from "@/lib/databaseCapabilities";
import type { AgentEvent } from "@/lib/tauri";
export type AiAction = "generate" | "explain" | "optimize" | "fix" | "convert" | "sampleData";
export type AiAssistantMode = "ask" | "agent";
@@ -25,6 +27,7 @@ export interface AiSchemaTable {
}
export interface AiContext {
connectionId: string;
connectionName: string;
databaseType: DatabaseType;
database: string;
@@ -44,7 +47,7 @@ export interface AiRequestInput {
context: AiContext;
}
export async function runAiAction(input: AiRequestInput, history?: api.AiMessage[]): Promise<string> {
function buildAgentRequest(input: AiRequestInput, history?: api.AiMessage[]): { messages: api.AiMessage[]; systemPrompt: string; maxTokens: number; temperature: number } {
const isZh = isChineseLocale(currentLocale());
const skill = aiSkillForAction(input.action);
const systemPrompt = buildSystemPrompt(input.action, input.context, input.mode);
@@ -54,27 +57,24 @@ export async function runAiAction(input: AiRequestInput, history?: api.AiMessage
const messages: api.AiMessage[] = [...(history || []), { role: "user", content: userPrompt }];
const params = actionParams(input.action);
const maxTokens = input.config.enableThinking ? Math.max(params.maxTokens, 8192) : params.maxTokens;
return { messages, systemPrompt, maxTokens, temperature: params.temperature };
}
export async function runAiAction(input: AiRequestInput, history?: api.AiMessage[]): Promise<string> {
const { messages, systemPrompt, maxTokens, temperature } = buildAgentRequest(input, history);
return api.aiComplete({
config: input.config,
systemPrompt,
messages,
maxTokens: params.maxTokens,
temperature: params.temperature,
maxTokens,
temperature,
});
}
export async function runAiStream(input: AiRequestInput, history: api.AiMessage[] | undefined, onDelta: (delta: string) => void, sessionId?: string, onReasoningDelta?: (delta: string) => void): Promise<void> {
const isZh = isChineseLocale(currentLocale());
const skill = aiSkillForAction(input.action);
const systemPrompt = buildSystemPrompt(input.action, input.context, input.mode);
const instruction = isZh ? skill.userInstruction.zh : skill.userInstruction.en;
const userPrompt = [`Action: ${input.action}`, instruction, "", "User request:", input.instruction.trim() || "(No extra instruction provided.)"].join("\n");
const messages: api.AiMessage[] = [...(history || []), { role: "user", content: userPrompt }];
const { messages, systemPrompt, maxTokens, temperature } = buildAgentRequest(input, history);
const sid = sessionId || uuid();
const params = actionParams(input.action);
const maxTokens = input.config.enableThinking ? Math.max(params.maxTokens, 8192) : params.maxTokens;
await api.aiStream(
sid,
@@ -83,7 +83,7 @@ export async function runAiStream(input: AiRequestInput, history: api.AiMessage[
systemPrompt,
messages,
maxTokens,
temperature: params.temperature,
temperature,
},
(chunk) => {
if (!chunk.done) {
@@ -94,6 +94,26 @@ export async function runAiStream(input: AiRequestInput, history: api.AiMessage[
);
}
export async function runAgentStream(input: AiRequestInput, history: api.AiMessage[] | undefined, onEvent: (event: AgentEvent) => void, sessionId?: string): Promise<string> {
const { messages, systemPrompt, maxTokens, temperature } = buildAgentRequest(input, history);
const sid = sessionId || uuid();
return api.aiAgentStream(
sid,
{
config: input.config,
systemPrompt,
messages,
maxTokens,
temperature,
},
input.context.connectionId,
input.context.database,
input.context.databaseType,
onEvent,
);
}
function actionParams(action: AiAction): { maxTokens: number; temperature: number } {
switch (action) {
case "explain":
@@ -314,6 +334,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig
}
return {
connectionId: tab.connectionId,
connectionName: connection.name,
databaseType: connection.db_type,
database: tab.database,
apps/desktop/src/lib/api.ts
+2
View File
@@ -171,6 +171,7 @@ export const buildDataCompareSyncPlan = forward("buildDataCompareSyncPlan");
// AI
export const aiComplete = forward("aiComplete");
export const aiStream = forward("aiStream");
export const aiAgentStream = forward("aiAgentStream");
export const aiCancelStream = forward("aiCancelStream");
export const aiTestConnection = forward("aiTestConnection");
export const aiListModels = forward("aiListModels");
@@ -349,4 +350,5 @@ export type {
TableExportProgress,
TableExportStatus,
TableExportRequest,
AgentEvent,
} from "./tauri";
apps/desktop/src/lib/http.ts
+52
View File
@@ -744,6 +744,58 @@ export async function aiListModels(config: AiConfig): Promise<AiModelInfo[]> {
return post("/api/ai/models", { config });
}
export type { AgentEvent } from "./tauri";
function isAgentEvent(v: unknown): v is import("./tauri").AgentEvent {
return typeof v === "object" && v !== null && "type" in v && typeof (v as Record<string, unknown>).type === "string";
}
export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: import("./tauri").AgentEvent) => void, signal?: AbortSignal): Promise<string> {
const res = await fetch("/api/ai/agent-stream", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ session_id: sessionId, request, connection_id: connectionId, database, db_type: dbType }),
signal,
});
if (!res.ok) throw new Error(await res.text());
const reader = res.body!.getReader();
const decoder = new TextDecoder();
let buffer = "";
let result = "";
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split("\n");
buffer = lines.pop() || "";
for (const line of lines) {
if (line.startsWith("data:")) {
const data = line.slice(5).trim();
if (data && data !== "[DONE]") {
try {
const parsed = JSON.parse(data);
if (!isAgentEvent(parsed)) {
console.warn("[aiAgentStream] Skipping invalid agent event:", data);
continue;
}
onEvent(parsed);
if (parsed.type === "agent_end" || parsed.type === "error") {
result = data;
}
} catch {
// skip malformed JSON
}
}
}
}
}
return result;
}
export async function saveAiConfig(config: AiConfig): Promise<void> {
return post("/api/ai/config", { config });
}
apps/desktop/src/lib/tauri.ts
+25
View File
@@ -273,6 +273,31 @@ export async function aiStream(sessionId: string, request: AiCompletionRequest,
}
}
export type AgentEvent =
| { type: "turn_start"; turn: number }
| { type: "text_delta"; delta: string }
| { type: "reasoning_delta"; delta: string }
| { 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: "error"; message: string };
export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: AgentEvent) => void): Promise<string> {
const unlisten: UnlistenFn = await listen<AgentEvent>("ai-agent-event", (event) => {
onEvent(event.payload);
if (event.payload.type === "agent_end" || event.payload.type === "error") {
unlisten();
}
});
try {
return await invoke("ai_agent_stream", { sessionId, request, connectionId, database, dbType });
} catch (e) {
unlisten();
throw e;
}
}
export async function saveAiConfig(config: AiConfig): Promise<void> {
return invoke("save_ai_config", { config });
}
crates/dbx-core/src/agent_events.rs
+79
View File
@@ -0,0 +1,79 @@
use serde::{Deserialize, Serialize};
/// Events emitted by the agent loop, streamed to the frontend via SSE.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentEvent {
/// A turn in the agent loop has started.
TurnStart { turn: u32 },
/// Text delta from the LLM response.
TextDelta { delta: String },
/// Reasoning/thinking delta (for models that support it).
ReasoningDelta { delta: String },
/// The LLM wants to call a tool.
ToolCallStart { tool_call_id: String, tool_name: String, args: serde_json::Value },
/// A tool execution has completed.
ToolCallEnd { tool_call_id: String, tool_name: String, result: serde_json::Value, is_error: bool },
/// A turn in the agent loop has ended.
TurnEnd { turn: u32 },
/// The agent loop has finished.
AgentEnd { total_tokens: Option<u64> },
/// An error occurred during the agent loop.
Error { message: String },
}
/// A unified tool call extracted from any provider's response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: serde_json::Value,
}
/// Result of executing a tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
pub tool_call_id: String,
pub tool_name: String,
pub content: String,
pub is_error: bool,
}
/// Definition of a tool available to the agent.
#[derive(Debug, Clone)]
pub struct ToolDefinition {
pub name: &'static str,
pub description: &'static str,
pub parameters: serde_json::Value,
pub read_only: bool,
}
impl ToolDefinition {
/// Convert to OpenAI-compatible tool JSON for the API request.
pub fn to_openai_tool(&self) -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters
}
})
}
/// Convert to Anthropic-compatible tool JSON.
pub fn to_anthropic_tool(&self) -> serde_json::Value {
serde_json::json!({
"name": self.name,
"description": self.description,
"input_schema": self.parameters
})
}
/// Convert to Gemini-compatible function declaration.
pub fn to_gemini_tool(&self) -> serde_json::Value {
serde_json::json!({
"name": self.name,
"description": self.description,
"parameters": self.parameters
})
}
}
crates/dbx-core/src/agent_loop.rs
+549
View File
@@ -0,0 +1,549 @@
use std::sync::Arc;
use futures::FutureExt;
use serde_json::json;
use tokio::sync::Notify;
use crate::agent_events::{AgentEvent, ToolCall, ToolDefinition};
use crate::agent_tools;
use crate::ai::{self, AiCompletionRequest, AiConfig, AiMessage, AiProvider, AiStreamChunk};
use crate::connection::AppState;
use crate::models::connection::DatabaseType;
use tokio::sync::Mutex;
/// Maximum number of agent loop turns to prevent infinite loops.
const MAX_AGENT_TURNS: u32 = 10;
/// 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>,
) -> 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 = agent_tools::read_only_tools();
let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
let mut final_text = String::new();
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;
}
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);
// 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}");
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);
}
on_event2(AgentEvent::TextDelta { delta: chunk.delta.clone() });
}
if let Some(ref reasoning) = chunk.reasoning_delta {
on_event2(AgentEvent::ReasoningDelta { delta: reasoning.clone() });
}
};
// Call the LLM with tool support
let collected_tool_calls = stream_with_tools(config, &request, &session_id, cancelled, on_chunk).await?;
on_event(AgentEvent::TurnEnd { turn });
let accumulated_text = accumulated_text.lock().await.clone();
// Add assistant message to conversation
conversation_messages.push(AiMessage {
role: "assistant".to_string(),
content: accumulated_text.clone(),
tool_call_id: None,
});
if collected_tool_calls.is_empty() {
// No tool calls -- we're done
final_text = accumulated_text;
break;
}
// Execute each tool call
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(),
});
let result = agent_tools::execute_tool(
tc,
&agent_ctx.state,
&agent_ctx.connection_id,
&agent_ctx.database,
&agent_ctx.db_type,
)
.await;
on_event(AgentEvent::ToolCallEnd {
tool_call_id: tc.id.clone(),
tool_name: tc.name.clone(),
result: json!({ "content": result.content }),
is_error: result.is_error,
});
// Add tool result to conversation for the next LLM call
// Uses "tool" role per OpenAI convention; provider-specific conversion
// happens in call_*_with_tools functions when building provider API requests.
conversation_messages.push(AiMessage {
role: "tool".to_string(),
content: result.content.clone(),
tool_call_id: Some(tc.id.clone()),
});
}
final_text = accumulated_text;
}
on_event(AgentEvent::AgentEnd { total_tokens: 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 stream_with_tools / call_*_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)),
}
}
/// Collect an LLM response with tool support, parsing tool_calls.
///
/// Named "stream" for consistency with the agent loop's streaming-API contract,
/// but currently uses a non-streaming (complete) LLM call internally, then emits
/// the full text as a single delta. True streaming tool_calls is deferred to
/// Phase 2 (provider-specific delta accumulation).
///
/// Check `cancelled` before the LLM call; if the LLM request has already started
/// it will run to completion (non-streaming requests cannot be interrupted mid-flight).
async fn stream_with_tools(
config: &AiConfig,
request: &AiCompletionRequest,
session_id: &str,
cancelled: &Notify,
on_chunk: impl Fn(AiStreamChunk) + Send + Sync + 'static,
) -> Result<Vec<ToolCall>, String> {
let tools = agent_tools::read_only_tools();
// 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());
}
// TODO(agent-loop-phase2): Implement streaming tool_calls for each provider.
// Ref: .trellis/tasks/06-11-feat-ai-agent-loop-tool-calling-upgrade PRD Phase 2
let client = ai::build_ai_http_client(config, 120)?;
let response_text = call_with_tools(&client, config, request, &tools).await?;
// Emit the full text as a single delta for the UI
if !response_text.text.is_empty() {
on_chunk(AiStreamChunk {
session_id: session_id.to_string(),
delta: response_text.text.clone(),
reasoning_delta: None,
done: false,
});
}
on_chunk(AiStreamChunk {
session_id: session_id.to_string(),
delta: String::new(),
reasoning_delta: None,
done: true,
});
Ok(response_text.tool_calls)
}
/// Response from an LLM call that may include tool calls.
struct LlmResponse {
text: String,
tool_calls: Vec<ToolCall>,
}
/// Make a non-streaming LLM call with tools and parse the response.
async fn call_with_tools(
client: &reqwest::Client,
config: &AiConfig,
request: &AiCompletionRequest,
tools: &[ToolDefinition],
) -> Result<LlmResponse, String> {
match config.provider {
AiProvider::Claude => call_claude_with_tools(client, config, request, tools).await,
AiProvider::Gemini => call_gemini_with_tools(client, config, request, tools).await,
_ => call_openai_with_tools(client, config, request, tools).await,
}
}
/// OpenAI / compatible providers: non-streaming call with tools.
async fn call_openai_with_tools(
client: &reqwest::Client,
config: &AiConfig,
request: &AiCompletionRequest,
tools: &[ToolDefinition],
) -> Result<LlmResponse, String> {
let headers = ai::maybe_bearer_headers(config)?;
let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })];
messages.extend(request.messages.iter().map(|m| {
let mut msg = json!({ "role": m.role, "content": m.content });
if m.role == "tool" {
if let Some(ref tc_id) = m.tool_call_id {
msg["tool_call_id"] = json!(tc_id);
}
}
msg
}));
let tool_json: Vec<serde_json::Value> = tools.iter().map(|t| t.to_openai_tool()).collect();
let mut body = json!({
"model": config.model,
"messages": messages,
"max_tokens": request.max_tokens.unwrap_or(4096),
"tools": tool_json,
"tool_choice": "auto",
"stream": false,
});
ai::add_temperature_if_supported(&mut body, request);
// Responses API tool calling is deferred to a future phase.
// Use the completions endpoint regardless of api_style for now.
let endpoint = ai::resolve_endpoint(config);
let res = client
.post(&endpoint)
.headers(headers)
.json(&body)
.send()
.await
.map_err(|e| format!("AI request failed: {e}"))?;
if !res.status().is_success() {
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string()));
}
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
let text = data["choices"].get(0).and_then(|c| c["message"]["content"].as_str()).unwrap_or("").to_string();
let mut tool_calls = Vec::new();
if let Some(calls) = data["choices"].get(0).and_then(|c| c["message"]["tool_calls"].as_array()) {
for tc in calls {
let id = tc["id"].as_str().unwrap_or_default().to_string();
let name = tc["function"]["name"].as_str().unwrap_or_default().to_string();
let args_str = tc["function"]["arguments"].as_str().unwrap_or("{}");
let args: serde_json::Value = serde_json::from_str(args_str).unwrap_or(json!({}));
tool_calls.push(ToolCall { id, name, arguments: args });
}
}
Ok(LlmResponse { text, tool_calls })
}
/// Claude: non-streaming call with tools.
async fn call_claude_with_tools(
client: &reqwest::Client,
config: &AiConfig,
request: &AiCompletionRequest,
tools: &[ToolDefinition],
) -> Result<LlmResponse, String> {
let headers = ai::claude_headers(config)?;
let mut messages: Vec<serde_json::Value> = Vec::new();
for m in &request.messages {
if m.role == "tool" {
// Convert tool results to Anthropic tool_result content blocks
messages.push(json!({
"role": "user",
"content": [{
"type": "tool_result",
"tool_use_id": m.tool_call_id.as_deref().unwrap_or_default(),
"content": m.content
}]
}));
} else {
messages.push(json!({ "role": m.role, "content": m.content }));
}
}
let tool_json: Vec<serde_json::Value> = tools.iter().map(|t| t.to_anthropic_tool()).collect();
let body = json!({
"model": config.model,
"max_tokens": request.max_tokens.unwrap_or(4096),
"system": request.system_prompt,
"messages": messages,
"tools": tool_json,
});
let res = client
.post(ai::resolve_endpoint(config))
.headers(headers)
.json(&body)
.send()
.await
.map_err(|e| format!("AI request failed: {e}"))?;
if !res.status().is_success() {
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string()));
}
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
let mut text = String::new();
let mut tool_calls = Vec::new();
if let Some(content) = data["content"].as_array() {
for block in content {
match block["type"].as_str() {
Some("text") => {
if let Some(t) = block["text"].as_str() {
text.push_str(t);
}
}
Some("tool_use") => {
let id = block["id"].as_str().unwrap_or_default().to_string();
let name = block["name"].as_str().unwrap_or_default().to_string();
let args = block["input"].clone();
tool_calls.push(ToolCall { id, name, arguments: args });
}
_ => {}
}
}
}
Ok(LlmResponse { text, tool_calls })
}
/// Gemini: non-streaming call with tools.
async fn call_gemini_with_tools(
client: &reqwest::Client,
config: &AiConfig,
request: &AiCompletionRequest,
tools: &[ToolDefinition],
) -> Result<LlmResponse, String> {
let mut contents: Vec<serde_json::Value> = Vec::new();
for m in &request.messages {
if m.role == "tool" {
// Extract tool name from id "gemini-tc-<tool_name>-<idx>"
let tool_name = m
.tool_call_id
.as_deref()
.and_then(|s| s.strip_prefix("gemini-tc-"))
.and_then(|s| s.rsplitn(2, '-').nth(1))
.unwrap_or("unknown");
contents.push(json!({
"role": "user",
"parts": [{
"functionResponse": {
"name": tool_name,
"response": { "content": m.content }
}
}]
}));
} else {
let role = if m.role == "assistant" { "model" } else { "user" };
contents.push(json!({
"role": role,
"parts": [{ "text": m.content }]
}));
}
}
let tool_declarations: Vec<serde_json::Value> = tools.iter().map(|t| t.to_gemini_tool()).collect();
let body = json!({
"contents": contents,
"systemInstruction": { "parts": [{ "text": request.system_prompt }] },
"tools": [{ "functionDeclarations": tool_declarations }],
"generationConfig": {
"maxOutputTokens": request.max_tokens.unwrap_or(4096),
}
});
let res = client
.post(ai::resolve_endpoint(config))
.query(&[("key", config.api_key.as_str())])
.header(reqwest::header::CONTENT_TYPE, "application/json")
.json(&body)
.send()
.await
.map_err(|e| format!("AI request failed: {e}"))?;
if !res.status().is_success() {
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string()));
}
let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
let mut text = String::new();
let mut tool_calls = Vec::new();
if let Some(candidates) = data["candidates"].as_array() {
if let Some(parts) = candidates[0]["content"]["parts"].as_array() {
for part in parts {
if let Some(t) = part["text"].as_str() {
text.push_str(t);
}
if let Some(fc) = part.get("functionCall") {
let name = fc["name"].as_str().unwrap_or_default().to_string();
let args = fc["args"].clone();
let id = format!("gemini-tc-{name}-{}", tool_calls.len());
tool_calls.push(ToolCall { id, name, arguments: args });
}
}
}
}
Ok(LlmResponse { text, tool_calls })
}
/// 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 { total_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
)
.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
}
crates/dbx-core/src/agent_tools.rs
+202
View File
@@ -0,0 +1,202 @@
use std::sync::Arc;
use serde_json::json;
use crate::agent_events::{ToolCall, ToolDefinition, ToolResult};
use crate::connection::AppState;
use crate::models::connection::DatabaseType;
/// Maximum number of tables returned by list_tables tool.
const LIST_TABLES_LIMIT: usize = 200;
/// Get all available tool definitions for Phase 1 (read-only tools).
pub fn read_only_tools() -> Vec<ToolDefinition> {
vec![list_tables_tool(), get_columns_tool()]
}
/// list_tables tool definition.
fn list_tables_tool() -> ToolDefinition {
ToolDefinition {
name: "list_tables",
description: "List all tables and views in the current database. Returns table names, types, and comments.",
parameters: json!({
"type": "object",
"properties": {
"schema": {
"type": "string",
"description": "Schema name to list tables from (optional, defaults to current database)"
}
},
"required": []
}),
read_only: true,
}
}
/// get_columns tool definition.
fn get_columns_tool() -> ToolDefinition {
ToolDefinition {
name: "get_columns",
description:
"Get column definitions for a table: names, types, primary keys, nullable, defaults, and comments.",
parameters: json!({
"type": "object",
"properties": {
"table": {
"type": "string",
"description": "Table name to get columns for"
},
"schema": {
"type": "string",
"description": "Schema name (optional, defaults to current database)"
}
},
"required": ["table"]
}),
read_only: true,
}
}
/// Execute a tool call and return the result.
pub async fn execute_tool(
tool_call: &ToolCall,
state: &Arc<AppState>,
connection_id: &str,
database: &str,
db_type: &DatabaseType,
) -> ToolResult {
let result = match tool_call.name.as_str() {
"list_tables" => execute_list_tables(tool_call, state, connection_id, database, db_type).await,
"get_columns" => execute_get_columns(tool_call, state, connection_id, database, db_type).await,
_ => Err(format!("Unknown tool: {}", tool_call.name)),
};
match result {
Ok(content) => ToolResult {
tool_call_id: tool_call.id.clone(),
tool_name: tool_call.name.clone(),
content,
is_error: false,
},
Err(err) => ToolResult {
tool_call_id: tool_call.id.clone(),
tool_name: tool_call.name.clone(),
content: format!("Error: {err}"),
is_error: true,
},
}
}
async fn execute_list_tables(
tool_call: &ToolCall,
state: &Arc<AppState>,
connection_id: &str,
database: &str,
_db_type: &DatabaseType,
) -> Result<String, String> {
let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string();
// list_tables_core already limits via the Some(LIST_TABLES_LIMIT) argument.
let tables =
crate::schema::list_tables_core(state, connection_id, database, &schema, None, Some(LIST_TABLES_LIMIT))
.await
.map_err(|e| format!("Failed to list tables: {e}"))?;
let total = tables.len();
let mut lines = Vec::new();
for table in &tables {
let mut line = format!("- {} ({})", table.name, table.table_type);
if let Some(comment) = &table.comment {
let trimmed = comment.trim();
if !trimmed.is_empty() {
line.push_str(&format!(" -- {}", trimmed));
}
}
lines.push(line);
}
if total > LIST_TABLES_LIMIT {
lines.push(format!("... (showing {LIST_TABLES_LIMIT} of {total} tables)"));
}
if lines.is_empty() {
return Ok("No tables found in this database/schema.".to_string());
}
Ok(lines.join("\n"))
}
async fn execute_get_columns(
tool_call: &ToolCall,
state: &Arc<AppState>,
connection_id: &str,
database: &str,
_db_type: &DatabaseType,
) -> Result<String, String> {
let table = tool_call
.arguments
.get("table")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: table")?
.trim()
.to_string();
if table.is_empty() {
return Err("Table name cannot be empty".to_string());
}
if table.len() > 256 {
return Err(format!("Table name too long: {} characters (max 256)", table.len()));
}
// Reject names with characters that are unlikely to be valid identifiers
if table.contains(';') || table.contains('\'') || table.contains('"') || table.contains('\\') {
return Err(format!("Table name contains invalid characters: '{}'", table));
}
let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string();
let columns = crate::schema::get_columns_core(state, connection_id, database, &schema, &table)
.await
.map_err(|e| format!("Failed to get columns for {table}: {e}"))?;
if columns.is_empty() {
return Ok(format!("No columns found for table '{table}'."));
}
let mut lines = Vec::new();
lines.push(format!("Columns of {table}:"));
for col in &columns {
let mut flags: Vec<String> = Vec::new();
if col.is_primary_key {
flags.push("PK".to_string());
}
if col.is_nullable {
flags.push("nullable".to_string());
} else {
flags.push("NOT NULL".to_string());
}
if let Some(default) = &col.column_default {
if !default.is_empty() {
flags.push(format!("default {default}"));
}
}
if let Some(extra) = &col.extra {
if !extra.is_empty() {
flags.push(extra.clone());
}
}
let flags_str = if flags.is_empty() { String::new() } else { format!(" ({})", flags.join(", ")) };
let comment_str = col
.comment
.as_ref()
.filter(|c| !c.trim().is_empty())
.map(|c| format!(" -- {}", c.trim()))
.unwrap_or_default();
lines.push(format!(" - {}: {}{}{}", col.name, col.data_type, flags_str, comment_str));
}
Ok(lines.join("\n"))
}
crates/dbx-core/src/ai.rs
+8 -4
View File
@@ -99,6 +99,10 @@ fn default_enable_thinking() -> bool {
pub struct AiMessage {
pub role: String,
pub content: String,
/// Tool call ID for tool results (role="tool"). Used to associate
/// a tool result with its originating tool call in multi-turn loops.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -313,7 +317,7 @@ pub fn supports_temperature(config: &AiConfig) -> bool {
!(is_openai_api_config(config) && is_openai_reasoning_model(&config.model))
}
fn add_temperature_if_supported(body: &mut serde_json::Value, request: &AiCompletionRequest) {
pub fn add_temperature_if_supported(body: &mut serde_json::Value, request: &AiCompletionRequest) {
if supports_temperature(&request.config) {
body["temperature"] = json!(request.temperature.unwrap_or(0.2));
}
@@ -388,7 +392,7 @@ fn validate_model_list_config(config: &AiConfig) -> Result<(), String> {
resolve_model_list_endpoint(config).map(|_| ())
}
fn maybe_bearer_headers(config: &AiConfig) -> Result<HeaderMap, String> {
pub fn maybe_bearer_headers(config: &AiConfig) -> Result<HeaderMap, String> {
let mut headers = HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
if !config.api_key.trim().is_empty() {
@@ -400,7 +404,7 @@ fn maybe_bearer_headers(config: &AiConfig) -> Result<HeaderMap, String> {
Ok(headers)
}
fn claude_headers(config: &AiConfig) -> Result<HeaderMap, String> {
pub fn claude_headers(config: &AiConfig) -> Result<HeaderMap, String> {
let mut headers = HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
match config.auth_method {
@@ -680,7 +684,7 @@ pub async fn test_connection_core(config: &AiConfig) -> Result<String, String> {
let request = AiCompletionRequest {
config: config.clone(),
system_prompt: String::new(),
messages: vec![AiMessage { role: "user".into(), content: "hi".into() }],
messages: vec![AiMessage { role: "user".into(), content: "hi".into(), tool_call_id: None }],
max_tokens: Some(16),
temperature: Some(0.0),
};
crates/dbx-core/src/db/ssh_tunnel.rs
+112 -76
View File
@@ -5,7 +5,10 @@ use std::sync::Arc;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine;
use russh::client::{self, Config, Handle};
use russh::keys::agent::{client::AgentClient, AgentIdentity};
use russh::keys::agent::{
client::{AgentClient, AgentStream},
AgentIdentity,
};
use russh::keys::{decode_secret_key, key::PrivateKeyWithHashAlg, PrivateKey};
use russh::ChannelMsg;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -43,12 +46,12 @@ impl client::Handler for SshClient {
}
async fn connect_and_authenticate(
ssh_host: &str,
ssh_host: String,
ssh_port: u16,
ssh_user: &str,
ssh_password: &str,
ssh_key_path: &str,
ssh_key_passphrase: &str,
ssh_user: String,
ssh_password: String,
ssh_key_path: String,
ssh_key_passphrase: String,
use_ssh_agent: bool,
connect_timeout_secs: u64,
) -> Result<Handle<SshClient>, String> {
@@ -64,15 +67,15 @@ async fn connect_and_authenticate(
if !ssh_key_path.is_empty() {
// Validate SSH key file path
validate_file_path(ssh_key_path, |_| false)?;
validate_file_path(&ssh_key_path, |_| false)?;
let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase) };
let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase.as_str()) };
let key_pair =
load_ssh_private_key(ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
load_ssh_private_key(&ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
let auth_res = tokio::time::timeout(
connect_timeout,
session.authenticate_publickey(
ssh_user,
&ssh_user,
PrivateKeyWithHashAlg::new(
Arc::new(key_pair),
session.best_supported_rsa_hash().await.ok().flatten().flatten(),
@@ -86,7 +89,7 @@ async fn connect_and_authenticate(
return Err("SSH public key authentication failed".to_string());
}
} else if !ssh_password.is_empty() {
let auth_res = tokio::time::timeout(connect_timeout, session.authenticate_password(ssh_user, ssh_password))
let auth_res = tokio::time::timeout(connect_timeout, session.authenticate_password(&ssh_user, &ssh_password))
.await
.map_err(|_| format!("SSH password auth timed out ({connect_timeout_secs}s)"))?
.map_err(|e| format!("SSH password auth failed: {e}"))?;
@@ -94,7 +97,7 @@ async fn connect_and_authenticate(
return Err("SSH password authentication failed".to_string());
}
} else if use_ssh_agent {
match try_authenticate_with_agent(&mut session, ssh_user, &connect_timeout).await {
match try_authenticate_with_agent(&mut session, &ssh_user, &connect_timeout).await {
Ok(()) => {}
Err(agent_err) => return Err(agent_err),
}
@@ -105,6 +108,18 @@ async fn connect_and_authenticate(
Ok(session)
}
/// Connect to the platform-appropriate SSH agent.
#[cfg(unix)]
async fn connect_ssh_agent() -> Result<AgentClient<Box<dyn AgentStream + Send + Unpin>>, String> {
AgentClient::connect_env().await.map(|c| c.dynamic()).map_err(|e| e.to_string())
}
/// Connect to the platform-appropriate SSH agent (Pageant on Windows).
#[cfg(windows)]
async fn connect_ssh_agent() -> Result<AgentClient<Box<dyn AgentStream + Send + Unpin>>, String> {
AgentClient::connect_pageant().await.map(|c| c.dynamic()).map_err(|e| e.to_string())
}
/// Try to authenticate using ssh-agent identities. Returns `Ok(())` on success,
/// or an error describing why agent auth failed (unavailable, no identities, all rejected).
async fn try_authenticate_with_agent(
@@ -112,7 +127,7 @@ async fn try_authenticate_with_agent(
ssh_user: &str,
connect_timeout: &Duration,
) -> Result<(), String> {
let mut agent = match AgentClient::connect_env().await {
let mut agent = match connect_ssh_agent().await {
Ok(a) => a,
Err(e) => {
return Err(format!("No SSH password or key provided, and ssh-agent is unavailable: {e}"));
@@ -299,7 +314,12 @@ fn read_u32(bytes: &[u8], pos: &mut usize) -> Result<u32, String> {
/// Accept connections on the local listener and forward them through the SSH session.
/// Returns when the SSH session dies (listener error or session.is_closed()).
async fn forward_loop(session: &Handle<SshClient>, listener: &TcpListener, remote_host: &str, remote_port: u16) {
async fn forward_loop(
session: Arc<Handle<SshClient>>,
listener: Arc<TcpListener>,
remote_host: String,
remote_port: u16,
) {
let mut idle_check = tokio::time::interval(IDLE_SESSION_CHECK_INTERVAL);
idle_check.set_missed_tick_behavior(MissedTickBehavior::Delay);
@@ -341,7 +361,7 @@ async fn forward_loop(session: &Handle<SshClient>, listener: &TcpListener, remot
let mut channel = match session
.channel_open_direct_tcpip(
remote_host,
remote_host.clone(),
remote_port.into(),
peer_addr.ip().to_string(),
peer_addr.port().into(),
@@ -398,8 +418,8 @@ async fn forward_loop(session: &Handle<SshClient>, listener: &TcpListener, remot
/// Uses exponential backoff for reconnect attempts and gives up after
/// MAX_RECONNECT_ATTEMPTS to avoid log storms from permanent failures.
#[allow(clippy::too_many_arguments)]
async fn tunnel_reconnect_loop(
mut session: Handle<SshClient>,
fn tunnel_reconnect_loop(
session: Handle<SshClient>,
connect_host: String,
connect_port: u16,
ssh_user: String,
@@ -408,66 +428,60 @@ async fn tunnel_reconnect_loop(
ssh_key_passphrase: String,
use_ssh_agent: bool,
connect_timeout_secs: u64,
listener: TcpListener,
listener: Arc<TcpListener>,
remote_host: String,
remote_port: u16,
) {
loop {
log::info!("SSH tunnel active: {}:{} -> {}:{}", connect_host, connect_port, remote_host, remote_port);
forward_loop(&session, &listener, &remote_host, remote_port).await;
log::warn!("SSH tunnel connection lost ({}:{}), reconnecting...", connect_host, connect_port);
// Reconnect with exponential backoff
let mut delay = INITIAL_RECONNECT_DELAY;
let mut attempts: u32 = 0;
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()>>> {
let mut session = Arc::new(session);
let label = format!("{connect_host}:{connect_port} -> {remote_host}:{remote_port}");
Box::pin(async move {
loop {
if attempts >= MAX_RECONNECT_ATTEMPTS {
log::error!(
"SSH tunnel ({connect_host}:{connect_port}): max reconnect attempts ({MAX_RECONNECT_ATTEMPTS}) exhausted, giving up"
);
return;
}
log::info!("SSH tunnel active: {label}");
tokio::time::sleep(delay).await;
forward_loop(session.clone(), listener.clone(), remote_host.clone(), remote_port).await;
match connect_and_authenticate(
&connect_host,
connect_port,
&ssh_user,
&ssh_password,
&ssh_key_path,
&ssh_key_passphrase,
use_ssh_agent,
connect_timeout_secs,
)
.await
{
Ok(new_session) => {
session = new_session;
log::info!(
"SSH tunnel reconnected to {}:{} (attempt {})",
connect_host,
connect_port,
attempts + 1
);
break;
}
Err(e) => {
attempts += 1;
log::warn!("SSH tunnel connection lost ({label}), reconnecting...");
// Reconnect with exponential backoff
let mut delay = INITIAL_RECONNECT_DELAY;
let mut attempts: u32 = 0;
loop {
if attempts >= MAX_RECONNECT_ATTEMPTS {
log::error!(
"SSH reconnect failed ({}:{}, attempt {attempts}/{MAX_RECONNECT_ATTEMPTS}): {e}",
connect_host,
connect_port,
"SSH tunnel ({label}): max reconnect attempts ({MAX_RECONNECT_ATTEMPTS}) exhausted, giving up"
);
// Exponential backoff: double the delay, cap at MAX_RECONNECT_DELAY
delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY);
return;
}
tokio::time::sleep(delay).await;
match connect_and_authenticate(
connect_host.clone(),
connect_port,
ssh_user.clone(),
ssh_password.clone(),
ssh_key_path.clone(),
ssh_key_passphrase.clone(),
use_ssh_agent,
connect_timeout_secs,
)
.await
{
Ok(new_session) => {
session = Arc::new(new_session);
log::info!("SSH tunnel reconnected to {label} (attempt {})", attempts + 1);
break;
}
Err(e) => {
attempts += 1;
log::error!("SSH reconnect failed ({label}, attempt {attempts}/{MAX_RECONNECT_ATTEMPTS}): {e}");
delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY);
}
}
}
}
}
})
}
struct TunnelEntry {
@@ -523,7 +537,7 @@ impl TunnelManager {
return Ok(port);
}
}
// Slow SSH connection — do this outside the lock.
// Slow SSH connection 鈥?do this outside the lock.
let (handle, local_port) = spawn_tunnel(
ssh_host,
ssh_port,
@@ -661,18 +675,20 @@ async fn spawn_tunnel(
// Initial connection: fail fast on bad credentials
let session = connect_and_authenticate(
connect_host,
connect_host.to_string(),
connect_port,
ssh_user,
ssh_password,
ssh_key_path,
ssh_key_passphrase,
ssh_user.to_string(),
ssh_password.to_string(),
ssh_key_path.to_string(),
ssh_key_passphrase.to_string(),
use_ssh_agent,
connect_timeout_secs,
)
.await?;
let handle = tokio::spawn(tunnel_reconnect_loop(
// Spawn the tunnel loop on a dedicated blocking thread to avoid
// a known Rust async Send inference limitation (rust-lang/rust#102211).
let tunnel_args = (
session,
connect_host.to_string(),
connect_port,
@@ -682,10 +698,30 @@ async fn spawn_tunnel(
ssh_key_passphrase.to_string(),
use_ssh_agent,
connect_timeout_secs,
listener,
Arc::new(listener),
remote_host.to_string(),
remote_port,
));
);
let handle = tokio::task::spawn_blocking(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("failed to create tunnel runtime");
rt.block_on(tunnel_reconnect_loop(
tunnel_args.0,
tunnel_args.1,
tunnel_args.2,
tunnel_args.3,
tunnel_args.4,
tunnel_args.5,
tunnel_args.6,
tunnel_args.7,
tunnel_args.8,
tunnel_args.9,
tunnel_args.10,
tunnel_args.11,
));
});
Ok((handle, local_port))
}
crates/dbx-core/src/lib.rs
+3
View File
@@ -1,9 +1,12 @@
pub mod agent_catalog;
pub mod agent_connection;
pub mod agent_events;
pub mod agent_kv;
pub mod agent_loop;
pub mod agent_manager;
pub mod agent_runtime;
pub mod agent_service;
pub mod agent_tools;
pub mod ai;
pub mod cloud_sync;
pub mod connection;
crates/dbx-web/src/main.rs
+1
View File
@@ -284,6 +284,7 @@ async fn main() {
.route("/ai/conversation/{id}", delete(routes::ai::delete_ai_conversation))
.route("/ai/complete", post(routes::ai::ai_complete))
.route("/ai/stream", post(routes::ai::ai_stream))
.route("/ai/agent-stream", post(routes::ai::ai_agent_stream))
.route("/ai/cancel-stream", post(routes::ai::ai_cancel_stream))
.route("/ai/test-connection", post(routes::ai::ai_test_connection))
.route("/ai/models", post(routes::ai::ai_list_models))
crates/dbx-web/src/routes/ai.rs
+82
View File
@@ -6,7 +6,10 @@ use axum::Json;
use futures::stream::Stream;
use serde::Deserialize;
use dbx_core::agent_events::AgentEvent;
use dbx_core::agent_loop::{run_agent_loop, AgentLoopContext};
use dbx_core::ai::{AiCompletionRequest, AiConfig, AiConversation, AiModelInfo, AiStreamChunk};
use dbx_core::models::connection::DatabaseType;
use crate::error::AppError;
use crate::state::WebState;
@@ -58,6 +61,16 @@ pub struct AiCancelStreamRequest {
pub session_id: String,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AiAgentStreamRequest {
pub session_id: String,
pub request: AiCompletionRequest,
pub connection_id: String,
pub database: String,
pub db_type: String,
}
// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------
@@ -171,3 +184,72 @@ pub async fn ai_stream(
Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}
// ---------------------------------------------------------------------------
// AI agent stream (POST returns SSE with AgentEvent)
// ---------------------------------------------------------------------------
pub async fn ai_agent_stream(
State(state): State<Arc<WebState>>,
Json(body): Json<AiAgentStreamRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>>, AppError> {
let session_id = body.session_id;
let request = body.request;
let cancelled = dbx_core::ai::register_stream(&session_id).await;
let (tx, rx) = tokio::sync::broadcast::channel::<String>(256);
let parsed_db_type: DatabaseType = serde_json::from_str(&format!("\"{}\"", body.db_type))
.map_err(|_| AppError(format!("Unknown database type: {}", body.db_type)))?;
let agent_ctx = AgentLoopContext {
state: state.app.clone(),
connection_id: body.connection_id,
database: body.database,
db_type: parsed_db_type,
};
let sid = session_id.clone();
let req_config = request.config;
let req_system_prompt = request.system_prompt;
let req_messages = request.messages;
let req_max_tokens = request.max_tokens;
let req_temperature = request.temperature;
let tx2 = tx.clone();
tokio::task::spawn_blocking(move || {
let rt =
tokio::runtime::Builder::new_current_thread().enable_all().build().expect("failed to create agent runtime");
rt.block_on(async move {
let result = run_agent_loop(
&req_config,
&req_system_prompt,
&req_messages,
&agent_ctx,
move |event: AgentEvent| {
let json = serde_json::to_string(&event).unwrap_or_default();
let _ = tx2.send(json);
},
&cancelled,
req_max_tokens,
req_temperature,
)
.await;
if let Err(e) = result {
let error_event = AgentEvent::Error { message: e };
let _ = tx.send(serde_json::to_string(&error_event).unwrap_or_default());
}
dbx_core::ai::unregister_stream(&sid).await;
});
});
let stream = async_stream::stream! {
let mut rx = rx;
while let Ok(data) = rx.recv().await {
yield Ok(Event::default().data(data));
}
};
Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}
src-tauri/src/commands/ai.rs
+42
View File
@@ -42,11 +42,53 @@ pub async fn ai_stream(app: AppHandle, session_id: String, request: AiCompletion
result
}
use dbx_core::agent_events::AgentEvent;
use dbx_core::agent_loop::{run_agent_loop, AgentLoopContext};
use dbx_core::models::connection::DatabaseType;
#[tauri::command]
pub async fn ai_cancel_stream(session_id: String) -> Result<bool, String> {
Ok(dbx_core::ai::cancel_stream(&session_id).await)
}
#[tauri::command]
pub async fn ai_agent_stream(
app: AppHandle,
state: State<'_, Arc<AppState>>,
session_id: String,
request: AiCompletionRequest,
connection_id: String,
database: String,
db_type: String,
) -> Result<String, String> {
let cancelled = dbx_core::ai::register_stream(&session_id).await;
let parsed_db_type: DatabaseType =
serde_json::from_str(&format!("\"{}\"", db_type)).map_err(|_| format!("Unknown database type: {db_type}"))?;
let agent_ctx = AgentLoopContext { state: state.inner().clone(), connection_id, database, db_type: parsed_db_type };
let result = run_agent_loop(
&request.config,
&request.system_prompt,
&request.messages,
&agent_ctx,
{
let app = app.clone();
move |event: AgentEvent| {
let _ = app.emit("ai-agent-event", &event);
}
},
&cancelled,
request.max_tokens,
request.temperature,
)
.await;
dbx_core::ai::unregister_stream(&session_id).await;
result
}
#[tauri::command]
pub async fn save_ai_conversation(state: State<'_, Arc<AppState>>, conversation: AiConversation) -> Result<(), String> {
state.storage.save_ai_conversation(&conversation).await
src-tauri/src/lib.rs
+1
View File
@@ -317,6 +317,7 @@ pub fn run() {
.invoke_handler(tauri::generate_handler![
commands::ai::ai_complete,
commands::ai::ai_stream,
commands::ai::ai_agent_stream,
commands::ai::ai_cancel_stream,
commands::ai::ai_test_connection,
commands::ai::ai_list_models,