From c0e39a395c8de141b474d73e7ff558d575adea11 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:27:15 +0300 Subject: [PATCH] =?UTF-8?q?perf(ai-chat):=20deferred-=D0=B0=D0=BA=D1=82?= =?UTF-8?q?=D0=B8=D0=B2=D0=B0=D1=86=D0=B8=D1=8F=20=D1=82=D1=83=D0=BB=D0=BE?= =?UTF-8?q?=D0=B2=20=D0=B2=20metadata=20=D1=87=D0=B0=D1=82=D0=B0=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Активированный set сбрасывался каждый ход → модель заново гоняла loadTools, чтобы переактивировать те же тулы (лишний round-trip на каждом ходу). Теперь набор персистится в metadata чата и сидируется на следующем ходу. - Миграция: jsonb-колонка metadata на ai_chats (default '{}'); db.d.ts дополнен вручную (AiChats.metadata: Generated). - seedActivatedTools(metadata, validDeferredNames): читает сохранённый набор, ПЕРЕСЕКАЯ с актуальными validDeferredNames — смена allowlist/ролей не воскресит несуществующий тул (иначе prepareAgentStep получил бы фантомное активное имя). Сид только при deferredEnabled. - Персист на завершении хода (once-guard, во всех терминальных ветках рядом со snapshotTurnEnd): детерминированно отсортированный набор, merge в существующий bag (другие ключи сохраняются), запись пропускается если ничего нового не активировано (обычный ход не даёт лишней записи). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 40 +++++++++ .../src/core/ai-chat/ai-chat.service.ts | 84 ++++++++++++++++++- .../20260707T120000-ai-chat-metadata.ts | 24 ++++++ apps/server/src/database/types/db.d.ts | 3 + 4 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 70293eb0..095bb34a 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -31,6 +31,7 @@ import { OUTPUT_DEGENERATION_ERROR, lastAssistantContextTokens, lastAssistantReplayOverflow, + seedActivatedTools, } from './ai-chat.service'; import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types'; import { buildSystemPrompt } from './ai-chat.prompt'; @@ -439,6 +440,45 @@ describe('lastAssistantContextTokens', () => { }); }); +// #490 deferred-tool activation persisted across turns. +describe('seedActivatedTools', () => { + const valid = new Set(['Search_web', 'getPageJson', 'diffPageVersions']); + + it('seeds from persisted metadata, intersected with current valid names', () => { + expect( + seedActivatedTools( + { activatedTools: ['Search_web', 'getPageJson'] }, + valid, + ), + ).toEqual(['Search_web', 'getPageJson']); + }); + + it('drops a stored tool that is no longer valid (allowlist/role changed)', () => { + // 'Habr_publish' was activated before but is not in the current allowlist. + expect( + seedActivatedTools({ activatedTools: ['Search_web', 'Habr_publish'] }, valid), + ).toEqual(['Search_web']); + }); + + it('is empty/robust for missing, non-array, or unknown-shaped metadata', () => { + expect(seedActivatedTools(undefined, valid)).toEqual([]); + expect(seedActivatedTools({}, valid)).toEqual([]); + expect(seedActivatedTools({ activatedTools: 'nope' }, valid)).toEqual([]); + expect( + seedActivatedTools({ activatedTools: [1, 'getPageJson', null] }, valid), + ).toEqual(['getPageJson']); + }); + + it('de-duplicates stored names', () => { + expect( + seedActivatedTools( + { activatedTools: ['getPageJson', 'getPageJson'] }, + valid, + ), + ).toEqual(['getPageJson']); + }); +}); + describe('lastAssistantReplayOverflow', () => { const row = ( role: string, diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index b4c65333..80078827 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -936,10 +936,17 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // supplied or the supplied one does not belong to this workspace. let isNewChat = false; let chatId = body.chatId; + // Persisted chat-level metadata bag (#490): read once here so the deferred-tool + // activation set can be seeded from the previous turn. Undefined for a new chat. + let chatMetadata: Record | undefined; if (chatId) { const existing = await this.aiChatRepo.findById(chatId, workspace.id); if (!existing) { chatId = undefined; + } else { + chatMetadata = (existing.metadata ?? undefined) as + | Record + | undefined; } } // The open page the client sent is attacker-controllable — BOTH its id and @@ -1393,10 +1400,19 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // tools + ALL external MCP tools), computed from the ACTUAL toolset so an // external tool is loadable by its namespaced name. loadTools rejects any // name outside this set. - const activatedTools = new Set(); const validDeferredNames = new Set( Object.keys(baseTools).filter((k) => !CORE_TOOL_SET.has(k)), ); + // #490: seed the activation set from the chat's PERSISTED set so the model + // does not re-run loadTools every turn to re-activate the same tools. Only + // when deferred loading is enabled, and ALWAYS intersected with the CURRENT + // valid deferred names — an allowlist/role change must never resurrect a tool + // that no longer exists (prepareAgentStep would get a phantom active name). + const activatedTools = new Set( + deferredEnabled + ? seedActivatedTools(chatMetadata, validDeferredNames) + : [], + ); // Add the loadTools meta-tool ONLY when the feature is enabled; when off the // toolset and behavior are exactly as before. const tools = deferredEnabled @@ -1406,6 +1422,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { } : baseTools; + // #490: persist the (deterministically ordered) activation set back onto the + // chat metadata at turn end, so the NEXT turn seeds from it. Once-guarded and + // skipped when nothing new was activated (the set equals its seed) so an + // ordinary turn adds no extra write. Preserves other metadata keys. + let activatedToolsPersisted = false; + const persistActivatedTools = async (): Promise => { + if (!deferredEnabled || activatedToolsPersisted || !chatId) return; + activatedToolsPersisted = true; + const current = [...activatedTools].sort(); + const seeded = seedActivatedTools(chatMetadata, validDeferredNames).sort(); + if (current.length === 0 || current.join('') === seeded.join('')) { + return; // nothing new activated -> no write + } + try { + await this.aiChatRepo.update( + chatId, + { + metadata: { + ...(chatMetadata ?? {}), + activatedTools: current, + }, + } as never, + workspace.id, + ); + } catch (err) { + this.logger.warn( + `Failed to persist activated tools (chat ${chatId}): ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + }; + // Accumulate the turn's streamed output so a provider error / disconnect can // persist the PARTIAL answer the user already saw — the SDK's onError/onAbort // callbacks don't hand us the in-progress text. `capturedSteps` holds finished @@ -1750,6 +1799,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // own edits are baked in — and this also SEEDS the snapshot on the first // turn. Runs once across every terminal path (see snapshotTurnEnd). await snapshotTurnEnd(); + // #490: persist the deferred-tool activation set for the next turn. + await persistActivatedTools(); // Generate the chat title for a freshly created chat AFTER the stream's // provider call has completed — NOT concurrently with it. The z.ai coding @@ -1812,6 +1863,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // committed before the error must be baked into the snapshot, or the // next turn would mis-report it as a user edit. await snapshotTurnEnd(); + // #490: persist the deferred-tool activation set for the next turn. + await persistActivatedTools(); }, onAbort: async ({ steps }) => { // #444: distinguish a degeneration abort (our internal controller) from @@ -1837,6 +1890,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { ); await closeExternalClients(); await snapshotTurnEnd(); + // #490: persist the deferred-tool activation set for the next turn. + await persistActivatedTools(); return; } const partialChars = @@ -1872,6 +1927,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // committed before the client disconnect / stop() must be baked into the // snapshot, or the next turn would mis-report it as a user edit. await snapshotTurnEnd(); + // #490: persist the deferred-tool activation set for the next turn. + await persistActivatedTools(); }, }); @@ -2202,6 +2259,31 @@ export function lastAssistantContextTokens( return undefined; } +/** + * Seed the per-turn deferred-tool activation set from a chat's persisted metadata + * (#490), INTERSECTED with the current valid deferred names. Persisting the set + * across turns saves the model re-running loadTools every turn to re-activate the + * same tools; intersecting on load means a changed allowlist / role can never + * resurrect a tool that no longer exists (which would hand prepareAgentStep a + * phantom active name). Tolerant of any stored shape — a non-array is ignored. + */ +export function seedActivatedTools( + metadata: Record | undefined, + validDeferredNames: ReadonlySet, +): string[] { + const stored = metadata?.activatedTools; + if (!Array.isArray(stored)) return []; + const seen = new Set(); + const out: string[] = []; + for (const name of stored) { + if (typeof name === 'string' && validDeferredNames.has(name) && !seen.has(name)) { + seen.add(name); + out.push(name); + } + } + return out; +} + /** * Whether the most recent assistant turn was rejected for CONTEXT OVERFLOW * (#490): its row carries `metadata.replayOverflow` (stamped by the stream's diff --git a/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts b/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts new file mode 100644 index 00000000..429a6bd9 --- /dev/null +++ b/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts @@ -0,0 +1,24 @@ +import { type Kysely, sql } from 'kysely'; + +export async function up(db: Kysely): Promise { + // Chat-level metadata bag (#490). First use: the deferred-tool ACTIVATION set + // (`activatedTools`) is persisted here so it survives across turns — previously + // the set was reset every turn, forcing the model to re-run loadTools and pay a + // fresh round-trip to re-activate the same tools each turn. On load the stored + // set is intersected with the current valid deferred names, so an allowlist / + // role change can never inject a now-nonexistent tool. + // + // jsonb, defaulted to '{}' so every row (incl. pre-migration ones, backfilled + // by the default) is a readable object — the app never has to null-guard the + // bag itself, only individual keys. + await db.schema + .alterTable('ai_chats') + .addColumn('metadata', 'jsonb', (col) => + col.notNull().defaultTo(sql`'{}'::jsonb`), + ) + .execute(); +} + +export async function down(db: Kysely): Promise { + await db.schema.alterTable('ai_chats').dropColumn('metadata').execute(); +} diff --git a/apps/server/src/database/types/db.d.ts b/apps/server/src/database/types/db.d.ts index bf5a9663..a17d628c 100644 --- a/apps/server/src/database/types/db.d.ts +++ b/apps/server/src/database/types/db.d.ts @@ -606,6 +606,9 @@ export interface AiChats { // The document the chat was created in (open page at first message). NULL => // started outside any document. ON DELETE SET NULL on the page FK. pageId: string | null; + // Chat-level metadata bag (#490). jsonb, defaulted to '{}'. First key: + // `activatedTools` — the deferred-tool activation set persisted across turns. + metadata: Generated; createdAt: Generated; updatedAt: Generated; deletedAt: Timestamp | null;