feat(ai-tools): 新增历史查询与片段探针

- 新增 inspect_saved_queries 与 inspect_sql_snippets 内置工具
- 拆出本地 SQL 资产快照 helper,并补齐执行器与测试覆盖
- 补充工具目录展示、系统提示和执行状态文案
This commit is contained in:
Syngnat
2026-06-08 19:11:23 +08:00
parent 8aa2932878
commit d7879d9ef0
12 changed files with 471 additions and 3 deletions

View File

@@ -876,6 +876,8 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
mcpTools,
toolContextMap: toolContextMapRef.current,
sqlLogs: useStore.getState().sqlLogs,
savedQueries: useStore.getState().savedQueries,
sqlSnippets: useStore.getState().sqlSnippets,
});
const toolResultMsg: AIChatMessage = buildToolResultMessage({
id: genId(),

View File

@@ -6,7 +6,7 @@ import AIBuiltinToolsCatalog from './AIBuiltinToolsCatalog';
import { buildOverlayWorkbenchTheme } from '../../utils/overlayWorkbenchTheme';
describe('AIBuiltinToolsCatalog', () => {
it('renders the AI-context flow, workspace-tab flow, sql log replay flow, and both snapshot tools', () => {
it('renders the workspace flows, snapshot tools, and local saved-sql discovery tools', () => {
const markup = renderToStaticMarkup(
<AIBuiltinToolsCatalog
darkMode={false}
@@ -36,6 +36,10 @@ describe('AIBuiltinToolsCatalog', () => {
expect(markup).toContain('inspect_workspace_tabs');
expect(markup).toContain('回看最近执行记录');
expect(markup).toContain('inspect_recent_sql_logs');
expect(markup).toContain('复用历史 SQL');
expect(markup).toContain('inspect_saved_queries');
expect(markup).toContain('查找模板片段');
expect(markup).toContain('inspect_sql_snippets');
expect(markup).toContain('理解样例数据');
expect(markup).toContain('preview_table_rows');
});

View File

@@ -62,6 +62,16 @@ const BUILTIN_TOOL_FLOWS = [
steps: 'inspect_recent_sql_logs → get_columns / get_indexes / execute_sql',
description: '适合追查刚刚执行失败的 SQL、慢查询耗时或基于真实执行历史继续让 AI 给解释和优化建议。',
},
{
title: '复用历史 SQL',
steps: 'inspect_saved_queries → get_columns / execute_sql',
description: '适合先找本地保存过的查询脚本,再核对字段和只读验证,避免把之前写过的 SQL 重新手打一遍。',
},
{
title: '查找模板片段',
steps: 'inspect_sql_snippets',
description: '适合先找团队已有的 SQL 片段模板、补全前缀和常用骨架,再决定是否继续改写。',
},
{
title: '理解样例数据',
steps: 'get_columns → preview_table_rows',

View File

@@ -469,6 +469,86 @@ describe('aiLocalToolExecutor', () => {
expect(result.content).not.toContain('SELECT * FROM users LIMIT 10');
});
it('returns local saved queries so the model can reuse historical sql scripts', async () => {
const result = await executeLocalAIToolCall({
toolCall: buildToolCall('inspect_saved_queries', {
keyword: '支付',
connectionId: 'conn-1',
}),
connections: [buildConnection()],
mcpTools: [],
toolContextMap: new Map(),
savedQueries: [
{
id: 'saved-1',
name: '支付订单核对',
sql: 'SELECT * FROM orders WHERE status = \'paid\'',
connectionId: 'conn-1',
dbName: 'crm',
createdAt: 2,
},
{
id: 'saved-2',
name: '用户列表',
sql: 'SELECT * FROM users',
connectionId: 'conn-1',
dbName: 'crm',
createdAt: 1,
},
],
runtime: {
getDatabases: vi.fn(),
getTables: vi.fn(),
},
});
expect(result.success).toBe(true);
expect(result.content).toContain('"totalMatched":1');
expect(result.content).toContain('支付订单核对');
expect(result.content).toContain('"connectionName":"主库"');
expect(result.content).toContain('status = \'paid\'');
});
it('returns sql snippets so the model can inspect local query templates', async () => {
const result = await executeLocalAIToolCall({
toolCall: buildToolCall('inspect_sql_snippets', {
keyword: '支付',
}),
connections: [buildConnection()],
mcpTools: [],
toolContextMap: new Map(),
sqlSnippets: [
{
id: 'snippet-1',
prefix: 'sel',
name: 'SELECT 模板',
body: 'SELECT * FROM ${1:table};',
isBuiltin: true,
createdAt: 1,
},
{
id: 'snippet-2',
prefix: 'pay',
name: '支付模板',
description: '支付对账',
body: 'SELECT * FROM pay_orders WHERE created_at >= ${1:start};',
isBuiltin: false,
createdAt: 2,
},
],
runtime: {
getDatabases: vi.fn(),
getTables: vi.fn(),
},
});
expect(result.success).toBe(true);
expect(result.content).toContain('"totalMatched":1');
expect(result.content).toContain('"prefix":"pay"');
expect(result.content).toContain('"customCount":1');
expect(result.content).toContain('pay_orders');
});
it('returns a database overview bundle with per-table column previews in one tool call', async () => {
const result = await executeLocalAIToolCall({
toolCall: buildToolCall('inspect_database_bundle', {

View File

@@ -1,13 +1,26 @@
import { DBGetAllColumns, DBGetDatabases, DBGetTables } from '../../../wailsjs/go/app/App';
import type { SqlLog } from '../../store';
import type { AIChatMessage, AIContextItem, AIMCPToolDescriptor, AIToolCall, SavedConnection, TabData } from '../../types';
import type {
AIChatMessage,
AIContextItem,
AIMCPToolDescriptor,
AIToolCall,
SavedConnection,
SavedQuery,
SqlSnippet,
TabData,
} from '../../types';
import { buildRpcConnectionConfig } from '../../utils/connectionRpcConfig';
import { buildAIReadonlyPreviewSQL } from '../../utils/aiSqlLimit';
import { buildPaginatedSelectSQL, quoteQualifiedIdent } from '../../utils/sql';
import { resolveAITableSchemaToolResult } from '../../utils/aiTableSchemaTool';
import { buildAIContextSnapshot } from './aiContextInsights';
import { buildCurrentConnectionSnapshot } from './aiConnectionInsights';
import {
buildSavedQueriesSnapshot,
buildSqlSnippetsSnapshot,
} from './aiSavedSqlInsights';
import {
buildActiveTabSnapshot,
buildRecentSqlLogsSnapshot,
@@ -44,6 +57,8 @@ export interface ExecuteLocalAIToolCallOptions {
mcpTools: AIMCPToolDescriptor[];
toolContextMap: Map<string, AIToolContextEntry>;
sqlLogs?: SqlLog[];
savedQueries?: SavedQuery[];
sqlSnippets?: SqlSnippet[];
runtime?: Partial<AILocalToolRuntime>;
}
@@ -171,6 +186,8 @@ export async function executeLocalAIToolCall({
mcpTools,
toolContextMap,
sqlLogs = [],
savedQueries = [],
sqlSnippets = [],
runtime,
}: ExecuteLocalAIToolCallOptions): Promise<ExecuteLocalAIToolCallResult> {
const mergedRuntime = { ...buildDefaultRuntime(), ...(runtime || {}) };
@@ -645,6 +662,37 @@ export async function executeLocalAIToolCall({
}
break;
}
case 'inspect_saved_queries': {
try {
content = JSON.stringify(buildSavedQueriesSnapshot({
savedQueries,
connections,
keyword: args.keyword,
connectionId: args.connectionId,
dbName: args.dbName,
limit: args.limit,
includeSql: args.includeSql !== false,
}));
success = true;
} catch (error: any) {
content = `读取已保存查询失败: ${error?.message || error}`;
}
break;
}
case 'inspect_sql_snippets': {
try {
content = JSON.stringify(buildSqlSnippetsSnapshot({
sqlSnippets,
keyword: args.keyword,
limit: args.limit,
includeBody: args.includeBody !== false,
}));
success = true;
} catch (error: any) {
content = `读取 SQL 片段失败: ${error?.message || error}`;
}
break;
}
case 'preview_table_rows': {
const connection = findConnection(connections, args.connectionId);
if (!connection) {

View File

@@ -0,0 +1,95 @@
import { describe, expect, it } from 'vitest';
import type { SavedConnection, SavedQuery, SqlSnippet } from '../../types';
import { buildSavedQueriesSnapshot, buildSqlSnippetsSnapshot } from './aiSavedSqlInsights';
const connections: SavedConnection[] = [
{
id: 'conn-1',
name: '本地开发库',
config: {
type: 'mysql',
host: '127.0.0.1',
port: 3306,
user: 'root',
},
},
];
describe('aiSavedSqlInsights', () => {
it('filters saved queries by keyword and returns sql previews with connection metadata', () => {
const savedQueries: SavedQuery[] = [
{
id: 'query-1',
name: '支付订单查询',
sql: 'SELECT * FROM orders WHERE status = "paid"',
connectionId: 'conn-1',
dbName: 'crm',
createdAt: 2,
},
{
id: 'query-2',
name: '用户清单',
sql: 'SELECT id, name FROM users',
connectionId: 'conn-1',
dbName: 'crm',
createdAt: 1,
},
];
const snapshot = buildSavedQueriesSnapshot({
savedQueries,
connections,
keyword: '支付',
});
expect(snapshot.totalMatched).toBe(1);
expect(snapshot.queries).toHaveLength(1);
expect(snapshot.queries[0]).toMatchObject({
id: 'query-1',
connectionName: '本地开发库',
connectionType: 'mysql',
dbName: 'crm',
});
expect(snapshot.queries[0].sqlPreview).toContain('status = "paid"');
});
it('filters sql snippets by keyword and keeps builtin/custom counts', () => {
const sqlSnippets: SqlSnippet[] = [
{
id: 'snippet-1',
prefix: 'sel',
name: 'SELECT 模板',
description: '快速生成 select',
body: 'SELECT * FROM ${1:table};',
isBuiltin: true,
createdAt: 1,
},
{
id: 'snippet-2',
prefix: 'pay',
name: '支付对账',
description: '支付结果核对模板',
body: 'SELECT * FROM pay_orders WHERE created_at >= ${1:start};',
isBuiltin: false,
createdAt: 2,
},
];
const snapshot = buildSqlSnippetsSnapshot({
sqlSnippets,
keyword: '支付',
});
expect(snapshot.totalMatched).toBe(1);
expect(snapshot.returnedSnippets).toBe(1);
expect(snapshot.builtinCount).toBe(0);
expect(snapshot.customCount).toBe(1);
expect(snapshot.snippets[0]).toMatchObject({
id: 'snippet-2',
prefix: 'pay',
isBuiltin: false,
});
expect(snapshot.snippets[0].bodyPreview).toContain('pay_orders');
});
});

View File

@@ -0,0 +1,151 @@
import type { SavedConnection, SavedQuery, SqlSnippet } from '../../types';
const SAVED_QUERY_SQL_PREVIEW_LIMIT = 4000;
const SQL_SNIPPET_BODY_PREVIEW_LIMIT = 2000;
const normalizeLimit = (input: unknown, fallback: number, max: number): number => {
const value = Math.floor(Number(input) || fallback);
if (value < 1) return 1;
if (value > max) return max;
return value;
};
const normalizeKeyword = (input: unknown): string => String(input || '').trim().toLowerCase();
const matchesKeyword = (keyword: string, fields: Array<string | undefined>): boolean => {
if (!keyword) {
return true;
}
return fields.some((field) => String(field || '').toLowerCase().includes(keyword));
};
export const buildSavedQueriesSnapshot = (params: {
savedQueries?: SavedQuery[];
connections: SavedConnection[];
keyword?: unknown;
connectionId?: unknown;
dbName?: unknown;
limit?: unknown;
includeSql?: unknown;
}) => {
const {
savedQueries = [],
connections,
keyword,
connectionId,
dbName,
limit,
includeSql = true,
} = params;
const safeKeyword = normalizeKeyword(keyword);
const safeConnectionId = String(connectionId || '').trim();
const safeDbName = String(dbName || '').trim();
const safeLimit = normalizeLimit(limit, 12, 50);
const shouldIncludeSql = includeSql !== false;
const filteredQueries = [...savedQueries]
.sort((left, right) => Number(right.createdAt || 0) - Number(left.createdAt || 0))
.filter((query) => {
if (safeConnectionId && query.connectionId !== safeConnectionId) {
return false;
}
if (safeDbName && query.dbName !== safeDbName) {
return false;
}
const connection = connections.find((item) => item.id === query.connectionId);
return matchesKeyword(safeKeyword, [
query.name,
query.sql,
query.dbName,
connection?.name,
connection?.config?.type,
]);
});
const visibleQueries = filteredQueries.slice(0, safeLimit).map((query) => {
const connection = connections.find((item) => item.id === query.connectionId);
const sqlText = String(query.sql || '').trim();
const sqlPreview = shouldIncludeSql ? sqlText.slice(0, SAVED_QUERY_SQL_PREVIEW_LIMIT) : '';
return {
id: query.id,
name: query.name,
connectionId: query.connectionId,
connectionName: connection?.name || '',
connectionType: connection?.config?.type || '',
dbName: query.dbName || '',
createdAt: Number(query.createdAt || 0),
sqlCharCount: sqlText.length,
sqlTruncated: shouldIncludeSql && sqlText.length > sqlPreview.length,
sqlPreview,
};
});
return {
keyword: safeKeyword,
connectionId: safeConnectionId,
dbName: safeDbName,
includeSql: shouldIncludeSql,
limit: safeLimit,
totalMatched: filteredQueries.length,
returnedQueries: visibleQueries.length,
truncated: filteredQueries.length > visibleQueries.length,
queries: visibleQueries,
};
};
export const buildSqlSnippetsSnapshot = (params: {
sqlSnippets?: SqlSnippet[];
keyword?: unknown;
limit?: unknown;
includeBody?: unknown;
}) => {
const {
sqlSnippets = [],
keyword,
limit,
includeBody = true,
} = params;
const safeKeyword = normalizeKeyword(keyword);
const safeLimit = normalizeLimit(limit, 20, 80);
const shouldIncludeBody = includeBody !== false;
const filteredSnippets = [...sqlSnippets]
.sort((left, right) => left.prefix.localeCompare(right.prefix))
.filter((snippet) =>
matchesKeyword(safeKeyword, [
snippet.prefix,
snippet.name,
snippet.description,
snippet.body,
]));
const visibleSnippets = filteredSnippets.slice(0, safeLimit).map((snippet) => {
const bodyText = String(snippet.body || '').trim();
const bodyPreview = shouldIncludeBody ? bodyText.slice(0, SQL_SNIPPET_BODY_PREVIEW_LIMIT) : '';
return {
id: snippet.id,
prefix: snippet.prefix,
name: snippet.name,
description: snippet.description || '',
isBuiltin: snippet.isBuiltin === true,
createdAt: Number(snippet.createdAt || 0),
bodyCharCount: bodyText.length,
bodyTruncated: shouldIncludeBody && bodyText.length > bodyPreview.length,
bodyPreview,
};
});
return {
keyword: safeKeyword,
includeBody: shouldIncludeBody,
limit: safeLimit,
totalMatched: filteredSnippets.length,
returnedSnippets: visibleSnippets.length,
truncated: filteredSnippets.length > visibleSnippets.length,
builtinCount: visibleSnippets.filter((snippet) => snippet.isBuiltin).length,
customCount: visibleSnippets.filter((snippet) => !snippet.isBuiltin).length,
snippets: visibleSnippets,
};
};

View File

@@ -68,7 +68,7 @@ describe('buildAISystemContextMessages', () => {
connections: [connections[0]],
tabs: [],
activeTabId: null,
availableToolNames: ['inspect_workspace_tabs', 'inspect_ai_context', 'inspect_current_connection', 'get_columns'],
availableToolNames: ['inspect_workspace_tabs', 'inspect_ai_context', 'inspect_current_connection', 'inspect_saved_queries', 'inspect_sql_snippets', 'get_columns'],
skills,
userPromptSettings,
});
@@ -77,6 +77,8 @@ describe('buildAISystemContextMessages', () => {
expect(joined).toContain('inspect_workspace_tabs 盘点当前工作区');
expect(joined).toContain('inspect_ai_context 读取当前挂载的表结构上下文');
expect(joined).toContain('inspect_current_connection');
expect(joined).toContain('inspect_saved_queries');
expect(joined).toContain('inspect_sql_snippets');
expect(joined).toContain('当前连接');
expect(joined).toContain('以下是当前用户的自定义补充提示词(全局)');
expect(joined).toContain('以下是当前用户的自定义补充提示词(数据库会话)');

View File

@@ -315,6 +315,18 @@ SELECT * FROM users WHERE status = 1;
content: '如果用户提到“当前连接”“当前数据源”“我现在连的是哪个库/地址”“这个连接走没走 SSH/代理”,优先调用 inspect_current_connection 读取当前活动连接摘要,不要凭界面或记忆猜测。',
});
}
if (availableToolNames.includes('inspect_saved_queries')) {
systemMessages.push({
role: 'system',
content: '如果用户提到“保存过的查询”“历史 SQL”“之前写过的语句”“帮我找以前那条脚本”优先调用 inspect_saved_queries 读取本地已保存查询,再决定是否继续核对字段或复用 SQL。',
});
}
if (availableToolNames.includes('inspect_sql_snippets')) {
systemMessages.push({
role: 'system',
content: '如果用户提到“SQL 片段”“snippet”“模板前缀”“常用模板”优先调用 inspect_sql_snippets 读取本地 SQL 片段库,不要凭记忆编造现有模板。',
});
}
appendCustomPromptGroup(systemMessages, ['database'], userPromptSettings);
appendSkillPromptGroup(systemMessages, ['database'], skills, availableToolNames);

View File

@@ -38,6 +38,8 @@ const TOOL_ACTION_LABELS: Record<string, string> = {
inspect_active_tab: '读取当前活动页签',
inspect_workspace_tabs: '盘点当前工作区页签',
inspect_recent_sql_logs: '回看最近 SQL 执行日志',
inspect_saved_queries: '检索本地已保存查询',
inspect_sql_snippets: '读取 SQL 片段模板',
preview_table_rows: '预览真实样例数据',
execute_sql: '执行只读 SQL 验证',
};

View File

@@ -10,6 +10,16 @@ describe('aiToolRegistry', () => {
expect(info?.tool.function.description).toContain('SSH/代理/HTTP 隧道状态');
});
it('registers the saved-query and sql-snippet inspectors as builtin tools', () => {
const savedQueryTool = BUILTIN_AI_TOOL_INFO.find((item) => item.name === 'inspect_saved_queries');
const snippetTool = BUILTIN_AI_TOOL_INFO.find((item) => item.name === 'inspect_sql_snippets');
expect(savedQueryTool?.desc).toContain('已保存的 SQL 查询');
expect(savedQueryTool?.tool.function.description).toContain('历史查询');
expect(snippetTool?.desc).toContain('SQL 片段模板');
expect(snippetTool?.tool.function.description).toContain('片段模板');
});
it('keeps builtin tools and MCP tools in the unified runtime tool chain', () => {
const tools = buildAvailableAIChatTools([{
alias: 'custom_probe',
@@ -27,6 +37,8 @@ describe('aiToolRegistry', () => {
}]);
expect(tools.some((item) => item.function.name === 'inspect_current_connection')).toBe(true);
expect(tools.some((item) => item.function.name === 'inspect_saved_queries')).toBe(true);
expect(tools.some((item) => item.function.name === 'inspect_sql_snippets')).toBe(true);
expect(tools.some((item) => item.function.name === 'custom_probe')).toBe(true);
});
});

View File

@@ -421,6 +421,56 @@ export const BUILTIN_AI_TOOL_INFO: AIBuiltinToolInfo[] = [
},
},
},
{
name: "inspect_saved_queries",
icon: "💾",
desc: "查看本地已保存的 SQL 查询",
detail:
"可按关键词、连接或数据库过滤,返回保存查询的名称、所属连接、数据库和 SQL 预览。适合用户提到“我之前保存过的查询”“帮我找那条历史 SQL”时先从真实本地收藏里检索。",
params: "keyword?, connectionId?, dbName?, limit?, includeSql?(默认 true)",
tool: {
type: "function",
function: {
name: "inspect_saved_queries",
description:
"读取本地已保存的 SQL 查询列表,可按关键词、连接和数据库过滤,并返回每条查询的名称、所属连接、数据库与 SQL 预览。适用于用户想找历史查询、复用旧 SQL、核对保存脚本时先读取真实本地记录。",
parameters: {
type: "object",
properties: {
keyword: { type: "string", description: "可选按查询名称、SQL 文本、连接名或数据库名做关键词筛选" },
connectionId: { type: "string", description: "可选,只看某个连接下保存的查询" },
dbName: { type: "string", description: "可选,只看某个数据库下保存的查询" },
limit: { type: "number", description: "可选,最多返回多少条,默认 12最大 50" },
includeSql: { type: "boolean", description: "可选,是否附带 SQL 预览,默认 true" },
},
},
},
},
},
{
name: "inspect_sql_snippets",
icon: "🧩",
desc: "查看 SQL 片段模板",
detail:
"返回本地 SQL 片段的 prefix、名称、说明和模板预览可按关键词过滤。适合用户想找现成模板、补全片段、团队约定 SQL 模板时先读取真实片段库。",
params: "keyword?, limit?, includeBody?(默认 true)",
tool: {
type: "function",
function: {
name: "inspect_sql_snippets",
description:
"读取本地 SQL 片段模板列表,可按关键词过滤,并返回 prefix、名称、说明和模板预览。适用于用户想找 snippet、复用模板、核对 SQL 片段配置时,先读取真实本地片段库。",
parameters: {
type: "object",
properties: {
keyword: { type: "string", description: "可选,按 prefix、名称、描述或模板内容做关键词筛选" },
limit: { type: "number", description: "可选,最多返回多少条,默认 20最大 80" },
includeBody: { type: "boolean", description: "可选,是否附带模板内容预览,默认 true" },
},
},
},
},
},
{
name: "execute_sql",
icon: "▶️",