diff --git a/.env.example b/.env.example index d32b67c7..d3d63309 100644 --- a/.env.example +++ b/.env.example @@ -202,6 +202,13 @@ MCP_DOCMOST_PASSWORD= # Default 900000 (15 min). # AI_MCP_CALL_TIMEOUT_MS=900000 +# Deferred tool loading for the in-app AI chat (#332). Default ON: the agent sees +# a compact and only CORE tools + a loadTools meta-tool are active +# each step; deferred tools (the fat/rare ones + all external MCP tools) load on +# demand. Set AI_CHAT_DEFERRED_TOOLS=false to restore the old "all tools always +# active" behavior. +# AI_CHAT_DEFERRED_TOOLS=true + # --- Anonymous public-share AI assistant --- # Opt-in per workspace (AI settings -> "public share assistant"; off by default). # When enabled, anonymous visitors of a published share can ask an AI about that diff --git a/apps/server/src/core/ai-chat/tools/tool-tiers.spec.ts b/apps/server/src/core/ai-chat/tools/tool-tiers.spec.ts index 367f4f4b..8aa039bd 100644 --- a/apps/server/src/core/ai-chat/tools/tool-tiers.spec.ts +++ b/apps/server/src/core/ai-chat/tools/tool-tiers.spec.ts @@ -13,6 +13,12 @@ import { // SHARED_TOOL_SPECS contract spec) so the tier metadata is checked against // exactly what @docmost/mcp ships. import { SHARED_TOOL_SPECS } from '../../../../../../packages/mcp/src/tool-specs'; +// For the live-toolset partition test (F3): the REAL adapter, so the catalog is +// checked against the tools AiChatToolsService.forUser() actually builds — not a +// static list that could drift from it. +import { AiChatToolsService } from './ai-chat-tools.service'; +import * as loader from './docmost-client.loader'; +import type { DocmostClientLike } from './docmost-client.loader'; /** * #332 deferred tool loading — tier metadata, catalog assembly, and the @@ -68,15 +74,94 @@ describe('buildInAppDeferredCatalog (#332)', () => { expect(names).not.toContain('editPageText'); // shared core }); - it('lists all 28 deferred tools (16 inline + 12 shared)', () => { - expect(catalog).toHaveLength(28); - // Each entry is a "name — purpose" line. + it('renders every entry as a "name — purpose" line', () => { + // Non-empty catalog (the length is pinned structurally by the live-toolset + // partition test below, not by a magic constant that rots on every new tool). + expect(catalog.length).toBeGreaterThan(0); for (const entry of catalog) { expect(entry.catalogLine).toMatch(/ — /); } }); }); +/** + * F3 — the deferred is built from STATIC metadata (INLINE_TOOL_TIERS + * + SHARED_TOOL_SPECS), but the loadable-by-name set is derived at RUNTIME from the + * actual toolset (`Object.keys(baseTools)` in ai-chat.service.ts). Those two must + * agree or a tool becomes loadable-but-invisible (agent thinks it doesn't exist) or + * catalogued-but-phantom. INLINE_TOOL_TIERS is a plain hand-maintained Record with + * no compile-time link to the tools AiChatToolsService.forUser() builds, so nothing + * else catches that drift. This test uses forUser()'s LIVE keys as the source of + * truth (mirroring ai-chat-tools.service.spec.ts's loader mock) and asserts a + * two-way partition against buildInAppDeferredCatalog — replacing the old magic + * toHaveLength(28), so a tool added to forUser() without a catalog line (or a + * catalog line without a real tool) fails the suite instead of silently vanishing. + */ +describe('deferred catalog ↔ live forUser() toolset partition (#332, F3)', () => { + let toolKeys: string[]; + const catalogNames = buildInAppDeferredCatalog(SHARED_TOOL_SPECS as never).map( + (e) => e.name, + ); + + beforeAll(async () => { + // Intercept the ESM loader so forUser() builds against the TS-source shared + // specs (no @docmost/mcp build) and never touches the network. + jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue({ + DocmostClient: function () { + return {} as DocmostClientLike; + } as unknown as loader.DocmostClientCtor, + sharedToolSpecs: SHARED_TOOL_SPECS as Record, + }); + const service = new AiChatToolsService( + { + generateAccessToken: jest.fn().mockResolvedValue('access-token'), + generateCollabToken: jest.fn().mockResolvedValue('collab-token'), + } as never, + {} as never, // aiService — not exercised while merely BUILDING the tools + {} as never, // pageEmbeddingRepo + {} as never, // spaceMemberRepo + {} as never, // pagePermissionRepo + // sandboxStore: forUser() eagerly calls asSink() to wire the stash tool. + { + asSink: () => ({ put: jest.fn(), has: jest.fn(), evict: jest.fn() }), + } as never, + ); + const tools = await service.forUser( + { id: 'user-1', email: 'u@example.com', workspaceId: 'ws-1' } as never, + 'session-1', + 'ws-1', + 'chat-1', + ); + toolKeys = Object.keys(tools); + }); + + afterAll(() => { + jest.restoreAllMocks(); + }); + + it('exposes a non-trivial toolset (sanity: the mock actually built tools)', () => { + expect(toolKeys.length).toBeGreaterThan(20); + }); + + it('every non-core live tool is present in the catalog (no capability silently hidden)', () => { + // forUser() does not itself add loadTools (ai-chat.service does), but guard + // anyway. Every remaining non-core key MUST have a catalog line. + const catalogSet = new Set(catalogNames); + const missing = toolKeys.filter( + (k) => !CORE_TOOL_SET.has(k) && k !== LOAD_TOOLS_NAME && !catalogSet.has(k), + ); + expect(missing).toEqual([]); + }); + + it('every catalog entry corresponds to a real, non-core live tool (no phantom)', () => { + const liveSet = new Set(toolKeys); + const phantom = catalogNames.filter( + (n) => !liveSet.has(n) || CORE_TOOL_SET.has(n), + ); + expect(phantom).toEqual([]); + }); +}); + describe('buildExternalToolCatalog + shortenForCatalog (#332)', () => { it('derives a short "name — purpose" line from each external tool description', () => { const catalog = buildExternalToolCatalog({ diff --git a/apps/server/test/integration/ai-chat-stream.int-spec.ts b/apps/server/test/integration/ai-chat-stream.int-spec.ts index 17deac63..a226e864 100644 --- a/apps/server/test/integration/ai-chat-stream.int-spec.ts +++ b/apps/server/test/integration/ai-chat-stream.int-spec.ts @@ -1,5 +1,7 @@ import * as http from 'node:http'; import { Kysely } from 'kysely'; +import { tool } from 'ai'; +import { z } from 'zod'; import { MockLanguageModelV3, convertArrayToReadableStream } from 'ai/test'; import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; @@ -318,4 +320,174 @@ describe('AiChatService.stream [integration]', () => { true, ); }); + + /** + * #332 deferred tool loading, the ON path. The riskiest property is that the + * per-turn `activatedTools` Set is created FRESH inside each stream() call, so a + * tool a previous turn activated via loadTools is NOT still active when the next + * turn starts — the new turn begins "cold" (CORE + loadTools only). The unit + * tests only exercise pure prepareAgentStep with hand-fed Sets; this pins the + * real wiring end-to-end (loadTools.execute -> activatedTools -> prepareStep -> + * per-step activeTools) against the real streamText loop, and proves there is no + * cross-turn leak. We drive a MockLanguageModelV3 whose step 1 calls + * loadTools(['createPage']) and assert, via the model's recorded per-step + * CallOptions.tools (the AI SDK filters the provider tool list by activeTools), + * that the deferred tool becomes active on the SAME turn's next step but NOT on a + * fresh turn's first step. + */ + describe('deferred tool loading ON — per-turn activation, no leak (#332)', () => { + // A stub deferred (non-core) tool the agent can activate. Its execute is never + // called — the model only needs to SEE it become active — but it must be a + // valid AI-SDK tool so the SDK includes it in a step's tool list once active. + const createPageStub = tool({ + description: 'create a new page', + inputSchema: z.object({ title: z.string() }), + execute: async () => ({ id: 'p-stub' }), + }); + + // A CORE tool in the toolset, so a cold step shows CORE tools ARE active while + // the deferred createPage is not. `searchPages` is in CORE_TOOL_SET. + const searchPagesStub = tool({ + description: 'search the wiki', + inputSchema: z.object({ query: z.string() }), + execute: async () => [], + }); + + // Same lifecycle harness as buildService() above, but with deferred loading ON + // and a toolset that exposes exactly one deferred tool (createPage) so it is + // catalogued + loadable-by-name. Kept separate so the OFF scenarios are + // untouched. + function buildDeferredService(): AiChatService { + return new AiChatService( + { getChatModel: async () => null } as any, + aiChatRepo, + msgRepo, + {} as any, + { resolve: async () => null } as any, + { + forUser: async () => ({ + searchPages: searchPagesStub, + createPage: createPageStub, + }), + getInAppDeferredCatalog: async () => [ + { name: 'createPage', catalogLine: 'createPage — create a new page.' }, + ], + } as any, + mcpClients as any, + {} as any, + {} as any, + {} as any, + // #332: deferred tool loading ON — the property under test. + { isAiChatDeferredToolsEnabled: () => true } as any, + ); + } + + // Drive ONE stream() turn against `model` and wait for the assistant row to + // settle (mirrors runStream, but builds the deferred-ON service). + async function runDeferredTurn( + model: MockLanguageModelV3, + chatId: string, + body: any, + ): Promise { + closeCalls = 0; + const service = buildDeferredService(); + const { res, cleanup } = await makeRealResponse(); + try { + await service.stream({ + user: { id: userId, workspaceId } as any, + workspace: { id: workspaceId, name: 'WS' } as any, + sessionId: 'sess-1', + body, + res: { raw: res } as any, + signal: new AbortController().signal, + model: model as any, + role: null, + } as any); + await waitFor(async () => { + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + return rows.some( + (r) => + r.role === 'assistant' && + ['completed', 'error', 'aborted'].includes(r.status as string), + ); + }); + await waitFor(() => closeCalls > 0, { timeoutMs: 5_000 }); + } finally { + await cleanup(); + } + } + + // Tool names the provider actually received for a recorded step (activeTools + // filters this list, so it reflects what was active that step). + const toolNames = (call: any): string[] => + ((call?.tools ?? []) as any[]).map((t) => t?.name).filter(Boolean); + + // A model that, on step 1, calls loadTools(['createPage']); on step 2, answers. + function loadThenAnswerModel(): MockLanguageModelV3 { + let step = 0; + return new MockLanguageModelV3({ + doStream: async () => { + const n = step++; + if (n === 0) { + return { + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { + type: 'tool-call', + toolCallId: 'lt1', + toolName: 'loadTools', + input: JSON.stringify({ names: ['createPage'] }), + }, + { + type: 'finish', + finishReason: 'tool-calls', + usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 }, + }, + ] as any), + }; + } + return { stream: successStream() }; + }, + } as any); + } + + it('activates a deferred tool for the SAME turn, and a NEW turn starts cold (no leak)', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + + // --- Turn 1: loadTools(createPage) on step 1, then answer on step 2. --- + const model1 = loadThenAnswerModel(); + await runDeferredTurn(model1, chatId, { + chatId, + messages: [userUiMessage('Make me a page')], + }); + + // The turn ran at least two steps (the load round-trip + the answer). + expect(model1.doStreamCalls.length).toBeGreaterThanOrEqual(2); + const step1Tools = toolNames(model1.doStreamCalls[0]); + const step2Tools = toolNames(model1.doStreamCalls[1]); + + // Step 1 starts cold: CORE tools + the loadTools meta-tool are active, but + // the deferred createPage is NOT yet. + expect(step1Tools).toContain('loadTools'); + expect(step1Tools).toContain('searchPages'); // a CORE tool, always active + expect(step1Tools).not.toContain('createPage'); + // Step 2 of the SAME turn sees the just-activated deferred tool. + expect(step2Tools).toContain('createPage'); + + // --- Turn 2 on the SAME chat: must start cold again. --- + const model2 = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + await runDeferredTurn(model2, chatId, { + chatId, + messages: [userUiMessage('And another thing')], + }); + + const nextTurnFirstStep = toolNames(model2.doStreamCalls[0]); + expect(nextTurnFirstStep).toContain('loadTools'); + // The activated set is per-turn: the prior turn's createPage did NOT leak, + // so the fresh turn's first step sees it deferred again. + expect(nextTurnFirstStep).not.toContain('createPage'); + }); + }); });