test(ai-chat): document AI_CHAT_DEFERRED_TOOLS + pin ON-path & catalog completeness (#341 review F1-F3)
- F1: document AI_CHAT_DEFERRED_TOOLS in .env.example (AI_* section) — default ON = deferred loading (compact catalog + loadTools), =false restores the old "all tools always active" behavior. - F2: integration test of the ON path in ai-chat-stream.int-spec.ts — a deferred tool activated via loadTools is active on the SAME turn's next step but a fresh turn starts cold (CORE + loadTools only), proving the per-turn activatedTools Set does not leak across turns/chats. Drives the real streamText loop with a MockLanguageModelV3 and inspects recorded per-step activeTools-filtered tools. - F3: replace the magic toHaveLength(28) in tool-tiers.spec.ts with a two-way partition against the LIVE in-app toolset (AiChatToolsService.forUser keys): every non-core tool must appear in buildInAppDeferredCatalog and every catalog entry must map to a real non-core tool — so a future tool forgotten in INLINE_TOOL_TIERS fails the suite instead of silently vanishing from the agent. No production logic change (mechanism was already reviewed correct). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
agent_coder
@@ -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 <tool_catalog> 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
|
||||
|
||||
@@ -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 <tool_catalog> 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<string, loader.SharedToolSpec>,
|
||||
});
|
||||
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({
|
||||
|
||||
@@ -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<void> {
|
||||
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');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user