feat(plugins): add AI tools and consent-gated data access

This commit is contained in:
t8y2 committed 2026-09-25 01:56:26 +08:00
1 parent 1bc6f70cd8
commit 26a015cd86
71 files changed
+4613 -129

No files matched your search

+38 -2
View File
@@ -1,3 +1,5 @@
use std::borrow::Cow;
use serde::{Deserialize, Serialize};
/// Events emitted by the agent loop, streamed to the frontend via SSE.
@@ -18,6 +20,25 @@ pub enum AgentEvent {
ReasoningDelta { delta: String },
/// The LLM wants to call a tool.
ToolCallStart { tool_call_id: String, tool_name: String, args: serde_json::Value },
/// A plugin tool call that may change state waits for the user's explicit
/// approval. The run stays paused until the client answers through the
/// approval command, the approval times out, or the run is cancelled.
ToolApprovalRequired {
approval_id: String,
tool_call_id: String,
tool_name: String,
plugin_id: String,
plugin_name: String,
/// The tool name as the plugin declared it (without the DBX prefix).
plugin_tool: String,
connection_id: String,
connection_name: String,
/// Exactly the arguments DBX will forward if the user approves.
args: serde_json::Value,
timeout_secs: u64,
},
/// The approval requested by `ToolApprovalRequired` has been settled.
ToolApprovalResolved { approval_id: String, tool_call_id: String, outcome: ToolApprovalOutcome },
/// 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.
@@ -41,6 +62,18 @@ pub enum AgentEvent {
Error { message: String },
}
/// How a tool approval request ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolApprovalOutcome {
Approved,
Denied,
/// Nobody answered before the approval deadline.
TimedOut,
/// The run was cancelled while the approval was pending.
Cancelled,
}
/// A unified tool call extracted from any provider's response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
@@ -67,10 +100,13 @@ pub struct ToolResult {
}
/// Definition of a tool available to the agent.
///
/// Built-in tools use static names and descriptions; tools contributed by
/// plugins are discovered at run time, so both fields accept owned text too.
#[derive(Debug, Clone)]
pub struct ToolDefinition {
pub name: &'static str,
pub description: &'static str,
pub name: Cow<'static, str>,
pub description: Cow<'static, str>,
pub parameters: serde_json::Value,
pub read_only: bool,
/// Whether this tool can be executed in parallel with other tools.
+6 -6
View File
@@ -5698,8 +5698,8 @@ mod tests {
let (endpoint, server) = spawn_json_capture_server("text/event-stream", response).await;
let request = anthropic_compatible_test_request(endpoint);
let tools = [crate::agent_events::ToolDefinition {
name: "get_tables",
description: "List tables",
name: "get_tables".into(),
description: "List tables".into(),
parameters: serde_json::json!({
"type": "object",
"properties": { "schema": { "type": "string" } }
@@ -5754,8 +5754,8 @@ mod tests {
let request = claude_http_test_request(endpoint);
let client = build_ai_http_client(&request.config, 10).unwrap();
let tools = [crate::agent_events::ToolDefinition {
name: "get_tables",
description: "List tables",
name: "get_tables".into(),
description: "List tables".into(),
parameters: serde_json::json!({ "type": "object", "properties": {} }),
read_only: true,
parallel_ok: true,
@@ -7330,8 +7330,8 @@ mod tests {
assert_eq!(input[3]["call_id"], "call_1");
let tool = crate::agent_events::ToolDefinition {
name: "list_tables",
description: "List tables",
name: "list_tables".into(),
description: "List tables".into(),
parameters: serde_json::json!({"type": "object"}),
read_only: true,
parallel_ok: true,
@@ -0,0 +1,166 @@
//! End-to-end check of plugin MCP tools in the built-in AI agent, run against
//! the packaged example plugin: install it, open a plugin connection, opt the
//! plugin in, discover its tools through the real sidecar, and call them the
//! way an agent run does.
use std::path::PathBuf;
use std::sync::Arc;
use dbx_core::agent_events::ToolCall;
use dbx_core::connection::AppState;
use dbx_core::models::connection::ConnectionConfig;
use dbx_core::plugin_tools::{self, PluginToolSet};
use dbx_core::plugins::{PluginInstallPolicy, PluginPackageInstaller};
use dbx_core::storage::Storage;
use serde_json::{json, Value};
const PLUGIN_ID: &str = "dbx.example.hello";
const APP_VERSION: &str = "0.5.68";
const CONNECTION_ID: &str = "hello-ai-connection";
#[tokio::main]
async fn main() {
if let Err(error) = run().await {
eprintln!("plugin AI tools smoke failed: {error}");
std::process::exit(1);
}
}
async fn run() -> Result<(), String> {
let package = std::env::args().nth(1).map(PathBuf::from).ok_or("Usage: plugin_ai_tools_smoke <package.dbxp>")?;
let root = tempfile::tempdir().map_err(|error| error.to_string())?;
let plugin_dir = root.path().join("plugins");
PluginPackageInstaller::new(plugin_dir.clone(), APP_VERSION.to_string())?
.install_file(&package, PluginInstallPolicy::LocalDevelopment)?;
let storage = Storage::open(&root.path().join("dbx.db")).await?;
let state = Arc::new(AppState::new_with_plugin_dir_and_app_version(storage, plugin_dir, APP_VERSION));
let config: ConnectionConfig = serde_json::from_value(json!({
"id": CONNECTION_ID,
"name": "Hello AI connection",
"db_type": "plugin",
"driver_profile": "plugin",
"host": "localhost",
"port": 22,
"username": "",
"password": "",
"database": null,
"external_config": { "greeting": "Hello" },
"plugin_id": PLUGIN_ID,
"plugin_connection_provider": "dbx.example.hello.connection",
"plugin_connection_type": "hello",
"connection_secrets": { "api_token": "smoke-secret" }
}))
.map_err(|error| error.to_string())?;
state.configs.write().await.insert(CONNECTION_ID.to_string(), config);
// Nothing reaches the model before the user opts the plugin in, and
// nothing for a connection that is not open.
if !plugin_tools::discover_plugin_tools(&state, None).await.is_empty() {
return Err("tools were exposed before the plugin was enabled for the AI".to_string());
}
state.storage.set_ai_plugin_tool_plugin_enabled(PLUGIN_ID, true).await?;
if !plugin_tools::discover_plugin_tools(&state, None).await.is_empty() {
return Err("tools were exposed without an open plugin connection".to_string());
}
state.get_or_create_pool(CONNECTION_ID, None).await?;
let tools = plugin_tools::discover_plugin_tools(&state, None).await;
let definitions = tools.definitions();
let names = definitions.iter().map(|definition| definition.name.as_ref()).collect::<Vec<_>>();
if names != ["hello__hello_greet", "hello__hello_set_greeting"] {
return Err(format!("unexpected tool names: {names:?}"));
}
let read_only = definitions.iter().map(|definition| definition.read_only).collect::<Vec<_>>();
if read_only != [true, false] {
return Err(format!("readOnlyHint was not honored: {read_only:?}"));
}
if definitions.iter().any(|definition| definition.parameters.pointer("/properties/connectionId").is_some()) {
return Err("the host-bound connectionId leaked into a model-facing schema".to_string());
}
expect_tool_text(
&state,
&tools,
"hello__hello_greet",
json!({ "name": "Agent" }),
"Hello, Agent, from Hello AI connection!",
)
.await?;
// A real run pauses here for the user's approval; this call stands in for
// an approved one.
expect_tool_text(
&state,
&tools,
"hello__hello_set_greeting",
json!({ "greeting": "Hi" }),
"Hello AI connection now greets with \"Hi\".",
)
.await?;
expect_tool_text(
&state,
&tools,
"hello__hello_greet",
json!({ "name": "Agent" }),
"Hi, Agent, from Hello AI connection!",
)
.await?;
let preview = plugin_tools::preview_plugin_tools(&state, None, PLUGIN_ID).await?;
let preview_names = preview.tools.iter().map(|tool| (tool.name.as_str(), tool.read_only)).collect::<Vec<_>>();
if preview_names != [("hello_greet", true), ("hello_set_greeting", false)]
|| preview.open_connections.iter().map(|connection| connection.id.as_str()).collect::<Vec<_>>()
!= [CONNECTION_ID]
{
return Err(format!("unexpected preview: {preview:?}"));
}
let call = ToolCall {
id: "revoked-call".to_string(),
name: "hello__hello_set_greeting".to_string(),
arguments: json!({ "greeting": "Must not run" }),
provider_payload: None,
};
let prepared = tools.prepare_call(&call).ok_or("missing write tool")??;
state.storage.set_ai_plugin_tool_plugin_enabled(PLUGIN_ID, false).await?;
let revoked = plugin_tools::execute_plugin_tool(&state, None, &call, &prepared).await;
if !revoked.is_error || !revoked.content.contains("has been disabled") {
return Err(format!("revoked AI access was not enforced: {revoked:?}"));
}
state.storage.set_ai_plugin_tool_plugin_enabled(PLUGIN_ID, true).await?;
expect_tool_text(
&state,
&tools,
"hello__hello_greet",
json!({ "name": "Agent" }),
"Hi, Agent, from Hello AI connection!",
)
.await?;
state.remove_connection_pools(CONNECTION_ID).await;
let closed = plugin_tools::execute_plugin_tool(&state, None, &call, &prepared).await;
if !closed.is_error || !closed.content.contains("no longer open") {
return Err(format!("closed plugin connection was not rejected: {closed:?}"));
}
println!("plugin AI tools smoke passed: {names:?}");
Ok(())
}
async fn expect_tool_text(
state: &Arc<AppState>,
tools: &PluginToolSet,
name: &str,
arguments: Value,
expected: &str,
) -> Result<(), String> {
let call = ToolCall { id: format!("call-{name}"), name: name.to_string(), arguments, provider_payload: None };
let prepared = tools.prepare_call(&call).ok_or_else(|| format!("{name} is not a plugin tool"))??;
if prepared.connection_id != CONNECTION_ID || prepared.arguments.get("connectionId") != Some(&json!(CONNECTION_ID))
{
return Err(format!("{name} was not bound to the open connection: {prepared:?}"));
}
let result = plugin_tools::execute_plugin_tool(state, None, &call, &prepared).await;
if result.is_error || result.content != expected {
return Err(format!("{name} returned {:?} (error: {})", result.content, result.is_error));
}
Ok(())
}
+269 -17
View File
@@ -6,13 +6,15 @@ use futures::FutureExt;
use serde_json::json;
use tokio::sync::Notify;
use crate::agent_events::{AgentEvent, ToolCall, ToolDefinition, ToolResult};
use crate::agent_events::{AgentEvent, ToolApprovalOutcome, ToolCall, ToolDefinition, ToolResult};
use crate::agent_tools;
use crate::ai::{self, AiCompletionRequest, AiConfig, AiMessage, AiProvider, AiStreamChunk, AiTaskContract};
use crate::ai_cli_agent::CliAgentCommandSpec;
use crate::connection::AppState;
use crate::models::connection::DatabaseType;
use crate::plugin_tools::{self, PluginToolSet, PreparedPluginToolCall};
use crate::token_usage::TokenUsage;
use crate::tool_approval::{self, ToolApprovalWait};
/// Default number of agent loop turns to prevent infinite loops.
/// Users can raise the limit in Settings → AI; it is clamped to
@@ -88,6 +90,12 @@ pub struct AgentLoopContext {
/// Stable per-conversation key forwarded to providers that support prompt
/// caching (OpenAI Responses API). `None` disables the field entirely.
pub prompt_cache_key: Option<String>,
/// Client AI session id. Plugin tool calls that need the user's approval
/// are routed to the client through it; without one they are refused.
pub session_id: Option<String>,
/// Runtime that owns plugin sidecar sessions. The web server runs each
/// agent loop on a runtime of its own, so plugin calls are spawned here.
pub host_runtime: Option<tokio::runtime::Handle>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -281,11 +289,21 @@ async fn run_agent_loop_inner(
.await;
}
let mut sql_permissions = agent_ctx.sql_permissions.clone();
let tools = if is_agent_mode {
let mut tools = if is_agent_mode {
agent_tools::all_tools(agent_ctx.db_type, sql_permissions.clone())
} else {
agent_tools::read_only_tools(agent_ctx.db_type)
};
// Plugin tools join agent runs only; ask mode keeps its database-only
// read tools.
let plugin_tool_set = if is_agent_mode {
plugin_tools::discover_plugin_tools(&agent_ctx.state, agent_ctx.host_runtime.as_ref()).await
} else {
PluginToolSet::default()
};
tools.extend(plugin_tool_set.definitions());
let plugin_system_prompt = plugin_tool_set.prompt_section().map(|section| format!("{system_prompt}\n\n{section}"));
let system_prompt = plugin_system_prompt.as_deref().unwrap_or(system_prompt);
let task_contract = task_contract.cloned();
let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
let mut final_text = String::new();
@@ -542,6 +560,34 @@ async fn run_agent_loop_inner(
});
}
// Plugin tool calls are validated first; a call not declared read-only
// runs only after the user approved it. Nothing has executed yet, so a
// cancellation here stops the run without side effects.
let mut plugin_calls = match gate_plugin_tool_calls(
&plugin_tool_set,
&collected_tool_calls,
agent_ctx,
&on_event,
cancelled,
)
.await
{
Some(plugin_calls) => plugin_calls,
None => {
for tc in &collected_tool_calls {
on_event(AgentEvent::ToolCallEnd {
tool_call_id: tc.id.clone(),
tool_name: tc.name.clone(),
result: json!({ "content": "Cancelled before the tool ran." }),
is_error: true,
});
}
final_text = accumulated_text;
loop_exit = LoopExit::Cancelled;
break;
}
};
// Execute tool calls: parallel for read tools, sequential for execute_query
let state2 = Arc::clone(&agent_ctx.state);
let conn2 = agent_ctx.connection_id.clone();
@@ -552,7 +598,7 @@ async fn run_agent_loop_inner(
// Split by index into parallel and sequential groups using tool metadata
let tool_parallel_map: std::collections::HashMap<&str, bool> =
tools.iter().map(|t| (t.name, t.parallel_ok)).collect();
tools.iter().map(|t| (t.name.as_ref(), t.parallel_ok)).collect();
let (parallel_indices, sequential_indices): (Vec<usize>, Vec<usize>) = (0..collected_tool_calls.len())
.partition(|&i| *tool_parallel_map.get(collected_tool_calls[i].name.as_str()).unwrap_or(&false));
@@ -585,19 +631,26 @@ async fn run_agent_loop_inner(
let mut sequential_results = Vec::with_capacity(sequential_indices.len());
for &i in &sequential_indices {
let tc = make_tc(&collected_tool_calls[i]);
let execution_permissions = sequential_tool_permissions(&tc, db_type, &mut sql_permissions);
sequential_results.push(
agent_tools::execute_tool(
&tc,
&state2,
&conn2,
&db2,
schema2.as_deref(),
&db_type,
execution_permissions,
)
.await,
);
let result = match plugin_calls.remove(&i) {
Some(Ok(prepared)) => {
plugin_tools::execute_plugin_tool(&state2, agent_ctx.host_runtime.as_ref(), &tc, &prepared).await
}
Some(Err(not_executed)) => not_executed,
None => {
let execution_permissions = sequential_tool_permissions(&tc, db_type, &mut sql_permissions);
agent_tools::execute_tool(
&tc,
&state2,
&conn2,
&db2,
schema2.as_deref(),
&db_type,
execution_permissions,
)
.await
}
};
sequential_results.push(result);
}
// Merge results back into original order
@@ -1095,7 +1148,7 @@ fn estimate_tool_schema_tokens(tools: &[ToolDefinition]) -> u32 {
.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
estimate_text_tokens(&tool.name) + estimate_text_tokens(&tool.description) + schema_tokens + 16
})
.sum()
}
@@ -1290,6 +1343,98 @@ fn write_attempt_response(sql: &str, targets_production: bool) -> WriteAttemptRe
}
}
const PLUGIN_TOOL_DECLINED: &str =
"The user declined this tool call, so it did not run. Do not retry it; continue without it.";
const PLUGIN_TOOL_APPROVAL_TIMED_OUT: &str =
"Nobody approved this tool call in time, so it did not run. Do not retry it; tell the user what you intended to run.";
const PLUGIN_TOOL_APPROVAL_UNAVAILABLE: &str =
"This tool may change state and needs the user's approval, which this session cannot ask for. It did not run.";
/// Per-call outcome of the plugin tool gate: run the prepared call, or reply
/// with the ready result explaining why it did not run.
type PluginCallGate = std::collections::HashMap<usize, Result<PreparedPluginToolCall, ToolResult>>;
/// Validates this turn's plugin tool calls and asks the user to approve every
/// call that is not declared read-only, one at a time in call order. Returns
/// `None` when the run was cancelled while an approval was pending.
async fn gate_plugin_tool_calls<F>(
plugin_tool_set: &PluginToolSet,
tool_calls: &[ToolCall],
agent_ctx: &AgentLoopContext,
on_event: &F,
cancelled: &Notify,
) -> Option<PluginCallGate>
where
F: Fn(AgentEvent) + Sync,
{
let mut gate = PluginCallGate::new();
if plugin_tool_set.is_empty() {
return Some(gate);
}
for (index, tool_call) in tool_calls.iter().enumerate() {
let Some(prepared) = plugin_tool_set.prepare_call(tool_call) else {
continue;
};
let outcome = match prepared {
Err(error) => Err(plugin_tools::not_executed_result(tool_call, error)),
Ok(prepared) if prepared.read_only => Ok(prepared),
Ok(prepared) => match agent_ctx.session_id.as_deref() {
None => Err(plugin_tools::not_executed_result(tool_call, PLUGIN_TOOL_APPROVAL_UNAVAILABLE)),
Some(session_id) => match request_tool_approval(session_id, tool_call, &prepared, on_event, cancelled)
.await
{
ToolApprovalWait::Approved => Ok(prepared),
ToolApprovalWait::Denied => Err(plugin_tools::not_executed_result(tool_call, PLUGIN_TOOL_DECLINED)),
ToolApprovalWait::TimedOut => {
Err(plugin_tools::not_executed_result(tool_call, PLUGIN_TOOL_APPROVAL_TIMED_OUT))
}
ToolApprovalWait::Cancelled => return None,
},
},
};
gate.insert(index, outcome);
}
Some(gate)
}
async fn request_tool_approval<F>(
session_id: &str,
tool_call: &ToolCall,
prepared: &PreparedPluginToolCall,
on_event: &F,
cancelled: &Notify,
) -> ToolApprovalWait
where
F: Fn(AgentEvent) + Sync,
{
let approval = tool_approval::register_tool_approval(session_id);
let approval_id = approval.id().to_string();
on_event(AgentEvent::ToolApprovalRequired {
approval_id: approval_id.clone(),
tool_call_id: tool_call.id.clone(),
tool_name: tool_call.name.clone(),
plugin_id: prepared.plugin_id.clone(),
plugin_name: prepared.plugin_name.clone(),
plugin_tool: prepared.plugin_tool.clone(),
connection_id: prepared.connection_id.clone(),
connection_name: prepared.connection_name.clone(),
args: prepared.arguments.clone(),
timeout_secs: tool_approval::TOOL_APPROVAL_TIMEOUT.as_secs(),
});
let wait = approval.wait(tool_approval::TOOL_APPROVAL_TIMEOUT, cancelled).await;
on_event(AgentEvent::ToolApprovalResolved {
approval_id,
tool_call_id: tool_call.id.clone(),
outcome: match wait {
ToolApprovalWait::Approved => ToolApprovalOutcome::Approved,
ToolApprovalWait::Denied => ToolApprovalOutcome::Denied,
ToolApprovalWait::TimedOut => ToolApprovalOutcome::TimedOut,
ToolApprovalWait::Cancelled => ToolApprovalOutcome::Cancelled,
},
});
wait
}
fn sequential_tool_permissions(
tool_call: &ToolCall,
db_type: DatabaseType,
@@ -1517,6 +1662,113 @@ fn summarize_message_content(content: &str) -> String {
mod tests {
use super::*;
fn plugin_tool_set() -> PluginToolSet {
PluginToolSet::from_listings_for_tests(
&[("io.dbx.kafka", "Kafka Studio")],
vec![(
plugin_tools::OpenPluginConnection {
connection_id: "k1".to_string(),
connection_name: "prod-kafka".to_string(),
plugin_id: "io.dbx.kafka".to_string(),
},
json!({ "tools": [
{ "name": "kafka_topics_list", "annotations": { "readOnlyHint": true }, "inputSchema": { "type": "object", "properties": {} } },
{ "name": "kafka_topics_delete", "inputSchema": { "type": "object", "properties": { "topics": { "type": "array", "items": { "type": "string" } } } } }
] }),
)],
)
}
fn plugin_tool_call(id: &str, name: &str) -> ToolCall {
ToolCall {
id: id.to_string(),
name: name.to_string(),
arguments: json!({ "topics": ["orders"] }),
provider_payload: None,
}
}
async fn gate_context(session_id: Option<&str>) -> (tempfile::TempDir, AgentLoopContext) {
let temp_dir = tempfile::tempdir().unwrap();
let storage = crate::storage::Storage::open(&temp_dir.path().join("storage.db")).await.unwrap();
let ctx = AgentLoopContext {
state: Arc::new(AppState::new(storage)),
connection_id: "db".to_string(),
database: String::new(),
selected_databases: Vec::new(),
schema: None,
db_type: DatabaseType::Postgres,
cli_mcp_server_command: None,
sql_permissions: agent_tools::AgentSqlPermissions::default(),
max_agent_turns: DEFAULT_MAX_AGENT_TURNS,
prompt_cache_key: None,
session_id: session_id.map(str::to_string),
host_runtime: None,
};
(temp_dir, ctx)
}
#[tokio::test]
async fn plugin_tool_gate_runs_read_only_calls_and_asks_before_the_rest() {
let (_dir, ctx) = gate_context(Some("session-gate")).await;
let set = plugin_tool_set();
let calls = vec![
plugin_tool_call("c1", "kafka__kafka_topics_list"),
plugin_tool_call("c2", "kafka__kafka_topics_delete"),
plugin_tool_call("c3", "list_tables"),
];
let events = Arc::new(Mutex::new(Vec::new()));
let recorded = events.clone();
// The "user" approves as soon as the question appears.
let on_event = move |event: AgentEvent| {
if let AgentEvent::ToolApprovalRequired { approval_id, .. } = &event {
assert!(tool_approval::resolve_tool_approval("session-gate", approval_id, true));
}
recorded.lock().unwrap().push(event);
};
let gate = gate_plugin_tool_calls(&set, &calls, &ctx, &on_event, &Notify::new()).await.unwrap();
assert!(gate[&0].as_ref().is_ok_and(|prepared| prepared.read_only));
let approved = gate[&1].as_ref().unwrap();
assert_eq!(approved.connection_id, "k1");
assert_eq!(approved.arguments, json!({ "topics": ["orders"] }));
assert!(!gate.contains_key(&2), "database tools bypass the plugin gate");
let events = events.lock().unwrap();
assert_eq!(events.len(), 2, "only the non-read-only call asks: {events:?}");
assert!(
matches!(&events[0], AgentEvent::ToolApprovalRequired { tool_call_id, plugin_tool, .. } if tool_call_id == "c2" && plugin_tool == "kafka_topics_delete")
);
assert!(matches!(&events[1], AgentEvent::ToolApprovalResolved { outcome: ToolApprovalOutcome::Approved, .. }));
}
#[tokio::test]
async fn plugin_tool_gate_refuses_declined_and_unaskable_calls_and_stops_on_cancel() {
let set = plugin_tool_set();
let calls = vec![plugin_tool_call("c1", "kafka__kafka_topics_delete")];
let (_dir, ctx) = gate_context(Some("session-deny")).await;
let deny = |event: AgentEvent| {
if let AgentEvent::ToolApprovalRequired { approval_id, .. } = &event {
tool_approval::resolve_tool_approval("session-deny", approval_id, false);
}
};
let gate = gate_plugin_tool_calls(&set, &calls, &ctx, &deny, &Notify::new()).await.unwrap();
let declined = gate[&0].as_ref().unwrap_err();
assert!(declined.is_error);
assert_eq!(declined.content, PLUGIN_TOOL_DECLINED);
// No client session: nobody can be asked, so the call never runs.
let (_dir, headless) = gate_context(None).await;
let gate = gate_plugin_tool_calls(&set, &calls, &headless, &|_event| {}, &Notify::new()).await.unwrap();
assert_eq!(gate[&0].as_ref().unwrap_err().content, PLUGIN_TOOL_APPROVAL_UNAVAILABLE);
let (_dir, ctx) = gate_context(Some("session-cancel")).await;
let cancelled = Notify::new();
cancelled.notify_one();
assert!(gate_plugin_tool_calls(&set, &calls, &ctx, &|_event| {}, &cancelled).await.is_none());
}
#[test]
fn compaction_prompt_preserves_attachment_trust_boundary() {
assert!(COMPACT_SYSTEM_PROMPT.contains("<attached-text-data>"));
+52 -31
View File
@@ -280,13 +280,14 @@ pub fn is_vector_db(db_type: DatabaseType) -> bool {
/// the current UTC time plus a caller-provided local offset.
fn get_current_time_tool() -> ToolDefinition {
ToolDefinition {
name: "get_current_time",
name: "get_current_time".into(),
description: "Get the current date and time with timezone information. \
Pass the client UTC offset from the system prompt; when \
omitted, local time safely falls back to UTC. Use this to resolve \
relative time expressions like \"last 7 days\", \
\"yesterday\", \"this month\" into concrete dates \
for constructing SQL queries.",
for constructing SQL queries."
.into(),
parameters: json!({
"type": "object",
"properties": {
@@ -341,7 +342,9 @@ fn execute_get_current_time(tool_call: &ToolCall) -> Result<String, String> {
/// Get read-only tool definitions for the given database type.
/// Returns vector tools for vector DBs, SQL tools otherwise.
pub fn read_only_tools(db_type: DatabaseType) -> Vec<ToolDefinition> {
if is_vector_db(db_type) {
if db_type == DatabaseType::Plugin {
vec![get_current_time_tool()]
} else if is_vector_db(db_type) {
vec![list_collections_tool(), get_current_time_tool()]
} else {
vec![list_databases_tool(), list_tables_tool(), get_columns_tool(db_type), get_current_time_tool()]
@@ -352,6 +355,9 @@ pub fn read_only_tools(db_type: DatabaseType) -> Vec<ToolDefinition> {
/// Includes read-only tools plus execute_query, get_sample_data, and
/// explain_query for database types that support them.
pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) -> Vec<ToolDefinition> {
if db_type == DatabaseType::Plugin {
return vec![get_current_time_tool()];
}
if is_vector_db(db_type) {
return vec![list_collections_tool(), browse_collection_tool(), get_current_time_tool()];
}
@@ -375,8 +381,8 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) ->
fn list_databases_tool() -> ToolDefinition {
ToolDefinition {
name: "list_databases",
description: "List databases available through the current connection. If more than one database is returned, cross-database read queries can use fully qualified names such as database.table (or database.schema.table for SQL Server).",
name: "list_databases".into(),
description: "List databases available through the current connection. If more than one database is returned, cross-database read queries can use fully qualified names such as database.table (or database.schema.table for SQL Server).".into(),
parameters: json!({"type": "object", "properties": {}, "required": []}),
read_only: true,
parallel_ok: true,
@@ -385,8 +391,8 @@ fn list_databases_tool() -> ToolDefinition {
fn mongo_execute_query_tool(_sql_permissions: AgentSqlPermissions) -> ToolDefinition {
ToolDefinition {
name: "execute_query",
description: "Execute a read-only MongoDB shell command and return results (max 50 rows). Use commands such as db.collection.find({}), db.collection.findOne({}), db.collection.aggregate([]), or db.collection.countDocuments({}). Write commands are not available to the MongoDB Agent.",
name: "execute_query".into(),
description: "Execute a read-only MongoDB shell command and return results (max 50 rows). Use commands such as db.collection.find({}), db.collection.findOne({}), db.collection.aggregate([]), or db.collection.countDocuments({}). Write commands are not available to the MongoDB Agent.".into(),
parameters: json!({
"type": "object",
"properties": {
@@ -423,8 +429,8 @@ fn solr_execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefiniti
request in one ```sql code block and ask for confirmation."
};
ToolDefinition {
name: "execute_query",
description,
name: "execute_query".into(),
description: description.into(),
parameters: json!({
"type": "object",
"properties": {
@@ -461,8 +467,9 @@ fn solr_execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefiniti
/// 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.",
name: "list_tables".into(),
description: "List all tables and views in the current database. Returns table names, types, and comments."
.into(),
parameters: json!({
"type": "object",
"properties": {
@@ -485,15 +492,17 @@ fn list_tables_tool() -> ToolDefinition {
/// get_columns tool definition.
fn get_columns_tool(db_type: DatabaseType) -> ToolDefinition {
ToolDefinition {
name: "get_columns",
name: "get_columns".into(),
description: if db_type == DatabaseType::MongoDb {
"Sample up to 100 documents from a MongoDB collection and infer up to 512 top-level field names and types. \
The sample may be smaller and is not a complete schema or a guarantee of required fields. \
Nested documents and arrays remain object and array fields; numeric BSON types are reported as number."
.into()
} else {
"Get column definitions for a table: names, types, primary keys, nullable, defaults, and comments. \
Use this when the user asks about table structure, column details, or field information — \
even if some schema context was provided, this tool returns the authoritative and complete column list."
.into()
},
parameters: json!({
"type": "object",
@@ -529,8 +538,8 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition {
"Execute a read-only SQL query and return results (default 50 rows, up to 1000 with the limit argument). Cross-database reads may use fully qualified names such as database.table (or database.schema.table for SQL Server) without switching the current database. This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run."
};
ToolDefinition {
name: "execute_query",
description,
name: "execute_query".into(),
description: description.into(),
parameters: json!({
"type": "object",
"properties": {
@@ -571,8 +580,8 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition {
/// get_sample_data tool definition.
fn get_sample_data_tool() -> ToolDefinition {
ToolDefinition {
name: "get_sample_data",
description: "Get sample rows from a table to understand its data. Returns up to 20 rows.",
name: "get_sample_data".into(),
description: "Get sample rows from a table to understand its data. Returns up to 20 rows.".into(),
parameters: json!({
"type": "object",
"properties": {
@@ -603,11 +612,12 @@ fn get_sample_data_tool() -> ToolDefinition {
/// explain_query tool definition (Phase 3).
fn explain_query_tool() -> ToolDefinition {
ToolDefinition {
name: "explain_query",
name: "explain_query".into(),
description: "Get the execution plan for a SQL query using EXPLAIN. \
Shows how the database will execute the query (scan type, indexes, cost). \
Only read-only queries (SELECT, WITH, SHOW, DESCRIBE, EXPLAIN) are allowed. \
Use this to analyze query performance and suggest index optimizations.",
Use this to analyze query performance and suggest index optimizations."
.into(),
parameters: json!({
"type": "object",
"properties": {
@@ -626,8 +636,9 @@ fn explain_query_tool() -> ToolDefinition {
/// list_collections tool definition (vector databases).
fn list_collections_tool() -> ToolDefinition {
ToolDefinition {
name: "list_collections",
description: "List all collections in the current vector database. Returns collection names and dimensions.",
name: "list_collections".into(),
description: "List all collections in the current vector database. Returns collection names and dimensions."
.into(),
parameters: json!({
"type": "object",
"properties": {},
@@ -641,8 +652,8 @@ fn list_collections_tool() -> ToolDefinition {
/// browse_collection tool definition (vector databases).
fn browse_collection_tool() -> ToolDefinition {
ToolDefinition {
name: "browse_collection",
description: "Browse documents in a collection. Returns up to 20 items with payload/metadata (vectors excluded for compactness). For ChromaDB, use the collection id (UUID from list_collections) instead of the collection name.",
name: "browse_collection".into(),
description: "Browse documents in a collection. Returns up to 20 items with payload/metadata (vectors excluded for compactness). For ChromaDB, use the collection id (UUID from list_collections) instead of the collection name.".into(),
parameters: json!({
"type": "object",
"properties": {
@@ -1697,7 +1708,7 @@ for line in sys.stdin:
#[test]
fn vector_read_only_tools_do_not_include_collection_browsing() {
let tools = read_only_tools(DatabaseType::Qdrant);
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"list_collections"));
assert!(!names.contains(&"browse_collection"));
@@ -1707,17 +1718,27 @@ for line in sys.stdin:
#[test]
fn vector_agent_tools_include_collection_browsing() {
let tools = all_tools(DatabaseType::Qdrant, AgentSqlPermissions::default());
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"list_collections"));
assert!(names.contains(&"browse_collection"));
assert!(names.contains(&"get_current_time"));
}
#[test]
fn plugin_connections_only_receive_database_independent_builtin_tools() {
for tools in
[read_only_tools(DatabaseType::Plugin), all_tools(DatabaseType::Plugin, AgentSqlPermissions::default())]
{
let names = tools.iter().map(|tool| tool.name.as_ref()).collect::<Vec<_>>();
assert_eq!(names, ["get_current_time"]);
}
}
#[test]
fn solr_agent_registers_rest_execute_query_tool() {
let tools = all_tools(DatabaseType::Solr, AgentSqlPermissions::default());
let names = tools.iter().map(|tool| tool.name).collect::<Vec<_>>();
let names = tools.iter().map(|tool| tool.name.as_ref()).collect::<Vec<_>>();
// Solr gets the shared metadata tools plus a REST `execute_query`; SQL
// sample/explain helpers are not meaningful for a non-SQL backend.
assert!(names.contains(&"list_tables"));
@@ -1746,7 +1767,7 @@ for line in sys.stdin:
#[tokio::test]
async fn mongodb_agent_registers_shell_query_tool_and_routes_find_one_as_read_only() {
let tools = all_tools(DatabaseType::MongoDb, AgentSqlPermissions::default());
let names = tools.iter().map(|tool| tool.name).collect::<Vec<_>>();
let names = tools.iter().map(|tool| tool.name.as_ref()).collect::<Vec<_>>();
assert!(names.contains(&"execute_query"));
assert!(!names.contains(&"get_sample_data"));
assert!(!names.contains(&"explain_query"));
@@ -1906,7 +1927,7 @@ for line in sys.stdin:
#[test]
fn oracle_agent_tools_include_explain_query() {
let tools = all_tools(DatabaseType::Oracle, AgentSqlPermissions::default());
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"explain_query"));
}
@@ -2456,28 +2477,28 @@ for line in sys.stdin:
#[test]
fn get_current_time_is_in_all_tools_postgres() {
let tools = all_tools(DatabaseType::Postgres, AgentSqlPermissions::default());
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"get_current_time"), "get_current_time missing from all_tools(Postgres)");
}
#[test]
fn get_current_time_is_in_read_only_tools_postgres() {
let tools = read_only_tools(DatabaseType::Postgres);
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"get_current_time"), "get_current_time missing from read_only_tools(Postgres)");
}
#[test]
fn get_current_time_is_in_all_tools_qdrant() {
let tools = all_tools(DatabaseType::Qdrant, AgentSqlPermissions::default());
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"get_current_time"), "get_current_time missing from all_tools(Qdrant)");
}
#[test]
fn get_current_time_is_in_read_only_tools_qdrant() {
let tools = read_only_tools(DatabaseType::Qdrant);
let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect();
let names: Vec<&str> = tools.iter().map(|tool| tool.name.as_ref()).collect();
assert!(names.contains(&"get_current_time"), "get_current_time missing from read_only_tools(Qdrant)");
}
+2
View File
@@ -5,4 +5,6 @@ pub mod agent_kv;
pub mod agent_loop;
pub mod agent_tools;
pub mod mcp_policy;
pub mod plugin_tools;
pub mod prompt_template;
pub mod tool_approval;
File diff suppressed because it is too large. Load diff
+158
View File
@@ -0,0 +1,158 @@
//! Interactive approvals for agent tool calls.
//!
//! A tool call that may change state outside DBX (today: plugin MCP tools not
//! declared read-only) must not run on the model's word alone. The agent loop
//! registers a pending approval, emits `AgentEvent::ToolApprovalRequired`, and
//! suspends until the client answers through [`resolve_tool_approval`], the
//! deadline passes, or the run is cancelled.
//!
//! The registry is process-wide and deliberately runtime-agnostic (a std mutex
//! plus `oneshot` channels): the web server runs each agent loop on its own
//! runtime while the approval arrives through an HTTP handler on another one.
//! An approval is bound to the AI session that requested it, so a client can
//! only answer questions its own run asked.
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use std::time::Duration;
use tokio::sync::{oneshot, Notify};
/// How long a pending approval waits for the user before it counts as denied.
pub const TOOL_APPROVAL_TIMEOUT: Duration = Duration::from_secs(300);
struct PendingApproval {
session_id: String,
responder: oneshot::Sender<bool>,
}
fn pending_approvals() -> &'static Mutex<HashMap<String, PendingApproval>> {
static PENDING: OnceLock<Mutex<HashMap<String, PendingApproval>>> = OnceLock::new();
PENDING.get_or_init(|| Mutex::new(HashMap::new()))
}
/// How a wait for approval ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolApprovalWait {
Approved,
Denied,
TimedOut,
Cancelled,
}
/// A registered approval. Dropping it removes the entry, so an abandoned wait
/// (cancelled future, early return) can never leave a stale approval behind
/// for a later answer to hit.
pub struct PendingToolApproval {
approval_id: String,
receiver: Option<oneshot::Receiver<bool>>,
}
impl PendingToolApproval {
pub fn id(&self) -> &str {
&self.approval_id
}
/// Waits for the user's answer. A dropped responder (the entry was purged)
/// counts as a denial: without an answer there is no approval.
pub async fn wait(mut self, timeout: Duration, cancelled: &Notify) -> ToolApprovalWait {
let Some(receiver) = self.receiver.take() else {
return ToolApprovalWait::Denied;
};
tokio::select! {
biased;
_ = cancelled.notified() => ToolApprovalWait::Cancelled,
answer = tokio::time::timeout(timeout, receiver) => match answer {
Ok(Ok(true)) => ToolApprovalWait::Approved,
Ok(Ok(false)) | Ok(Err(_)) => ToolApprovalWait::Denied,
Err(_) => ToolApprovalWait::TimedOut,
},
}
}
}
impl Drop for PendingToolApproval {
fn drop(&mut self) {
pending_approvals().lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&self.approval_id);
}
}
/// Registers a pending approval for `session_id`.
pub fn register_tool_approval(session_id: &str) -> PendingToolApproval {
let approval_id = uuid::Uuid::new_v4().to_string();
let (responder, receiver) = oneshot::channel();
pending_approvals()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.insert(approval_id.clone(), PendingApproval { session_id: session_id.to_string(), responder });
PendingToolApproval { approval_id, receiver: Some(receiver) }
}
/// Answers a pending approval. Returns `false` when no approval with this id is
/// pending for `session_id` (already answered, timed out, or never asked).
pub fn resolve_tool_approval(session_id: &str, approval_id: &str, approved: bool) -> bool {
let mut pending = pending_approvals().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
let belongs_to_session = pending.get(approval_id).is_some_and(|entry| entry.session_id == session_id);
if !belongs_to_session {
return false;
}
let Some(entry) = pending.remove(approval_id) else {
return false;
};
entry.responder.send(approved).is_ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn resolved_approval_reports_the_user_answer() {
let cancelled = Notify::new();
let approval = register_tool_approval("session-a");
let id = approval.id().to_string();
assert!(resolve_tool_approval("session-a", &id, true));
assert_eq!(approval.wait(Duration::from_secs(5), &cancelled).await, ToolApprovalWait::Approved);
let denied = register_tool_approval("session-a");
let id = denied.id().to_string();
assert!(resolve_tool_approval("session-a", &id, false));
assert_eq!(denied.wait(Duration::from_secs(5), &cancelled).await, ToolApprovalWait::Denied);
}
#[tokio::test]
async fn another_session_cannot_answer_an_approval() {
let cancelled = Notify::new();
let approval = register_tool_approval("owner");
let id = approval.id().to_string();
assert!(!resolve_tool_approval("intruder", &id, true));
assert!(resolve_tool_approval("owner", &id, false));
assert_eq!(approval.wait(Duration::from_secs(5), &cancelled).await, ToolApprovalWait::Denied);
}
#[tokio::test]
async fn unanswered_approval_times_out_and_is_forgotten() {
let cancelled = Notify::new();
let approval = register_tool_approval("session-b");
let id = approval.id().to_string();
assert_eq!(approval.wait(Duration::from_millis(20), &cancelled).await, ToolApprovalWait::TimedOut);
assert!(!resolve_tool_approval("session-b", &id, true), "a timed-out approval must not accept a late answer");
}
#[tokio::test]
async fn cancellation_ends_the_wait() {
let cancelled = Notify::new();
cancelled.notify_one();
let approval = register_tool_approval("session-c");
assert_eq!(approval.wait(Duration::from_secs(5), &cancelled).await, ToolApprovalWait::Cancelled);
}
#[tokio::test]
async fn cancellation_wins_over_an_answer_that_arrived_at_the_same_time() {
let cancelled = Notify::new();
let approval = register_tool_approval("session-race");
assert!(resolve_tool_approval("session-race", approval.id(), true));
cancelled.notify_one();
assert_eq!(approval.wait(Duration::from_secs(5), &cancelled).await, ToolApprovalWait::Cancelled);
}
}
+2
View File
@@ -17,6 +17,8 @@ pub use ai::agent_explain;
pub use ai::agent_kv;
pub use ai::agent_loop;
pub use ai::agent_tools;
pub use ai::plugin_tools;
pub use ai::tool_approval;
pub use dbx_drivers::agent_catalog;
pub use dbx_drivers::agent_connection;
pub use dbx_drivers::agent_manager;
+195 -2
View File
@@ -56,6 +56,14 @@ const MCP_HTTP_SERVER_SETTINGS_KEY: &str = "mcp_http_server_settings";
const MAX_RETRIES_KEY: &str = "max_retries";
const HISTORY_RETENTION_LIMIT_KEY: &str = "history_retention_limit";
const SQL_FILE_UPLOAD_MAX_MB_KEY: &str = "sql_file_upload_max_mb";
/// Plugin ids whose MCP tools the built-in AI agent may call.
const AI_PLUGIN_TOOL_PLUGINS_KEY: &str = "ai_plugin_tool_plugins";
/// `{ pluginId: [connectionId, ...] }` — connections a plugin may read through
/// the `host.data:read` Host API. Written only after an explicit user consent.
const PLUGIN_DATA_GRANTS_KEY: &str = "plugin_data_grants";
/// Upper bound for one plugin's persisted data grants; consent is per
/// connection, so a legitimate plugin never approaches it.
const MAX_PLUGIN_DATA_GRANTS_PER_PLUGIN: usize = 512;
const APP_STATE_AI_GLOBAL_INSTRUCTIONS_KEY: &str = "ai_global_custom_instructions";
const APP_STATE_AI_CHAT_SELECTION_KEY: &str = "ai_chat_selection_v1";
const SNIPPET_SYNC_IDS_KEY: &str = "snippet_sync_ids";
@@ -3132,6 +3140,56 @@ fn load_history_retention_limit_from_conn(conn: &Connection) -> Result<u32, Stri
Ok(history_retention_limit_from_settings(&settings))
}
fn app_settings_map_from_conn(conn: &Connection) -> Result<serde_json::Map<String, serde_json::Value>, String> {
let current: Option<String> = conn
.query_row("SELECT settings_json FROM app_settings WHERE id = 1", [], |row| row.get(0))
.optional()
.map_err(|e| e.to_string())?;
match current {
Some(json) => serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(&json)
.map_err(|e| format!("invalid app settings JSON: {e}")),
None => Ok(serde_json::Map::new()),
}
}
fn write_app_settings_map(
conn: &Connection,
settings: &serde_json::Map<String, serde_json::Value>,
) -> Result<(), String> {
let json = serde_json::to_string(settings).map_err(|e| e.to_string())?;
conn.execute("INSERT OR REPLACE INTO app_settings (id, settings_json) VALUES (1, ?1)", [json])
.map(|_| ())
.map_err(|e| e.to_string())
}
/// Reads a JSON string array setting, dropping blanks and duplicates so a
/// hand-edited or partially written value can never widen a permission list.
fn normalized_string_list(value: Option<&serde_json::Value>) -> Vec<String> {
let mut values: Vec<String> = value
.and_then(serde_json::Value::as_array)
.map(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.map(str::trim)
.filter(|item| !item.is_empty())
.map(str::to_string)
.collect()
})
.unwrap_or_default();
values.sort();
values.dedup();
values
}
fn normalized_setting_id(value: &str, label: &str) -> Result<String, String> {
let trimmed = value.trim();
if trimmed.is_empty() || trimmed.chars().count() > 256 {
return Err(format!("A valid {label} is required"));
}
Ok(trimmed.to_string())
}
impl Storage {
pub async fn save_history_entry(&self, entry: &HistoryEntry) -> Result<(), String> {
let entry = entry.clone();
@@ -3956,8 +4014,14 @@ impl Storage {
.query_row("SELECT settings_json FROM app_settings WHERE id = 1", [], |row| row.get(0))
.optional()
.map_err(|e| e.to_string())?;
let dedicated_keys =
[MCP_GLOBAL_POLICY_KEY, MAX_RETRIES_KEY, SQL_FILE_UPLOAD_MAX_MB_KEY, HISTORY_RETENTION_LIMIT_KEY];
let dedicated_keys = [
MCP_GLOBAL_POLICY_KEY,
MAX_RETRIES_KEY,
SQL_FILE_UPLOAD_MAX_MB_KEY,
HISTORY_RETENTION_LIMIT_KEY,
AI_PLUGIN_TOOL_PLUGINS_KEY,
PLUGIN_DATA_GRANTS_KEY,
];
for key in dedicated_keys {
settings.remove(key);
}
@@ -4067,6 +4131,101 @@ impl Storage {
.map_err(|e| format!("MCP_POLICY_UNAVAILABLE: {e}"))
}
/// Plugin ids whose MCP tools the built-in AI agent may call. Empty until
/// the user enables a plugin in the Plugin Center.
pub async fn load_ai_plugin_tool_plugin_ids(&self) -> Result<Vec<String>, String> {
let settings = self.load_app_settings_json().await?;
Ok(normalized_string_list(settings.get(AI_PLUGIN_TOOL_PLUGINS_KEY)))
}
/// Enables or disables built-in AI access to one plugin's tools and returns
/// the updated list. The read-modify-write runs inside one connection
/// closure so concurrent settings saves cannot drop the change.
pub async fn set_ai_plugin_tool_plugin_enabled(
&self,
plugin_id: &str,
enabled: bool,
) -> Result<Vec<String>, String> {
let plugin_id = normalized_setting_id(plugin_id, "plugin id")?;
self.with_conn(move |conn| {
let mut settings = app_settings_map_from_conn(conn)?;
let mut plugin_ids = normalized_string_list(settings.get(AI_PLUGIN_TOOL_PLUGINS_KEY));
plugin_ids.retain(|candidate| candidate != &plugin_id);
if enabled {
plugin_ids.push(plugin_id);
plugin_ids.sort();
}
settings.insert(AI_PLUGIN_TOOL_PLUGINS_KEY.to_string(), serde_json::json!(plugin_ids));
write_app_settings_map(conn, &settings)?;
Ok(plugin_ids)
})
.await
}
/// Connections the plugin may read through `host.data:read`.
pub async fn load_plugin_data_grants(&self, plugin_id: &str) -> Result<Vec<String>, String> {
let settings = self.load_app_settings_json().await?;
Ok(normalized_string_list(settings.get(PLUGIN_DATA_GRANTS_KEY).and_then(|grants| grants.get(plugin_id.trim()))))
}
/// Records or revokes the user's consent for one plugin to read one
/// connection and returns the plugin's updated grant list.
pub async fn set_plugin_data_grant(
&self,
plugin_id: &str,
connection_id: &str,
granted: bool,
) -> Result<Vec<String>, String> {
let plugin_id = normalized_setting_id(plugin_id, "plugin id")?;
let connection_id = normalized_setting_id(connection_id, "connection id")?;
self.with_conn(move |conn| {
let mut settings = app_settings_map_from_conn(conn)?;
let mut all_grants = match settings.remove(PLUGIN_DATA_GRANTS_KEY) {
Some(serde_json::Value::Object(grants)) => grants,
_ => serde_json::Map::new(),
};
let mut connection_ids = normalized_string_list(all_grants.get(&plugin_id));
connection_ids.retain(|candidate| candidate != &connection_id);
if granted {
if connection_ids.len() >= MAX_PLUGIN_DATA_GRANTS_PER_PLUGIN {
return Err(format!(
"A plugin can hold at most {MAX_PLUGIN_DATA_GRANTS_PER_PLUGIN} data access grants"
));
}
connection_ids.push(connection_id);
connection_ids.sort();
}
if connection_ids.is_empty() {
all_grants.remove(&plugin_id);
} else {
all_grants.insert(plugin_id, serde_json::json!(connection_ids));
}
if !all_grants.is_empty() {
settings.insert(PLUGIN_DATA_GRANTS_KEY.to_string(), serde_json::Value::Object(all_grants));
}
write_app_settings_map(conn, &settings)?;
Ok(connection_ids)
})
.await
}
/// Drops everything the user granted to a plugin (AI tool access and data
/// grants). Called after uninstall so a reinstall starts without consent.
pub async fn forget_plugin_permissions(&self, plugin_id: &str) -> Result<(), String> {
let plugin_id = normalized_setting_id(plugin_id, "plugin id")?;
self.with_conn(move |conn| {
let mut settings = app_settings_map_from_conn(conn)?;
let mut plugin_ids = normalized_string_list(settings.get(AI_PLUGIN_TOOL_PLUGINS_KEY));
plugin_ids.retain(|candidate| candidate != &plugin_id);
settings.insert(AI_PLUGIN_TOOL_PLUGINS_KEY.to_string(), serde_json::json!(plugin_ids));
if let Some(serde_json::Value::Object(grants)) = settings.get_mut(PLUGIN_DATA_GRANTS_KEY) {
grants.remove(&plugin_id);
}
write_app_settings_map(conn, &settings)
})
.await
}
pub async fn load_mcp_http_server_settings(&self) -> Result<McpHttpServerSettings, String> {
let settings = self.load_app_settings_json().await?;
match settings.get(MCP_HTTP_SERVER_SETTINGS_KEY) {
@@ -11092,6 +11251,40 @@ mod tests {
assert_eq!(storage.load_max_retries().await.unwrap(), crate::ai::MAX_MAX_RETRIES);
}
#[tokio::test]
async fn plugin_ai_tools_and_data_grants_persist_and_survive_legacy_settings_saves() {
let path = temp_db_path("plugin-permissions");
let storage = Storage::open(&path).await.unwrap();
assert!(storage.load_ai_plugin_tool_plugin_ids().await.unwrap().is_empty());
assert!(storage.load_plugin_data_grants("io.dbx.chart").await.unwrap().is_empty());
assert_eq!(storage.set_ai_plugin_tool_plugin_enabled("io.dbx.ssh", true).await.unwrap(), ["io.dbx.ssh"]);
storage.set_ai_plugin_tool_plugin_enabled("io.dbx.kafka", true).await.unwrap();
storage.set_ai_plugin_tool_plugin_enabled("io.dbx.ssh", true).await.unwrap();
assert_eq!(storage.load_ai_plugin_tool_plugin_ids().await.unwrap(), ["io.dbx.kafka", "io.dbx.ssh"]);
storage.set_plugin_data_grant("io.dbx.chart", "conn-b", true).await.unwrap();
assert_eq!(storage.set_plugin_data_grant("io.dbx.chart", "conn-a", true).await.unwrap(), ["conn-a", "conn-b"]);
storage.set_plugin_data_grant("io.dbx.other", "conn-a", true).await.unwrap();
// A settings writer that rebuilds the whole JSON map must not drop the
// dedicated permission keys.
storage.save_password_hash("hash").await.unwrap();
storage.save_app_settings_json(&serde_json::Map::new()).await.unwrap();
assert_eq!(storage.load_ai_plugin_tool_plugin_ids().await.unwrap(), ["io.dbx.kafka", "io.dbx.ssh"]);
assert_eq!(storage.load_plugin_data_grants("io.dbx.chart").await.unwrap(), ["conn-a", "conn-b"]);
assert_eq!(storage.set_plugin_data_grant("io.dbx.chart", "conn-b", false).await.unwrap(), ["conn-a"]);
assert_eq!(storage.set_ai_plugin_tool_plugin_enabled("io.dbx.kafka", false).await.unwrap(), ["io.dbx.ssh"]);
assert!(storage.set_plugin_data_grant("io.dbx.chart", " ", true).await.is_err());
storage.set_ai_plugin_tool_plugin_enabled("io.dbx.chart", true).await.unwrap();
storage.forget_plugin_permissions("io.dbx.chart").await.unwrap();
assert_eq!(storage.load_ai_plugin_tool_plugin_ids().await.unwrap(), ["io.dbx.ssh"]);
assert!(storage.load_plugin_data_grants("io.dbx.chart").await.unwrap().is_empty());
assert_eq!(storage.load_plugin_data_grants("io.dbx.other").await.unwrap(), ["conn-a"]);
}
#[tokio::test]
async fn max_retries_survives_stale_app_settings_save() {
let path = temp_db_path("max-retries-stale-save");
+1
View File
@@ -3,6 +3,7 @@ pub mod document_ops;
pub mod hbase_ops;
pub mod mongo_ops;
pub mod object_cache;
pub mod plugin_data;
pub mod plugin_plan;
pub mod query_cancel;
pub mod redis_ops;
+473
View File
@@ -0,0 +1,473 @@
//! Read-only data access for plugins (`host.data:read`, Host API 1.4).
//!
//! A plugin UI can ask the host to run one read-only SQL statement on a DBX
//! connection. The plugin never receives a driver, a pool, a credential, or a
//! connection string, and the host owns every safety decision:
//!
//! * the plugin must be installed, compatible, and declare `host.data:read`;
//! * the user must have granted this plugin access to this connection — the
//! grant is persisted per (plugin, connection) and revocable in the Plugin
//! Center, and a plugin cannot grant itself anything;
//! * the connection must already be open: like the plan and schema metadata
//! APIs, DBX never connects on a plugin's behalf;
//! * the SQL must be a single statement the shared risk classifier rates
//! read-only (the same gate as MCP read-only access and the AI agent), and
//! may not switch the session's database, because the statement runs on the
//! connection's shared pool;
//! * execution uses the host's own connection, timeout ceiling, and row and
//! byte caps.
use std::time::Instant;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::plugin_plan::plugin_plan_timeout_secs;
use super::{execute_sql_statement_with_options, QueryExecutionOptions};
use crate::connection::AppState;
use crate::db::QueryResult;
use crate::models::connection::{ConnectionConfig, DatabaseType};
use crate::sql_risk::SqlRisk;
pub const PLUGIN_DATA_READ_PERMISSION: &str = "host.data:read";
/// Error code prefix for a missing (or revoked) user grant, so the host UI can
/// ask for consent instead of showing a failure.
pub const PLUGIN_DATA_ACCESS_NOT_GRANTED: &str = "PLUGIN_DATA_ACCESS_NOT_GRANTED";
pub const DEFAULT_PLUGIN_DATA_MAX_ROWS: usize = 500;
pub const MAX_PLUGIN_DATA_MAX_ROWS: usize = 5_000;
/// Serialized row payload cap; well inside what the plugin bridge carries.
pub const MAX_PLUGIN_DATA_RESULT_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_PLUGIN_DATA_SQL_CHARS: usize = 100_000;
const MAX_PLUGIN_DATA_NAME_CHARS: usize = 256;
/// One read request. There is deliberately no way to pass a session id, an
/// execution mode, or several statements.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginDataQueryRequest {
pub connection_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub database: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub schema: Option<String>,
pub sql: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_rows: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timeout_ms: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginDataColumn {
pub name: String,
/// Database type name when the driver reports one.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub data_type: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginDataQueryResult {
pub db_type: String,
pub columns: Vec<PluginDataColumn>,
pub rows: Vec<Vec<Value>>,
/// True when rows were cut by the row cap, the byte cap, or the driver.
pub truncated: bool,
pub elapsed_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginDataGrant {
pub connection_id: String,
/// Absent when the granted connection no longer exists.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub connection_name: Option<String>,
}
/// Runs one read-only statement for `plugin_id`. `plugin_id` must come from
/// the host (the plugin bridge binds it), never from plugin-supplied data.
pub async fn query_plugin_data(
state: &AppState,
plugin_id: &str,
request: PluginDataQueryRequest,
) -> Result<PluginDataQueryResult, String> {
let request = validate_request(request)?;
require_plugin_permission(state, plugin_id)?;
// The grant is checked before the connection is resolved, so an unknown id
// and an ungranted one look the same to the plugin.
require_grant(state, plugin_id, &request.connection_id).await?;
let config = connection_config(state, &request.connection_id).await?;
if !state.is_connection_open(&request.connection_id).await {
return Err("Connection is not open".to_string());
}
require_sql_connection(&config)?;
require_single_read_only_statement(&request.sql, config.db_type)?;
let max_rows = request.max_rows.unwrap_or(DEFAULT_PLUGIN_DATA_MAX_ROWS).clamp(1, MAX_PLUGIN_DATA_MAX_ROWS);
let started = Instant::now();
let result = execute_sql_statement_with_options(
state,
&request.connection_id,
request.database.as_deref().unwrap_or_default(),
&request.sql,
request.schema.as_deref(),
None,
QueryExecutionOptions {
max_rows: Some(max_rows),
max_result_bytes: Some(MAX_PLUGIN_DATA_RESULT_BYTES),
timeout_secs: Some(plugin_plan_timeout_secs(request.timeout_ms, &config)),
..Default::default()
},
)
.await?;
let shaped = shape_result(config.db_type, result, max_rows, started.elapsed().as_millis() as u64);
// Audit trail without SQL text or values: literals can carry user data.
log::info!(
"[plugin-data] plugin={plugin_id} connection={} rows={} truncated={}",
request.connection_id,
shaped.rows.len(),
shaped.truncated
);
Ok(shaped)
}
/// The connections the user granted to `plugin_id`, with current names.
pub async fn list_plugin_data_grants(state: &AppState, plugin_id: &str) -> Result<Vec<PluginDataGrant>, String> {
let connection_ids = state.storage.load_plugin_data_grants(plugin_id).await?;
let configs = state.configs.read().await;
Ok(connection_ids
.into_iter()
.map(|connection_id| PluginDataGrant {
connection_name: configs.get(&connection_id).map(|config| config.name.clone()),
connection_id,
})
.collect())
}
/// Records the user's answer to a consent prompt, or a revocation. Granting
/// requires a plugin that declared the permission and an existing SQL
/// connection, so a grant can never name something the API would refuse.
pub async fn set_plugin_data_grant(
state: &AppState,
plugin_id: &str,
connection_id: &str,
granted: bool,
) -> Result<Vec<PluginDataGrant>, String> {
if granted {
require_plugin_permission(state, plugin_id)?;
let config = connection_config(state, connection_id.trim()).await?;
require_sql_connection(&config)?;
}
state.storage.set_plugin_data_grant(plugin_id, connection_id, granted).await?;
list_plugin_data_grants(state, plugin_id).await
}
fn validate_request(mut request: PluginDataQueryRequest) -> Result<PluginDataQueryRequest, String> {
request.connection_id = required_name(&request.connection_id, "connectionId")?;
let sql = request.sql.trim();
if sql.is_empty() {
return Err("Plugin data queries need a non-empty sql".to_string());
}
if sql.chars().count() > MAX_PLUGIN_DATA_SQL_CHARS {
return Err(format!("Plugin data query sql exceeds {MAX_PLUGIN_DATA_SQL_CHARS} characters"));
}
request.sql = sql.to_string();
request.database = optional_name(request.database, "database")?;
request.schema = optional_name(request.schema, "schema")?;
Ok(request)
}
fn required_name(value: &str, label: &str) -> Result<String, String> {
let trimmed = value.trim();
if trimmed.is_empty() || trimmed.chars().count() > MAX_PLUGIN_DATA_NAME_CHARS {
return Err(format!("Plugin data queries need a valid {label}"));
}
Ok(trimmed.to_string())
}
fn optional_name(value: Option<String>, label: &str) -> Result<Option<String>, String> {
match value.as_deref().map(str::trim) {
None | Some("") => Ok(None),
Some(trimmed) if trimmed.chars().count() > MAX_PLUGIN_DATA_NAME_CHARS => {
Err(format!("Plugin data query {label} is too long"))
}
Some(trimmed) => Ok(Some(trimmed.to_string())),
}
}
fn require_plugin_permission(state: &AppState, plugin_id: &str) -> Result<(), String> {
let plugin =
state.plugins.find_plugin(plugin_id)?.ok_or_else(|| format!("Plugin '{plugin_id}' is not installed"))?;
if !plugin.compatibility.compatible {
return Err(format!("Plugin '{plugin_id}' is not compatible with this DBX version"));
}
if !plugin.manifest.permissions.iter().any(|permission| permission == PLUGIN_DATA_READ_PERMISSION) {
return Err(format!("Plugin has not declared permission '{PLUGIN_DATA_READ_PERMISSION}'"));
}
Ok(())
}
async fn require_grant(state: &AppState, plugin_id: &str, connection_id: &str) -> Result<(), String> {
let grants = state.storage.load_plugin_data_grants(plugin_id).await?;
if grants.iter().any(|granted| granted == connection_id) {
return Ok(());
}
Err(format!("{PLUGIN_DATA_ACCESS_NOT_GRANTED}: the user has not granted this plugin access to the connection"))
}
async fn connection_config(state: &AppState, connection_id: &str) -> Result<ConnectionConfig, String> {
state.configs.read().await.get(connection_id).cloned().ok_or_else(|| "Connection config not found".to_string())
}
fn require_sql_connection(config: &ConnectionConfig) -> Result<(), String> {
let sql_capable = crate::query_execution_sql::supports_sql_query(config.db_type)
&& !matches!(config.db_type, DatabaseType::Plugin | DatabaseType::MessageQueue | DatabaseType::Mqtt);
if sql_capable {
return Ok(());
}
Err(format!("Plugin data queries are not available for '{}' connections", config.db_type.as_str()))
}
fn require_single_read_only_statement(sql: &str, db_type: DatabaseType) -> Result<(), String> {
if crate::sql::split_sql_statements_for_database(sql, db_type).len() != 1 {
return Err("Plugin data queries accept exactly one SQL statement".to_string());
}
// Messages stay generic so a plugin cannot use them to probe the gate.
match crate::sql_risk::classify_sql_risk_for_database(sql, db_type) {
Ok(SqlRisk::ReadOnly) => {}
Ok(_) => return Err("Plugin data queries are read-only; the statement was rejected".to_string()),
Err(_) => return Err("The statement could not be verified as read-only".to_string()),
}
if crate::sql_risk::mcp_sql_has_forbidden_database_switch(sql, db_type) {
return Err("Plugin data queries cannot switch the session database; pass `database` instead".to_string());
}
Ok(())
}
fn shape_result(db_type: DatabaseType, result: QueryResult, max_rows: usize, elapsed_ms: u64) -> PluginDataQueryResult {
let columns = result
.columns
.iter()
.enumerate()
.map(|(index, name)| PluginDataColumn {
name: name.clone(),
data_type: result.column_types.get(index).filter(|data_type| !data_type.is_empty()).cloned(),
})
.collect();
let mut truncated = result.truncated || result.has_more || result.rows.len() > max_rows;
let mut rows = Vec::new();
let mut bytes = 0usize;
for row in result.rows.into_iter().take(max_rows) {
let row_bytes = serde_json::to_vec(&row).map_or(0, |encoded| encoded.len());
if bytes + row_bytes > MAX_PLUGIN_DATA_RESULT_BYTES {
truncated = true;
break;
}
bytes += row_bytes;
rows.push(row);
}
PluginDataQueryResult { db_type: db_type.as_str().to_string(), columns, rows, truncated, elapsed_ms }
}
#[cfg(test)]
mod tests {
use super::*;
fn request(sql: &str) -> PluginDataQueryRequest {
PluginDataQueryRequest {
connection_id: " conn-1 ".to_string(),
database: Some(" ".to_string()),
schema: None,
sql: sql.to_string(),
max_rows: None,
timeout_ms: None,
}
}
#[test]
fn requests_are_trimmed_and_bounded() {
let validated = validate_request(request(" SELECT 1 ")).unwrap();
assert_eq!(validated.connection_id, "conn-1");
assert_eq!(validated.database, None);
assert_eq!(validated.sql, "SELECT 1");
assert!(validate_request(request(" ")).is_err());
assert!(
validate_request(PluginDataQueryRequest { connection_id: " ".to_string(), ..request("SELECT 1") }).is_err()
);
assert!(validate_request(request(&"x".repeat(MAX_PLUGIN_DATA_SQL_CHARS + 1))).is_err());
}
#[test]
fn only_a_single_read_only_statement_passes() {
assert!(require_single_read_only_statement("SELECT * FROM users WHERE id = 1", DatabaseType::Postgres).is_ok());
assert!(require_single_read_only_statement("WITH t AS (SELECT 1) SELECT * FROM t", DatabaseType::Mysql).is_ok());
for rejected in [
"DELETE FROM users",
"UPDATE users SET name = 'x' WHERE id = 1",
"DROP TABLE users",
"SELECT * FROM users FOR UPDATE",
"SELECT 1; SELECT 2",
"SELECT 1; DELETE FROM users",
] {
assert!(require_single_read_only_statement(rejected, DatabaseType::Postgres).is_err(), "{rejected}");
}
assert!(require_single_read_only_statement("USE other_db", DatabaseType::Mysql).is_err());
}
#[test]
fn non_sql_connections_are_rejected() {
let config = |db_type: &str| -> ConnectionConfig {
serde_json::from_value(serde_json::json!({
"id": "c", "name": "c", "db_type": db_type, "host": "h", "port": 1, "username": "", "password": ""
}))
.unwrap()
};
assert!(require_sql_connection(&config("postgres")).is_ok());
for db_type in ["redis", "mongodb", "plugin"] {
assert!(require_sql_connection(&config(db_type)).is_err(), "{db_type}");
}
}
#[test]
fn results_respect_the_row_cap_and_keep_column_types() {
let result: QueryResult = serde_json::from_value(serde_json::json!({
"columns": ["id", "name"],
"column_types": ["int4", ""],
"rows": (0..5).map(|index| serde_json::json!([index, "n"])).collect::<Vec<_>>(),
"affected_rows": 0,
"execution_time_ms": 0
}))
.unwrap();
let shaped = shape_result(DatabaseType::Postgres, result, 3, 7);
assert_eq!(shaped.rows.len(), 3);
assert!(shaped.truncated);
assert_eq!(shaped.columns[0].data_type.as_deref(), Some("int4"));
assert_eq!(shaped.columns[1].data_type, None);
assert_eq!(shaped.db_type, "postgres");
assert_eq!(shaped.elapsed_ms, 7);
}
/// A UI-only plugin package declaring `permissions`, in the `.dbxp` layout
/// the installer verifies (every file covered by `checksums.json`).
fn ui_plugin_package(id: &str, permissions: &[&str]) -> Vec<u8> {
use sha2::{Digest, Sha256};
use std::io::Write;
let manifest = serde_json::to_vec_pretty(&serde_json::json!({
"manifest_version": 1,
"id": id,
"name": "Data sample",
"version": "1.0.0",
"publisher": "sample",
"engines": { "dbx": ">=0.5.0", "host_api": "^1.4" },
"permissions": permissions,
"entrypoints": { "ui": { "root": "ui", "entry": "ui/index.html" } },
"contributions": [{ "type": "workbench", "id": format!("{id}.main"), "label": "Data" }]
}))
.unwrap();
let files = [("manifest.json", manifest), ("ui/index.html", b"<!doctype html>".to_vec())];
let checksums = serde_json::to_vec_pretty(&serde_json::json!({
"algorithm": "sha256",
"files": files.iter().map(|(path, bytes)| (path.to_string(), format!("{:x}", Sha256::digest(bytes)))).collect::<std::collections::BTreeMap<_, _>>()
}))
.unwrap();
let mut output = std::io::Cursor::new(Vec::new());
{
let mut zip = zip::ZipWriter::new(&mut output);
let options = zip::write::SimpleFileOptions::default().unix_permissions(0o644);
for (path, bytes) in files
.iter()
.map(|(path, bytes)| (*path, bytes.as_slice()))
.chain([("checksums.json", checksums.as_slice())])
{
zip.start_file(path, options).unwrap();
zip.write_all(bytes).unwrap();
}
zip.finish().unwrap();
}
output.into_inner()
}
#[tokio::test]
async fn granted_plugins_read_an_open_sqlite_connection_and_nothing_else() {
use crate::plugins::{PluginInstallPolicy, PluginPackageInstaller};
const PLUGIN: &str = "sample.data";
let root = tempfile::tempdir().unwrap();
let plugin_dir = root.path().join("plugins");
let installer = PluginPackageInstaller::new(plugin_dir.clone(), "0.6.0").unwrap();
installer
.install_bytes(&ui_plugin_package(PLUGIN, &["host.data:read"]), PluginInstallPolicy::LocalDevelopment)
.unwrap();
installer
.install_bytes(&ui_plugin_package("sample.nodata", &[]), PluginInstallPolicy::LocalDevelopment)
.unwrap();
let storage = crate::storage::Storage::open(&root.path().join("dbx.db")).await.unwrap();
let state = AppState::new_with_plugin_dir_and_app_version(storage, plugin_dir, "0.6.0");
let path = root.path().join("orders.db");
crate::db::sqlite::connect_path_create_if_missing(path.to_str().unwrap()).await.unwrap();
let config: ConnectionConfig = serde_json::from_value(serde_json::json!({
"id": "orders", "name": "Local orders", "db_type": "sqlite", "host": path.to_str().unwrap(),
"port": 0, "username": "", "password": "", "database": null,
"one_time": false, "save_password": false, "read_only": false
}))
.unwrap();
state.configs.write().await.insert("orders".to_string(), config);
for sql in [
"CREATE TABLE orders(id INTEGER PRIMARY KEY, status TEXT)",
"INSERT INTO orders(status) VALUES ('paid'), ('paid'), ('refunded')",
] {
execute_sql_statement_with_options(&state, "orders", "main", sql, None, None, Default::default())
.await
.unwrap();
}
let request = |sql: &str| PluginDataQueryRequest {
connection_id: "orders".to_string(),
database: Some("main".to_string()),
schema: None,
sql: sql.to_string(),
max_rows: None,
timeout_ms: None,
};
let refused = query_plugin_data(&state, PLUGIN, request("SELECT status FROM orders")).await.unwrap_err();
assert!(refused.starts_with(PLUGIN_DATA_ACCESS_NOT_GRANTED), "{refused}");
assert!(set_plugin_data_grant(&state, "sample.nodata", "orders", true).await.is_err());
assert!(set_plugin_data_grant(&state, PLUGIN, "missing-connection", true).await.is_err());
let grants = set_plugin_data_grant(&state, PLUGIN, "orders", true).await.unwrap();
assert_eq!(
grants,
[PluginDataGrant {
connection_id: "orders".to_string(),
connection_name: Some("Local orders".to_string())
}]
);
let rows = query_plugin_data(
&state,
PLUGIN,
PluginDataQueryRequest { max_rows: Some(2), ..request("SELECT status FROM orders ORDER BY id") },
)
.await
.unwrap();
assert_eq!(rows.columns[0].name, "status");
assert_eq!(rows.rows, [[serde_json::json!("paid")], [serde_json::json!("paid")]]);
assert!(rows.truncated);
// The gate holds on a granted connection: nothing but one read runs.
for rejected in ["DELETE FROM orders", "SELECT 1; DELETE FROM orders", "DROP TABLE orders"] {
assert!(query_plugin_data(&state, PLUGIN, request(rejected)).await.is_err(), "{rejected}");
}
let count = query_plugin_data(&state, PLUGIN, request("SELECT count(*) FROM orders")).await.unwrap();
assert_eq!(count.rows[0][0], serde_json::json!(3));
set_plugin_data_grant(&state, PLUGIN, "orders", false).await.unwrap();
let revoked = query_plugin_data(&state, PLUGIN, request("SELECT 1")).await.unwrap_err();
assert!(revoked.starts_with(PLUGIN_DATA_ACCESS_NOT_GRANTED), "{revoked}");
set_plugin_data_grant(&state, PLUGIN, "orders", true).await.unwrap();
state.remove_connection_pools("orders").await;
assert_eq!(query_plugin_data(&state, PLUGIN, request("SELECT 1")).await.unwrap_err(), "Connection is not open");
}
}
@@ -53,6 +53,15 @@ impl PluginConnectionHandle {
self.session.as_ref().is_none_or(|session| session.status().state == PluginSessionState::Running)
}
/// The lifecycle payload this connection was opened with: provider,
/// hydrated connection, and the runtime endpoint *after* DBX transport
/// layers. Host-initiated calls that act on an open connection (for
/// example plugin MCP tools) must reuse it instead of rebuilding a payload
/// from the saved config, which would bypass the tunnel.
pub fn lifecycle_params(&self) -> &serde_json::Value {
&self.params
}
pub async fn disconnect(&self) -> Result<(), String> {
let Some(session) = &self.session else {
return Ok(());
@@ -13,10 +13,12 @@ pub const SUPPORTED_PLUGIN_MANIFEST_VERSION: u32 = 1;
/// schema metadata (`host.getTableMetadata` behind `host.schema:read`) and the
/// plugin-initiated clipboard Host API (`host.clipboardRead` behind the
/// `host.clipboard:read` permission; clipboard writes reuse the existing
/// ungated `host.copy`). All are additive: older plugins keep working, and a
/// plugin that wants a capability must check the advertised version (or the
/// matching `capabilities` / `host.features` entry) before calling it.
pub const SUPPORTED_PLUGIN_HOST_API_VERSION: &str = "1.3.0";
/// ungated `host.copy`). 1.4 adds consent-gated read-only data queries
/// (`host.queryData` behind `host.data:read`). All are additive: older plugins
/// keep working, and a plugin that wants a capability must check the
/// advertised version (or the matching `capabilities` / `host.features` entry)
/// before calling it.
pub const SUPPORTED_PLUGIN_HOST_API_VERSION: &str = "1.4.0";
/// Capabilities the host advertises to a plugin backend at `plugin/initialize`.
pub const SUPPORTED_PLUGIN_HOST_FEATURES: &[&str] = &["host.requestUserInput"];
pub const SUPPORTED_PLUGIN_PROTOCOL_VERSION: u32 = 1;
@@ -34,6 +36,7 @@ pub const SUPPORTED_PLUGIN_PERMISSIONS: &[&str] = &[
"host.storage",
"host.ai",
"host.clipboard:read",
"host.data:read",
];
/// Cap the number of `host.network:<origin>` entries so a manifest cannot bloat
@@ -2240,27 +2243,27 @@ mod tests {
assert_eq!(declared, SUPPORTED_PLUGIN_PERMISSIONS.iter().map(|value| value.to_string()).collect::<Vec<_>>());
}
/// The reason for the 1.3.0 bump: `engines.host_api` is how a plugin states
/// "I need the schema metadata API" or "I need clipboard reads", so the
/// advertised version has to satisfy `^1.3` while a floor this host cannot
/// meet stays rejected.
/// `engines.host_api` is how a plugin states "I need the schema metadata
/// API" / "I need clipboard reads" (1.3) or "I need data queries" (1.4), so
/// the advertised version has to satisfy each floor while a floor this host
/// cannot meet stays rejected.
#[test]
fn host_api_advertises_the_floor_a_schema_metadata_plugin_declares() {
let advertised = semver::Version::parse(SUPPORTED_PLUGIN_HOST_API_VERSION)
.expect("the advertised Host API version must be semver");
assert!(
semver::VersionReq::parse("^1.3").unwrap().matches(&advertised),
"the host must satisfy the schema metadata API floor it asks plugins to declare"
"the host must satisfy the schema metadata and clipboard-read floor it asks plugins to declare"
);
assert!(
semver::VersionReq::parse("^1.3").unwrap().matches(&advertised),
"the host must satisfy the clipboard-read floor it asks plugins to declare"
semver::VersionReq::parse("^1.4").unwrap().matches(&advertised),
"the host must satisfy the data-query floor it asks plugins to declare"
);
for requirement in ["^1.0", "^1.1", "^1.2", "^1.3", ">=1.1.0, <2.0.0"] {
for requirement in ["^1.0", "^1.1", "^1.2", "^1.3", "^1.4", ">=1.1.0, <2.0.0"] {
assert!(host_api_requirement_errors(requirement).is_empty(), "{requirement} must be satisfiable");
}
for requirement in [">=1.4.0", "^2.0"] {
for requirement in [">=1.5.0", "^2.0"] {
assert!(!host_api_requirement_errors(requirement).is_empty(), "{requirement} must be rejected");
}
}
+5
View File
@@ -55,6 +55,8 @@ const BLOCKED_EXACT: &[&str] = &[
"app-settings/config/decrypt",
// 读取服务器本地 ~/.ssh/config
"ssh/config-hosts",
// 插件数据访问授权写入(插件读取已存连接数据的许可)
"plugin/data/grant",
];
/// 非 GET 请求整族拦截的前缀:插件/JDBC/驱动安装与运行时控制、云同步外发通道。
@@ -76,6 +78,9 @@ const BLOCKED_AI_EXACT: &[&str] = &[
"ai/test-connection",
"ai/models",
"ai/model-effort",
// 允许内置 AI 调用插件工具的开关;预览会拉起插件 sidecar
"ai/plugin-tools/plugins",
"ai/plugin-tools/preview",
];
pub(crate) fn is_demo_blocked(method: &Method, suffix: &str) -> bool {
+9
View File
@@ -733,6 +733,9 @@ async fn serve() {
.route("/query/get-explain-info", post(routes::query::get_explain_info))
.route("/query/plugin-plan-capabilities", post(routes::query::get_plugin_plan_capabilities))
.route("/query/plugin-estimated-plan", post(routes::query::get_plugin_estimated_plan))
.route("/plugin/data/query", post(routes::query::query_plugin_data))
.route("/plugin/data/grants", post(routes::query::get_plugin_data_grants))
.route("/plugin/data/grant", post(routes::query::set_plugin_data_grant))
.route("/query/build-create-user-sql", post(routes::query::build_create_user_sql))
.route("/query/build-table-select-sql", post(routes::query::build_table_select_sql))
.route("/query/build-database-search-sql", post(routes::query::build_database_search_sql))
@@ -1211,6 +1214,12 @@ async fn serve() {
.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/tool-approval", post(routes::ai::ai_resolve_tool_approval))
.route(
"/ai/plugin-tools/plugins",
get(routes::ai::get_ai_plugin_tool_plugins).post(routes::ai::set_ai_plugin_tool_plugin_enabled),
)
.route("/ai/plugin-tools/preview", post(routes::ai::preview_plugin_ai_tools))
.route("/ai/test-connection", post(routes::ai::ai_test_connection))
.route("/ai/models", post(routes::ai::ai_list_models))
.route("/ai/model-effort", post(routes::ai::ai_resolve_model_effort))
+56
View File
@@ -344,6 +344,58 @@ pub async fn ai_cancel_stream(Json(body): Json<AiCancelStreamRequest>) -> Result
Ok(Json(result))
}
// ---------------------------------------------------------------------------
// Plugin tools for the built-in agent
// ---------------------------------------------------------------------------
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AiResolveToolApprovalRequest {
pub session_id: String,
pub approval_id: String,
pub approved: bool,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AiPluginToolPluginRequest {
pub plugin_id: String,
#[serde(default)]
pub enabled: bool,
}
/// Answers a pending plugin tool approval of the agent run `session_id`.
pub async fn ai_resolve_tool_approval(Json(body): Json<AiResolveToolApprovalRequest>) -> Result<Json<bool>, AppError> {
Ok(Json(dbx_core::tool_approval::resolve_tool_approval(&body.session_id, &body.approval_id, body.approved)))
}
pub async fn get_ai_plugin_tool_plugins(State(state): State<Arc<WebState>>) -> Result<Json<Vec<String>>, AppError> {
Ok(Json(state.app.storage.load_ai_plugin_tool_plugin_ids().await.map_err(AppError::from)?))
}
pub async fn set_ai_plugin_tool_plugin_enabled(
State(state): State<Arc<WebState>>,
Json(body): Json<AiPluginToolPluginRequest>,
) -> Result<Json<Vec<String>>, AppError> {
let plugin_ids = state
.app
.storage
.set_ai_plugin_tool_plugin_enabled(&body.plugin_id, body.enabled)
.await
.map_err(AppError::bad_request)?;
Ok(Json(plugin_ids))
}
pub async fn preview_plugin_ai_tools(
State(state): State<Arc<WebState>>,
Json(body): Json<AiPluginToolPluginRequest>,
) -> Result<Json<dbx_core::plugin_tools::PluginToolPreview>, AppError> {
let preview = dbx_core::plugin_tools::preview_plugin_tools(&state.app, None, &body.plugin_id)
.await
.map_err(AppError::bad_request)?;
Ok(Json(preview))
}
// ---------------------------------------------------------------------------
// AI stream (POST returns SSE directly)
// ---------------------------------------------------------------------------
@@ -461,6 +513,10 @@ pub async fn ai_agent_stream(
sql_permissions,
max_agent_turns,
prompt_cache_key: request.prompt_cache_key.clone(),
session_id: Some(session_id.clone()),
// The loop below runs on its own current-thread runtime; plugin
// sidecar calls must stay on the server runtime that owns the sessions.
host_runtime: Some(tokio::runtime::Handle::current()),
};
let sid = session_id.clone();
+5
View File
@@ -359,12 +359,17 @@ pub async fn uninstall_plugin(
}
let root_dir = state.app.plugins.root_dir().to_path_buf();
let app_version = state.app.plugins.app_version().to_string();
let plugin_id = request.plugin_id.clone();
tokio::task::spawn_blocking(move || {
PluginPackageInstaller::new(root_dir, app_version)?.uninstall(&request.plugin_id)
})
.await
.map_err(|error| AppError::internal(error.to_string()))?
.map_err(AppError::bad_request)?;
// A reinstall must ask for AI tool access and data grants again.
if let Err(error) = state.app.storage.forget_plugin_permissions(&plugin_id).await {
log::warn!("Failed to clear permissions of uninstalled plugin {plugin_id}: {error}");
}
list_plugins(State(state)).await
}
+53
View File
@@ -886,6 +886,59 @@ pub async fn get_plugin_estimated_plan(
Ok(Json(result))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct QueryPluginDataRequest {
pub plugin_id: String,
pub request: dbx_core::query::plugin_data::PluginDataQueryRequest,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginDataGrantRequest {
pub plugin_id: String,
#[serde(default)]
pub connection_id: String,
#[serde(default)]
pub granted: bool,
}
/// Plugin Host API: consent-gated read-only data query (`host.data:read`).
pub async fn query_plugin_data(
State(state): State<Arc<WebState>>,
Json(body): Json<QueryPluginDataRequest>,
) -> Result<Json<dbx_core::query::plugin_data::PluginDataQueryResult>, AppError> {
let result = dbx_core::query::plugin_data::query_plugin_data(&state.app, &body.plugin_id, body.request)
.await
.map_err(AppError::from)?;
Ok(Json(result))
}
pub async fn get_plugin_data_grants(
State(state): State<Arc<WebState>>,
Json(body): Json<PluginDataGrantRequest>,
) -> Result<Json<Vec<dbx_core::query::plugin_data::PluginDataGrant>>, AppError> {
let grants = dbx_core::query::plugin_data::list_plugin_data_grants(&state.app, &body.plugin_id)
.await
.map_err(AppError::from)?;
Ok(Json(grants))
}
pub async fn set_plugin_data_grant(
State(state): State<Arc<WebState>>,
Json(body): Json<PluginDataGrantRequest>,
) -> Result<Json<Vec<dbx_core::query::plugin_data::PluginDataGrant>>, AppError> {
let grants = dbx_core::query::plugin_data::set_plugin_data_grant(
&state.app,
&body.plugin_id,
&body.connection_id,
body.granted,
)
.await
.map_err(AppError::bad_request)?;
Ok(Json(grants))
}
pub async fn build_create_user_sql(Json(req): Json<BuildCreateUserSqlRequest>) -> Result<Json<String>, AppError> {
Ok(Json(dbx_core::db_admin_sql::build_create_user_sql(&req.username, &req.password, &req.tablespace)))
}