|
| 1 | +/** |
| 2 | + * MCP-Native Evolution Handler — policy evolution via the MCP tool protocol. |
| 3 | + * |
| 4 | + * Instead of blocking on stdin (which conflicts with MCP stdio transport), |
| 5 | + * this handler works in two asynchronous steps: |
| 6 | + * |
| 7 | + * 1. On deny: returns a structured denial response with a suggestion ID. |
| 8 | + * The agent presents the suggestion to the user in chat. |
| 9 | + * 2. On approve: the agent calls `policy_evolution_approve` with the |
| 10 | + * suggestion ID and the user's decision. The handler applies the change. |
| 11 | + * |
| 12 | + * The agent then retries the original tool call. |
| 13 | + */ |
| 14 | + |
| 15 | +import { nanoid } from 'nanoid'; |
| 16 | +import type { ActionRequest, Policy, Session } from '../types.js'; |
| 17 | +import type { EvolutionDecision, PolicySuggestion } from './types.js'; |
| 18 | +import { suggestPolicyChange } from './suggestion.js'; |
| 19 | +import { applyPolicyChange, writePolicyToFile } from './writer.js'; |
| 20 | + |
| 21 | +// --------------------------------------------------------------------------- |
| 22 | +// Types |
| 23 | +// --------------------------------------------------------------------------- |
| 24 | + |
| 25 | +interface PendingSuggestion { |
| 26 | + suggestion: PolicySuggestion; |
| 27 | + action: ActionRequest; |
| 28 | + sessionId: string; |
| 29 | + createdAt: number; |
| 30 | +} |
| 31 | + |
| 32 | +interface McpToolContent { |
| 33 | + type: 'text'; |
| 34 | + text: string; |
| 35 | +} |
| 36 | + |
| 37 | +interface McpToolResponse { |
| 38 | + [key: string]: unknown; |
| 39 | + content: McpToolContent[]; |
| 40 | + isError?: boolean; |
| 41 | +} |
| 42 | + |
| 43 | +// --------------------------------------------------------------------------- |
| 44 | +// Tool definition |
| 45 | +// --------------------------------------------------------------------------- |
| 46 | + |
| 47 | +const TOOL_NAME = 'policy_evolution_approve'; |
| 48 | + |
| 49 | +const TOOL_DEFINITION = { |
| 50 | + name: TOOL_NAME, |
| 51 | + description: |
| 52 | + 'Approve or deny a suggested policy change after a tool call was denied. ' + |
| 53 | + 'Present the suggestion to the user first and relay their decision.', |
| 54 | + inputSchema: { |
| 55 | + type: 'object' as const, |
| 56 | + properties: { |
| 57 | + suggestion_id: { |
| 58 | + type: 'string', |
| 59 | + description: 'The suggestion ID from the denial message', |
| 60 | + }, |
| 61 | + decision: { |
| 62 | + type: 'string', |
| 63 | + enum: ['add-to-policy', 'allow-once', 'deny'], |
| 64 | + description: |
| 65 | + 'The user\'s decision: "add-to-policy" persists to YAML, ' + |
| 66 | + '"allow-once" applies for this session only, "deny" keeps the block', |
| 67 | + }, |
| 68 | + }, |
| 69 | + required: ['suggestion_id', 'decision'], |
| 70 | + }, |
| 71 | +}; |
| 72 | + |
| 73 | +// --------------------------------------------------------------------------- |
| 74 | +// Handler class |
| 75 | +// --------------------------------------------------------------------------- |
| 76 | + |
| 77 | +export class McpEvolutionHandler { |
| 78 | + private pendingSuggestions = new Map<string, PendingSuggestion>(); |
| 79 | + private readonly policyPath: string; |
| 80 | + |
| 81 | + constructor(policyPath: string) { |
| 82 | + this.policyPath = policyPath; |
| 83 | + } |
| 84 | + |
| 85 | + /** The MCP tool name used by this handler. */ |
| 86 | + static readonly TOOL_NAME = TOOL_NAME; |
| 87 | + |
| 88 | + /** |
| 89 | + * Return the MCP tool definition for `policy_evolution_approve`. |
| 90 | + */ |
| 91 | + getToolDefinition(): typeof TOOL_DEFINITION { |
| 92 | + return TOOL_DEFINITION; |
| 93 | + } |
| 94 | + |
| 95 | + /** |
| 96 | + * Build a denial response that includes a policy-change suggestion. |
| 97 | + * |
| 98 | + * Called by the proxy when a tool call is denied and evolution is enabled. |
| 99 | + * Returns `null` if the denial is not suggestible (e.g. budget exceeded), |
| 100 | + * so the proxy should fall back to a plain denial message. |
| 101 | + */ |
| 102 | + buildDenialResponse( |
| 103 | + action: ActionRequest, |
| 104 | + reasons: string[], |
| 105 | + policy: Policy, |
| 106 | + sessionId: string, |
| 107 | + ): McpToolResponse | null { |
| 108 | + const suggestion = suggestPolicyChange(action, reasons, policy); |
| 109 | + if (!suggestion) { |
| 110 | + return null; |
| 111 | + } |
| 112 | + |
| 113 | + const suggestionId = nanoid(12); |
| 114 | + this.pendingSuggestions.set(suggestionId, { |
| 115 | + suggestion, |
| 116 | + action, |
| 117 | + sessionId, |
| 118 | + createdAt: Date.now(), |
| 119 | + }); |
| 120 | + |
| 121 | + const text = [ |
| 122 | + `Action denied by policy: ${reasons.join('; ')}`, |
| 123 | + '', |
| 124 | + `[Policy Evolution] Suggested change: ${suggestion.description}`, |
| 125 | + `Suggestion ID: ${suggestionId}`, |
| 126 | + '', |
| 127 | + 'Present this suggestion to the user and ask for their decision:', |
| 128 | + ' - "add-to-policy" — permanently add to the policy file', |
| 129 | + ' - "allow-once" — allow for this session only', |
| 130 | + ' - "deny" — keep the restriction', |
| 131 | + '', |
| 132 | + `Then call the "${TOOL_NAME}" tool with the suggestion_id and their decision.`, |
| 133 | + ].join('\n'); |
| 134 | + |
| 135 | + return { |
| 136 | + content: [{ type: 'text', text }], |
| 137 | + isError: true, |
| 138 | + }; |
| 139 | + } |
| 140 | + |
| 141 | + /** |
| 142 | + * Handle a call to `policy_evolution_approve`. |
| 143 | + * |
| 144 | + * Validates the suggestion ID, applies the policy change, and returns |
| 145 | + * a success or error message. |
| 146 | + */ |
| 147 | + handleApproval( |
| 148 | + args: Record<string, unknown>, |
| 149 | + session: Session, |
| 150 | + ): McpToolResponse { |
| 151 | + const suggestionId = args.suggestion_id as string | undefined; |
| 152 | + const decision = args.decision as EvolutionDecision | undefined; |
| 153 | + |
| 154 | + if (!suggestionId || !decision) { |
| 155 | + return { |
| 156 | + content: [{ type: 'text', text: 'Missing required fields: suggestion_id, decision' }], |
| 157 | + isError: true, |
| 158 | + }; |
| 159 | + } |
| 160 | + |
| 161 | + const validDecisions: EvolutionDecision[] = ['add-to-policy', 'allow-once', 'deny']; |
| 162 | + if (!validDecisions.includes(decision)) { |
| 163 | + return { |
| 164 | + content: [{ |
| 165 | + type: 'text', |
| 166 | + text: `Invalid decision "${decision}". Must be one of: ${validDecisions.join(', ')}`, |
| 167 | + }], |
| 168 | + isError: true, |
| 169 | + }; |
| 170 | + } |
| 171 | + |
| 172 | + const pending = this.pendingSuggestions.get(suggestionId); |
| 173 | + if (!pending) { |
| 174 | + return { |
| 175 | + content: [{ |
| 176 | + type: 'text', |
| 177 | + text: `Suggestion "${suggestionId}" not found or already resolved.`, |
| 178 | + }], |
| 179 | + isError: true, |
| 180 | + }; |
| 181 | + } |
| 182 | + |
| 183 | + // Clean up the pending suggestion |
| 184 | + this.pendingSuggestions.delete(suggestionId); |
| 185 | + |
| 186 | + if (decision === 'deny') { |
| 187 | + return { |
| 188 | + content: [{ type: 'text', text: 'Policy change denied. The restriction remains in place.' }], |
| 189 | + }; |
| 190 | + } |
| 191 | + |
| 192 | + // Apply the change |
| 193 | + try { |
| 194 | + const updated = applyPolicyChange(session.policy, pending.suggestion); |
| 195 | + Object.assign(session.policy, updated); |
| 196 | + |
| 197 | + if (decision === 'add-to-policy') { |
| 198 | + writePolicyToFile(updated, this.policyPath); |
| 199 | + return { |
| 200 | + content: [{ |
| 201 | + type: 'text', |
| 202 | + text: `Policy updated and saved to disk. ${pending.suggestion.description} You can now retry the original action.`, |
| 203 | + }], |
| 204 | + }; |
| 205 | + } |
| 206 | + |
| 207 | + // allow-once |
| 208 | + return { |
| 209 | + content: [{ |
| 210 | + type: 'text', |
| 211 | + text: `Policy updated for this session only (not saved to disk). ${pending.suggestion.description} You can now retry the original action.`, |
| 212 | + }], |
| 213 | + }; |
| 214 | + } catch (err) { |
| 215 | + return { |
| 216 | + content: [{ |
| 217 | + type: 'text', |
| 218 | + text: `Failed to apply policy change: ${(err as Error).message}`, |
| 219 | + }], |
| 220 | + isError: true, |
| 221 | + }; |
| 222 | + } |
| 223 | + } |
| 224 | + |
| 225 | + /** |
| 226 | + * Check whether a tool name is handled by this evolution handler. |
| 227 | + */ |
| 228 | + isEvolutionTool(toolName: string): boolean { |
| 229 | + return toolName === TOOL_NAME; |
| 230 | + } |
| 231 | + |
| 232 | + /** |
| 233 | + * Get the number of pending suggestions (useful for testing). |
| 234 | + */ |
| 235 | + getPendingCount(): number { |
| 236 | + return this.pendingSuggestions.size; |
| 237 | + } |
| 238 | +} |
0 commit comments