fix(mcp): harden agent write paths

This commit is contained in:
Saravana Achu Mac 2026-05-05 10:08:58 -07:00
parent 3c5b9986c9
commit 8a53dfd489
2 changed files with 306 additions and 30 deletions

View File

@ -1,11 +1,28 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { NOTES_MCP_TOOL_NAMES } from './note-tool-contracts.js';
const { listNotesMock, getNoteMock, createNoteMock, createNoteAgentActionMock } = vi.hoisted(() => ({
const {
listNotesMock,
getNoteMock,
createNoteMock,
updateNoteMock,
createRelationshipMock,
createNoteTaskMock,
createNoteArtifactMock,
createNoteAgentActionMock,
listNoteAgentActionsMock,
getWorkspaceMock,
} = vi.hoisted(() => ({
listNotesMock: vi.fn(),
getNoteMock: vi.fn(),
createNoteMock: vi.fn(),
updateNoteMock: vi.fn(),
createRelationshipMock: vi.fn(),
createNoteTaskMock: vi.fn(),
createNoteArtifactMock: vi.fn(),
createNoteAgentActionMock: vi.fn(),
listNoteAgentActionsMock: vi.fn(),
getWorkspaceMock: vi.fn(),
}));
vi.mock('../lib/product-config.js', () => ({ PRODUCT_ID: 'notelett' }));
@ -13,9 +30,23 @@ vi.mock('../modules/notes/repository.js', () => ({
listNotes: listNotesMock,
getNote: getNoteMock,
createNote: createNoteMock,
updateNote: updateNoteMock,
}));
vi.mock('../modules/note-agent-actions/repository.js', () => ({
createNoteAgentAction: createNoteAgentActionMock,
listNoteAgentActions: listNoteAgentActionsMock,
}));
vi.mock('../modules/note-relationships/repository.js', () => ({
createRelationship: createRelationshipMock,
}));
vi.mock('../modules/note-tasks/repository.js', () => ({
createNoteTask: createNoteTaskMock,
}));
vi.mock('../modules/note-artifacts/repository.js', () => ({
createNoteArtifact: createNoteArtifactMock,
}));
vi.mock('../modules/workspaces/repository.js', () => ({
getWorkspace: getWorkspaceMock,
}));
import { NotesExecutableMcpTools, getNotesExecutableMcpTool } from './note-tools.js';
@ -35,6 +66,18 @@ const req = {
describe('note executable MCP tools', () => {
beforeEach(() => {
vi.clearAllMocks();
listNoteAgentActionsMock.mockResolvedValue({ items: [], total: 0 });
getWorkspaceMock.mockResolvedValue({
id: 'ws_1',
productId: 'notelett',
userId: 'user_1',
name: 'Workspace',
members: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T00:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
});
});
it('exposes executable tools for each core note contract', () => {
@ -164,6 +207,7 @@ describe('note executable MCP tools', () => {
expect(createNoteMock).not.toHaveBeenCalled();
expect(createNoteAgentActionMock).not.toHaveBeenCalled();
expect(listNoteAgentActionsMock).not.toHaveBeenCalled();
expect(result).toMatchObject({ dryRun: true, state: 'proposed' });
});
@ -208,4 +252,185 @@ describe('note executable MCP tools', () => {
correlationId: 'corr_1',
});
});
it('rejects admin write tools for non-admin MCP callers', async () => {
const tool = getNotesExecutableMcpTool(NOTES_MCP_TOOL_NAMES.createDraft);
await expect(tool?.execute(
tool.inputSchema.parse({
workspaceId: 'ws_1',
title: 'Draft title',
body: 'Draft body',
}),
{ ...req, jwtPayload: { ...req.jwtPayload, role: 'viewer' } }
)).rejects.toThrow('MCP tool requires admin role');
expect(createNoteMock).not.toHaveBeenCalled();
});
it('rejects create_draft outside a user-owned product workspace', async () => {
getWorkspaceMock.mockResolvedValueOnce({
id: 'ws_1',
productId: 'other-product',
userId: 'user_1',
name: 'Foreign Workspace',
members: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T00:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
});
const tool = getNotesExecutableMcpTool(NOTES_MCP_TOOL_NAMES.createDraft);
await expect(tool?.execute(
tool.inputSchema.parse({
workspaceId: 'ws_1',
title: 'Draft title',
body: 'Draft body',
}),
req
)).rejects.toThrow('Workspace not found');
expect(createNoteMock).not.toHaveBeenCalled();
});
it('rejects duplicate idempotency keys before creating another draft', async () => {
listNoteAgentActionsMock.mockResolvedValueOnce({
items: [{ toolName: NOTES_MCP_TOOL_NAMES.createDraft }],
total: 1,
});
const tool = getNotesExecutableMcpTool(NOTES_MCP_TOOL_NAMES.createDraft);
await expect(tool?.execute(
tool.inputSchema.parse({
workspaceId: 'ws_1',
title: 'Draft title',
body: 'Draft body',
idempotencyKey: 'idem_1',
}),
req
)).rejects.toThrow('Duplicate idempotency key for MCP tool');
expect(createNoteMock).not.toHaveBeenCalled();
expect(createNoteAgentActionMock).not.toHaveBeenCalled();
});
it('updates a scoped note and records audit metadata', async () => {
getNoteMock.mockResolvedValue({
id: 'note_1',
productId: 'notelett',
workspaceId: 'ws_1',
userId: 'user_1',
title: 'Before',
body: 'Body',
status: 'draft',
tags: [],
links: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T00:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
});
updateNoteMock.mockResolvedValue({
id: 'note_1',
productId: 'notelett',
workspaceId: 'ws_1',
userId: 'user_1',
title: 'After',
body: 'Body',
status: 'active',
tags: ['ready'],
links: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T01:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
});
const tool = getNotesExecutableMcpTool(NOTES_MCP_TOOL_NAMES.updateNote);
const result = await tool?.execute(
tool.inputSchema.parse({
noteId: 'note_1',
workspaceId: 'ws_1',
title: 'After',
status: 'active',
tags: ['ready'],
agentId: 'agent_1',
idempotencyKey: 'idem_update',
}),
req
);
expect(updateNoteMock).toHaveBeenCalledWith('note_1', 'ws_1', expect.objectContaining({
title: 'After',
status: 'active',
tags: ['ready'],
updatedBy: 'user_1',
}));
expect(createNoteAgentActionMock).toHaveBeenCalledWith(expect.objectContaining({
productId: 'notelett',
workspaceId: 'ws_1',
userId: 'user_1',
noteId: 'note_1',
actorId: 'agent_1',
actorType: 'agent',
toolName: NOTES_MCP_TOOL_NAMES.updateNote,
actionType: 'update',
state: 'applied',
idempotencyKey: 'idem_update',
correlationId: 'req_1',
}));
expect(result).toMatchObject({ dryRun: false, idempotencyKey: 'idem_update' });
});
it('requires both linked notes to belong to the caller and product scope', async () => {
getNoteMock
.mockResolvedValueOnce({
id: 'from',
productId: 'notelett',
workspaceId: 'ws_1',
userId: 'user_1',
title: 'From',
body: 'Body',
status: 'active',
tags: [],
links: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T00:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
})
.mockResolvedValueOnce({
id: 'to',
productId: 'other-product',
workspaceId: 'ws_1',
userId: 'user_1',
title: 'To',
body: 'Body',
status: 'active',
tags: [],
links: [],
createdAt: '2026-03-10T00:00:00.000Z',
updatedAt: '2026-03-10T00:00:00.000Z',
createdBy: 'user_1',
updatedBy: 'user_1',
});
const tool = getNotesExecutableMcpTool(NOTES_MCP_TOOL_NAMES.linkNotes);
await expect(tool?.execute(
tool.inputSchema.parse({
workspaceId: 'ws_1',
fromNoteId: 'from',
toNoteId: 'to',
relationshipType: 'related',
}),
req
)).rejects.toThrow('Note not found');
expect(createRelationshipMock).not.toHaveBeenCalled();
expect(createNoteAgentActionMock).not.toHaveBeenCalled();
});
});

View File

@ -1,9 +1,10 @@
import { randomUUID } from 'node:crypto';
import type { ZodTypeAny } from 'zod';
import { ConflictError, ForbiddenError } from '@bytelyst/errors';
import { PRODUCT_ID } from '../lib/product-config.js';
import { createNote, getNote, listNotes, updateNote } from '../modules/notes/repository.js';
import type { NoteDoc } from '../modules/notes/types.js';
import { createNoteAgentAction } from '../modules/note-agent-actions/repository.js';
import { createNoteAgentAction, listNoteAgentActions } from '../modules/note-agent-actions/repository.js';
import type { NoteAgentActionDoc } from '../modules/note-agent-actions/types.js';
import { createRelationship } from '../modules/note-relationships/repository.js';
import type { NoteRelationshipDoc } from '../modules/note-relationships/types.js';
@ -11,6 +12,7 @@ import { createNoteTask } from '../modules/note-tasks/repository.js';
import type { NoteTaskDoc } from '../modules/note-tasks/types.js';
import { createNoteArtifact } from '../modules/note-artifacts/repository.js';
import type { NoteArtifactDoc } from '../modules/note-artifacts/types.js';
import { getWorkspace } from '../modules/workspaces/repository.js';
import {
AttachArtifactToolOutputSchema,
CreateNoteDraftToolOutputSchema,
@ -84,6 +86,55 @@ function requireUserId(req: NotesMcpRequest): string {
return userId;
}
function requireToolRole(req: NotesMcpRequest, requiredRole: NoteToolRole): void {
if (requiredRole === 'viewer') return;
const role = req.jwtPayload?.role;
if (requiredRole === 'admin' && (role === 'admin' || role === 'owner' || role === 'super_admin')) {
return;
}
if (requiredRole === 'super_admin' && role === 'super_admin') {
return;
}
throw new ForbiddenError(`MCP tool requires ${requiredRole} role`);
}
async function requireWorkspaceAccess(workspaceId: string, userId: string): Promise<void> {
const workspace = await getWorkspace(workspaceId, userId);
if (!workspace || workspace.productId !== PRODUCT_ID) {
throw new Error('Workspace not found');
}
}
async function requireScopedNote(noteId: string, workspaceId: string, userId: string): Promise<NoteDoc> {
const note = await getNote(noteId, workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
return note;
}
async function assertUnusedIdempotencyKey(
userId: string,
workspaceId: string,
toolName: string,
idempotencyKey?: string,
): Promise<void> {
if (!idempotencyKey) return;
const { items } = await listNoteAgentActions(userId, PRODUCT_ID, {
workspaceId,
idempotencyKey,
limit: 10,
offset: 0,
});
if (items.some(action => action.toolName === toolName)) {
throw new ConflictError('Duplicate idempotency key for MCP tool');
}
}
function mapNoteSummary(note: NoteDoc) {
return {
id: note.id,
@ -113,6 +164,7 @@ function mapNote(note: NoteDoc) {
}
async function executeListNotes(args: ListNotesToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.list.requiredRole);
const userId = requireUserId(req);
const result = await listNotes(userId, PRODUCT_ID, args);
return ListNotesToolOutputSchema.parse({
@ -124,11 +176,9 @@ async function executeListNotes(args: ListNotesToolInput, req: NotesMcpRequest)
}
async function executeGetNote(args: GetNoteToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.get.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
return mapNote(note);
}
@ -148,6 +198,7 @@ function getMatchFields(note: NoteDoc, query: string): Array<'title' | 'body' |
}
async function executeSearchNotes(args: SearchNotesToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.search.requiredRole);
const userId = requireUserId(req);
const result = await listNotes(userId, PRODUCT_ID, {
workspaceId: args.workspaceId,
@ -222,7 +273,9 @@ async function recordCreateDraftAction(
}
async function executeCreateDraft(args: CreateNoteDraftToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.createDraft.requiredRole);
const userId = requireUserId(req);
await requireWorkspaceAccess(args.workspaceId, userId);
const draft = buildDraftNote(args, userId);
if (args.dryRun) {
@ -235,6 +288,7 @@ async function executeCreateDraft(args: CreateNoteDraftToolInput, req: NotesMcpR
});
}
await assertUnusedIdempotencyKey(userId, args.workspaceId, NOTES_MCP_TOOL_NAMES.createDraft, args.idempotencyKey);
const created = await createNote(draft);
await recordCreateDraftAction(created, args, req, userId);
@ -248,11 +302,9 @@ async function executeCreateDraft(args: CreateNoteDraftToolInput, req: NotesMcpR
}
async function executeUpdateNote(args: UpdateNoteToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.updateNote.requiredRole);
const userId = requireUserId(req);
const existing = await getNote(args.noteId, args.workspaceId);
if (!existing || existing.userId !== userId || existing.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const existing = await requireScopedNote(args.noteId, args.workspaceId, userId);
const updates: Partial<NoteDoc> = {
updatedAt: new Date().toISOString(),
@ -273,6 +325,7 @@ async function executeUpdateNote(args: UpdateNoteToolInput, req: NotesMcpRequest
});
}
await assertUnusedIdempotencyKey(userId, args.workspaceId, NOTES_MCP_TOOL_NAMES.updateNote, args.idempotencyKey);
const updated = await updateNote(args.noteId, args.workspaceId, updates);
if (!updated) throw new Error('Update failed');
@ -308,7 +361,11 @@ async function executeUpdateNote(args: UpdateNoteToolInput, req: NotesMcpRequest
}
async function executeLinkNotes(args: LinkNotesToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.linkNotes.requiredRole);
const userId = requireUserId(req);
await requireScopedNote(args.fromNoteId, args.workspaceId, userId);
await requireScopedNote(args.toNoteId, args.workspaceId, userId);
await assertUnusedIdempotencyKey(userId, args.workspaceId, NOTES_MCP_TOOL_NAMES.linkNotes, args.idempotencyKey);
const now = new Date().toISOString();
const relationship: NoteRelationshipDoc = {
@ -374,11 +431,9 @@ function extractTasksFromBody(body: string): Array<{ title: string; description?
}
async function executeExtractTasks(args: ExtractTasksToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.extractTasks.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
const extracted = extractTasksFromBody(note.body);
const now = new Date().toISOString();
@ -400,6 +455,7 @@ async function executeExtractTasks(args: ExtractTasksToolInput, req: NotesMcpReq
}));
if (!args.dryRun) {
await assertUnusedIdempotencyKey(userId, args.workspaceId, NOTES_MCP_TOOL_NAMES.extractTasks, args.idempotencyKey);
for (const task of taskDocs) {
await createNoteTask(task);
}
@ -444,7 +500,9 @@ async function executeExtractTasks(args: ExtractTasksToolInput, req: NotesMcpReq
}
async function executeAttachArtifact(args: AttachArtifactToolInput, req: NotesMcpRequest) {
requireToolRole(req, NotesMcpToolDefinitions.attachArtifact.requiredRole);
const userId = requireUserId(req);
await requireScopedNote(args.noteId, args.workspaceId, userId);
const now = new Date().toISOString();
const artifactDoc: NoteArtifactDoc = {
@ -481,6 +539,7 @@ async function executeAttachArtifact(args: AttachArtifactToolInput, req: NotesMc
});
}
await assertUnusedIdempotencyKey(userId, args.workspaceId, NOTES_MCP_TOOL_NAMES.attachArtifact, args.idempotencyKey);
const created = await createNoteArtifact(artifactDoc);
const action: NoteAgentActionDoc = {
@ -524,11 +583,9 @@ async function executeAttachArtifact(args: AttachArtifactToolInput, req: NotesMc
// ── Run Prompt MCP tool implementation ────────────────────────
async function executeRunPrompt(args: RunPromptToolInput, req: NotesMcpRequest) {
requireToolRole(req, RunPromptMcpToolDefinition.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
const { executePrompt } = await import('../modules/note-prompts/runner.js');
const promptRepo = await import('../modules/note-prompts/repository.js');
@ -581,11 +638,9 @@ async function executeRunPrompt(args: RunPromptToolInput, req: NotesMcpRequest)
// ── Smart Action MCP tool implementations ─────────────────────
async function executeSuggestTags(args: SuggestTagsToolInput, req: NotesMcpRequest) {
requireToolRole(req, SmartActionMcpToolDefinitions.suggestTags.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
const plain = stripHtmlForEmbedding(note.body ?? '');
const provider = llm();
@ -608,11 +663,9 @@ async function executeSuggestTags(args: SuggestTagsToolInput, req: NotesMcpReque
}
async function executeCheckDuplicates(args: CheckDuplicatesToolInput, req: NotesMcpRequest) {
requireToolRole(req, SmartActionMcpToolDefinitions.checkDuplicates.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
const plain = stripHtmlForEmbedding(note.body ?? '');
const noteEmbedding = await embedText(plain);
@ -647,11 +700,9 @@ async function executeCheckDuplicates(args: CheckDuplicatesToolInput, req: Notes
}
async function executeSuggestLinks(args: SuggestLinksToolInput, req: NotesMcpRequest) {
requireToolRole(req, SmartActionMcpToolDefinitions.suggestLinks.requiredRole);
const userId = requireUserId(req);
const note = await getNote(args.noteId, args.workspaceId);
if (!note || note.userId !== userId || note.productId !== PRODUCT_ID) {
throw new Error('Note not found');
}
const note = await requireScopedNote(args.noteId, args.workspaceId, userId);
const plain = stripHtmlForEmbedding(note.body ?? '');
const noteEmbedding = await embedText(plain);