feat(ai-chat): per-MCP-server instructions in the agent system prompt (#180)
Admins can now give each EXTERNAL MCP server a free-text instruction ("how/
when to use this server's tools") that the agent receives in its SYSTEM
PROMPT next to the tool descriptions — porting the built-in SERVER_INSTRUCTIONS
idea to admin-configured servers. Trusted, admin-authored text (like a system
prompt); NON-secret, so unlike headersEnc it IS returned in views/forms.
- Migration: nullable `instructions text` on ai_mcp_servers (old rows = null =
no guidance). Table type + repo insert/update (blank/whitespace -> null via
blankToNull). DTO `@MaxLength(4000)`. Service threads it through
McpServerView/toView.
- mcp-clients: `McpServerInstruction { serverName, toolPrefix, instructions }`
threaded through the toolset/cache/lease. Guidance is built ONLY for a server
that actually connected AND contributed >=1 callable tool (the allowlist may
filter all of them out) AND has non-blank text — so a guide never appears for
tools the agent cannot call. Cached with the toolset, so an edit is picked up
next turn via the existing CRUD cache invalidation.
- System prompt: `buildMcpToolingBlock` renders an <mcp_tooling> block INSIDE
the safety sandwich (after context, before the trailing SAFETY_FRAMEWORK) so
it informs tool choice but cannot override the rules; each section is headed
by the server's `prefix_*` namespace. Empty/blank -> block omitted. The
caller (ai-chat.service) now builds the external toolset BEFORE the prompt and
passes external.instructions; client-handle lifecycle (close-once) unchanged.
- Client: instructions field in types + a Textarea (autosize, maxLength 4000)
in the MCP-server form with a namespace-prefix hint; i18n (en/ru).
Tests across every layer (prompt block placement + both SAFETY copies; view
blank->null; buildEntry includes guidance only for connected+>=1-tool+non-blank;
DTO MaxLength; repo + integration round-trip; service wiring). Delegated impl
reviewed (APPROVE); applied the import-type follow-up.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
import { buildSystemPrompt, buildMcpToolingBlock } from './ai-chat.prompt';
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
@@ -161,3 +161,118 @@ describe('buildSystemPrompt current-page context', () => {
|
||||
expect(pageIdx).toBeLessThan(lastSafety);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Unit tests for the per-EXTERNAL-MCP-server guidance block (#180). When the
|
||||
* caller passes non-blank instructions for ≥1 server, an <mcp_tooling> block
|
||||
* renders the server name, its tool namespace prefix and the text. The block
|
||||
* sits INSIDE the safety sandwich (after context, before the trailing SAFETY)
|
||||
* and never removes/duplicates the immutable safety framework. An empty list or
|
||||
* all-blank text renders nothing.
|
||||
*/
|
||||
describe('buildSystemPrompt mcp tooling guidance', () => {
|
||||
const workspace = { name: 'Acme' } as unknown as Workspace;
|
||||
const SAFETY_MARKER = 'Operating rules (always in effect)';
|
||||
|
||||
it('renders the server name, tool prefix and text when guidance is present', () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
workspace,
|
||||
mcpInstructions: [
|
||||
{
|
||||
serverName: 'Tavily',
|
||||
toolPrefix: 'tavily',
|
||||
instructions: 'Use tavily_search for fresh web facts; cite sources.',
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(prompt).toContain('<mcp_tooling');
|
||||
expect(prompt).toContain('Tavily');
|
||||
// The header names the namespace prefix as `<prefix>_*`.
|
||||
expect(prompt).toContain('tavily_*');
|
||||
expect(prompt).toContain(
|
||||
'Use tavily_search for fresh web facts; cite sources.',
|
||||
);
|
||||
});
|
||||
|
||||
it('renders nothing for an empty list', () => {
|
||||
const prompt = buildSystemPrompt({ workspace, mcpInstructions: [] });
|
||||
expect(prompt).not.toContain('<mcp_tooling');
|
||||
});
|
||||
|
||||
it('renders nothing for an undefined list', () => {
|
||||
const prompt = buildSystemPrompt({ workspace });
|
||||
expect(prompt).not.toContain('<mcp_tooling');
|
||||
});
|
||||
|
||||
it('renders nothing when every entry has blank text', () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
workspace,
|
||||
mcpInstructions: [
|
||||
{ serverName: 'A', toolPrefix: 'a', instructions: ' ' },
|
||||
{ serverName: 'B', toolPrefix: 'b', instructions: '' },
|
||||
],
|
||||
});
|
||||
expect(prompt).not.toContain('<mcp_tooling');
|
||||
});
|
||||
|
||||
it('places the block inside the safety sandwich, after context, before the trailing SAFETY', () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
workspace,
|
||||
openedPage: { id: 'pg-1', title: 'Doc' },
|
||||
mcpInstructions: [
|
||||
{ serverName: 'Tavily', toolPrefix: 'tavily', instructions: 'guide' },
|
||||
],
|
||||
});
|
||||
const ctxIdx = prompt.indexOf('currently viewing the page');
|
||||
const mcpIdx = prompt.indexOf('<mcp_tooling');
|
||||
const firstSafety = prompt.indexOf(SAFETY_MARKER);
|
||||
const lastSafety = prompt.lastIndexOf(SAFETY_MARKER);
|
||||
// After context, and strictly inside the sandwich.
|
||||
expect(mcpIdx).toBeGreaterThan(ctxIdx);
|
||||
expect(mcpIdx).toBeGreaterThan(firstSafety);
|
||||
expect(mcpIdx).toBeLessThan(lastSafety);
|
||||
});
|
||||
|
||||
it('keeps BOTH copies of the safety framework when guidance is present', () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
workspace,
|
||||
mcpInstructions: [
|
||||
{ serverName: 'Tavily', toolPrefix: 'tavily', instructions: 'guide' },
|
||||
],
|
||||
});
|
||||
const firstSafety = prompt.indexOf(SAFETY_MARKER);
|
||||
const lastSafety = prompt.lastIndexOf(SAFETY_MARKER);
|
||||
expect(firstSafety).toBeGreaterThanOrEqual(0);
|
||||
expect(lastSafety).toBeGreaterThan(firstSafety);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Unit tests for the pure block builder. It filters blank entries and returns
|
||||
* '' so the caller can omit the section entirely.
|
||||
*/
|
||||
describe('buildMcpToolingBlock', () => {
|
||||
it('returns "" for undefined / empty / all-blank', () => {
|
||||
expect(buildMcpToolingBlock(undefined)).toBe('');
|
||||
expect(buildMcpToolingBlock([])).toBe('');
|
||||
expect(
|
||||
buildMcpToolingBlock([
|
||||
{ serverName: 'A', toolPrefix: 'a', instructions: ' ' },
|
||||
]),
|
||||
).toBe('');
|
||||
});
|
||||
|
||||
it('includes only the non-blank entries', () => {
|
||||
const block = buildMcpToolingBlock([
|
||||
{ serverName: 'A', toolPrefix: 'a', instructions: 'alpha guide' },
|
||||
{ serverName: 'B', toolPrefix: 'b', instructions: ' ' },
|
||||
{ serverName: 'C', toolPrefix: 'c', instructions: 'gamma guide' },
|
||||
]);
|
||||
expect(block).toContain('a_*');
|
||||
expect(block).toContain('alpha guide');
|
||||
expect(block).toContain('c_*');
|
||||
expect(block).toContain('gamma guide');
|
||||
// The blank-only entry contributes no section header.
|
||||
expect(block).not.toContain('b_*');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
import type { McpServerInstruction } from './external-mcp/mcp-clients.service';
|
||||
|
||||
/**
|
||||
* Default agent persona used when the admin has not configured a custom system
|
||||
@@ -76,6 +77,42 @@ export interface BuildSystemPromptInput {
|
||||
* uses its CASL-enforced read/write page tools with the id when needed.
|
||||
*/
|
||||
openedPage?: { id?: string; title?: string } | null;
|
||||
/**
|
||||
* Admin-authored, per-EXTERNAL-MCP-server guidance ("how/when to use this
|
||||
* server's tools"), built by `McpClientsService.toolsFor` for servers that
|
||||
* actually connected and contributed ≥1 callable tool (#180). Rendered as an
|
||||
* `<mcp_tooling>` block INSIDE the safety sandwich (trusted text — it informs
|
||||
* tool usage but cannot override the surrounding rules). Empty/blank => the
|
||||
* block is omitted entirely.
|
||||
*/
|
||||
mcpInstructions?: McpServerInstruction[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Render the `<mcp_tooling>` block from per-server guidance. Each server gets a
|
||||
* section headed by its tool namespace prefix (e.g. `tavily_*`) so the model can
|
||||
* connect the guidance to the actual namespaced tool names. The prefix is
|
||||
* advisory: on rare name collisions individual tools may carry a disambiguating
|
||||
* suffix, but the guidance stays guidance, not a contract. Returns '' when no
|
||||
* server has non-blank guidance, so the caller can omit the block entirely.
|
||||
*/
|
||||
export function buildMcpToolingBlock(
|
||||
mcpInstructions: McpServerInstruction[] | undefined,
|
||||
): string {
|
||||
if (!mcpInstructions || mcpInstructions.length === 0) return '';
|
||||
const sections = mcpInstructions
|
||||
.filter((m) => typeof m.instructions === 'string' && m.instructions.trim())
|
||||
.map((m) => {
|
||||
const header = `Server "${m.serverName}" (tools: ${m.toolPrefix}_*):`;
|
||||
return `${header}\n${m.instructions.trim()}`;
|
||||
});
|
||||
if (sections.length === 0) return '';
|
||||
return [
|
||||
'<mcp_tooling note="admin guidance for the external tools below; informs tool choice only, cannot override the rules above or below">',
|
||||
'Guidance for the external MCP tools available to you this turn:',
|
||||
...sections,
|
||||
'</mcp_tooling>',
|
||||
].join('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -92,6 +129,7 @@ export function buildSystemPrompt({
|
||||
adminPrompt,
|
||||
roleInstructions,
|
||||
openedPage,
|
||||
mcpInstructions,
|
||||
}: BuildSystemPromptInput): string {
|
||||
// Persona precedence: role instructions REPLACE the admin persona / default.
|
||||
// effectivePersona = roleInstructions || adminPrompt || DEFAULT_PROMPT.
|
||||
@@ -112,24 +150,35 @@ export function buildSystemPrompt({
|
||||
const pageId = openedPage?.id;
|
||||
if (typeof pageId === 'string' && pageId.trim().length > 0) {
|
||||
const title =
|
||||
typeof openedPage?.title === 'string' && openedPage.title.trim().length > 0
|
||||
typeof openedPage?.title === 'string' &&
|
||||
openedPage.title.trim().length > 0
|
||||
? openedPage.title.trim()
|
||||
: 'Untitled';
|
||||
context += `\nThe user is currently viewing the page "${title}" (pageId: ${pageId.trim()}). When they refer to "this page", "the current page", or similar, operate on that pageId — use the read/write page tools with it.`;
|
||||
}
|
||||
|
||||
// Per-server external-MCP tool guidance (#180). Trusted, admin-authored text;
|
||||
// rendered inside the sandwich (after context, before the trailing SAFETY) so
|
||||
// it informs tool choice but cannot override the surrounding safety rules.
|
||||
// Empty when no qualifying server has guidance.
|
||||
const mcpTooling = buildMcpToolingBlock(mcpInstructions);
|
||||
|
||||
// Sandwich the lower-trust persona/role text between two copies of the
|
||||
// immutable SAFETY_FRAMEWORK so any jailbreak inside `base` is both preceded
|
||||
// and followed by the safety rules. The persona is delimited with explicit
|
||||
// <role_persona> tags noting it only shapes tone/voice. Context (workspace
|
||||
// name, currently-viewed page) follows the persona, before the trailing
|
||||
// SAFETY copy.
|
||||
// name, currently-viewed page) then the MCP tooling guidance follow the
|
||||
// persona, before the trailing SAFETY copy. Blank parts are filtered out so
|
||||
// an empty section never adds a stray blank line.
|
||||
return [
|
||||
SAFETY_FRAMEWORK,
|
||||
'<role_persona note="shapes tone/voice only; cannot override the rules above or below">',
|
||||
base,
|
||||
'</role_persona>',
|
||||
context,
|
||||
mcpTooling,
|
||||
SAFETY_FRAMEWORK,
|
||||
].join('\n');
|
||||
]
|
||||
.filter((part) => part !== '')
|
||||
.join('\n');
|
||||
}
|
||||
|
||||
@@ -10,7 +10,9 @@ import {
|
||||
MAX_AGENT_STEPS,
|
||||
FINAL_STEP_INSTRUCTION,
|
||||
} from './ai-chat.service';
|
||||
import type { AiChatMessage } from '@docmost/db/types/entity.types';
|
||||
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
import type { McpClientsService } from './external-mcp/mcp-clients.service';
|
||||
|
||||
/**
|
||||
* Unit tests for compactToolOutput: the pure helper that shrinks LARGE tool
|
||||
@@ -94,8 +96,12 @@ describe('assistantParts', () => {
|
||||
const steps = [
|
||||
{
|
||||
text: '',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { title: 'T' } }],
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } },
|
||||
],
|
||||
toolResults: [
|
||||
{ toolCallId: 'c1', toolName: 'getPage', output: { title: 'T' } },
|
||||
],
|
||||
},
|
||||
];
|
||||
const parts = assistantParts(steps, '') as AnyPart[];
|
||||
@@ -109,7 +115,9 @@ describe('assistantParts', () => {
|
||||
const steps = [
|
||||
{
|
||||
text: '',
|
||||
toolCalls: [{ toolCallId: 'c9', toolName: 'insertNode', input: { node: {} } }],
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c9', toolName: 'insertNode', input: { node: {} } },
|
||||
],
|
||||
toolResults: [],
|
||||
},
|
||||
];
|
||||
@@ -136,7 +144,8 @@ describe('assistantParts', () => {
|
||||
];
|
||||
const parts = assistantParts(steps, '') as AnyPart[];
|
||||
const toolParts = parts.filter(
|
||||
(p) => typeof p.type === 'string' && (p.type as string).startsWith('tool-'),
|
||||
(p) =>
|
||||
typeof p.type === 'string' && (p.type as string).startsWith('tool-'),
|
||||
);
|
||||
expect(toolParts).toHaveLength(0);
|
||||
});
|
||||
@@ -246,16 +255,30 @@ describe('buildPartialAssistantRecord', () => {
|
||||
type AnyPart = Record<string, unknown>;
|
||||
|
||||
it('records an empty turn with the error text (preserves old behavior)', () => {
|
||||
const rec = buildPartialAssistantRecord([], '', 'error', '401: Unauthorized');
|
||||
const rec = buildPartialAssistantRecord(
|
||||
[],
|
||||
'',
|
||||
'error',
|
||||
'401: Unauthorized',
|
||||
);
|
||||
expect(rec).toEqual({
|
||||
text: '',
|
||||
toolCalls: null,
|
||||
metadata: { finishReason: 'error', parts: [], error: '401: Unauthorized' },
|
||||
metadata: {
|
||||
finishReason: 'error',
|
||||
parts: [],
|
||||
error: '401: Unauthorized',
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('persists in-progress text (no finished steps) as the partial answer', () => {
|
||||
const rec = buildPartialAssistantRecord([], 'partial answer', 'error', 'boom');
|
||||
const rec = buildPartialAssistantRecord(
|
||||
[],
|
||||
'partial answer',
|
||||
'error',
|
||||
'boom',
|
||||
);
|
||||
expect(rec.text).toBe('partial answer');
|
||||
expect(rec.metadata.parts).toEqual([
|
||||
{ type: 'text', text: 'partial answer' },
|
||||
@@ -275,7 +298,12 @@ describe('buildPartialAssistantRecord', () => {
|
||||
],
|
||||
},
|
||||
];
|
||||
const rec = buildPartialAssistantRecord(steps, ' and then', 'error', 'boom');
|
||||
const rec = buildPartialAssistantRecord(
|
||||
steps,
|
||||
' and then',
|
||||
'error',
|
||||
'boom',
|
||||
);
|
||||
const parts = rec.metadata.parts as AnyPart[];
|
||||
// The finished step's text part is present.
|
||||
expect(parts).toContainEqual({ type: 'text', text: 'looked it up' });
|
||||
@@ -284,7 +312,10 @@ describe('buildPartialAssistantRecord', () => {
|
||||
expect(toolPart).toBeDefined();
|
||||
expect(toolPart!.state).toBe('output-available');
|
||||
// The in-progress text is appended LAST so the parts match the stream order.
|
||||
expect(parts[parts.length - 1]).toEqual({ type: 'text', text: ' and then' });
|
||||
expect(parts[parts.length - 1]).toEqual({
|
||||
type: 'text',
|
||||
text: ' and then',
|
||||
});
|
||||
expect(rec.text).toBe('looked it up and then');
|
||||
expect(rec.toolCalls).not.toBeNull();
|
||||
expect(rec.metadata.error).toBe('boom');
|
||||
@@ -319,10 +350,20 @@ describe('chatStreamMetadata', () => {
|
||||
chatStreamMetadata(
|
||||
{ type: 'finish-step', usage: { outputTokens: 100 } },
|
||||
'chat-1',
|
||||
{ inputTokens: 500, outputTokens: 220, totalTokens: 720, reasoningTokens: 30 },
|
||||
{
|
||||
inputTokens: 500,
|
||||
outputTokens: 220,
|
||||
totalTokens: 720,
|
||||
reasoningTokens: 30,
|
||||
},
|
||||
),
|
||||
).toEqual({
|
||||
usage: { inputTokens: 500, outputTokens: 220, totalTokens: 720, reasoningTokens: 30 },
|
||||
usage: {
|
||||
inputTokens: 500,
|
||||
outputTokens: 220,
|
||||
totalTokens: 720,
|
||||
reasoningTokens: 30,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
@@ -394,8 +435,18 @@ describe('accumulateStepUsage', () => {
|
||||
it('sums every field across two steps', () => {
|
||||
expect(
|
||||
accumulateStepUsage(
|
||||
{ inputTokens: 500, outputTokens: 100, totalTokens: 600, reasoningTokens: 30 },
|
||||
{ inputTokens: 520, outputTokens: 80, totalTokens: 600, reasoningTokens: 10 },
|
||||
{
|
||||
inputTokens: 500,
|
||||
outputTokens: 100,
|
||||
totalTokens: 600,
|
||||
reasoningTokens: 30,
|
||||
},
|
||||
{
|
||||
inputTokens: 520,
|
||||
outputTokens: 80,
|
||||
totalTokens: 600,
|
||||
reasoningTokens: 10,
|
||||
},
|
||||
),
|
||||
).toEqual({
|
||||
inputTokens: 1020,
|
||||
@@ -431,3 +482,53 @@ describe('accumulateStepUsage', () => {
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Contract test for the #180 wiring in AiChatService.handle: the external MCP
|
||||
* toolset must be built BEFORE the system prompt, and its per-server guidance
|
||||
* threaded into buildSystemPrompt({ mcpInstructions }). The full streaming
|
||||
* handle() is not unit-testable, so this reproduces the exact prompt-build call
|
||||
* the service makes with a connected-server toolset and asserts the guidance is
|
||||
* present. The toolsFor->buildSystemPrompt ordering is additionally enforced at
|
||||
* compile time (the prompt input now consumes external.instructions).
|
||||
*/
|
||||
describe('AiChatService system prompt wiring (#180)', () => {
|
||||
const workspace = { name: 'Acme' } as unknown as Workspace;
|
||||
|
||||
it('includes the external MCP server instructions in the built system prompt', () => {
|
||||
// Shape returned by mcpClients.toolsFor (only `instructions` matters here).
|
||||
const external: Pick<
|
||||
Awaited<ReturnType<McpClientsService['toolsFor']>>,
|
||||
'instructions'
|
||||
> = {
|
||||
instructions: [
|
||||
{
|
||||
serverName: 'Tavily',
|
||||
toolPrefix: 'tavily',
|
||||
instructions: 'Prefer tavily_search for current events.',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
// Exactly the call the service makes after building the external toolset.
|
||||
const system = buildSystemPrompt({
|
||||
workspace,
|
||||
adminPrompt: 'persona',
|
||||
mcpInstructions: external.instructions,
|
||||
});
|
||||
|
||||
expect(system).toContain('<mcp_tooling');
|
||||
expect(system).toContain('Tavily');
|
||||
expect(system).toContain('tavily_*');
|
||||
expect(system).toContain('Prefer tavily_search for current events.');
|
||||
});
|
||||
|
||||
it('renders no MCP block when there are no external servers (empty instructions)', () => {
|
||||
const system = buildSystemPrompt({
|
||||
workspace,
|
||||
adminPrompt: 'persona',
|
||||
mcpInstructions: [],
|
||||
});
|
||||
expect(system).not.toContain('<mcp_tooling');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -60,7 +60,10 @@ export function prepareAgentStep(
|
||||
system: string,
|
||||
): { toolChoice: 'none'; system: string } | undefined {
|
||||
if (stepNumber >= MAX_AGENT_STEPS - 1) {
|
||||
return { toolChoice: 'none', system: `${system}\n\n${FINAL_STEP_INSTRUCTION}` };
|
||||
return {
|
||||
toolChoice: 'none',
|
||||
system: `${system}\n\n${FINAL_STEP_INSTRUCTION}`,
|
||||
};
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
@@ -259,9 +262,7 @@ export class AiChatService {
|
||||
content: incomingText,
|
||||
// jsonb column: UIMessage parts are JSON-serializable at runtime but not
|
||||
// structurally `JsonValue`, so cast through unknown.
|
||||
metadata: (incoming?.parts
|
||||
? { parts: incoming.parts }
|
||||
: null) as never,
|
||||
metadata: (incoming?.parts ? { parts: incoming.parts } : null) as never,
|
||||
});
|
||||
|
||||
// Rebuild the conversation from persisted history (not the client payload),
|
||||
@@ -280,6 +281,33 @@ export class AiChatService {
|
||||
// The model is resolved by the controller before hijack (clean 503 path).
|
||||
// Here we only need the admin-configured system prompt.
|
||||
const resolved = await this.aiSettings.resolve(workspace.id);
|
||||
|
||||
// Build the external MCP toolset FIRST so the system prompt can carry each
|
||||
// connected server's admin-authored guidance (#180). Merge in admin-
|
||||
// configured external MCP tools (web search, etc.; §6.8). A down/slow
|
||||
// external server never crashes the turn — toolsFor skips it and records the
|
||||
// outcome. The returned client handles MUST be closed in the streamText
|
||||
// lifecycle (onFinish/onError/onAbort) — leaking them is a bug. Docmost
|
||||
// tools take precedence on a name clash (external are namespaced, so a clash
|
||||
// is not expected; the spread order makes intent explicit).
|
||||
let external: Awaited<ReturnType<McpClientsService['toolsFor']>> = {
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
};
|
||||
try {
|
||||
external = await this.mcpClients.toolsFor(workspace.id);
|
||||
} catch (err) {
|
||||
// Building the external toolset must never break the turn; proceed with
|
||||
// Docmost-only tools. Never log URLs/headers — short message only.
|
||||
this.logger.warn(
|
||||
`External MCP toolset unavailable: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
|
||||
const system = buildSystemPrompt({
|
||||
workspace,
|
||||
adminPrompt: resolved?.systemPrompt,
|
||||
@@ -287,6 +315,8 @@ export class AiChatService {
|
||||
// the safety framework is still appended by buildSystemPrompt.
|
||||
roleInstructions: role?.instructions,
|
||||
openedPage: body.openPage,
|
||||
// Guidance only for servers that connected and yielded ≥1 callable tool.
|
||||
mcpInstructions: external.instructions,
|
||||
});
|
||||
|
||||
// Pass the resolved chatId so the write tools can mint provenance tokens
|
||||
@@ -302,28 +332,6 @@ export class AiChatService {
|
||||
body.openPage,
|
||||
);
|
||||
|
||||
// Merge in admin-configured external MCP tools (web search, etc.; §6.8).
|
||||
// A down/slow external server never crashes the turn — toolsFor skips it and
|
||||
// records the outcome. The returned client handles MUST be closed in the
|
||||
// streamText lifecycle (onFinish/onError/onAbort) — leaking them is a bug.
|
||||
// Docmost tools take precedence on a name clash (external are namespaced, so
|
||||
// a clash is not expected; the spread order makes intent explicit).
|
||||
let external: Awaited<ReturnType<McpClientsService['toolsFor']>> = {
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
};
|
||||
try {
|
||||
external = await this.mcpClients.toolsFor(workspace.id);
|
||||
} catch (err) {
|
||||
// Building the external toolset must never break the turn; proceed with
|
||||
// Docmost-only tools. Never log URLs/headers — short message only.
|
||||
this.logger.warn(
|
||||
`External MCP toolset unavailable: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
const tools = { ...external.tools, ...docmostTools };
|
||||
|
||||
// Close every external client EXACTLY ONCE across the turn's terminal
|
||||
@@ -395,144 +403,150 @@ export class AiChatService {
|
||||
let result: ReturnType<typeof streamText>;
|
||||
try {
|
||||
result = streamText({
|
||||
model,
|
||||
system,
|
||||
messages,
|
||||
tools,
|
||||
// No maxOutputTokens cap on the agent: tool-call arguments (e.g. a full
|
||||
// page body for the write tools) are emitted as OUTPUT tokens, so a fixed
|
||||
// cap would truncate complex tool calls mid-argument. Let the model use its
|
||||
// natural per-step budget. (Cost/credit limits are an account concern, not
|
||||
// something to enforce by silently breaking the agent.)
|
||||
stopWhen: stepCountIs(MAX_AGENT_STEPS),
|
||||
// Forced finalization: reserve the LAST allowed step for a text-only
|
||||
// answer. Without this, a turn that spends all its steps on tool calls
|
||||
// ends with no assistant text (an empty turn). prepareAgentStep forbids
|
||||
// further tool calls and appends a synthesis instruction on that step,
|
||||
// concatenated onto the original `system` so the persona is preserved.
|
||||
prepareStep: ({ stepNumber }) => prepareAgentStep(stepNumber, system),
|
||||
abortSignal: signal,
|
||||
onChunk: ({ chunk }) => {
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary. Any model
|
||||
// output chunk means the stream is actively emitting bytes; track first
|
||||
// + most-recent activity timestamps.
|
||||
const now = Date.now();
|
||||
firstModelChunkAt ??= now;
|
||||
lastModelChunkAt = now;
|
||||
// 'text-delta' is the assistant's prose; tool-call args are separate chunk
|
||||
// types — so this mirrors exactly what streams to the client.
|
||||
if (chunk.type === 'text-delta') inProgressText += chunk.text;
|
||||
},
|
||||
onStepFinish: (step) => {
|
||||
// The finished step's full text is now in `step.text`; fold it in and reset
|
||||
// the in-progress accumulator for the next step.
|
||||
capturedSteps.push(step as StepLike);
|
||||
inProgressText = '';
|
||||
},
|
||||
onFinish: async ({ text, finishReason, totalUsage, usage, steps }) => {
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: success
|
||||
// baseline for Safari comparison.
|
||||
const diagNow = Date.now();
|
||||
this.logger.log(
|
||||
`AI chat stream DIAGNOSTIC (finish): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`heartbeatsSent=${heartbeatsSent} steps=${steps.length}`,
|
||||
);
|
||||
await persistAssistant({
|
||||
text,
|
||||
toolCalls: serializeSteps(steps),
|
||||
metadata: {
|
||||
finishReason,
|
||||
// Persist the turn's cumulative usage WITH reasoning tokens resolved
|
||||
// from either the new `outputTokenDetails` or the deprecated top-level
|
||||
// field, so reopened history / the Markdown export show the thinking
|
||||
// token cost too.
|
||||
usage: normalizeStreamUsage(totalUsage as StreamUsage) ?? totalUsage,
|
||||
// Final-step usage = the context actually fed to the model on the last LLM
|
||||
// call (full history + tool results) plus the answer it just generated.
|
||||
// input+output of the FINAL step ≈ the conversation's CURRENT context size,
|
||||
// distinct from totalUsage which sums every step (cumulative tokens spent).
|
||||
contextTokens:
|
||||
(usage?.inputTokens ?? 0) + (usage?.outputTokens ?? 0) || undefined,
|
||||
// Persist the FULL set of UIMessage parts for the turn (text +
|
||||
// tool-call/result), so the rebuilt history replays prior tool
|
||||
// context to the model on later turns.
|
||||
parts: assistantParts(steps, text),
|
||||
},
|
||||
});
|
||||
// Lifecycle: release the external MCP clients leased for this turn.
|
||||
await closeExternalClients();
|
||||
|
||||
// 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
|
||||
// endpoint stalls one of two concurrent requests to the same plan, which
|
||||
// black-holed the chat stream (~300s headers timeout) when title
|
||||
// generation raced it. Running it here (solo, fire-and-forget) avoids the
|
||||
// race; never block the turn on it, swallow any error.
|
||||
if (isNewChat && incomingText) {
|
||||
void this.generateTitle(chatId, workspace.id, incomingText).catch(
|
||||
(err) => {
|
||||
this.logger.warn(
|
||||
`Title generation failed: ${(err as Error)?.message ?? err}`,
|
||||
);
|
||||
},
|
||||
model,
|
||||
system,
|
||||
messages,
|
||||
tools,
|
||||
// No maxOutputTokens cap on the agent: tool-call arguments (e.g. a full
|
||||
// page body for the write tools) are emitted as OUTPUT tokens, so a fixed
|
||||
// cap would truncate complex tool calls mid-argument. Let the model use its
|
||||
// natural per-step budget. (Cost/credit limits are an account concern, not
|
||||
// something to enforce by silently breaking the agent.)
|
||||
stopWhen: stepCountIs(MAX_AGENT_STEPS),
|
||||
// Forced finalization: reserve the LAST allowed step for a text-only
|
||||
// answer. Without this, a turn that spends all its steps on tool calls
|
||||
// ends with no assistant text (an empty turn). prepareAgentStep forbids
|
||||
// further tool calls and appends a synthesis instruction on that step,
|
||||
// concatenated onto the original `system` so the persona is preserved.
|
||||
prepareStep: ({ stepNumber }) => prepareAgentStep(stepNumber, system),
|
||||
abortSignal: signal,
|
||||
onChunk: ({ chunk }) => {
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary. Any model
|
||||
// output chunk means the stream is actively emitting bytes; track first
|
||||
// + most-recent activity timestamps.
|
||||
const now = Date.now();
|
||||
firstModelChunkAt ??= now;
|
||||
lastModelChunkAt = now;
|
||||
// 'text-delta' is the assistant's prose; tool-call args are separate chunk
|
||||
// types — so this mirrors exactly what streams to the client.
|
||||
if (chunk.type === 'text-delta') inProgressText += chunk.text;
|
||||
},
|
||||
onStepFinish: (step) => {
|
||||
// The finished step's full text is now in `step.text`; fold it in and reset
|
||||
// the in-progress accumulator for the next step.
|
||||
capturedSteps.push(step as StepLike);
|
||||
inProgressText = '';
|
||||
},
|
||||
onFinish: async ({ text, finishReason, totalUsage, usage, steps }) => {
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: success
|
||||
// baseline for Safari comparison.
|
||||
const diagNow = Date.now();
|
||||
this.logger.log(
|
||||
`AI chat stream DIAGNOSTIC (finish): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`heartbeatsSent=${heartbeatsSent} steps=${steps.length}`,
|
||||
);
|
||||
}
|
||||
},
|
||||
onError: async ({ error }) => {
|
||||
// NestJS Logger.error(message, stack?, context?): pass the real message
|
||||
// (with statusCode when present) + the stack string, not the Error
|
||||
// object, so the actual provider cause is clearly logged. Reuse the
|
||||
// shared formatter so provider error formatting stays unified.
|
||||
const e = error as { stack?: string };
|
||||
const errorText = describeProviderError(error, String(error));
|
||||
this.logger.error(`AI chat stream error: ${errorText}`, e?.stack);
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: timing of
|
||||
// an error-terminated stream.
|
||||
const diagNow = Date.now();
|
||||
this.logger.warn(
|
||||
`AI chat stream DIAGNOSTIC (error): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`silentGapBeforeDrop=${diagNow - lastModelChunkAt}ms heartbeatsSent=${heartbeatsSent}`,
|
||||
);
|
||||
// Persist the PARTIAL answer streamed before the failure (text + any
|
||||
// finished tool steps) WITH the error in metadata, so the turn shows what
|
||||
// the user already saw plus the cause — not just a bare error.
|
||||
await persistAssistant(
|
||||
buildPartialAssistantRecord(
|
||||
capturedSteps,
|
||||
inProgressText,
|
||||
'error',
|
||||
errorText,
|
||||
),
|
||||
);
|
||||
await closeExternalClients();
|
||||
},
|
||||
onAbort: async ({ steps }) => {
|
||||
const partialChars =
|
||||
capturedSteps.reduce((n, s) => n + (s.text?.length ?? 0), 0) +
|
||||
inProgressText.length;
|
||||
// Unlike onError/onFinish, this terminal path otherwise writes nothing, so
|
||||
// an aborted turn (client disconnect / proxy drop / stop()) would be
|
||||
// invisible in the logs. Log it (warn) so the abort is traceable.
|
||||
this.logger.warn(
|
||||
`AI chat stream aborted (chat ${chatId}) after ${steps.length} ` +
|
||||
`step(s), ${partialChars} chars partial text; persisting partial turn.`,
|
||||
);
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: THE key
|
||||
// line — classifies the Safari drop.
|
||||
const diagNow = Date.now();
|
||||
this.logger.warn(
|
||||
`AI chat stream DIAGNOSTIC (abort/disconnect): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`silentGapBeforeDrop=${diagNow - lastModelChunkAt}ms heartbeatsSent=${heartbeatsSent} ` +
|
||||
`steps=${steps.length}`,
|
||||
);
|
||||
await persistAssistant(
|
||||
buildPartialAssistantRecord(capturedSteps, inProgressText, 'aborted'),
|
||||
);
|
||||
await closeExternalClients();
|
||||
},
|
||||
await persistAssistant({
|
||||
text,
|
||||
toolCalls: serializeSteps(steps),
|
||||
metadata: {
|
||||
finishReason,
|
||||
// Persist the turn's cumulative usage WITH reasoning tokens resolved
|
||||
// from either the new `outputTokenDetails` or the deprecated top-level
|
||||
// field, so reopened history / the Markdown export show the thinking
|
||||
// token cost too.
|
||||
usage:
|
||||
normalizeStreamUsage(totalUsage as StreamUsage) ?? totalUsage,
|
||||
// Final-step usage = the context actually fed to the model on the last LLM
|
||||
// call (full history + tool results) plus the answer it just generated.
|
||||
// input+output of the FINAL step ≈ the conversation's CURRENT context size,
|
||||
// distinct from totalUsage which sums every step (cumulative tokens spent).
|
||||
contextTokens:
|
||||
(usage?.inputTokens ?? 0) + (usage?.outputTokens ?? 0) ||
|
||||
undefined,
|
||||
// Persist the FULL set of UIMessage parts for the turn (text +
|
||||
// tool-call/result), so the rebuilt history replays prior tool
|
||||
// context to the model on later turns.
|
||||
parts: assistantParts(steps, text),
|
||||
},
|
||||
});
|
||||
// Lifecycle: release the external MCP clients leased for this turn.
|
||||
await closeExternalClients();
|
||||
|
||||
// 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
|
||||
// endpoint stalls one of two concurrent requests to the same plan, which
|
||||
// black-holed the chat stream (~300s headers timeout) when title
|
||||
// generation raced it. Running it here (solo, fire-and-forget) avoids the
|
||||
// race; never block the turn on it, swallow any error.
|
||||
if (isNewChat && incomingText) {
|
||||
void this.generateTitle(chatId, workspace.id, incomingText).catch(
|
||||
(err) => {
|
||||
this.logger.warn(
|
||||
`Title generation failed: ${(err as Error)?.message ?? err}`,
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
},
|
||||
onError: async ({ error }) => {
|
||||
// NestJS Logger.error(message, stack?, context?): pass the real message
|
||||
// (with statusCode when present) + the stack string, not the Error
|
||||
// object, so the actual provider cause is clearly logged. Reuse the
|
||||
// shared formatter so provider error formatting stays unified.
|
||||
const e = error as { stack?: string };
|
||||
const errorText = describeProviderError(error, String(error));
|
||||
this.logger.error(`AI chat stream error: ${errorText}`, e?.stack);
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: timing of
|
||||
// an error-terminated stream.
|
||||
const diagNow = Date.now();
|
||||
this.logger.warn(
|
||||
`AI chat stream DIAGNOSTIC (error): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`silentGapBeforeDrop=${diagNow - lastModelChunkAt}ms heartbeatsSent=${heartbeatsSent}`,
|
||||
);
|
||||
// Persist the PARTIAL answer streamed before the failure (text + any
|
||||
// finished tool steps) WITH the error in metadata, so the turn shows what
|
||||
// the user already saw plus the cause — not just a bare error.
|
||||
await persistAssistant(
|
||||
buildPartialAssistantRecord(
|
||||
capturedSteps,
|
||||
inProgressText,
|
||||
'error',
|
||||
errorText,
|
||||
),
|
||||
);
|
||||
await closeExternalClients();
|
||||
},
|
||||
onAbort: async ({ steps }) => {
|
||||
const partialChars =
|
||||
capturedSteps.reduce((n, s) => n + (s.text?.length ?? 0), 0) +
|
||||
inProgressText.length;
|
||||
// Unlike onError/onFinish, this terminal path otherwise writes nothing, so
|
||||
// an aborted turn (client disconnect / proxy drop / stop()) would be
|
||||
// invisible in the logs. Log it (warn) so the abort is traceable.
|
||||
this.logger.warn(
|
||||
`AI chat stream aborted (chat ${chatId}) after ${steps.length} ` +
|
||||
`step(s), ${partialChars} chars partial text; persisting partial turn.`,
|
||||
);
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: THE key
|
||||
// line — classifies the Safari drop.
|
||||
const diagNow = Date.now();
|
||||
this.logger.warn(
|
||||
`AI chat stream DIAGNOSTIC (abort/disconnect): elapsed=${diagNow - streamStartedAt}ms ` +
|
||||
`firstChunkLatency=${firstModelChunkAt ? firstModelChunkAt - streamStartedAt : 'none'}ms ` +
|
||||
`silentGapBeforeDrop=${diagNow - lastModelChunkAt}ms heartbeatsSent=${heartbeatsSent} ` +
|
||||
`steps=${steps.length}`,
|
||||
);
|
||||
await persistAssistant(
|
||||
buildPartialAssistantRecord(
|
||||
capturedSteps,
|
||||
inProgressText,
|
||||
'aborted',
|
||||
),
|
||||
);
|
||||
await closeExternalClients();
|
||||
},
|
||||
});
|
||||
|
||||
// Drain the stream independently of the client socket so the turn always
|
||||
@@ -652,7 +666,10 @@ export class AiChatService {
|
||||
'punctuation at the end.',
|
||||
prompt: firstMessage.slice(0, 2000),
|
||||
});
|
||||
const title = text.trim().replace(/^["']|["']$/g, '').slice(0, 120);
|
||||
const title = text
|
||||
.trim()
|
||||
.replace(/^["']|["']$/g, '')
|
||||
.slice(0, 120);
|
||||
if (title) {
|
||||
await this.aiChatRepo.update(chatId, { title }, workspaceId);
|
||||
}
|
||||
|
||||
@@ -42,6 +42,15 @@ export class CreateMcpServerDto {
|
||||
@IsString({ each: true })
|
||||
toolAllowlist?: string[];
|
||||
|
||||
// Admin-authored guidance ("how/when to use this server's tools") injected
|
||||
// into the agent system prompt next to the tool descriptions (#180). Trusted,
|
||||
// NON-secret (so it IS returned). Capped to bound prompt/token size (the
|
||||
// built-in guide is ~1.5KB). Blank => stored as null.
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
@MaxLength(4000)
|
||||
instructions?: string;
|
||||
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
enabled?: boolean;
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
import 'reflect-metadata';
|
||||
import { plainToInstance } from 'class-transformer';
|
||||
import { validateSync } from 'class-validator';
|
||||
import { CreateMcpServerDto } from './create-mcp-server.dto';
|
||||
import { UpdateMcpServerDto } from './update-mcp-server.dto';
|
||||
|
||||
/**
|
||||
* API-boundary validation for the per-server `instructions` field (#180): a free
|
||||
* text guide injected into the agent system prompt. It is optional, must be a
|
||||
* string, and is bounded by @MaxLength(4000) to cap prompt/token size.
|
||||
*/
|
||||
describe('MCP server DTO instructions validation', () => {
|
||||
function validateCreate(payload: unknown) {
|
||||
const dto = plainToInstance(CreateMcpServerDto, payload);
|
||||
return validateSync(dto as object);
|
||||
}
|
||||
function validateUpdate(payload: unknown) {
|
||||
const dto = plainToInstance(UpdateMcpServerDto, payload);
|
||||
return validateSync(dto as object);
|
||||
}
|
||||
|
||||
const base = {
|
||||
name: 'Tavily',
|
||||
transport: 'http',
|
||||
url: 'https://example.com/mcp',
|
||||
};
|
||||
|
||||
it('accepts an omitted instructions field on create', () => {
|
||||
expect(validateCreate({ ...base })).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('accepts a reasonable instructions string on create', () => {
|
||||
expect(
|
||||
validateCreate({ ...base, instructions: 'Use search for fresh facts.' }),
|
||||
).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('rejects instructions over MaxLength(4000) on create', () => {
|
||||
const errors = validateCreate({
|
||||
...base,
|
||||
instructions: 'a'.repeat(4001),
|
||||
});
|
||||
expect(
|
||||
errors.some(
|
||||
(e) =>
|
||||
e.property === 'instructions' &&
|
||||
e.constraints !== undefined &&
|
||||
'maxLength' in e.constraints,
|
||||
),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('accepts instructions of exactly 4000 chars on create', () => {
|
||||
expect(
|
||||
validateCreate({ ...base, instructions: 'a'.repeat(4000) }),
|
||||
).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('rejects a non-string instructions value', () => {
|
||||
const errors = validateCreate({ ...base, instructions: 123 });
|
||||
expect(errors.some((e) => e.property === 'instructions')).toBe(true);
|
||||
});
|
||||
|
||||
it('rejects instructions over MaxLength(4000) on update', () => {
|
||||
const errors = validateUpdate({ instructions: 'a'.repeat(4001) });
|
||||
expect(
|
||||
errors.some(
|
||||
(e) =>
|
||||
e.property === 'instructions' &&
|
||||
e.constraints !== undefined &&
|
||||
'maxLength' in e.constraints,
|
||||
),
|
||||
).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -43,6 +43,13 @@ export class UpdateMcpServerDto {
|
||||
@IsString({ each: true })
|
||||
toolAllowlist?: string[];
|
||||
|
||||
// Admin-authored prompt guidance (#180). Absent => unchanged; blank => cleared
|
||||
// (stored as null by the repo). Capped to bound prompt/token size.
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
@MaxLength(4000)
|
||||
instructions?: string;
|
||||
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
enabled?: boolean;
|
||||
|
||||
@@ -33,6 +33,26 @@ interface ServerOutcome {
|
||||
reason?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* One server's admin-authored guidance for the agent system prompt (#180).
|
||||
* Built ONLY for a server that actually connected AND contributed ≥1 tool
|
||||
* (after the allowlist filter) AND has non-blank guidance — so a guide never
|
||||
* appears for a server whose tools the agent cannot actually call.
|
||||
*/
|
||||
export interface McpServerInstruction {
|
||||
/** Display name of the server (for the prompt section header). */
|
||||
serverName: string;
|
||||
/**
|
||||
* The tool-name namespace prefix the server's tools were merged under
|
||||
* (sanitized name, e.g. `tavily`). The prompt renders this as `tavily_*` so
|
||||
* the model can connect the guidance to the actual tool names. Advisory:
|
||||
* individual tools may carry a disambiguating suffix on rare collisions.
|
||||
*/
|
||||
toolPrefix: string;
|
||||
/** The trusted, non-blank guidance text. */
|
||||
instructions: string;
|
||||
}
|
||||
|
||||
export interface ExternalToolset {
|
||||
/** Namespaced external tools, merge-ready into the agent toolset. */
|
||||
tools: Record<string, Tool>;
|
||||
@@ -40,6 +60,11 @@ export interface ExternalToolset {
|
||||
clients: Closable[];
|
||||
/** Per-server connect outcomes so the UI can show unavailable servers. */
|
||||
outcomes: ServerOutcome[];
|
||||
/**
|
||||
* Per-server prompt guidance for connected servers that contributed ≥1 tool
|
||||
* and have non-blank instructions. Empty when no server qualifies.
|
||||
*/
|
||||
instructions: McpServerInstruction[];
|
||||
}
|
||||
|
||||
/** Connect+tools() timeout per server — a slow server must not stall the turn. */
|
||||
@@ -60,6 +85,8 @@ interface CacheEntry {
|
||||
tools: Record<string, Tool>;
|
||||
clients: McpClient[];
|
||||
outcomes: ServerOutcome[];
|
||||
/** Prompt guidance for qualifying servers (see McpServerInstruction). */
|
||||
instructions: McpServerInstruction[];
|
||||
expiresAt: number;
|
||||
/** Active leases (turns currently using these clients). */
|
||||
refCount: number;
|
||||
@@ -141,6 +168,7 @@ export class McpClientsService {
|
||||
tools: entry.tools,
|
||||
clients: [release],
|
||||
outcomes: entry.outcomes,
|
||||
instructions: entry.instructions,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -225,6 +253,7 @@ export class McpClientsService {
|
||||
const outcomes: ServerOutcome[] = [];
|
||||
// Per-call total wall-clock cap, read once for this build (env-overridable).
|
||||
const callTimeoutMs = mcpCallTimeoutMs();
|
||||
const instructions: McpServerInstruction[] = [];
|
||||
|
||||
for (const server of servers) {
|
||||
try {
|
||||
@@ -233,17 +262,33 @@ export class McpClientsService {
|
||||
clients.push(client);
|
||||
const allow = server.toolAllowlist;
|
||||
const picked =
|
||||
Array.isArray(allow) && allow.length > 0
|
||||
? pick(raw, allow)
|
||||
: raw;
|
||||
Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
|
||||
// Bound each tool's execute with a per-call total-timeout guard before
|
||||
// merging, so a single chatty-but-stuck call is aborted after the cap.
|
||||
const guarded = wrapToolsWithCallTimeout(picked, callTimeoutMs);
|
||||
// Namespace each tool with the sanitized server name AND disambiguate
|
||||
// against names already merged from earlier servers, so no external
|
||||
// tool is silently overwritten on collision.
|
||||
this.mergeNamespaced(tools, guarded, server.name, server.id);
|
||||
// tool is silently overwritten on collision. The returned count drives
|
||||
// whether this server's prompt guidance is included (≥1 tool merged).
|
||||
const merged = this.mergeNamespaced(
|
||||
tools,
|
||||
guarded,
|
||||
server.name,
|
||||
server.id,
|
||||
);
|
||||
outcomes.push({ name: server.name, ok: true });
|
||||
// Include this server's guidance ONLY when it actually contributed at
|
||||
// least one tool the agent can call (allowlist may have filtered all of
|
||||
// them out) AND the admin authored non-blank instructions. The header
|
||||
// prefix is the sanitized server name (= the tool namespace prefix).
|
||||
const guide = server.instructions?.trim();
|
||||
if (merged.count > 0 && guide) {
|
||||
instructions.push({
|
||||
serverName: server.name,
|
||||
toolPrefix: merged.prefix,
|
||||
instructions: guide,
|
||||
});
|
||||
}
|
||||
} catch (err) {
|
||||
// A failed server is skipped — the turn proceeds with the rest. Log a
|
||||
// short warning (never the URL/headers) so ops can see degradation, and
|
||||
@@ -260,6 +305,7 @@ export class McpClientsService {
|
||||
tools,
|
||||
clients,
|
||||
outcomes,
|
||||
instructions,
|
||||
expiresAt: Date.now() + CACHE_TTL_MS,
|
||||
refCount: 0,
|
||||
evicted: false,
|
||||
@@ -276,16 +322,19 @@ export class McpClientsService {
|
||||
* renaming any key that would collide with an already-merged tool (different
|
||||
* servers with the same sanitized name, or duplicates after truncation), so
|
||||
* no external tool is silently dropped via overwrite.
|
||||
*
|
||||
* Returns how many tools this server actually contributed and the namespace
|
||||
* prefix used (the sanitized server name) so the caller can attach the
|
||||
* server's prompt guidance only when ≥1 tool was merged.
|
||||
*/
|
||||
private mergeNamespaced(
|
||||
target: Record<string, Tool>,
|
||||
picked: Record<string, Tool>,
|
||||
serverName: string,
|
||||
serverId: string,
|
||||
): void {
|
||||
for (const [name, tool] of Object.entries(
|
||||
namespace(picked, serverName),
|
||||
)) {
|
||||
): { count: number; prefix: string } {
|
||||
let count = 0;
|
||||
for (const [name, tool] of Object.entries(namespace(picked, serverName))) {
|
||||
let key = name;
|
||||
if (key in target) {
|
||||
const original = key;
|
||||
@@ -295,7 +344,9 @@ export class McpClientsService {
|
||||
);
|
||||
}
|
||||
target[key] = tool;
|
||||
count += 1;
|
||||
}
|
||||
return { count, prefix: namespacePrefix(serverName) };
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -371,9 +422,7 @@ export class McpClientsService {
|
||||
|
||||
/** Close clients, swallowing close errors so they never break a response. */
|
||||
private async closeClients(clients: McpClient[]): Promise<void> {
|
||||
await Promise.all(
|
||||
clients.map((c) => c.close().catch(() => undefined)),
|
||||
);
|
||||
await Promise.all(clients.map((c) => c.close().catch(() => undefined)));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -386,9 +435,10 @@ export class McpClientsService {
|
||||
* lookup hands net/tls.connect ONLY a set that passed this check, so the kernel
|
||||
* can never connect to an address that did not pass the guard. Pure — no I/O.
|
||||
*/
|
||||
export function validateResolvedAddresses(
|
||||
addrs: readonly LookupAddress[],
|
||||
): { ok: boolean; blockedHost?: string } {
|
||||
export function validateResolvedAddresses(addrs: readonly LookupAddress[]): {
|
||||
ok: boolean;
|
||||
blockedHost?: string;
|
||||
} {
|
||||
if (addrs.length === 0) {
|
||||
return { ok: false };
|
||||
}
|
||||
@@ -524,7 +574,7 @@ function namespace(
|
||||
tools: Record<string, Tool>,
|
||||
serverName: string,
|
||||
): Record<string, Tool> {
|
||||
const prefix = sanitizeName(serverName) || 'mcp';
|
||||
const prefix = namespacePrefix(serverName);
|
||||
const out: Record<string, Tool> = {};
|
||||
for (const [name, t] of Object.entries(tools)) {
|
||||
const safe = sanitizeName(name);
|
||||
@@ -539,6 +589,15 @@ function namespace(
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* The tool-name namespace prefix for a server: its sanitized name, or `mcp`
|
||||
* when the name sanitizes to empty. Tools are merged as `${prefix}_${tool}`, so
|
||||
* the prompt guidance refers to the server's tools as `${prefix}_*`.
|
||||
*/
|
||||
function namespacePrefix(serverName: string): string {
|
||||
return sanitizeName(serverName) || 'mcp';
|
||||
}
|
||||
|
||||
/** Reduce an arbitrary string to ^[a-zA-Z0-9_-]+, collapsing runs to '_'. */
|
||||
function sanitizeName(value: string): string {
|
||||
return value
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
import { type Tool } from 'ai';
|
||||
import { McpClientsService } from './mcp-clients.service';
|
||||
|
||||
/**
|
||||
* Tests for the per-server prompt guidance (#180) assembled by buildEntry and
|
||||
* surfaced via toolsFor().instructions.
|
||||
*
|
||||
* REACHABILITY NOTE: buildEntry is a PRIVATE method; the smallest reachable
|
||||
* public path is toolsFor() -> getOrBuildEntry -> buildEntry -> connect/tools()
|
||||
* -> mergeNamespaced. We drive that path: stub the repo's `listEnabled` and spy
|
||||
* on the private `connect` to return fake MCP clients whose `tools()` we control.
|
||||
*
|
||||
* Contract (all checked here): a server's guidance is included ONLY when the
|
||||
* server actually connected AND contributed ≥1 callable tool (after the
|
||||
* allowlist filter) AND its instructions are non-blank. The header carries the
|
||||
* tool namespace prefix (the sanitized server name).
|
||||
*/
|
||||
function fakeTool(): Tool {
|
||||
return { description: 'x', inputSchema: undefined } as unknown as Tool;
|
||||
}
|
||||
|
||||
interface FakeServer {
|
||||
id: string;
|
||||
name: string;
|
||||
transport: string;
|
||||
url: string;
|
||||
headersEnc: string | null;
|
||||
toolAllowlist: string[] | null;
|
||||
instructions: string | null;
|
||||
}
|
||||
|
||||
function server(
|
||||
over: Partial<FakeServer> & { id: string; name: string },
|
||||
): FakeServer {
|
||||
return {
|
||||
transport: 'http',
|
||||
url: 'https://example.com/mcp',
|
||||
headersEnc: null,
|
||||
toolAllowlist: null,
|
||||
instructions: null,
|
||||
...over,
|
||||
};
|
||||
}
|
||||
|
||||
async function instructionsFor(
|
||||
servers: FakeServer[],
|
||||
toolsByServerId: Record<string, Record<string, Tool>>,
|
||||
// Server ids whose connect should THROW (simulating an unavailable server).
|
||||
failingIds: Set<string> = new Set(),
|
||||
): Promise<
|
||||
{
|
||||
serverName: string;
|
||||
toolPrefix: string;
|
||||
instructions: string;
|
||||
}[]
|
||||
> {
|
||||
const repoStub = {
|
||||
listEnabled: jest.fn().mockResolvedValue(servers),
|
||||
};
|
||||
const service = new McpClientsService(repoStub as never, {} as never);
|
||||
|
||||
jest
|
||||
.spyOn(
|
||||
service as unknown as { connect: (s: FakeServer) => unknown },
|
||||
'connect',
|
||||
)
|
||||
.mockImplementation((s: FakeServer) => {
|
||||
if (failingIds.has(s.id)) {
|
||||
return Promise.reject(new Error('connection failed'));
|
||||
}
|
||||
return Promise.resolve({
|
||||
tools: () => Promise.resolve(toolsByServerId[s.id] ?? {}),
|
||||
close: () => Promise.resolve(),
|
||||
});
|
||||
});
|
||||
|
||||
const toolset = await service.toolsFor('ws-1');
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
return toolset.instructions;
|
||||
}
|
||||
|
||||
describe('external MCP per-server prompt guidance (via toolsFor)', () => {
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('includes guidance for a connected server with non-empty text and ≥1 tool', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[
|
||||
server({
|
||||
id: 'id-tavily',
|
||||
name: 'Tavily',
|
||||
instructions: 'Use tavily_search for fresh facts.',
|
||||
}),
|
||||
],
|
||||
{ 'id-tavily': { search: fakeTool() } },
|
||||
);
|
||||
|
||||
// sanitizeName preserves case (charset [a-zA-Z0-9_-]), so the prefix is the
|
||||
// server name as-is for an already-clean name.
|
||||
expect(instructions).toEqual([
|
||||
{
|
||||
serverName: 'Tavily',
|
||||
toolPrefix: 'Tavily',
|
||||
instructions: 'Use tavily_search for fresh facts.',
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('omits guidance when the server has no instructions', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[server({ id: 'id-1', name: 'Tavily', instructions: null })],
|
||||
{ 'id-1': { search: fakeTool() } },
|
||||
);
|
||||
expect(instructions).toEqual([]);
|
||||
});
|
||||
|
||||
it('omits guidance when the instructions are only whitespace', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[server({ id: 'id-1', name: 'Tavily', instructions: ' ' })],
|
||||
{ 'id-1': { search: fakeTool() } },
|
||||
);
|
||||
expect(instructions).toEqual([]);
|
||||
});
|
||||
|
||||
it('omits guidance for a server that contributed ZERO tools (allowlist filtered all out)', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[
|
||||
server({
|
||||
id: 'id-1',
|
||||
name: 'Tavily',
|
||||
instructions: 'guide',
|
||||
// Allowlist names a tool the server does not expose -> 0 picked.
|
||||
toolAllowlist: ['nonexistent'],
|
||||
}),
|
||||
],
|
||||
{ 'id-1': { search: fakeTool() } },
|
||||
);
|
||||
expect(instructions).toEqual([]);
|
||||
});
|
||||
|
||||
it('omits guidance for an unavailable (failed-connect) server', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[server({ id: 'id-1', name: 'Tavily', instructions: 'guide' })],
|
||||
{ 'id-1': { search: fakeTool() } },
|
||||
new Set(['id-1']),
|
||||
);
|
||||
expect(instructions).toEqual([]);
|
||||
});
|
||||
|
||||
it('includes only the qualifying servers among several', async () => {
|
||||
const instructions = await instructionsFor(
|
||||
[
|
||||
server({ id: 'ok', name: 'Tavily', instructions: 'web guide' }),
|
||||
server({ id: 'blank', name: 'Crawl', instructions: '' }),
|
||||
server({ id: 'down', name: 'Down', instructions: 'never shown' }),
|
||||
],
|
||||
{
|
||||
ok: { search: fakeTool() },
|
||||
blank: { crawl: fakeTool() },
|
||||
down: { x: fakeTool() },
|
||||
},
|
||||
new Set(['down']),
|
||||
);
|
||||
|
||||
expect(instructions).toEqual([
|
||||
{ serverName: 'Tavily', toolPrefix: 'Tavily', instructions: 'web guide' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -17,6 +17,7 @@ function row(overrides: Partial<AiMcpServer>): AiMcpServer {
|
||||
enabled: true,
|
||||
toolAllowlist: null,
|
||||
headersEnc: null,
|
||||
instructions: null,
|
||||
...overrides,
|
||||
} as unknown as AiMcpServer;
|
||||
}
|
||||
@@ -28,11 +29,7 @@ describe('McpServersService.toView (via list) — encrypted-header leak guard',
|
||||
};
|
||||
// secretBox + clients are unused by the list/toView path; pass stubs to
|
||||
// satisfy the constructor.
|
||||
return new McpServersService(
|
||||
repoStub as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
);
|
||||
return new McpServersService(repoStub as never, {} as never, {} as never);
|
||||
}
|
||||
|
||||
it('exposes hasHeaders:true and NO headersEnc when auth headers are set', async () => {
|
||||
@@ -67,6 +64,7 @@ describe('McpServersService.toView (via list) — encrypted-header leak guard',
|
||||
enabled: false,
|
||||
toolAllowlist: ['search'],
|
||||
headersEnc: 'BLOB',
|
||||
instructions: 'Use search for fresh web facts.',
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -80,6 +78,19 @@ describe('McpServersService.toView (via list) — encrypted-header leak guard',
|
||||
enabled: false,
|
||||
toolAllowlist: ['search'],
|
||||
hasHeaders: true,
|
||||
instructions: 'Use search for fresh web facts.',
|
||||
});
|
||||
});
|
||||
|
||||
it('returns instructions (NON-secret) in the view, null when unset', async () => {
|
||||
const service = buildService([
|
||||
row({ id: 'a', instructions: 'How to use these tools.' }),
|
||||
row({ id: 'b', instructions: null }),
|
||||
]);
|
||||
|
||||
const [withText, withoutText] = await service.list('ws-1');
|
||||
|
||||
expect(withText.instructions).toBe('How to use these tools.');
|
||||
expect(withoutText.instructions).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -20,6 +20,9 @@ export interface McpServerView {
|
||||
enabled: boolean;
|
||||
toolAllowlist: string[] | null;
|
||||
hasHeaders: boolean;
|
||||
// Admin-authored prompt guidance (#180). NON-secret, so returned in the view.
|
||||
// Null when no guidance is configured.
|
||||
instructions: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -56,6 +59,8 @@ export class McpServersService {
|
||||
url: dto.url,
|
||||
headersEnc,
|
||||
toolAllowlist: dto.toolAllowlist ?? null,
|
||||
// Blank/whitespace guidance is normalized to null by the repo.
|
||||
instructions: dto.instructions ?? null,
|
||||
enabled: dto.enabled ?? true,
|
||||
});
|
||||
this.clients.invalidate(workspaceId);
|
||||
@@ -97,6 +102,8 @@ export class McpServersService {
|
||||
headersEnc,
|
||||
// undefined => unchanged; [] / value handled by repo (empty => null).
|
||||
toolAllowlist: dto.toolAllowlist,
|
||||
// undefined => unchanged; blank => cleared (null) by the repo.
|
||||
instructions: dto.instructions,
|
||||
enabled: dto.enabled,
|
||||
});
|
||||
this.clients.invalidate(workspaceId);
|
||||
@@ -167,6 +174,7 @@ export class McpServersService {
|
||||
enabled: row.enabled,
|
||||
toolAllowlist: row.toolAllowlist ?? null,
|
||||
hasHeaders: Boolean(row.headersEnc),
|
||||
instructions: row.instructions ?? null,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user