From dc3bd82c90bcc56c1bdb5a94f9d2cdc534063f8b Mon Sep 17 00:00:00 2001 From: digua Date: Sat, 22 Aug 2026 12:45:35 +0800 Subject: [PATCH] fix(ai): enforce cross-chat execution budgets Count complete tool payloads with the runtime tokenizer so Chinese evidence respects the configured model context budget.\n\nResolve contact sources incrementally with cancellation checks and targeted member lookups to keep long cross-chat investigations responsive. --- apps/cli/src/ai/cross-chat-tool-adapter.ts | 3 +- .../main/ai/cross-chat-tool-adapter.ts | 3 +- packages/core/src/index.ts | 1 + .../query/__tests__/contact-queries.test.ts | 12 +++++ packages/core/src/query/contact-queries.ts | 8 ++- packages/core/src/query/index.ts | 1 + .../cross-chat-analysis/service.test.ts | 35 +++++++++++-- .../services/cross-chat-analysis/service.ts | 21 ++++++-- .../src/definitions/cross-chat-tools.test.ts | 49 +++++++++++++++++-- .../tools/src/definitions/cross-chat-tools.ts | 46 ++++++++++++++--- packages/tools/src/types.ts | 4 +- 11 files changed, 162 insertions(+), 21 deletions(-) diff --git a/apps/cli/src/ai/cross-chat-tool-adapter.ts b/apps/cli/src/ai/cross-chat-tool-adapter.ts index 7ba479f2..fb4a64f1 100644 --- a/apps/cli/src/ai/cross-chat-tool-adapter.ts +++ b/apps/cli/src/ai/cross-chat-tool-adapter.ts @@ -4,7 +4,7 @@ import type { PreprocessConfig, SessionRuntimeAdapter, } from '@openchatlab/node-runtime' -import { preprocessCrossChatMessages } from '@openchatlab/node-runtime' +import { countTokens, preprocessCrossChatMessages } from '@openchatlab/node-runtime' import { CROSS_CHAT_AGENT_TOOL_REGISTRY, executeToolForAgent, @@ -23,6 +23,7 @@ export function createCliCrossChatTools(options: { locale: options.locale, analysisService: options.analysisService, maxToolResultTokens: options.maxToolResultTokens, + countTokens, preprocessMessagesBySession: (sessionId, messages) => preprocessCrossChatMessages( options.sessionAdapter, diff --git a/apps/desktop/main/ai/cross-chat-tool-adapter.ts b/apps/desktop/main/ai/cross-chat-tool-adapter.ts index 5090987b..b3e71008 100644 --- a/apps/desktop/main/ai/cross-chat-tool-adapter.ts +++ b/apps/desktop/main/ai/cross-chat-tool-adapter.ts @@ -4,7 +4,7 @@ import type { PreprocessConfig, SessionRuntimeAdapter, } from '@openchatlab/node-runtime' -import { preprocessCrossChatMessages } from '@openchatlab/node-runtime' +import { countTokens, preprocessCrossChatMessages } from '@openchatlab/node-runtime' import { CROSS_CHAT_AGENT_TOOL_REGISTRY, executeToolForAgent, @@ -23,6 +23,7 @@ export function createElectronCrossChatTools(options: { locale: options.locale, analysisService: options.analysisService, maxToolResultTokens: options.maxToolResultTokens, + countTokens, preprocessMessagesBySession: (sessionId, messages) => preprocessCrossChatMessages( options.sessionAdapter, diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index fffb0bd0..db976062 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -62,6 +62,7 @@ export { getNonSystemMembersForContacts, getPrivateContactFacts, isValidContactPlatformId, + resolveContactMember, resolveOwnerMember, buildContactKey, shouldScopeContactToSession, diff --git a/packages/core/src/query/__tests__/contact-queries.test.ts b/packages/core/src/query/__tests__/contact-queries.test.ts index 634b4eee..2acf1d6b 100644 --- a/packages/core/src/query/__tests__/contact-queries.test.ts +++ b/packages/core/src/query/__tests__/contact-queries.test.ts @@ -13,6 +13,7 @@ import { getGroupRelationshipGraphFacts, getNonSystemMembersForContacts, getPrivateContactFacts, + resolveContactMember, resolveOwnerMember, } from '../contact-queries' import type { DatabaseAdapter, PreparedStatement, RunResult } from '../../interfaces' @@ -121,6 +122,17 @@ describe('contact query helpers', () => { }) }) + it('resolves one contact directly by stable platform id', () => { + assert.deepEqual(resolveContactMember(db, 'alice-pid'), { + id: 2, + platformId: 'alice-pid', + name: 'Alice G', + aliases: ['Ally', '小爱'], + avatar: 'alice.png', + }) + assert.equal(resolveContactMember(db, 'missing-pid'), null) + }) + it('returns saved member aliases for contact candidates', () => { const alice = getNonSystemMembersForContacts(db).find((member) => member.platformId === 'alice-pid') diff --git a/packages/core/src/query/contact-queries.ts b/packages/core/src/query/contact-queries.ts index 74eb80b5..2cabe33f 100644 --- a/packages/core/src/query/contact-queries.ts +++ b/packages/core/src/query/contact-queries.ts @@ -70,6 +70,12 @@ export function isValidContactPlatformId(platformId: string | null | undefined): export function resolveOwnerMember(db: DatabaseAdapter): ContactMemberRef | null { const meta = db.prepare('SELECT owner_id FROM meta LIMIT 1').get() as { owner_id: string | null } | undefined if (!isValidContactPlatformId(meta?.owner_id)) return null + + return resolveContactMember(db, meta.owner_id) +} + +export function resolveContactMember(db: DatabaseAdapter, platformId: string): ContactMemberRef | null { + if (!isValidContactPlatformId(platformId)) return null const aliasesSelect = hasColumn(db, 'member', 'aliases') ? 'aliases' : 'NULL as aliases' const row = db @@ -84,7 +90,7 @@ export function resolveOwnerMember(db: DatabaseAdapter): ContactMemberRef | null WHERE platform_id = ? AND ${nonSystemContactMemberCondition(db, 'm')} LIMIT 1` ) - .get(meta.owner_id) as ContactMemberRow | undefined + .get(platformId) as ContactMemberRow | undefined return row ? mapContactMemberRow(row) : null } diff --git a/packages/core/src/query/index.ts b/packages/core/src/query/index.ts index fe34cddb..0fc51b00 100644 --- a/packages/core/src/query/index.ts +++ b/packages/core/src/query/index.ts @@ -85,6 +85,7 @@ export { getNonSystemMembersForContacts, getPrivateContactFacts, isValidContactPlatformId, + resolveContactMember, resolveOwnerMember, } from './contact-queries' diff --git a/packages/node-runtime/src/services/cross-chat-analysis/service.test.ts b/packages/node-runtime/src/services/cross-chat-analysis/service.test.ts index 1e09654d..d73b48ee 100644 --- a/packages/node-runtime/src/services/cross-chat-analysis/service.test.ts +++ b/packages/node-runtime/src/services/cross-chat-analysis/service.test.ts @@ -312,11 +312,11 @@ test('contact lookup uses the all-history snapshot for typed names and candidate } }) -test('entity resolution uses contact keys and resolves per-session member ids without merging display names', () => { +test('entity resolution uses contact keys and resolves per-session member ids without merging display names', async () => { const { env, contactsService } = createFixture() try { const service = createCrossChatAnalysisService({ adapter: env.adapter, contactsService }) - const result = service.resolveEntities([ + const result = await service.resolveEntities([ { type: 'contact', contactKey: 'test:alice', displayName: 'Alice' }, { type: 'contact', contactKey: 'test:group-other:alice-other', displayName: 'Alice' }, ]) @@ -344,7 +344,7 @@ test('entity resolution uses contact keys and resolves per-session member ids wi } }) -test('entity resolution uses the all-history contact snapshot so older source sessions remain searchable', () => { +test('entity resolution uses the all-history contact snapshot so older source sessions remain searchable', async () => { const { env, contactsService: fixtureContactsService } = createFixture() try { const contactsService: Pick = { @@ -367,7 +367,7 @@ test('entity resolution uses the all-history contact snapshot so older source se } const service = createCrossChatAnalysisService({ adapter: env.adapter, contactsService }) - const result = service.resolveEntities([{ type: 'contact', contactKey: 'test:alice', displayName: 'Alice' }]) + const result = await service.resolveEntities([{ type: 'contact', contactKey: 'test:alice', displayName: 'Alice' }]) assert.deepEqual( result.contacts[0]?.sessions.map((session) => session.sessionId), @@ -380,6 +380,33 @@ test('entity resolution uses the all-history contact snapshot so older source se } }) +test('entity resolution honors interruption between contact source sessions', async () => { + const { env, contactsService } = createFixture() + try { + const openedSessions: string[] = [] + const adapter: SessionRuntimeAdapter = { + ...env.adapter, + openReadonly: (sessionId) => { + openedSessions.push(sessionId) + return env.adapter.openReadonly(sessionId) + }, + } + const controller = new AbortController() + const service = createCrossChatAnalysisService({ adapter, contactsService }) + const resolution = service.resolveEntities([{ type: 'contact', contactKey: 'test:alice', displayName: 'Alice' }], { + signal: controller.signal, + }) + + queueMicrotask(() => controller.abort()) + + await assert.rejects(resolution, { name: 'AbortError' }) + assert.ok(openedSessions.length > 0) + assert.ok(openedSessions.length < 2) + } finally { + env.cleanup() + } +}) + test('scoped search filters by resolved member ids and keeps compound evidence identity', async () => { const { env, contactsService } = createFixture() try { diff --git a/packages/node-runtime/src/services/cross-chat-analysis/service.ts b/packages/node-runtime/src/services/cross-chat-analysis/service.ts index 00e477ff..bdcc575c 100644 --- a/packages/node-runtime/src/services/cross-chat-analysis/service.ts +++ b/packages/node-runtime/src/services/cross-chat-analysis/service.ts @@ -4,6 +4,7 @@ import { getSearchMessageContext as getCoreSearchMessageContext, getSessionMeta, getSessionOverview, + resolveContactMember, resolveOwnerMember, searchMessagesByKeywords, type DatabaseAdapter, @@ -53,7 +54,7 @@ export interface CrossChatAnalysisServiceDeps { export interface CrossChatAnalysisService { lookupContact(query: string): CrossChatContactLookupResult - resolveEntities(refs: AIEntityRef[]): CrossChatEntityResolution + resolveEntities(refs: AIEntityRef[], options?: CrossChatOperationOptions): Promise searchMessages(request: CrossChatSearchRequest, options?: CrossChatOperationOptions): Promise getMessageContext(request: CrossChatMessageContextRequest): CrossChatMessageContextResult getOverview(request: CrossChatOverviewRequest, options?: CrossChatOperationOptions): Promise @@ -116,7 +117,11 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { } } - resolveEntities(refs: AIEntityRef[]): CrossChatEntityResolution { + async resolveEntities( + refs: AIEntityRef[], + options: CrossChatOperationOptions = {} + ): Promise { + throwIfAborted(options.signal) const contacts: CrossChatResolvedContact[] = [] const sessions: CrossChatResolvedSession[] = [] const unresolved: CrossChatUnresolvedEntity[] = [] @@ -125,6 +130,7 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { const failedSessionIds = new Set() for (const ref of refs) { + throwIfAborted(options.signal) if (ref.type === 'session') { candidateSessionIds.add(ref.sessionId) const descriptor = this.tryGetSessionDescriptor(ref.sessionId) @@ -136,6 +142,7 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { sessions.push({ ref, status: 'unresolved' }) unresolved.push({ ref, reason: 'session_not_found' }) } + await yieldToEventLoop() continue } @@ -157,6 +164,7 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { ref, reason: detail.cache.status === 'missing' ? 'contact_snapshot_missing' : 'contact_not_found', }) + await yieldToEventLoop() continue } @@ -167,7 +175,9 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { const unresolvedContactSessions: string[] = [] const failedContactSessions: string[] = [] + // 实体解析必须保留联系人全部精确来源；逐库让出事件循环，使长列表仍能及时响应取消。 for (const source of sourceSessions) { + throwIfAborted(options.signal) candidateSessionIds.add(source.id) try { const db = this.deps.adapter.openReadonly(source.id) @@ -177,7 +187,7 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { continue } const descriptor = getSessionDescriptor(source.id, db) - const member = getMembers(db).find((item) => item.platformId === contact.platformId) + const member = resolveContactMember(db, contact.platformId) if (!descriptor || !member) { unresolvedContactSessions.push(source.id) continue @@ -193,6 +203,8 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { failedContactSessions.push(source.id) failedSessionIds.add(source.id) appLogger.warn('cross-chat-analysis', `failed to resolve contact in source session: ${source.id}`, error) + } finally { + await yieldToEventLoop() } } @@ -211,8 +223,11 @@ class DefaultCrossChatAnalysisService implements CrossChatAnalysisService { failedSessionIds: failedContactSessions, }) if (status === 'unresolved') unresolved.push({ ref, reason: 'member_not_found' }) + if (sourceSessions.length === 0) await yieldToEventLoop() } + throwIfAborted(options.signal) + return { contacts, sessions, diff --git a/packages/tools/src/definitions/cross-chat-tools.test.ts b/packages/tools/src/definitions/cross-chat-tools.test.ts index d08d545f..63045845 100644 --- a/packages/tools/src/definitions/cross-chat-tools.test.ts +++ b/packages/tools/src/definitions/cross-chat-tools.test.ts @@ -44,7 +44,7 @@ function createContext(overrides: Partial = {}): C totalCandidates: 0, candidates: [], }), - resolveEntities: (_refs: AIEntityRef[]): CrossChatEntityResolution => ({ + resolveEntities: async (_refs: AIEntityRef[]): Promise => ({ contacts: [], sessions: [], unresolved: [], @@ -158,8 +158,9 @@ describe('cross-chat agent registry', () => { }, ], }), - resolveEntities: (refs) => { + resolveEntities: async (refs, options) => { resolvedRefs = refs + assert.equal(options?.signal, context.abortSignal) return { contacts: [], sessions: [], @@ -174,6 +175,8 @@ describe('cross-chat agent registry', () => { } }, } as Partial) + const controller = new AbortController() + context.abortSignal = controller.signal const tool = CROSS_CHAT_AGENT_TOOL_REGISTRY.find((item) => item.name === 'resolve_chat_entities') assert.ok(tool) @@ -207,7 +210,7 @@ describe('cross-chat agent registry', () => { }, ], }), - resolveEntities: (refs) => { + resolveEntities: async (refs) => { resolvedRefs = refs return { contacts: [], @@ -322,6 +325,46 @@ describe('cross-chat agent registry', () => { assert.equal(details.crossChatEvidence.sources[0]?.snippet, 'x'.repeat(500)) }) + it('uses the injected tokenizer when budgeting Chinese search payloads', async () => { + const context = createContext({ + searchMessages: async () => ({ + messages: [messageSource('a', 1, 1, '中文消息'.repeat(125))], + totalMatches: 1, + appliedFilters: { startTs: null, endTs: null, recentDays: null, sender: 'all' }, + coverage: { + candidateSessions: 1, + scannedSessions: 1, + matchedSessions: 1, + failedSessions: 0, + truncated: false, + truncatedReasons: [], + }, + }), + }) + context.maxToolResultTokens = 300 + const countTokens = (text: string) => { + let tokens = 0 + for (const char of text) tokens += /[\u3400-\u9fff]/u.test(char) ? 1 : 0.25 + return Math.ceil(tokens) + } + context.countTokens = countTokens + context.preprocessMessagesBySession = (_sessionId, messages) => messages + const tool = CROSS_CHAT_AGENT_TOOL_REGISTRY.find((item) => item.name === 'search_messages_globally') + assert.ok(tool) + + const result = await tool.handler({ keywords: ['中文'] }, context) + const details = result.data as { + returned: number + coverage: { truncatedReasons: string[] } + crossChatEvidence: { sources: unknown[] } + } + + assert.ok(countTokens(result.content) <= 300) + assert.equal(details.returned, 0) + assert.deepEqual(details.crossChatEvidence.sources, []) + assert.ok(details.coverage.truncatedReasons.includes('evidence_budget')) + }) + it('allows recent-message sampling only when explicit scopes are present', async () => { let captured: unknown const context = createContext({ diff --git a/packages/tools/src/definitions/cross-chat-tools.ts b/packages/tools/src/definitions/cross-chat-tools.ts index 5296e1f8..37f09703 100644 --- a/packages/tools/src/definitions/cross-chat-tools.ts +++ b/packages/tools/src/definitions/cross-chat-tools.ts @@ -100,7 +100,10 @@ const overviewSchema: JsonSchema = { required: ['scopes'], } -function resolveHandler(params: Record, context: CrossChatToolExecutionContext): ToolResult { +async function resolveHandler( + params: Record, + context: CrossChatToolExecutionContext +): Promise { const inputs = parseEntityInputs(params.entities) const refs: AIEntityRef[] = [] const contactLookups: CrossChatContactLookupResult[] = [] @@ -119,7 +122,9 @@ function resolveHandler(params: Record, context: CrossChatToolE } refs.push(input as AIEntityRef) } - const resolution = context.analysisService.resolveEntities(dedupeEntityRefs(refs)) + const resolution = await context.analysisService.resolveEntities(dedupeEntityRefs(refs), { + signal: context.abortSignal, + }) const data = { ...resolution, contactLookups } return { content: JSON.stringify(data), data } } @@ -169,7 +174,9 @@ async function searchHandler( coverage: buildCoverage(budgetTruncated), messages: messages.map(toModelMessage), }) - const limited = limitMessagesToBudget(safeMessages, context.maxToolResultTokens, buildModelData) + const limited = limitMessagesToBudget(safeMessages, context.maxToolResultTokens, buildModelData, { + countTokens: context.countTokens, + }) const modelData = buildModelData(limited.messages, limited.truncated) const evidence: CrossChatEvidencePayload = { version: 1, @@ -204,6 +211,7 @@ async function contextHandler( const limited = limitMessagesToBudget(safeMessages, context.maxToolResultTokens, buildModelData, { priorityIndexes: buildContextPriority(safeMessages, messageId), continueAfterOverflow: true, + countTokens: context.countTokens, }) const data = buildModelData(limited.messages, limited.truncated) return { content: JSON.stringify(data), data } @@ -354,7 +362,11 @@ function limitMessagesToBudget( messages: CrossChatMessageSource[], maxToolResultTokens: number | undefined, buildPayload: (messages: CrossChatMessageSource[], truncated: boolean) => unknown, - options: { priorityIndexes?: number[]; continueAfterOverflow?: boolean } = {} + options: { + priorityIndexes?: number[] + continueAfterOverflow?: boolean + countTokens?: (text: string) => number + } = {} ): { messages: CrossChatMessageSource[]; truncated: boolean } { const prepared = messages.map((message) => ({ ...message, @@ -364,7 +376,7 @@ function limitMessagesToBudget( return { messages: prepared, truncated: false } } - const maxChars = maxToolResultTokens * 4 + const countTokens = options.countTokens ?? estimatePayloadTokens const priorityIndexes = options.priorityIndexes ?? prepared.map((_, index) => index) const selectedIndexes = new Set() for (const index of priorityIndexes) { @@ -372,8 +384,8 @@ function limitMessagesToBudget( const candidateIndexes = [...selectedIndexes, index].sort((left, right) => left - right) const candidateMessages = candidateIndexes.map((candidateIndex) => prepared[candidateIndex]) const truncated = candidateMessages.length < prepared.length - const estimatedChars = JSON.stringify(buildPayload(candidateMessages, truncated)).length - if (estimatedChars > maxChars) { + const serializedCandidate = JSON.stringify(buildPayload(candidateMessages, truncated)) + if (countTokens(serializedCandidate) > maxToolResultTokens) { if (options.continueAfterOverflow) continue break } @@ -383,6 +395,26 @@ function limitMessagesToBudget( return { messages: limited, truncated: limited.length < prepared.length } } +function estimatePayloadTokens(text: string): number { + // Browser/MCP 等未注入 Node tokenizer 的调用方仍需避免按 ASCII 密度低估中文结果。 + let cjkChars = 0 + let otherChars = 0 + for (const char of text) { + const codePoint = char.codePointAt(0) ?? 0 + if ( + (codePoint >= 0x3400 && codePoint <= 0x4dbf) || + (codePoint >= 0x4e00 && codePoint <= 0x9fff) || + (codePoint >= 0x3040 && codePoint <= 0x30ff) || + (codePoint >= 0xac00 && codePoint <= 0xd7af) + ) { + cjkChars++ + } else { + otherChars++ + } + } + return Math.ceil(cjkChars / 1.6 + otherChars / 4) +} + function buildContextPriority(messages: CrossChatMessageSource[], anchorMessageId: number): number[] { const anchorIndex = messages.findIndex((message) => message.messageId === anchorMessageId) if (anchorIndex < 0) return messages.map((_, index) => index) diff --git a/packages/tools/src/types.ts b/packages/tools/src/types.ts index e8a80acc..a7fba279 100644 --- a/packages/tools/src/types.ts +++ b/packages/tools/src/types.ts @@ -357,7 +357,7 @@ export interface ToolExecutionContext { export interface CrossChatAnalysisToolService { lookupContact(query: string): CrossChatContactLookupResult - resolveEntities(refs: AIEntityRef[]): CrossChatEntityResolution + resolveEntities(refs: AIEntityRef[], options?: CrossChatOperationOptions): Promise searchMessages(request: CrossChatSearchRequest, options?: CrossChatOperationOptions): Promise getMessageContext(request: CrossChatMessageContextRequest): CrossChatMessageContextResult getOverview(request: CrossChatOverviewRequest, options?: CrossChatOperationOptions): Promise @@ -368,6 +368,8 @@ export interface CrossChatToolExecutionContext { abortSignal?: AbortSignal reportProgress?: (progress: ToolProgress) => void maxToolResultTokens?: number + /** 平台注入的统一 tokenizer；未注入时 handler 使用 CJK-aware 轻量估算。 */ + countTokens?: (text: string) => number analysisService: CrossChatAnalysisToolService preprocessMessagesBySession: ( sessionId: string,