diff --git a/src/mcp/built-in-servers.ts b/src/mcp/built-in-servers.ts index 91f99da..edef70e 100644 --- a/src/mcp/built-in-servers.ts +++ b/src/mcp/built-in-servers.ts @@ -24,6 +24,10 @@ export function getBuiltInMcpServers(env: Env): Record (ch === '*' ? '.*' : `\\${ch}`)); + return new RegExp(`^${escaped}$`).test(name); +} + +function matchesAnyGlob(name: string, patterns: string[] | undefined): boolean { + return !!patterns && patterns.some((p) => matchesGlob(name, p)); +} + /** * Convert a single JSON Schema property definition to a Zod type. * Handles the common scalar types returned by MCP servers. @@ -89,6 +103,7 @@ export function mcpToolToAITool( serverId: string, mcpTool: MCPToolDefinition, mcpClient: MCPClient, + continuationApprovalPatterns?: string[], ) { const toolName = `mcp__${serverId}__${mcpTool.name}`; const description = mcpTool.description ?? `MCP tool ${mcpTool.name} from server ${serverId}`; @@ -98,17 +113,27 @@ export function mcpToolToAITool( mcpTool.inputSchema as Record | undefined, ); + // Continuation-gated tools pause the agentic loop after they finish (see server.ts), + const isContinuationGated = matchesAnyGlob(mcpTool.name, continuationApprovalPatterns); + return { name: toolName, tool: tool({ description, inputSchema, + // Flag read by the stopWhen predicate in server.ts. `daAgent` is our own + // provider namespace and is ignored by the Bedrock provider. + ...(isContinuationGated + ? { providerOptions: { daAgent: { continuationApproval: true } } } + : {}), // Fail-closed gating for untrusted external MCP servers. Per the MCP spec // annotation defaults (readOnlyHint=false, destructiveHint=true) and the // fact that annotations are optional, we gate unless the tool tells us it // is safe: skip approval only when it is read-only OR explicitly // non-destructive. Everything else — including unannotated tools — - // requires approval. + // requires approval. This is independent of continuation gating: a tool + // can require both pre-execution approval and post-execution continuation + // approval. needsApproval: async () => { const { readOnlyHint, destructiveHint } = mcpTool.annotations ?? {}; return readOnlyHint !== true && destructiveHint !== false; @@ -146,6 +171,8 @@ export async function connectAndRegisterMCPTools( mcpConfig: { mcpServers: Record; toolAllowPatterns: string[]; + /** Per-server glob patterns for post-execution continuation approval. */ + continuationApprovalPatterns?: Record; }, options?: { headers?: Record; @@ -194,6 +221,7 @@ export async function connectAndRegisterMCPTools( const discoveredTools = await client.listTools(); clients.push(client); + const serverContinuationPatterns = mcpConfig.continuationApprovalPatterns?.[serverId]; let registeredCount = 0; for (const toolDef of discoveredTools) { try { @@ -201,6 +229,7 @@ export async function connectAndRegisterMCPTools( serverId, toolDef, client, + serverContinuationPatterns, ); tools[qualifiedName] = aiTool; registeredCount += 1; diff --git a/src/mcp/types.ts b/src/mcp/types.ts index 56552a2..6bd6d55 100644 --- a/src/mcp/types.ts +++ b/src/mcp/types.ts @@ -25,12 +25,18 @@ export type MCPServerConfig = StdioMCPServerConfig | RemoteMCPServerConfig; * `sendImsToken` — forward the user's IMS Bearer token in the Authorization header. * `apiKey` — static API key to send as x-api-key (omitted when undefined). * `instructions` — additional prompt instructions appended to the system prompt. + * `continuationApprovalPatterns` — glob patterns (e.g. `evaluate_*`) matched against + * the bare MCP tool name. Matching tools pause the agentic loop after they finish so + * the user can review results and decide whether to continue (post-execution + * "continuation approval" gate). This is independent of pre-execution approval — + * a matching tool may still require approval to run in the first place. */ export interface BuiltInMCPServerConfig { type: 'http' | 'sse'; url: string; sendImsToken?: boolean; instructions?: string; + continuationApprovalPatterns?: string[]; } export function isStdioConfig(config: MCPServerConfig): config is StdioMCPServerConfig { diff --git a/src/message-pipeline.ts b/src/message-pipeline.ts index f6c1b1a..2dd4526 100644 --- a/src/message-pipeline.ts +++ b/src/message-pipeline.ts @@ -426,3 +426,40 @@ export function expandLatestUserAttachmentsForModel( } /* eslint-enable @typescript-eslint/no-explicit-any */ + +/** A `data-continuation` transient stream part driving the client's Continue/Stop prompt. */ +export interface ContinuationPart { + type: 'data-continuation'; + transient: true; + data: { toolCallId: string; toolName: string }; +} + +interface StepLike { + toolCalls: Array<{ toolCallId: string; toolName: string }>; + toolResults: Array<{ toolCallId: string }>; +} + +/** + * The transient continuation parts to emit for the given (final) step: one per + * continuation-gated tool that actually produced a result in this step. A gated tool + * with no result (e.g. it was itself pre-execution-paused and never ran) is skipped so + * we never prompt "continue?" for a tool that hasn't finished. + */ +export function buildContinuationParts( + lastStep: StepLike | undefined, + requiresContinuationApproval: (toolName: string) => boolean, +): ContinuationPart[] { + if (!lastStep) return []; + const resultIds = new Set(lastStep.toolResults.map((r) => r.toolCallId)); + const parts: ContinuationPart[] = []; + for (const tc of lastStep.toolCalls) { + if (requiresContinuationApproval(tc.toolName) && resultIds.has(tc.toolCallId)) { + parts.push({ + type: 'data-continuation', + transient: true, + data: { toolCallId: tc.toolCallId, toolName: tc.toolName }, + }); + } + } + return parts; +} diff --git a/src/prompt-builder.ts b/src/prompt-builder.ts index 2b6390f..d407bfd 100644 --- a/src/prompt-builder.ts +++ b/src/prompt-builder.ts @@ -193,6 +193,31 @@ This is a critical issue. Use these blocks when they improve readability — for example, checklists for audits, alerts for important notes, toggle lists for detailed breakdowns. Do NOT overuse them for simple responses. +**Planning bracket** — for any operation involving 2 or more distinct steps or tool calls, use the planning bracket before executing anything: +1. Call \`enter_plan_mode\` — signals the start of planning (no side effects). +2. Reason about the steps needed. +3. Call \`exit_plan_mode\` with the full plan — the user reviews and clicks Run to approve. +4. After approval, execute all steps in order. + +**Task item** — after the user approves and you begin execution, emit \`:::task-item\` before and after each step: +\`\`\` +:::task-item +{ "label": "Same label as in exit_plan_mode", "status": "running" } +::: +\`\`\` +\`\`\` +:::task-item +{ "label": "Same label as in exit_plan_mode", "status": "done" } +::: +\`\`\` + +Rules: +- Always call \`enter_plan_mode\` first, then \`exit_plan_mode\` with ALL planned steps. +- Use the **exact same** \`label\` string in \`exit_plan_mode\` tasks and \`:::task-item\` directives — character-for-character identical. +- Do NOT use these for single-step or trivial responses — only for operations with 2+ distinct steps. +- After the user approves (clicks Run), for EVERY step: emit \`running\`, make the tool call, then emit \`done\` as the very first text after the tool result — before any commentary or prose. +- Never skip the \`done\` directive. Every step that started with \`running\` must end with \`done\`. + ## EDS HTML Content Rules ALL content you create or update via tools MUST be valid Edge Delivery Services (EDS) semantic HTML. Follow these rules strictly: diff --git a/src/server.ts b/src/server.ts index 0e4d900..6c4f844 100644 --- a/src/server.ts +++ b/src/server.ts @@ -26,6 +26,7 @@ import { ensureOrphanedToolResults, expandUserSelectionContextForModel, expandLatestUserAttachmentsForModel, + buildContinuationParts, } from './message-pipeline.js'; import { buildSystemPrompt } from './prompt-builder.js'; import { buildEarlyChatContext, resolveAsyncContext } from './chat-context.js'; @@ -72,7 +73,10 @@ export default { const url = new URL(request.url); if (url.pathname === '/chat') { if (request.method === 'HEAD') { - return new Response(null, { status: 200, headers: CORS_HEADERS }); + return new Response(null, { + status: 200, + headers: { ...CORS_HEADERS, 'Content-Length': '0' }, + }); } if (request.method === 'POST') { return handleChat(request, env); @@ -206,6 +210,14 @@ async function handleChat(request: Request, env: Env): Promise { const { allTools, mcpClients, mcpConfig, mcpErrors, generatedToolsIndex, builtInServers } = assembled; + // A tool declares a post-execution "continuation approval" gate via + // providerOptions.daAgent.continuationApproval (built-in tools set it inline; + // MCP tools get it during adaptation when they match a server pattern). When such + // a tool runs we halt the agentic loop after its result and prompt the user to + // continue — the LLM never decides whether to pause. + const requiresContinuationApproval = (toolName: string): boolean => + allTools[toolName]?.providerOptions?.daAgent?.continuationApproval === true; + console.log(`[da-agent:perf] early=${t1 - t0}ms parallel=${t2 - t1}ms pre-stream=${t2 - t0}ms`); const { messages, requestedSkills, imsToken, attachments = [], sessionId } = parsed.data; @@ -297,7 +309,7 @@ async function handleChat(request: Request, env: Env): Promise { }); const stream = createUIMessageStream({ - execute: ({ writer }) => { + execute: async ({ writer }) => { // Stream results for tools the user approved this round so the client can // move each approved card to its result state, before the model continues. for (const o of executedOutputs) { @@ -343,7 +355,15 @@ async function handleChat(request: Request, env: Env): Promise { system: systemPrompt, messages: modelMessages as ModelMessage[], tools: allTools, - stopWhen: stepCountIs(5), + // Halt after the normal step budget OR immediately after a step that ran a + // continuation-gated tool, so the user can review results before continuing. + stopWhen: [ + stepCountIs(5), + ({ steps }) => { + const last = steps.at(-1); + return !!last?.toolCalls?.some((tc) => requiresContinuationApproval(tc.toolName)); + }, + ], experimental_telemetry: { isEnabled: true, functionId: 'da-agent-chat', @@ -357,7 +377,32 @@ async function handleChat(request: Request, env: Env): Promise { }, }); - writer.merge(result.toUIMessageStream()); + // Merge the model stream manually so we can emit a transient `data-continuation` + // part after the model stream for any continuation-gated tool that just ran, while + // holding the terminal `finish` chunk so the ordering is + // `…tool-output-available, data-continuation, finish`. The transient part is + // delivered to the client but never merged into message history. (Outputs of tools + // approved this round were already streamed above from `executedOutputs`.) + const reader = result.toUIMessageStream().getReader(); + let finishChunk: Awaited>['value'] | null = null; + for (;;) { + // eslint-disable-next-line no-await-in-loop -- sequential stream consumption + const { done, value } = await reader.read(); + if (done) break; + if (value.type === 'finish') { + finishChunk = value; + } else { + writer.write(value); + } + } + + const continuationParts = buildContinuationParts( + (await result.steps).at(-1), + requiresContinuationApproval, + ); + for (const part of continuationParts) writer.write(part); + + if (finishChunk) writer.write(finishChunk); }, onError: (error) => { console.error('[da-agent] stream error:', formatErrorForLog(error)); diff --git a/src/tool-assembly.ts b/src/tool-assembly.ts index c4d3c2a..cb7a9d3 100644 --- a/src/tool-assembly.ts +++ b/src/tool-assembly.ts @@ -99,6 +99,7 @@ export async function assembleTools( } const builtInServers = getBuiltInMcpServers(env); + const continuationApprovalPatterns: Record = {}; for (const [id, builtIn] of Object.entries(builtInServers)) { const headers: Record = {}; @@ -111,6 +112,9 @@ export async function assembleTools( url: builtIn.url, ...(Object.keys(headers).length > 0 ? { headers } : {}), }; + if (builtIn.continuationApprovalPatterns?.length) { + continuationApprovalPatterns[id] = builtIn.continuationApprovalPatterns; + } } const mcpConfig = @@ -118,6 +122,7 @@ export async function assembleTools( ? { mcpServers: allMcpServers, toolAllowPatterns: Object.keys(allMcpServers).map((id) => `mcp__${id}__*`), + continuationApprovalPatterns, } : null; diff --git a/src/tools/tools.ts b/src/tools/tools.ts index 4444869..e07a254 100644 --- a/src/tools/tools.ts +++ b/src/tools/tools.ts @@ -461,7 +461,12 @@ export function createDATools( try { const content = await loadSkillBodyFromFolder(client, ctxOrg, ctxRepo, skillId); if (!content) return { error: `Skill "${skillId}" not found` }; - return { skillId, content }; + return { + skillId, + content, + _hint: + 'Skill loaded. Before executing any steps, call enter_plan_mode then exit_plan_mode with your planned tasks so the user can review and approve.', + }; } catch (e) { return { error: String(e) }; } @@ -557,6 +562,41 @@ export function createDATools( }, }); + // Planning bracket — mirrors AO's enter_plan_mode / exit_plan_mode built-in tools. + // enter_plan_mode: signals start of planning phase; no approval, no side effects. + tools.enter_plan_mode = tool({ + description: + 'Signal the start of a planning phase. Call this before reasoning about what steps to take ' + + 'for any operation involving 2 or more distinct steps or tool calls. ' + + 'No action is taken — this is a signal only. Follow it by calling exit_plan_mode with the full plan.', + inputSchema: z.object({}), + needsApproval: async () => false, + execute: async () => ({ planning: true }), + }); + + // exit_plan_mode: submits the plan for user review; requires approval before execution proceeds. + tools.exit_plan_mode = tool({ + description: + 'Submit the completed plan for the user to review before any actions are taken. ' + + 'Call this after enter_plan_mode, once you have determined all the steps. ' + + 'The user will see the plan card and click Run to approve execution. ' + + 'Use the same task labels later in :::task-item directives to report progress.', + inputSchema: z.object({ + title: z.string().describe('Short plan title (≤ 8 words)'), + description: z.string().optional().describe('One-line summary of what you are about to do'), + tasks: z + .array( + z.object({ + id: z.string().describe('Unique step identifier, e.g. "1", "2"'), + label: z.string().describe('Human-readable step description'), + }), + ) + .describe('Ordered list of steps to execute'), + }), + needsApproval: async () => true, + execute: async () => ({ approved: true }), + }); + // Memory tools write to internal agent metadata paths — no user approval needed. tools.write_project_memory = tool({ description: diff --git a/test/mcp/built-in-servers.test.ts b/test/mcp/built-in-servers.test.ts index 5dcf292..4716d05 100644 --- a/test/mcp/built-in-servers.test.ts +++ b/test/mcp/built-in-servers.test.ts @@ -33,4 +33,9 @@ describe('getBuiltInMcpServers', () => { ); expect(servers['governance-agent'].url).toBe('http://localhost:8000/mcp/'); }); + + it('gates evaluate_* tools behind a post-execution continuation approval', () => { + const servers = getBuiltInMcpServers(envWith()); + expect(servers['governance-agent'].continuationApprovalPatterns).toEqual(['evaluate_*']); + }); }); diff --git a/test/mcp/tool-adapter.test.ts b/test/mcp/tool-adapter.test.ts index d9e7fb7..248f928 100644 --- a/test/mcp/tool-adapter.test.ts +++ b/test/mcp/tool-adapter.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { mcpToolToAITool } from '../../src/mcp/tool-adapter.js'; +import { mcpToolToAITool, matchesGlob } from '../../src/mcp/tool-adapter.js'; import type { MCPClient, MCPToolDefinition } from '../../src/mcp/client.js'; const fakeClient = {} as MCPClient; @@ -9,6 +9,11 @@ function needsApproval(mcpTool: MCPToolDefinition): Promise return Promise.resolve(aiTool.needsApproval?.({}, { toolCallId: 'x', messages: [] })); } +// eslint-disable-next-line @typescript-eslint/no-explicit-any +function continuationFlag(aiTool: any): boolean { + return aiTool.providerOptions?.daAgent?.continuationApproval === true; +} + describe('mcpToolToAITool', () => { it('gates the tool behind approval when the server sets destructiveHint', async () => { expect( @@ -32,3 +37,60 @@ describe('mcpToolToAITool', () => { ).toBe(false); }); }); + +describe('matchesGlob', () => { + it('matches a trailing wildcard', () => { + expect(matchesGlob('evaluate_page', 'evaluate_*')).toBe(true); + expect(matchesGlob('evaluate_image', 'evaluate_*')).toBe(true); + }); + + it('does not match a different prefix', () => { + expect(matchesGlob('retrieve_page', 'evaluate_*')).toBe(false); + }); + + it('requires a full-string match (anchored)', () => { + expect(matchesGlob('pre_evaluate_page', 'evaluate_*')).toBe(false); + expect(matchesGlob('evaluate', 'evaluate_*')).toBe(false); + }); + + it('escapes regex metacharacters in the pattern', () => { + expect(matchesGlob('a.b', 'a.b')).toBe(true); + expect(matchesGlob('axb', 'a.b')).toBe(false); + }); +}); + +describe('mcpToolToAITool continuation approval', () => { + it('flags a matching tool for continuation approval independently of pre-exec approval', async () => { + const { tool: aiTool } = mcpToolToAITool( + 'governance-agent', + { name: 'evaluate_page' }, + fakeClient, + ['evaluate_*'], + ); + expect(continuationFlag(aiTool)).toBe(true); + // Continuation gating and pre-execution approval are independent: this + // unannotated tool still fails closed on the pre-exec gate while also being + // continuation-gated, so it can be both pre-gated and post-gated. + expect(await aiTool.needsApproval?.({}, { toolCallId: 'x', messages: [] })).toBe(true); + }); + + it('does not flag a non-matching tool and keeps annotation-based gating', async () => { + const { tool: aiTool } = mcpToolToAITool( + 'governance-agent', + { name: 'retrieve_brand_rules' }, + fakeClient, + ['evaluate_*'], + ); + expect(continuationFlag(aiTool)).toBe(false); + expect(await aiTool.needsApproval?.({}, { toolCallId: 'x', messages: [] })).toBe(true); + }); + + it('does not flag anything when no patterns are configured', () => { + const { tool: aiTool } = mcpToolToAITool( + 'governance-agent', + { name: 'evaluate_page' }, + fakeClient, + ); + expect(continuationFlag(aiTool)).toBe(false); + }); +}); diff --git a/test/message-pipeline.test.ts b/test/message-pipeline.test.ts index a174093..6794f82 100644 --- a/test/message-pipeline.test.ts +++ b/test/message-pipeline.test.ts @@ -7,6 +7,7 @@ import { ensureOrphanedToolResults, expandUserSelectionContextForModel, expandLatestUserAttachmentsForModel, + buildContinuationParts, TOOL_STATE, } from '../src/message-pipeline.js'; @@ -543,3 +544,65 @@ describe('toModelMessages', () => { expect(withOrphans[1].content[0].type).toBe('tool-result'); }); }); + +describe('buildContinuationParts', () => { + const gates = (name: string) => name === 'mcp__governance-agent__evaluate_page'; + + it('emits one transient part for a gated tool that produced a result', () => { + const parts = buildContinuationParts( + { + toolCalls: [{ toolCallId: 'call-a', toolName: 'mcp__governance-agent__evaluate_page' }], + toolResults: [{ toolCallId: 'call-a' }], + }, + gates, + ); + expect(parts).toEqual([ + { + type: 'data-continuation', + transient: true, + data: { toolCallId: 'call-a', toolName: 'mcp__governance-agent__evaluate_page' }, + }, + ]); + }); + + it('does not emit for a non-gated tool', () => { + const parts = buildContinuationParts( + { + toolCalls: [{ toolCallId: 'call-a', toolName: 'content_read' }], + toolResults: [{ toolCallId: 'call-a' }], + }, + gates, + ); + expect(parts).toEqual([]); + }); + + it('does not emit for a gated tool that produced no result', () => { + const parts = buildContinuationParts( + { + toolCalls: [{ toolCallId: 'call-a', toolName: 'mcp__governance-agent__evaluate_page' }], + toolResults: [], + }, + gates, + ); + expect(parts).toEqual([]); + }); + + it('emits one part per gated tool in a multi-tool step', () => { + const parts = buildContinuationParts( + { + toolCalls: [ + { toolCallId: 'call-a', toolName: 'mcp__governance-agent__evaluate_page' }, + { toolCallId: 'call-b', toolName: 'content_read' }, + { toolCallId: 'call-c', toolName: 'mcp__governance-agent__evaluate_page' }, + ], + toolResults: [{ toolCallId: 'call-a' }, { toolCallId: 'call-c' }], + }, + gates, + ); + expect(parts.map((p) => p.data.toolCallId)).toEqual(['call-a', 'call-c']); + }); + + it('returns nothing when there is no final step', () => { + expect(buildContinuationParts(undefined, gates)).toEqual([]); + }); +});