diff --git a/apps/server/src/services/llm/tools/attribute_tools.ts b/apps/server/src/services/llm/tools/attribute_tools.ts index f17799f7b2..ffb9b24bee 100644 --- a/apps/server/src/services/llm/tools/attribute_tools.ts +++ b/apps/server/src/services/llm/tools/attribute_tools.ts @@ -64,7 +64,7 @@ export const attributeTools = defineTools({ value: z.string().optional().describe("The attribute value (for relations, this must be a target noteId)") }), mutates: true, - execute: async ({ noteId, type, name, value = "" }) => { + execute: ({ noteId, type, name, value = "" }) => { const note = becca.getNote(noteId); if (!note) { return { error: "Note not found" }; @@ -98,7 +98,7 @@ export const attributeTools = defineTools({ attributeId: z.string().describe("The ID of the attribute to delete") }), mutates: true, - execute: async ({ noteId, attributeId }) => { + execute: ({ noteId, attributeId }) => { const attribute = becca.getAttribute(attributeId); if (!attribute) { return { error: "Attribute not found" }; diff --git a/apps/server/src/services/llm/tools/note_tools.ts b/apps/server/src/services/llm/tools/note_tools.ts index 3c2ece135d..1702eeab54 100644 --- a/apps/server/src/services/llm/tools/note_tools.ts +++ b/apps/server/src/services/llm/tools/note_tools.ts @@ -103,7 +103,7 @@ export const noteTools = defineTools({ content: z.string().describe("The new content for the note (Markdown for text notes, plain text for code notes)") }), mutates: true, - execute: async ({ noteId, content }) => { + execute: ({ noteId, content }) => { const note = becca.getNote(noteId); if (!note) { return { error: "Note not found" }; @@ -132,7 +132,7 @@ export const noteTools = defineTools({ content: z.string().describe("The content to append (Markdown for text notes, plain text for code notes)") }), mutates: true, - execute: async ({ noteId, content }) => { + execute: ({ noteId, content }) => { const note = becca.getNote(noteId); if (!note) { return { error: "Note not found" }; @@ -195,7 +195,7 @@ export const noteTools = defineTools({ mime: z.string().optional().describe("MIME type, REQUIRED for code notes (e.g. 'application/javascript;env=backend', 'text/jsx'). Ignored for other types.") }), mutates: true, - execute: async ({ parentNoteId, title, content, type, mime }) => { + execute: ({ parentNoteId, title, content, type, mime }) => { if (type === "code" && !mime) { return { error: "mime is required when creating code notes" }; } diff --git a/apps/server/src/services/llm/tools/tool_registry.ts b/apps/server/src/services/llm/tools/tool_registry.ts index 0df416037b..b8a1b6995e 100644 --- a/apps/server/src/services/llm/tools/tool_registry.ts +++ b/apps/server/src/services/llm/tools/tool_registry.ts @@ -12,14 +12,36 @@ import { tool } from "ai"; import type { z } from "zod"; import type { ToolSet } from "ai"; -export interface ToolDefinition { +import sql from "../../sql.js"; + +/** + * Type constraint that rejects Promises at compile time. + * Works by requiring `then` to be void if present - Promises have `then: Function`. + */ +type NotAPromise = T & { then?: void }; + +interface MutatingToolDefinition { description: string; inputSchema: z.ZodType; - /** Whether this tool modifies data (needs CLS + transaction wrapping). */ - mutates?: boolean; - execute: (args: any) => Promise; + /** Marks this tool as modifying data (needs CLS + transaction wrapping). */ + mutates: true; + /** + * Execute the tool synchronously. Must NOT be async because better-sqlite3 + * transactions are synchronous and would commit before awaits complete. + */ + execute: (args: any) => NotAPromise; } +interface ReadOnlyToolDefinition { + description: string; + inputSchema: z.ZodType; + mutates?: false; + /** Execute the tool. May be async for read-only operations. */ + execute: (args: any) => unknown; +} + +export type ToolDefinition = MutatingToolDefinition | ReadOnlyToolDefinition; + /** * A named collection of tool definitions that can be iterated or converted * to an AI SDK ToolSet. @@ -34,14 +56,20 @@ export class ToolRegistry implements Iterable<[string, ToolDefinition]> { /** * Convert to an AI SDK ToolSet for use with the LLM chat providers. + * Mutating tools are wrapped in a transaction for consistency with MCP. + * (CLS context is provided by the route handler.) */ toToolSet(): ToolSet { const set: ToolSet = {}; for (const [name, def] of this) { + const execute = def.mutates + ? (args: unknown) => sql.transactional(() => def.execute(args)) + : def.execute; + set[name] = tool({ description: def.description, inputSchema: def.inputSchema, - execute: def.execute + execute }); } return set; @@ -53,9 +81,13 @@ export class ToolRegistry implements Iterable<[string, ToolDefinition]> { * * ```ts * export const noteTools = defineTools({ - * search_notes: { description: "...", inputSchema: z.object({...}), execute: async (args) => {...} }, + * search_notes: { description: "...", inputSchema: z.object({...}), execute: (args) => {...} }, + * create_note: { description: "...", inputSchema: z.object({...}), mutates: true, execute: (args) => {...} }, * }); * ``` + * + * Note: Mutating tools (mutates: true) MUST have synchronous execute functions + * because better-sqlite3 transactions are synchronous. */ export function defineTools(tools: Record): ToolRegistry { return new ToolRegistry(tools); diff --git a/apps/server/src/services/mcp/mcp_server.ts b/apps/server/src/services/mcp/mcp_server.ts index 67cc8af3fc..382b5025c7 100644 --- a/apps/server/src/services/mcp/mcp_server.ts +++ b/apps/server/src/services/mcp/mcp_server.ts @@ -24,11 +24,15 @@ function registerTool(server: McpServer, name: string, def: ToolDefinition) { server.registerTool(name, { description: def.description, inputSchema: def.inputSchema - }, async (args: any): Promise => { - const run = () => def.execute(args); - const result = def.mutates - ? await cls.init(() => sql.transactional(run)) - : await run(); + }, (args: any): CallToolResult => { + const result = cls.init(() => { + cls.set("componentId", "mcp"); + + return def.mutates + ? sql.transactional(() => def.execute(args)) + : def.execute(args); + }); + return { content: [{ type: "text", text: JSON.stringify(result) }] }; }); }