From 28fedbd815ebb082302f01875de3776a36707b5a Mon Sep 17 00:00:00 2001 From: Leo <132990250+jxy0811@users.noreply.github.com> Date: Tue, 16 Jun 2026 11:30:02 +0800 Subject: [PATCH] fix(ai): recover from analysis tool timeouts (#229) Co-authored-by: Leo --- apps/desktop/main/ai/agent-stream-runner.ts | 9 +- apps/desktop/main/ai/tools/debug-executor.ts | 67 +++++++++ .../main/ai/tools/shared-tool-adapter.ts | 3 +- apps/desktop/main/ai/tools/types.ts | 1 + .../main/ai/tools/worker-data-provider.ts | 101 ++++++++------ apps/desktop/main/internal-api.ts | 2 + apps/desktop/main/worker/workerManager.ts | 105 +++++++++++--- .../main/worker/workerTimeoutPolicy.test.ts | 27 ++++ .../main/worker/workerTimeoutPolicy.ts | 48 +++++++ packages/core/src/schema/tables.ts | 5 +- packages/http-routes/src/context.ts | 20 +++ packages/http-routes/src/index.ts | 2 +- .../src/routes/web/ai-tools.test.ts | 132 ++++++++++++++++++ .../http-routes/src/routes/web/ai-tools.ts | 15 +- .../node-runtime/src/database-manager.test.ts | 75 +++++++++- .../src/migrations/chat-db-migrations.ts | 13 ++ .../src/definitions/analysis-tools.test.ts | 55 ++++++++ packages/tools/src/sql/definitions.ts | 2 +- packages/tools/src/types.ts | 1 + 19 files changed, 612 insertions(+), 71 deletions(-) create mode 100644 apps/desktop/main/ai/tools/debug-executor.ts create mode 100644 apps/desktop/main/worker/workerTimeoutPolicy.test.ts create mode 100644 apps/desktop/main/worker/workerTimeoutPolicy.ts create mode 100644 packages/http-routes/src/routes/web/ai-tools.test.ts diff --git a/apps/desktop/main/ai/agent-stream-runner.ts b/apps/desktop/main/ai/agent-stream-runner.ts index 2c4202d1..91489472 100644 --- a/apps/desktop/main/ai/agent-stream-runner.ts +++ b/apps/desktop/main/ai/agent-stream-runner.ts @@ -125,7 +125,13 @@ export function createElectronRunAgentStream(): ( round: 0, toolsUsed: 0, contextTokens: 0, - totalUsage: { promptTokens: 0, completionTokens: 0, totalTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0 }, + totalUsage: { + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }, updatedAt: Date.now(), } satisfies AgentRuntimeStatus, }) @@ -237,6 +243,7 @@ export function createElectronRunAgentStream(): ( preprocessConfig: params.preprocessConfig as ToolContext['preprocessConfig'], maxToolResultTokens, dataSnapshot, + abortSignal, } const agent = new Agent( diff --git a/apps/desktop/main/ai/tools/debug-executor.ts b/apps/desktop/main/ai/tools/debug-executor.ts new file mode 100644 index 00000000..a33a2db4 --- /dev/null +++ b/apps/desktop/main/ai/tools/debug-executor.ts @@ -0,0 +1,67 @@ +import { stripAvatarFields } from '@openchatlab/core' +import type { AiToolExecuteRequest, AiToolExecuteResult } from '@openchatlab/http-routes' +import { batchSegmentWithFrequency } from '@openchatlab/node-runtime' +import { AGENT_TOOL_REGISTRY } from '@openchatlab/tools' +import type { ToolExecutionContext } from '@openchatlab/tools' +import { t as i18nT } from '../../i18n' +import { WorkerDataProvider } from './worker-data-provider' + +const MAX_RESULT_CHARS = 500_000 + +function assertNotAborted(signal: AbortSignal): void { + if (signal.aborted) { + throw new Error('cancelled') + } +} + +export async function executeElectronAiTool(params: AiToolExecuteRequest): Promise { + const { toolName, sessionId, abortSignal } = params + const entry = AGENT_TOOL_REGISTRY.find((tool) => tool.name === toolName) + if (!entry) { + return { success: false, error: `Tool not found: ${toolName}` } + } + + try { + assertNotAborted(abortSignal) + const execCtx: ToolExecutionContext = { + sessionId, + abortSignal, + dataProvider: new WorkerDataProvider(sessionId, abortSignal), + segmentText: (texts, locale, options) => batchSegmentWithFrequency(texts, locale as any, options as any), + translateTemplate: (key: string) => { + const translated = i18nT(key) + return translated !== key ? translated : undefined + }, + } + + const startTime = Date.now() + const result = await entry.handler(params.params, execCtx) + const elapsed = Date.now() - startTime + assertNotAborted(abortSignal) + + let details = (result.data as Record | undefined) ?? undefined + let truncated = false + + if (details) { + stripAvatarFields(details) + const raw = JSON.stringify(details) + if (raw.length > MAX_RESULT_CHARS) { + truncated = true + details = { _truncated: true, _originalSize: raw.length, _preview: raw.slice(0, MAX_RESULT_CHARS) } + } + } + + return { + success: true, + elapsed, + content: [{ type: 'text', text: result.content }], + details, + truncated, + } + } catch (error) { + if (abortSignal.aborted) { + return { success: false, error: 'cancelled' } + } + return { success: false, error: error instanceof Error ? error.message : String(error) } + } +} diff --git a/apps/desktop/main/ai/tools/shared-tool-adapter.ts b/apps/desktop/main/ai/tools/shared-tool-adapter.ts index e70cbfd9..28caf4cd 100644 --- a/apps/desktop/main/ai/tools/shared-tool-adapter.ts +++ b/apps/desktop/main/ai/tools/shared-tool-adapter.ts @@ -19,10 +19,11 @@ interface AdaptOptions { function buildExecutionContext(ctx: ToolContext): ToolExecutionContext { return { - dataProvider: new WorkerDataProvider(ctx.sessionId), + dataProvider: new WorkerDataProvider(ctx.sessionId, ctx.abortSignal), sessionId: ctx.sessionId, locale: ctx.locale, timeFilter: ctx.timeFilter, + abortSignal: ctx.abortSignal, searchContextBefore: ctx.searchContextBefore, searchContextAfter: ctx.searchContextAfter, maxMessagesLimit: ctx.maxMessagesLimit, diff --git a/apps/desktop/main/ai/tools/types.ts b/apps/desktop/main/ai/tools/types.ts index 847fdf04..339b9941 100644 --- a/apps/desktop/main/ai/tools/types.ts +++ b/apps/desktop/main/ai/tools/types.ts @@ -62,6 +62,7 @@ export interface ToolContext { locale?: string /** 聊天记录预处理配置（全局） */ preprocessConfig?: PreprocessConfig + abortSignal?: AbortSignal /** 搜索结果上下文：向前取多少条（默认 3） */ searchContextBefore?: number /** 搜索结果上下文：向后取多少条（默认 3） */ diff --git a/apps/desktop/main/ai/tools/worker-data-provider.ts b/apps/desktop/main/ai/tools/worker-data-provider.ts index 68487ca0..d3409bc6 100644 --- a/apps/desktop/main/ai/tools/worker-data-provider.ts +++ b/apps/desktop/main/ai/tools/worker-data-provider.ts @@ -33,19 +33,37 @@ function mapSearchMessages(messages: workerManager.SearchMessageResult[]): RawMe } export class WorkerDataProvider implements ToolDataProvider { - constructor(private sessionId: string) {} + constructor( + private sessionId: string, + private abortSignal?: AbortSignal + ) {} + + private throwIfAborted(): void { + if (this.abortSignal?.aborted) { + throw new Error('cancelled') + } + } + + private async run(operation: () => Promise): Promise { + this.throwIfAborted() + const result = await operation() + this.throwIfAborted() + return result + } async searchMessages( keywords: string[], options?: { timeFilter?: TimeFilter; limit?: number; senderId?: number } ): Promise { - const result = await workerManager.searchMessages( - this.sessionId, - keywords, - options?.timeFilter, - options?.limit ?? 50, - 0, - options?.senderId + const result = await this.run(() => + workerManager.searchMessages( + this.sessionId, + keywords, + options?.timeFilter, + options?.limit ?? 50, + 0, + options?.senderId + ) ) return { messages: mapSearchMessages(result.messages), total: result.total } } @@ -54,13 +72,15 @@ export class WorkerDataProvider implements ToolDataProvider { keywords: string[], options?: { timeFilter?: TimeFilter; limit?: number; senderId?: number } ): Promise { - const result = await workerManager.deepSearchMessages( - this.sessionId, - keywords, - options?.timeFilter, - options?.limit ?? 50, - 0, - options?.senderId + const result = await this.run(() => + workerManager.deepSearchMessages( + this.sessionId, + keywords, + options?.timeFilter, + options?.limit ?? 50, + 0, + options?.senderId + ) ) return { messages: mapSearchMessages(result.messages), total: result.total } } @@ -70,31 +90,30 @@ export class WorkerDataProvider implements ToolDataProvider { contextBefore: number, contextAfter: number ): Promise { - const messages = await workerManager.getSearchMessageContext( - this.sessionId, - messageIds, - contextBefore, - contextAfter + const messages = await this.run(() => + workerManager.getSearchMessageContext(this.sessionId, messageIds, contextBefore, contextAfter) ) return mapSearchMessages(messages) } async getRecentMessages(options?: { timeFilter?: TimeFilter; limit?: number }): Promise { - const result = await workerManager.getRecentMessages(this.sessionId, options?.timeFilter, options?.limit ?? 50) + const result = await this.run(() => + workerManager.getRecentMessages(this.sessionId, options?.timeFilter, options?.limit ?? 50) + ) return { messages: mapSearchMessages(result.messages), total: result.total } } async getMessageContext(messageIds: number[], contextSize: number): Promise { - const messages = await workerManager.getMessageContext(this.sessionId, messageIds, contextSize) + const messages = await this.run(() => workerManager.getMessageContext(this.sessionId, messageIds, contextSize)) return mapSearchMessages(messages) } async getChatOverview(topN?: number): Promise { - return workerManager.getChatOverview(this.sessionId, topN) + return this.run(() => workerManager.getChatOverview(this.sessionId, topN)) } async getMembers(): Promise { - const members = await workerManager.getMembers(this.sessionId) + const members = await this.run(() => workerManager.getMembers(this.sessionId)) return members.map((m) => ({ id: m.id, platformId: m.platformId, @@ -107,7 +126,7 @@ export class WorkerDataProvider implements ToolDataProvider { async getMemberStats(options?: { timeFilter?: TimeFilter; top?: number }): Promise { const top = options?.top ?? 20 - const members = await workerManager.getMemberActivity(this.sessionId, options?.timeFilter) + const members = await this.run(() => workerManager.getMemberActivity(this.sessionId, options?.timeFilter)) return members.slice(0, top).map((m: any) => ({ name: m.name, messageCount: m.messageCount, @@ -116,19 +135,19 @@ export class WorkerDataProvider implements ToolDataProvider { } async getMemberNameHistory(memberId: number): Promise { - return workerManager.getMemberNameHistory(this.sessionId, memberId) + return this.run(() => workerManager.getMemberNameHistory(this.sessionId, memberId)) } async getTimeStats(type: 'hourly' | 'weekday' | 'daily', options?: { timeFilter?: TimeFilter }): Promise { const filter = options?.timeFilter switch (type) { case 'weekday': - return workerManager.getWeekdayActivity(this.sessionId, filter) + return this.run(() => workerManager.getWeekdayActivity(this.sessionId, filter)) case 'daily': - return workerManager.getDailyActivity(this.sessionId, filter) + return this.run(() => workerManager.getDailyActivity(this.sessionId, filter)) case 'hourly': default: - return workerManager.getHourlyActivity(this.sessionId, filter) + return this.run(() => workerManager.getHourlyActivity(this.sessionId, filter)) } } @@ -138,18 +157,20 @@ export class WorkerDataProvider implements ToolDataProvider { limit?: number, previewCount?: number ): Promise { - return workerManager.searchSegments(this.sessionId, keywords, timeFilter, limit, previewCount) + return this.run(() => workerManager.searchSegments(this.sessionId, keywords, timeFilter, limit, previewCount)) } async getSegmentMessages(segmentId: number, limit?: number): Promise { - return workerManager.getSegmentMessages(this.sessionId, segmentId, limit) + return this.run(() => workerManager.getSegmentMessages(this.sessionId, segmentId, limit)) } async getSegmentSummaries(options?: { limit?: number; timeFilter?: TimeFilter }): Promise { - return workerManager.getSegmentSummaries(this.sessionId, { - limit: options?.limit, - timeFilter: options?.timeFilter, - }) + return this.run(() => + workerManager.getSegmentSummaries(this.sessionId, { + limit: options?.limit, + timeFilter: options?.timeFilter, + }) + ) } async getConversationBetween( @@ -158,7 +179,9 @@ export class WorkerDataProvider implements ToolDataProvider { timeFilter?: TimeFilter, limit?: number ): Promise { - const result = await workerManager.getConversationBetween(this.sessionId, memberId1, memberId2, timeFilter, limit) + const result = await this.run(() => + workerManager.getConversationBetween(this.sessionId, memberId1, memberId2, timeFilter, limit) + ) return { messages: mapSearchMessages(result.messages), total: result.total, @@ -168,18 +191,18 @@ export class WorkerDataProvider implements ToolDataProvider { } async executeSql(sql: string): Promise { - return workerManager.executeRawSQL(this.sessionId, sql) + return this.run(() => workerManager.executeRawSQL(this.sessionId, sql)) } async executeParameterizedSql>( query: string, params: Record ): Promise { - return workerManager.pluginQuery(this.sessionId, query, params) + return this.run(() => workerManager.pluginQuery(this.sessionId, query, params)) } async getSchema(): Promise { - const tables = await workerManager.getSchema(this.sessionId) + const tables = await this.run(() => workerManager.getSchema(this.sessionId)) return tables.map((t) => ({ name: t.name, sql: t.columns.map((c) => `${c.name} ${c.type}${c.pk ? ' PK' : ''}${c.notnull ? ' NOT NULL' : ''}`).join(', '), diff --git a/apps/desktop/main/internal-api.ts b/apps/desktop/main/internal-api.ts index 8ff7fd7b..47be414e 100644 --- a/apps/desktop/main/internal-api.ts +++ b/apps/desktop/main/internal-api.ts @@ -41,6 +41,7 @@ import { getManager as getAIChatManager } from './ai/chats' import { getManager as getAssistantManager } from './ai/assistant/manager' import { getManager as getSkillManager } from './ai/skills/manager' import { createElectronRunAgentStream } from './ai/agent-stream-runner' +import { executeElectronAiTool } from './ai/tools/debug-executor' import { assertDesktopDataDirCompatible, getDesktopAppVersion } from './runtime-compat' export interface InternalEndpoint { @@ -153,6 +154,7 @@ export async function startInternalServer(pathProvider: PathProvider): Promise void reject: (error: Error) => void + timeout: ReturnType + restartOnTimeout: boolean onProgress?: (progress: ParseProgress) => void // 进度回调 } >() @@ -32,6 +40,49 @@ const pendingRequests = new Map< // 请求 ID 计数器 let requestIdCounter = 0 +function hasProgressRequest(): boolean { + for (const pending of pendingRequests.values()) { + if (pending.onProgress) return true + } + return false +} + +function rejectAllPending(error: Error): void { + for (const [id, pending] of pendingRequests.entries()) { + clearTimeout(pending.timeout) + pending.reject(error) + pendingRequests.delete(id) + } +} + +function hasNonRestartableRequest(): boolean { + for (const pending of pendingRequests.values()) { + if (!pending.restartOnTimeout) return true + } + return false +} + +function terminateWorkerAfterQueryTimeout(type: string, timedOutWorker: Worker): void { + if (!worker || worker !== timedOutWorker) return + + if (hasProgressRequest()) { + console.warn(`[WorkerManager] Query timed out while a progress task is active; keeping worker alive: ${type}`) + return + } + + if (hasNonRestartableRequest()) { + console.warn(`[WorkerManager] Query timed out while non-restartable work is active; keeping worker alive: ${type}`) + return + } + + console.warn(`[WorkerManager] Restarting worker after request timeout: ${type}`) + worker = null + rejectAllPending(new Error(`Worker restarted after request timeout: ${type}`)) + timedOutWorker.terminate().catch((error) => { + console.error('[WorkerManager] Failed to terminate timed-out worker:', error) + }) +} + /** * 获取数据库目录 */ @@ -75,7 +126,7 @@ export function initWorker(): void { console.log('[WorkerManager] Initializing worker at:', workerPath) try { - worker = new Worker(workerPath, { + const initializedWorker = new Worker(workerPath, { workerData: { dbDir: getDbDir(), cacheDir: getCacheDir(), @@ -84,9 +135,10 @@ export function initWorker(): void { appVersion: getDesktopAppVersion(app.getVersion()), }, }) + worker = initializedWorker // 监听 Worker 消息 - worker.on('message', (message) => { + initializedWorker.on('message', (message) => { const { id, type, success, result, error, payload } = message const pending = pendingRequests.get(id) @@ -102,6 +154,7 @@ export function initWorker(): void { // 处理完成或错误消息 pendingRequests.delete(id) + clearTimeout(pending.timeout) if (success) { pending.resolve(result) @@ -111,20 +164,21 @@ export function initWorker(): void { }) // 监听 Worker 错误 - worker.on('error', (error) => { + initializedWorker.on('error', (error) => { console.error('[WorkerManager] Worker error:', error) }) // 监听 Worker 退出 - worker.on('exit', (code) => { + initializedWorker.on('exit', (code) => { console.log('[WorkerManager] Worker exited with code:', code) + if (worker !== initializedWorker) { + console.log('[WorkerManager] Ignoring exit from replaced worker') + return + } worker = null // 拒绝所有等待中的请求 - for (const [id, pending] of pendingRequests.entries()) { - pending.reject(new Error('Worker exited unexpectedly')) - pendingRequests.delete(id) - } + rejectAllPending(new Error('Worker exited unexpectedly')) }) console.log('[WorkerManager] Worker initialized successfully') @@ -137,7 +191,7 @@ export function initWorker(): void { /** * 发送消息到 Worker 并等待响应 */ -function sendToWorker(type: string, payload: any, timeoutMs: number = 60000): Promise { +function sendToWorker(type: string, payload: any, options: number | WorkerRequestOptions = {}): Promise { return new Promise((resolve, reject) => { if (!worker) { try { @@ -148,19 +202,27 @@ function sendToWorker(type: string, payload: any, timeoutMs: number = 60000): } } + const timeoutMs = typeof options === 'number' ? options : (options.timeoutMs ?? 60000) + const restartOnTimeout = + typeof options === 'number' + ? isRestartableReadOnlyRequestType(type) + : (options.restartOnTimeout ?? isRestartableReadOnlyRequestType(type)) + const requestWorker = worker! const id = `req_${++requestIdCounter}` - pendingRequests.set(id, { resolve, reject }) - - worker!.postMessage({ id, type, payload }) - - // 设置超时 - setTimeout(() => { + const timeout = setTimeout(() => { if (pendingRequests.has(id)) { pendingRequests.delete(id) reject(new Error(`Worker request timeout: ${type}`)) + if (restartOnTimeout) { + terminateWorkerAfterQueryTimeout(type, requestWorker) + } } }, timeoutMs) + + pendingRequests.set(id, { resolve, reject, timeout, restartOnTimeout }) + + requestWorker.postMessage({ id, type, payload }) }) } @@ -186,17 +248,16 @@ function sendToWorkerWithProgress( const id = `req_${++requestIdCounter}` - pendingRequests.set(id, { resolve, reject, onProgress }) - - worker!.postMessage({ id, type, payload }) - - // 设置超时 - setTimeout(() => { + const timeout = setTimeout(() => { if (pendingRequests.has(id)) { pendingRequests.delete(id) reject(new Error(`Worker request timeout: ${type}`)) } }, timeoutMs) + + pendingRequests.set(id, { resolve, reject, timeout, restartOnTimeout: false, onProgress }) + + worker!.postMessage({ id, type, payload }) }) } @@ -210,6 +271,7 @@ export function closeWorker(): void { worker.terminate() worker = null + rejectAllPending(new Error('Worker terminated')) console.log('[WorkerManager] Worker terminated') } } @@ -230,6 +292,7 @@ export async function closeWorkerAsync(): Promise { worker.terminate() worker = null + rejectAllPending(new Error('Worker terminated')) console.log('[WorkerManager] Worker terminated (async)') } } diff --git a/apps/desktop/main/worker/workerTimeoutPolicy.test.ts b/apps/desktop/main/worker/workerTimeoutPolicy.test.ts new file mode 100644 index 00000000..c9aea90c --- /dev/null +++ b/apps/desktop/main/worker/workerTimeoutPolicy.test.ts @@ -0,0 +1,27 @@ +import { describe, it } from 'node:test' +import assert from 'node:assert/strict' +import { isRestartableReadOnlyRequestType } from './workerTimeoutPolicy' + +describe('worker timeout policy', () => { + it('allows timed-out read-only AI/debug requests to restart the worker', () => { + assert.equal(isRestartableReadOnlyRequestType('pluginQuery'), true) + assert.equal(isRestartableReadOnlyRequestType('executeRawSQL'), true) + assert.equal(isRestartableReadOnlyRequestType('getSchema'), true) + assert.equal(isRestartableReadOnlyRequestType('getChatOverview'), true) + assert.equal(isRestartableReadOnlyRequestType('searchMessages'), true) + assert.equal(isRestartableReadOnlyRequestType('getSegmentSummaries'), true) + }) + + it('keeps mutating, indexing, import, export, and unknown requests non-restartable', () => { + assert.equal(isRestartableReadOnlyRequestType('generateSessions'), false) + assert.equal(isRestartableReadOnlyRequestType('generateIncrementalSessions'), false) + assert.equal(isRestartableReadOnlyRequestType('clearSessions'), false) + assert.equal(isRestartableReadOnlyRequestType('updateMemberAliases'), false) + assert.equal(isRestartableReadOnlyRequestType('mergeMembers'), false) + assert.equal(isRestartableReadOnlyRequestType('deleteMember'), false) + assert.equal(isRestartableReadOnlyRequestType('streamImport'), false) + assert.equal(isRestartableReadOnlyRequestType('incrementalImport'), false) + assert.equal(isRestartableReadOnlyRequestType('exportFilterResultToFile'), false) + assert.equal(isRestartableReadOnlyRequestType('unknownFutureWorkerRequest'), false) + }) +}) diff --git a/apps/desktop/main/worker/workerTimeoutPolicy.ts b/apps/desktop/main/worker/workerTimeoutPolicy.ts new file mode 100644 index 00000000..0f7a323f --- /dev/null +++ b/apps/desktop/main/worker/workerTimeoutPolicy.ts @@ -0,0 +1,48 @@ +const RESTARTABLE_READ_ONLY_REQUEST_TYPES = new Set([ + 'pluginQuery', + 'executeRawSQL', + 'getSchema', + 'getAvailableYears', + 'getMemberActivity', + 'getHourlyActivity', + 'getDailyActivity', + 'getWeekdayActivity', + 'getMonthlyActivity', + 'getYearlyActivity', + 'getMessageLengthDistribution', + 'getMessageTypeDistribution', + 'getTimeRange', + 'getMemberNameHistory', + 'getCatchphraseAnalysis', + 'getLanguagePreferenceAnalysis', + 'getMentionAnalysis', + 'getMentionGraph', + 'getLaughAnalysis', + 'getClusterGraph', + 'getRelationshipStats', + 'getAllSessions', + 'getSession', + 'getChatOverview', + 'getMembers', + 'getMembersPaginated', + 'searchMessages', + 'deepSearchMessages', + 'getMessageContext', + 'getSearchMessageContext', + 'getRecentMessages', + 'getAllRecentMessages', + 'getConversationBetween', + 'hasSessionIndex', + 'getSessionStats', + 'getAllIndexStats', + 'getSessions', + 'getSessionsByTimeRange', + 'getRecentChatSessions', + 'searchSegments', + 'getSegmentMessages', + 'getSegmentSummaries', +]) + +export function isRestartableReadOnlyRequestType(type: string): boolean { + return RESTARTABLE_READ_ONLY_REQUEST_TYPES.has(type) +} diff --git a/packages/core/src/schema/tables.ts b/packages/core/src/schema/tables.ts index e202afa3..50353e65 100644 --- a/packages/core/src/schema/tables.ts +++ b/packages/core/src/schema/tables.ts @@ -8,7 +8,7 @@ */ /** 当前 Schema 版本（最新迁移的版本号） */ -export const CURRENT_SCHEMA_VERSION = 7 +export const CURRENT_SCHEMA_VERSION = 8 /** * Table DDL only (no indexes). Used by bulk-import workflows that defer @@ -82,6 +82,9 @@ export const CHAT_DB_TABLES = ` export const CHAT_DB_INDEXES = ` CREATE INDEX IF NOT EXISTS idx_message_ts ON message(ts); CREATE INDEX IF NOT EXISTS idx_message_sender ON message(sender_id); + CREATE INDEX IF NOT EXISTS idx_message_sender_ts ON message(sender_id, ts); + CREATE INDEX IF NOT EXISTS idx_message_type_ts ON message(type, ts); + CREATE INDEX IF NOT EXISTS idx_message_reply_to ON message(reply_to_message_id); CREATE INDEX IF NOT EXISTS idx_message_platform_id ON message(platform_message_id); CREATE INDEX IF NOT EXISTS idx_member_name_history_member_id ON member_name_history(member_id); CREATE INDEX IF NOT EXISTS idx_segment_time ON segment(start_ts, end_ts); diff --git a/packages/http-routes/src/context.ts b/packages/http-routes/src/context.ts index f03d8f8e..25f6ad88 100644 --- a/packages/http-routes/src/context.ts +++ b/packages/http-routes/src/context.ts @@ -69,6 +69,26 @@ export interface HttpRouteContext { onEvent: (chunk: AgentStreamChunk) => void, abortSignal: AbortSignal ) => Promise + + /** AI tool debug execution - platform-specific so Electron can keep DB work in its worker. */ + executeAiTool?: (params: AiToolExecuteRequest) => Promise +} + +export interface AiToolExecuteRequest { + testId: string + toolName: string + params: Record + sessionId: string + abortSignal: AbortSignal +} + +export interface AiToolExecuteResult { + success: boolean + elapsed?: number + content?: Array<{ type: 'text'; text: string }> + details?: unknown + truncated?: boolean + error?: string } export interface AgentStreamRequest { diff --git a/packages/http-routes/src/index.ts b/packages/http-routes/src/index.ts index 4bbdc505..f88b22f1 100644 --- a/packages/http-routes/src/index.ts +++ b/packages/http-routes/src/index.ts @@ -1,4 +1,4 @@ -export type { HttpRouteContext, AgentStreamRequest } from './context' +export type { HttpRouteContext, AgentStreamRequest, AiToolExecuteRequest, AiToolExecuteResult } from './context' export { registerSharedRoutes } from './register' export type { SharedRouteOptions } from './register' export { setAuthToken, setRequireAuth, authHook } from './auth' diff --git a/packages/http-routes/src/routes/web/ai-tools.test.ts b/packages/http-routes/src/routes/web/ai-tools.test.ts new file mode 100644 index 00000000..937e78d6 --- /dev/null +++ b/packages/http-routes/src/routes/web/ai-tools.test.ts @@ -0,0 +1,132 @@ +import { describe, it, afterEach } from 'node:test' +import assert from 'node:assert/strict' +import Fastify, { type FastifyInstance } from 'fastify' +import type { HttpRouteContext } from '../../context' +import { registerAiToolRoutes } from './ai-tools' + +function createContext(overrides: Partial = {}): HttpRouteContext { + return { + dbManager: { + open: () => null, + }, + ...overrides, + } as unknown as HttpRouteContext +} + +describe('ai tool debug routes', () => { + let app: FastifyInstance | null = null + + afterEach(async () => { + if (app) { + await app.close() + app = null + } + }) + + it('delegates tool execution to the platform hook when provided', async () => { + const calls: Array<{ toolName: string; sessionId: string; params: Record }> = [] + app = Fastify() + registerAiToolRoutes( + app, + createContext({ + executeAiTool: async ({ toolName, sessionId, params }) => { + calls.push({ toolName, sessionId, params }) + return { + success: true, + elapsed: 12, + content: [{ type: 'text', text: 'ok' }], + details: { delegated: true }, + truncated: false, + } + }, + }) + ) + await app.ready() + + const resp = await app.inject({ + method: 'POST', + url: '/_web/ai/tools/execute', + payload: { + testId: 'test-1', + toolName: 'get_chat_overview', + sessionId: 'chat-1', + params: { top_n: 3 }, + }, + }) + + assert.equal(resp.statusCode, 200) + assert.deepEqual(resp.json(), { + success: true, + elapsed: 12, + content: [{ type: 'text', text: 'ok' }], + details: { delegated: true }, + truncated: false, + }) + assert.deepEqual(calls, [{ toolName: 'get_chat_overview', sessionId: 'chat-1', params: { top_n: 3 } }]) + }) + + it('aborts an active platform tool execution when cancelled', async () => { + let startedResolve: (() => void) | undefined + const started = new Promise((resolve) => { + startedResolve = resolve + }) + + app = Fastify() + registerAiToolRoutes( + app, + createContext({ + executeAiTool: async ({ abortSignal }) => { + startedResolve?.() + return new Promise((resolve) => { + abortSignal.addEventListener('abort', () => resolve({ success: false, error: 'cancelled' }), { once: true }) + }) + }, + }) + ) + await app.ready() + + const executePromise = app.inject({ + method: 'POST', + url: '/_web/ai/tools/execute', + payload: { + testId: 'test-cancel', + toolName: 'get_chat_overview', + sessionId: 'chat-1', + params: {}, + }, + }) + await started + + const cancelResp = await app.inject({ + method: 'POST', + url: '/_web/ai/tools/cancel', + payload: { testId: 'test-cancel' }, + }) + const executeResp = await executePromise + + assert.equal(cancelResp.statusCode, 200) + assert.deepEqual(cancelResp.json(), { success: true }) + assert.equal(executeResp.statusCode, 200) + assert.deepEqual(executeResp.json(), { success: false, error: 'cancelled' }) + }) + + it('keeps the CoreDataProvider fallback when no platform hook is provided', async () => { + app = Fastify() + registerAiToolRoutes(app, createContext()) + await app.ready() + + const resp = await app.inject({ + method: 'POST', + url: '/_web/ai/tools/execute', + payload: { + testId: 'test-fallback', + toolName: 'get_chat_overview', + sessionId: 'missing-session', + params: {}, + }, + }) + + assert.equal(resp.statusCode, 404) + assert.deepEqual(resp.json(), { success: false, error: 'Session not found: missing-session' }) + }) +}) diff --git a/packages/http-routes/src/routes/web/ai-tools.ts b/packages/http-routes/src/routes/web/ai-tools.ts index 78b001a3..156a676f 100644 --- a/packages/http-routes/src/routes/web/ai-tools.ts +++ b/packages/http-routes/src/routes/web/ai-tools.ts @@ -32,19 +32,24 @@ export function registerAiToolRoutes(server: FastifyInstance, ctx: HttpRouteCont return reply.code(404).send({ success: false, error: `Tool not found: ${toolName}` }) } - const db = ctx.dbManager.open(sessionId) - if (!db) { - return reply.code(404).send({ success: false, error: `Session not found: ${sessionId}` }) - } - const abortController = new AbortController() activeToolTests.set(testId, abortController) try { + if (ctx.executeAiTool) { + return await ctx.executeAiTool({ testId, toolName, params, sessionId, abortSignal: abortController.signal }) + } + + const db = ctx.dbManager.open(sessionId) + if (!db) { + return reply.code(404).send({ success: false, error: `Session not found: ${sessionId}` }) + } + const execCtx: ToolExecutionContext = { sessionId, db, dataProvider: new CoreDataProvider(db), + abortSignal: abortController.signal, } const startTime = Date.now() diff --git a/packages/node-runtime/src/database-manager.test.ts b/packages/node-runtime/src/database-manager.test.ts index f2a9d963..f476788e 100644 --- a/packages/node-runtime/src/database-manager.test.ts +++ b/packages/node-runtime/src/database-manager.test.ts @@ -5,7 +5,7 @@ import path from 'node:path' import test from 'node:test' import Database from 'better-sqlite3' import type { PathProvider } from '@openchatlab/core' -import { CURRENT_SCHEMA_VERSION, getSessionInfo } from '@openchatlab/core' +import { CHAT_DB_SCHEMA, CURRENT_SCHEMA_VERSION, getSessionInfo } from '@openchatlab/core' import { DataDirCompatibilityError, readDataDirCompatibilityMeta } from './data-dir-compat' import { DatabaseManager } from './database-manager' @@ -30,6 +30,20 @@ function createPathProvider(root: string): PathProvider { } } +function getIndexNames(db: Database.Database): string[] { + const rows = db.prepare("SELECT name FROM sqlite_master WHERE type = 'index' ORDER BY name").all() as Array<{ + name: string + }> + return rows.map((row) => row.name) +} + +function assertAnalysisToolIndexes(db: Database.Database): void { + const indexes = getIndexNames(db) + assert.ok(indexes.includes('idx_message_reply_to')) + assert.ok(indexes.includes('idx_message_sender_ts')) + assert.ok(indexes.includes('idx_message_type_ts')) +} + test('constructor rejects missing runtime unless the test-only bypass is explicit', () => { const root = makeTempDir() @@ -154,6 +168,65 @@ test('open backfills FTS index when migrating legacy sessions', () => { manager.closeAll() }) +test('open migrates v7 databases to include analysis tool indexes', () => { + const root = makeTempDir() + const dbDir = path.join(root, 'data', 'databases') + fs.mkdirSync(dbDir, { recursive: true }) + const dbPath = path.join(dbDir, 'v7-analysis-indexes.db') + + const rawDb = new Database(dbPath, { nativeBinding }) + rawDb.exec(` + CREATE TABLE meta ( + name TEXT NOT NULL, + platform TEXT NOT NULL, + type TEXT NOT NULL, + imported_at INTEGER NOT NULL, + schema_version INTEGER DEFAULT 7 + ); + INSERT INTO meta (name, platform, type, imported_at, schema_version) + VALUES ('V7 Analysis Indexes', 'qq', 'group', 1000, 7); + + CREATE TABLE member ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + platform_id TEXT NOT NULL UNIQUE, + account_name TEXT + ); + INSERT INTO member (platform_id, account_name) VALUES ('u1', 'Alice'); + + CREATE TABLE message ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + sender_id INTEGER NOT NULL, + ts INTEGER NOT NULL, + type INTEGER NOT NULL, + content TEXT, + reply_to_message_id TEXT DEFAULT NULL, + platform_message_id TEXT DEFAULT NULL + ); + INSERT INTO message (sender_id, ts, type, content) VALUES (1, 1000, 0, 'hello indexed tools'); + `) + rawDb.close() + + const manager = new DatabaseManager(createPathProvider(root), { nativeBinding, allowMissingRuntimeForTests: true }) + const db = manager.open('v7-analysis-indexes') + assert.ok(db) + + const version = db.prepare('SELECT schema_version FROM meta LIMIT 1').get() as { schema_version: number } + assert.equal(version.schema_version, CURRENT_SCHEMA_VERSION) + assertAnalysisToolIndexes(db as unknown as Database.Database) + + manager.closeAll() +}) + +test('CHAT_DB_SCHEMA creates analysis tool indexes for new databases', () => { + const db = new Database(':memory:', { nativeBinding }) + try { + db.exec(CHAT_DB_SCHEMA) + assertAnalysisToolIndexes(db) + } finally { + db.close() + } +}) + test('open migrates v2 chat_session schema to current segment schema', () => { const root = makeTempDir() const dbDir = path.join(root, 'data', 'databases') diff --git a/packages/node-runtime/src/migrations/chat-db-migrations.ts b/packages/node-runtime/src/migrations/chat-db-migrations.ts index 0168b978..0002914d 100644 --- a/packages/node-runtime/src/migrations/chat-db-migrations.ts +++ b/packages/node-runtime/src/migrations/chat-db-migrations.ts @@ -395,5 +395,18 @@ export function getChatDbMigrations(deps?: MigrationDeps): CoreMigration[] { } }, }, + { + version: 8, + description: 'Add analysis tool performance indexes', + up: (db: DatabaseAdapter) => { + addColumnIfMissing(db, 'message', 'reply_to_message_id', 'TEXT DEFAULT NULL') + + db.exec(` + CREATE INDEX IF NOT EXISTS idx_message_sender_ts ON message(sender_id, ts); + CREATE INDEX IF NOT EXISTS idx_message_type_ts ON message(type, ts); + CREATE INDEX IF NOT EXISTS idx_message_reply_to ON message(reply_to_message_id); + `) + }, + }, ] } diff --git a/packages/tools/src/definitions/analysis-tools.test.ts b/packages/tools/src/definitions/analysis-tools.test.ts index b6f29da3..c3830072 100644 --- a/packages/tools/src/definitions/analysis-tools.test.ts +++ b/packages/tools/src/definitions/analysis-tools.test.ts @@ -6,6 +6,7 @@ import { getSegmentSummariesTool } from './get-segment-summaries' import { searchMessagesTool } from './search-messages' import { searchSegmentsTool } from './search-segments' import { schemaTool, sqlQueryTool } from './sql-query' +import { SQL_TOOL_DEFS, createSqlToolDefinition } from '../sql' import type { RawMessage, ToolDataProvider, ToolExecutionContext, TimeFilter } from '../types' function createContext( @@ -20,6 +21,12 @@ function createContext( } } +function createSqlTool(name: string) { + const def = SQL_TOOL_DEFS.find((tool) => tool.name === name) + assert.ok(def, `Expected SQL tool definition for ${name}`) + return { def, tool: createSqlToolDefinition(def) } +} + describe('high-risk analysis tool definitions', () => { it('search_messages passes filters to the provider and returns expanded context messages', async () => { const contextFilter: TimeFilter = { startTs: 1710000000, endTs: 1710000100 } @@ -232,4 +239,52 @@ describe('high-risk analysis tool definitions', () => { assert.deepEqual(result.data, schema) assert.deepEqual(JSON.parse(result.content), { tables: schema }) }) + + it('mutual_interaction_pairs uses adjacent-message windows instead of a message self-join', async () => { + const { def } = createSqlTool('mutual_interaction_pairs') + + assert.doesNotMatch(def.execution.query, /JOIN message b ON b\.sender_id != a\.sender_id/) + assert.match(def.execution.query, /LAG\(sender_id\) OVER/) + assert.match(def.execution.query, /ordered_messages/) + }) + + it('mutual_interaction_pairs formats adjacent interaction rows from the provider', async () => { + const { tool } = createSqlTool('mutual_interaction_pairs') + const calls: Array<{ query: string; params: Record }> = [] + const context = createContext({ + async executeParameterizedSql>(query: string, params: Record) { + calls.push({ query, params }) + return [{ member_a: 'Alice', member_b: 'Bob', interaction_count: 2 }] as T[] + }, + }) + + const result = await tool.handler({ days: 30, limit: 5 }, context) + + assert.equal(calls.length, 1) + assert.equal(calls[0]?.params.days, 30) + assert.equal(calls[0]?.params.limit, 5) + assert.match(result.content, /Alice/) + assert.match(result.content, /Bob/) + assert.deepEqual((result.data as { rows: unknown[]; rowCount: number }).rowCount, 1) + }) + + it('reply_interaction_ranking still ranks explicit reply relationships', async () => { + const { def, tool } = createSqlTool('reply_interaction_ranking') + const context = createContext({ + async executeParameterizedSql>(query: string, params: Record) { + assert.equal(query, def.execution.query) + assert.deepEqual(params, { days: 14, limit: 3 }) + return [{ replier_name: 'Bob', original_name: 'Alice', reply_count: 4 }] as T[] + }, + }) + + assert.match(def.execution.query, /reply_to_message_id/) + assert.match(def.execution.query, /ORDER BY reply_count DESC/) + + const result = await tool.handler({ days: 14, limit: 3 }, context) + + assert.match(result.content, /Bob/) + assert.match(result.content, /Alice/) + assert.match(result.content, /4/) + }) }) diff --git a/packages/tools/src/sql/definitions.ts b/packages/tools/src/sql/definitions.ts index 1ba8d5f2..26bc941a 100644 --- a/packages/tools/src/sql/definitions.ts +++ b/packages/tools/src/sql/definitions.ts @@ -128,7 +128,7 @@ export const SQL_TOOL_DEFS: SqlToolDef[] = [ execution: { type: 'sqlite', query: - "SELECT COALESCE(m1.group_nickname, m1.account_name) AS member_a, COALESCE(m2.group_nickname, m2.account_name) AS member_b, COUNT(*) AS interaction_count FROM message a JOIN message b ON b.sender_id != a.sender_id AND b.ts > a.ts AND b.ts <= a.ts + 300 JOIN member m1 ON a.sender_id = m1.id JOIN member m2 ON b.sender_id = m2.id WHERE a.sender_id < b.sender_id AND a.ts > unixepoch('now', '-' || @days || ' days') AND a.type = 0 AND b.type = 0 GROUP BY a.sender_id, b.sender_id ORDER BY interaction_count DESC LIMIT @limit", + "WITH ordered_messages AS (SELECT id, sender_id, ts, LAG(sender_id) OVER (ORDER BY ts, id) AS prev_sender_id, LAG(ts) OVER (ORDER BY ts, id) AS prev_ts FROM message WHERE type = 0 AND ts > unixepoch('now', '-' || @days || ' days')), pairs AS (SELECT CASE WHEN prev_sender_id < sender_id THEN prev_sender_id ELSE sender_id END AS member_a_id, CASE WHEN prev_sender_id < sender_id THEN sender_id ELSE prev_sender_id END AS member_b_id FROM ordered_messages WHERE prev_sender_id IS NOT NULL AND prev_sender_id != sender_id AND prev_ts IS NOT NULL AND ts <= prev_ts + 300) SELECT COALESCE(m1.group_nickname, m1.account_name) AS member_a, COALESCE(m2.group_nickname, m2.account_name) AS member_b, COUNT(*) AS interaction_count FROM pairs JOIN member m1 ON pairs.member_a_id = m1.id JOIN member m2 ON pairs.member_b_id = m2.id GROUP BY member_a_id, member_b_id ORDER BY interaction_count DESC LIMIT @limit", rowTemplate: '{member_a} ↔ {member_b}：{interaction_count} 次互动', summaryTemplate: '互动最频繁的 {rowCount} 对好友：', fallback: '该时间范围内没有检测到明显的互动关系', diff --git a/packages/tools/src/types.ts b/packages/tools/src/types.ts index 2458ad9f..905b81aa 100644 --- a/packages/tools/src/types.ts +++ b/packages/tools/src/types.ts @@ -232,6 +232,7 @@ export interface ToolExecutionContext { sessionId: string locale?: string timeFilter?: TimeFilter + abortSignal?: AbortSignal /** 抽象查询接口 */ dataProvider?: ToolDataProvider /** @deprecated 逐步迁移到 dataProvider 后移除；Electron 端不提供此字段 */