From 5d315fff21357698fc57171544b38e66817cc5b2 Mon Sep 17 00:00:00 2001 From: Hong Jiarong Date: Thu, 9 Jul 2026 20:48:03 +0800 Subject: [PATCH] fix: enforce role tool whitelist --- hub/src/agent/models.ts | 12 ++- hub/src/agent/roleTools.ts | 110 +++++++++++++++++++++++++++ hub/src/agent/runner.ts | 12 ++- hub/src/feishu/fileDeliveryTool.ts | 59 ++++++++++---- hub/src/feishu/trigger.ts | 9 ++- hub/test/integration/trigger.test.ts | 16 ++++ hub/test/unit/runner.test.ts | 44 +++++++++++ 7 files changed, 242 insertions(+), 20 deletions(-) create mode 100644 hub/src/agent/roleTools.ts diff --git a/hub/src/agent/models.ts b/hub/src/agent/models.ts index 700e3a5..d10b27c 100644 --- a/hub/src/agent/models.ts +++ b/hub/src/agent/models.ts @@ -17,10 +17,11 @@ * * A role bundles everything that distinguishes one agent persona from another: * model, system prompt, and the tool surface (files / cph / feishu / skills / - * mcps). Per-run tool whitelisting is enforced by {@link ToolRegistry.subset}; - * the model never sees tools outside its role's whitelist, even if it tries to - * call them — security does not rely on the system prompt. + * mcps). Per-run tool whitelisting is enforced by the Claude SDK runner and + * the cph_hub MCP server builder; security does not rely on the system prompt. */ +import { assertSupportedRoleTools } from "./roleTools.js"; + export interface RoleEntry { readonly id: string; /** Human label for the teacher-side switcher UX. */ @@ -36,7 +37,7 @@ export interface RoleEntry { /** * Tool names this role may use (whitelist). `undefined` ⇒ the full registered * set (back-compat / unrestricted roles). An empty array ⇒ no tools at all. - * Names not present in the registry are silently dropped at run setup. + * Invalid names fail fast when settings are loaded or the run is set up. */ readonly tools?: readonly string[] | undefined; } @@ -72,6 +73,9 @@ export class InMemoryModelRegistry implements ModelRegistry { constructor(models: readonly ModelEntry[], roles: readonly RoleEntry[] = []) { this.models = models; + for (const role of roles) { + if (role.tools !== undefined) assertSupportedRoleTools(role.tools); + } this.roles = new Map(roles.map((r) => [r.id, r])); } diff --git a/hub/src/agent/roleTools.ts b/hub/src/agent/roleTools.ts new file mode 100644 index 0000000..a157698 --- /dev/null +++ b/hub/src/agent/roleTools.ts @@ -0,0 +1,110 @@ +export const DEFAULT_CLAUDE_BUILT_IN_TOOLS = ["Read", "Write", "Bash", "Glob", "Grep"] as const; + +export const CPH_HUB_MCP_SERVER_NAME = "cph_hub"; +export const CPH_HUB_MCP_TOOL_IDS = ["send_file", "feishu_read_context", "request_approval"] as const; + +export type CphHubMcpToolId = (typeof CPH_HUB_MCP_TOOL_IDS)[number]; + +export interface ClaudeSdkToolConfig { + readonly tools: readonly string[]; + readonly allowedTools: readonly string[]; +} + +const ROLE_TOOL_TO_CLAUDE_BUILT_INS = new Map([ + ["read_file", ["Read"]], + ["write_file", ["Write"]], + ["list_files", ["Glob"]], + ["search_files", ["Grep"]], + ["bash", ["Bash"]], + // ADR-0017 replaced cph custom tools with Bash commands. Granting either + // cph role tool therefore exposes the SDK Bash tool; cph-only Bash narrowing + // would need a separate command-policy layer. + ["cph_check", ["Bash"]], + ["cph_build", ["Bash"]], + ["Read", ["Read"]], + ["Write", ["Write"]], + ["Bash", ["Bash"]], + ["Glob", ["Glob"]], + ["Grep", ["Grep"]], +]); + +const ROLE_TOOL_TO_CPH_HUB_MCP_TOOL = new Map([ + ["send_file", "send_file"], + ["feishu_read_context", "feishu_read_context"], + ["request_approval", "request_approval"], + ["mcp__cph_hub__send_file", "send_file"], + ["mcp__cph_hub__feishu_read_context", "feishu_read_context"], + ["mcp__cph_hub__request_approval", "request_approval"], +]); + +const SUPPORTED_ROLE_TOOLS = new Set([ + ...ROLE_TOOL_TO_CLAUDE_BUILT_INS.keys(), + ...ROLE_TOOL_TO_CPH_HUB_MCP_TOOL.keys(), +]); + +export function claudeSdkToolConfigForRole(roleTools: readonly string[] | undefined): ClaudeSdkToolConfig { + if (roleTools === undefined) { + const mcpTools = CPH_HUB_MCP_TOOL_IDS.map(claudeMcpToolName); + return { + tools: [...DEFAULT_CLAUDE_BUILT_IN_TOOLS], + allowedTools: [...DEFAULT_CLAUDE_BUILT_IN_TOOLS, ...mcpTools], + }; + } + + const builtIns: string[] = []; + const allowedTools: string[] = []; + for (const roleTool of roleTools) { + assertSupportedRoleTool(roleTool); + for (const tool of ROLE_TOOL_TO_CLAUDE_BUILT_INS.get(roleTool) ?? []) { + pushUnique(builtIns, tool); + pushUnique(allowedTools, tool); + } + + const mcpTool = ROLE_TOOL_TO_CPH_HUB_MCP_TOOL.get(roleTool); + if (mcpTool !== undefined) { + pushUnique(allowedTools, claudeMcpToolName(mcpTool)); + } + } + + return { tools: builtIns, allowedTools }; +} + +export function cphHubMcpToolsForRole(roleTools: readonly string[] | undefined): readonly CphHubMcpToolId[] { + if (roleTools === undefined) return [...CPH_HUB_MCP_TOOL_IDS]; + + const tools: CphHubMcpToolId[] = []; + for (const roleTool of roleTools) { + assertSupportedRoleTool(roleTool); + const mcpTool = ROLE_TOOL_TO_CPH_HUB_MCP_TOOL.get(roleTool); + if (mcpTool !== undefined) pushUnique(tools, mcpTool); + } + return tools; +} + +export function roleToolsAllow(roleTools: readonly string[] | undefined, roleTool: string): boolean { + if (roleTools === undefined) return true; + for (const configured of roleTools) { + assertSupportedRoleTool(configured); + if (configured === roleTool) return true; + if (ROLE_TOOL_TO_CPH_HUB_MCP_TOOL.get(configured) === roleTool) return true; + } + return false; +} + +export function assertSupportedRoleTools(roleTools: readonly string[]): void { + for (const roleTool of roleTools) assertSupportedRoleTool(roleTool); +} + +function claudeMcpToolName(tool: CphHubMcpToolId): string { + return `mcp__${CPH_HUB_MCP_SERVER_NAME}__${tool}`; +} + +function assertSupportedRoleTool(roleTool: string): void { + if (!SUPPORTED_ROLE_TOOLS.has(roleTool)) { + throw new Error(`unknown role tool id: ${roleTool}`); + } +} + +function pushUnique(items: T[], item: T): void { + if (!items.includes(item)) items.push(item); +} diff --git a/hub/src/agent/runner.ts b/hub/src/agent/runner.ts index a14f66a..82b259d 100644 --- a/hub/src/agent/runner.ts +++ b/hub/src/agent/runner.ts @@ -30,6 +30,7 @@ import { homedir } from "node:os"; import { query, type McpServerConfig, type SDKMessage, type SDKAssistantMessage, type SDKUserMessage, type SDKResultMessage, type SDKPartialAssistantMessage } from "@anthropic-ai/claude-agent-sdk"; import type { PrismaClient } from "@prisma/client"; +import { claudeSdkToolConfigForRole } from "./roleTools.js"; export interface ProjectContext { readonly projectId: string; @@ -54,6 +55,13 @@ export interface RunRequest { readonly systemPrompt: string | undefined; readonly providerEnv?: Record | undefined; readonly resumeSessionId?: string | undefined; + /** + * RoleEntry.tools from admin/runtime settings. Values are Hub role tool ids + * (for example read_file/cph_check/send_file), mapped to Claude SDK tool + * names at run setup. `undefined` means the full supported tool surface; [] + * means no tools. + */ + readonly tools?: readonly string[] | undefined; readonly mcpServers?: Record | undefined; readonly maxTurns?: number; readonly runId: string; @@ -131,10 +139,12 @@ export async function runAgent(req: RunRequest): Promise { let error: string | undefined; try { await persistAgentMessage(req, "user", req.prompt); + const toolConfig = claudeSdkToolConfigForRole(req.tools); const options: Parameters[0]["options"] = { cwd: req.project.workspaceDir, - allowedTools: ["Read", "Write", "Bash", "Glob", "Grep"], + tools: [...toolConfig.tools], + allowedTools: [...toolConfig.allowedTools], maxTurns: cap, includePartialMessages: true, // ADR-0018: bypass interactive prompts (headless server); the sandbox diff --git a/hub/src/feishu/fileDeliveryTool.ts b/hub/src/feishu/fileDeliveryTool.ts index 76ae653..2467525 100644 --- a/hub/src/feishu/fileDeliveryTool.ts +++ b/hub/src/feishu/fileDeliveryTool.ts @@ -1,9 +1,10 @@ -import { createSdkMcpServer, tool, type McpSdkServerConfigWithInstance } from "@anthropic-ai/claude-agent-sdk"; +import { createSdkMcpServer, tool, type McpSdkServerConfigWithInstance, type SdkMcpToolDefinition } from "@anthropic-ai/claude-agent-sdk"; import { z } from "zod"; import { sendApprovalCard, sendFile, type FeishuRuntime, type SendMessageOptions } from "./client.js"; import { resolveDeliverableFile } from "./fileDelivery.js"; import { readFeishuContext } from "./read.js"; import type { ApprovalManager } from "./approval.js"; +import { CPH_HUB_MCP_TOOL_IDS, type CphHubMcpToolId } from "../agent/roleTools.js"; export interface FileDeliveryToolOptions { readonly rt: FeishuRuntime; @@ -14,20 +15,15 @@ export interface FileDeliveryToolOptions { readonly sendOptions?: SendMessageOptions | undefined; readonly approvalManager: ApprovalManager; readonly onDelivered?: (path: string) => void; + readonly tools?: readonly CphHubMcpToolId[] | undefined; } export function createFileDeliveryMcpServer(options: FileDeliveryToolOptions): McpSdkServerConfigWithInstance { - return createSdkMcpServer({ - name: "cph_hub", - version: "0.0.0", - alwaysLoad: true, - instructions: - "Use send_file when the user asks to receive, resend, download, or attach a file. " + - "Do not claim a file was sent unless send_file returns success. " + - "If a generated file path is uncertain, list or inspect the workspace first, then call send_file with the actual path. " + - "Use feishu_read_context when the trigger context contains a reply, thread, or message anchor and more Feishu context is needed. " + - "Use request_approval when explicit human approval or confirmation is required before continuing.", - tools: [ + const enabledTools = new Set(options.tools ?? CPH_HUB_MCP_TOOL_IDS); + const tools: Array> = []; + + if (enabledTools.has("send_file")) { + tools.push( tool( "send_file", "Upload an existing project/repository file to the current Feishu chat. The path must point to a concrete file, for example build/student.pdf or README.md.", @@ -59,6 +55,11 @@ export function createFileDeliveryMcpServer(options: FileDeliveryToolOptions): M }, { alwaysLoad: true }, ), + ); + } + + if (enabledTools.has("feishu_read_context")) { + tools.push( tool( "feishu_read_context", "Read Feishu context for the current project's bound chat. Use anchor ids from the trigger context; this tool never reads arbitrary chats.", @@ -81,6 +82,11 @@ export function createFileDeliveryMcpServer(options: FileDeliveryToolOptions): M }, { alwaysLoad: true }, ), + ); + } + + if (enabledTools.has("request_approval")) { + tools.push( tool( "request_approval", "Send an interactive Feishu approval/confirmation card to the current chat and wait for a button click.", @@ -128,6 +134,33 @@ export function createFileDeliveryMcpServer(options: FileDeliveryToolOptions): M }, { alwaysLoad: true }, ), - ], + ); + } + + const instructions = mcpInstructions(enabledTools); + return createSdkMcpServer({ + name: "cph_hub", + version: "0.0.0", + alwaysLoad: true, + ...(instructions === "" ? {} : { instructions }), + tools, }); } + +function mcpInstructions(enabledTools: ReadonlySet): string { + const instructions: string[] = []; + if (enabledTools.has("send_file")) { + instructions.push( + "Use send_file when the user asks to receive, resend, download, or attach a file.", + "Do not claim a file was sent unless send_file returns success.", + "If a generated file path is uncertain, list or inspect the workspace first, then call send_file with the actual path.", + ); + } + if (enabledTools.has("feishu_read_context")) { + instructions.push("Use feishu_read_context when the trigger context contains a reply, thread, or message anchor and more Feishu context is needed."); + } + if (enabledTools.has("request_approval")) { + instructions.push("Use request_approval when explicit human approval or confirmation is required before continuing."); + } + return instructions.join(" "); +} diff --git a/hub/src/feishu/trigger.ts b/hub/src/feishu/trigger.ts index dfccd44..200cecb 100644 --- a/hub/src/feishu/trigger.ts +++ b/hub/src/feishu/trigger.ts @@ -41,6 +41,7 @@ import { ApprovalManager } from "./approval.js"; import { SenderNameCache } from "./senderCache.js"; import { TriggerQueue, triggerQueue as defaultTriggerQueue, type QueuedTrigger } from "./triggerQueue.js"; import { createSlashCommandRegistry, formatSlashHelpTarget, parseSlashHelpSubcommand, parseSlashInvocation } from "./slashCommands.js"; +import { cphHubMcpToolsForRole, roleToolsAllow } from "../agent/roleTools.js"; export { ApprovalManager } from "./approval.js"; export type { ApprovalResult, PendingApproval } from "./approval.js"; @@ -183,7 +184,8 @@ export function makeTriggerHandler(deps: TriggerDeps): TriggerHandler { // Per-role bundle (ADR-0017): systemPrompt seeds persona; tools whitelist // is enforced inside runAgent when it builds the SDK tools record. `tools: // undefined` means the full registered set (unrestricted role). - const systemPrompt = withFileDeliveryInstructions(role?.systemPrompt); + const roleTools = role?.tools; + const systemPrompt = withFileDeliveryInstructions(role?.systemPrompt, roleTools); const feishuTriggerContext = await buildFeishuTriggerContext({ msg, chatId, @@ -299,6 +301,7 @@ export function makeTriggerHandler(deps: TriggerDeps): TriggerHandler { workspaceDir: project.workspaceDir, sendOptions, approvalManager, + tools: cphHubMcpToolsForRole(roleTools), onDelivered: (path) => { deliveredFiles.push(path); }, @@ -315,6 +318,7 @@ export function makeTriggerHandler(deps: TriggerDeps): TriggerHandler { systemPrompt, providerEnv: provider.sdkEnv, resumeSessionId: sessionMetadata.claudeSessionId, + tools: roleTools, mcpServers: { cph_hub: fileDeliveryMcpServer }, maxTurns: runPolicy.maxTurns, runId: run.id, @@ -963,7 +967,8 @@ async function senderAuditMetadata(rt: FeishuRuntime, openId: string): Promise

{ expect(await prisma.agentRun.findMany()).toHaveLength(0); }); + it("passes the selected role tool whitelist into the runner", async () => { + await seedProject("proj-role-tools", "chat-role-tools"); + const trigger = makeTriggerHandler({ prisma, settings, logger: silentLogger, runAgent, messageBatcherOptions: { maxMessages: 1 } }); + + await trigger(makeEvent("chat-role-tools", "@_user_1 /review 检查讲义"), rt); + + await vi.waitFor(async () => { + const runs = await prisma.agentRun.findMany(); + expect(runs).toHaveLength(1); + expect(runs[0]?.status).toBe("COMPLETED"); + }); + + expect(runAgentCalls).toHaveLength(1); + expect(runAgentCalls[0]?.tools).toEqual(["read_file"]); + }); + it("control commands support a help subcommand", async () => { await seedProject("proj-new-help", "chat-new-help"); const trigger = makeTriggerHandler({ prisma, settings, logger: silentLogger, runAgent, messageBatcherOptions: { maxMessages: 1 } }); diff --git a/hub/test/unit/runner.test.ts b/hub/test/unit/runner.test.ts index af94ee6..591d84d 100644 --- a/hub/test/unit/runner.test.ts +++ b/hub/test/unit/runner.test.ts @@ -104,6 +104,50 @@ describe("runAgent", () => { expect(call?.options).not.toHaveProperty("resume"); }); + it("maps role tool ids to the Claude SDK tool whitelist", async () => { + queryMock.mockReturnValue(messages(assistantMessage("ok"), resultMessage("sdk-session-1"))); + + await runAgent({ + prompt: "审校一下", + model: undefined, + project: { projectId: "p", boundChatId: "c", workspaceDir: "/tmp/ws" }, + systemPrompt: undefined, + tools: ["read_file", "cph_check", "send_file"], + runId: "run-1", + sessionId: "hub-session-1", + prisma: stubPrisma, + }); + + expect(queryMock.mock.calls[0]?.[0]).toMatchObject({ + options: { + tools: ["Read", "Bash"], + allowedTools: ["Read", "Bash", "mcp__cph_hub__send_file"], + }, + }); + }); + + it("disables all SDK tools for an empty role tool whitelist", async () => { + queryMock.mockReturnValue(messages(assistantMessage("ok"), resultMessage("sdk-session-1"))); + + await runAgent({ + prompt: "只回答", + model: undefined, + project: { projectId: "p", boundChatId: "c", workspaceDir: "/tmp/ws" }, + systemPrompt: undefined, + tools: [], + runId: "run-1", + sessionId: "hub-session-1", + prisma: stubPrisma, + }); + + expect(queryMock.mock.calls[0]?.[0]).toMatchObject({ + options: { + tools: [], + allowedTools: [], + }, + }); + }); + it("returns SDK-reported cost when present", async () => { queryMock.mockReturnValue(messages(assistantMessage("ok"), resultMessage("sdk-session-1", 0.0042)));