diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts index 5add9494..753a3b28 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts @@ -120,21 +120,26 @@ describe('AiChatToolsService deletePage guardrail (H4)', () => { const tools = await buildTools(); const deletePage = tools.deletePage; - // The Zod input schema only allows `pageId`; parsing strips/ignores extra - // keys, so a permanent/force flag is never part of the validated input. + // inputSchema is now an AI SDK `Schema` (not a raw zod object). Its + // `validate` runs the same zod safeParse and forwards the STRIPPED data, so + // a permanent/force flag is never part of the validated input the SDK then + // hands to execute. const schema = (deletePage as unknown as { inputSchema: unknown }) .inputSchema as { - parse: (v: unknown) => Record; + validate: ( + v: unknown, + ) => Promise<{ success: boolean; value?: Record }>; }; - const parsed = schema.parse({ + const result = await schema.validate({ pageId: 'page-789', permanentlyDelete: true, forceDelete: true, }); - expect(parsed).toHaveProperty('pageId', 'page-789'); - expect(parsed).not.toHaveProperty('permanentlyDelete'); - expect(parsed).not.toHaveProperty('forceDelete'); + expect(result.success).toBe(true); + expect(result.value).toHaveProperty('pageId', 'page-789'); + expect(result.value).not.toHaveProperty('permanentlyDelete'); + expect(result.value).not.toHaveProperty('forceDelete'); }); }); @@ -207,21 +212,25 @@ describe('AiChatToolsService expanded toolset guardrails', () => { const tools = await buildTools(); const transformPage = tools.transformPage; - // The Zod input schema only allows pageId/transformJs/dryRun; parsing - // strips unknown keys, so deleteComments can never reach the client. + // inputSchema is now an AI SDK `Schema`; its `validate` runs the same zod + // safeParse, which only allows pageId/transformJs/dryRun and strips unknown + // keys — so deleteComments can never reach the client. const schema = (transformPage as unknown as { inputSchema: unknown }) .inputSchema as { - parse: (v: unknown) => Record; + validate: ( + v: unknown, + ) => Promise<{ success: boolean; value?: Record }>; }; - const parsed = schema.parse({ + const result = await schema.validate({ pageId: 'p', transformJs: '(d)=>d', dryRun: true, deleteComments: true, }); - expect(parsed).toHaveProperty('pageId', 'p'); - expect(parsed).not.toHaveProperty('deleteComments'); + expect(result.success).toBe(true); + expect(result.value).toHaveProperty('pageId', 'p'); + expect(result.value).not.toHaveProperty('deleteComments'); }); }); diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts index f6d38487..c0916fa9 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts @@ -15,6 +15,7 @@ import { } from './docmost-client.loader'; import { resolveCurrentPageResult } from './current-page.util'; import { parseNodeArg } from './parse-node-arg'; +import { modelFriendlyInput } from './model-friendly-input'; /** * Per-user, per-request adapter that exposes Docmost READ operations to the @@ -102,9 +103,9 @@ export class AiChatToolsService { ): Tool => tool({ description: spec.description, - inputSchema: spec.buildShape - ? z.object(spec.buildShape(z) as z.ZodRawShape) - : z.object({}), + inputSchema: modelFriendlyInput( + spec.buildShape ? (spec.buildShape(z) as z.ZodRawShape) : {}, + ), execute, }); @@ -118,7 +119,7 @@ export class AiChatToolsService { 'and entities), not a full sentence. If the first results look weak ' + 'or incomplete, search again with different wording or synonyms ' + 'before answering.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ query: z.string().describe('The search query.'), limit: z .number() @@ -227,7 +228,7 @@ export class AiChatToolsService { '"the current page", or "here" refers to. Returns the page id and title, ' + 'or null if the user is not currently on a page. Call this first whenever ' + 'the user refers to the current page without giving an explicit id.', - inputSchema: z.object({}), + inputSchema: modelFriendlyInput({}), execute: async () => resolveCurrentPageResult(openedPage), }), @@ -235,7 +236,7 @@ export class AiChatToolsService { description: 'Fetch a single page as Markdown by its page id. Returns the page ' + 'title and its Markdown content.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id (or slugId) of the page.'), }), execute: async ({ pageId }) => { @@ -259,7 +260,7 @@ export class AiChatToolsService { 'Create a new page with a Markdown body in a space, optionally under ' + 'a parent page. Returns the new page id and title. Reversible: a page ' + 'can be moved to trash later.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ title: z.string().describe('The title of the new page.'), content: z .string() @@ -294,7 +295,7 @@ export class AiChatToolsService { description: "Replace a page's body with new Markdown content (and optionally its " + 'title). Reversible: the previous version is kept in page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to update.'), content: z.string().describe('The new page body as Markdown.'), title: z @@ -316,7 +317,7 @@ export class AiChatToolsService { description: "Rename a page (change its title only; the body is untouched). " + 'Reversible: rename back at any time.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to rename.'), title: z.string().describe('The new title.'), }), @@ -331,7 +332,7 @@ export class AiChatToolsService { description: 'Move a page under a new parent page, or to the space root when no ' + 'parent is given. Reversible: move it back at any time.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to move.'), parentPageId: z .string() @@ -353,7 +354,7 @@ export class AiChatToolsService { description: 'Move a page to the trash (SOFT delete only — fully reversible; the ' + 'page can be restored from trash). This NEVER permanently deletes.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to move to trash.'), }), // GUARDRAIL (§14 H4): the only field ever passed to the client is @@ -379,7 +380,7 @@ export class AiChatToolsService { '"selection not found" error, retry with a corrected EXACT selection ' + 'copied verbatim from a single paragraph/block. Reversible via the ' + 'comment UI.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to comment on.'), content: z.string().describe('The comment body as Markdown.'), selection: z @@ -428,7 +429,7 @@ export class AiChatToolsService { description: 'Resolve or reopen a top-level comment thread (reversible — toggle ' + 'the resolved flag). Only top-level comments can be resolved.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ commentId: z .string() .describe('The id of the top-level comment to resolve/reopen.'), @@ -460,7 +461,7 @@ export class AiChatToolsService { 'List the most recent pages, optionally scoped to a single space. ' + 'Returns a bounded list (default 50, max 100). Pass tree:true (with ' + "spaceId) to instead get the space's full page hierarchy as a nested tree.", - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ spaceId: z .string() .optional() @@ -488,7 +489,7 @@ export class AiChatToolsService { 'List sidebar pages for a space. With no pageId, returns the ' + "space's ROOT pages; with a pageId, returns that page's direct " + 'CHILDREN.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ spaceId: z.string().describe('The id of the space.'), pageId: z .string() @@ -520,7 +521,7 @@ export class AiChatToolsService { description: 'Read a table as a matrix of cell texts (plus a parallel cellIds ' + 'matrix so cells can be addressed for rich edits).', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), tableRef: z .string() @@ -536,7 +537,7 @@ export class AiChatToolsService { listComments: tool({ description: 'List all comments on a page (content as Markdown).', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), }), execute: async ({ pageId }) => await client.listComments(pageId), @@ -544,7 +545,7 @@ export class AiChatToolsService { getComment: tool({ description: 'Fetch a single comment by id (content as Markdown).', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ commentId: z.string().describe('The id of the comment.'), }), execute: async ({ commentId }) => await client.getComment(commentId), @@ -554,7 +555,7 @@ export class AiChatToolsService { description: 'Find new comments across a space (optionally scoped to a subtree) ' + 'created after a given timestamp.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ spaceId: z.string().describe('The id of the space to scan.'), since: z .string() @@ -586,7 +587,7 @@ export class AiChatToolsService { description: 'Fetch a single page-history version including its lossless ' + 'ProseMirror content.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ historyId: z.string().describe('The id of the history version.'), }), execute: async ({ historyId }) => @@ -604,7 +605,7 @@ export class AiChatToolsService { 'Export a page to a single self-contained Docmost-flavoured ' + 'Markdown file (meta + body + comment threads). Lossless round-trip ' + 'with importPageMarkdown.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to export.'), }), execute: async ({ pageId }) => { @@ -630,7 +631,7 @@ export class AiChatToolsService { '{"type":"text","text":"x","marks":[{"type":"bold"}]}. The node arg ' + 'may be a JSON object or a JSON string (both accepted). Reversible: ' + 'the previous version is kept in page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), nodeId: z .string() @@ -663,7 +664,7 @@ export class AiChatToolsService { '{"type":"text","text":"x","marks":[{"type":"bold"}]}. The node arg ' + 'may be a JSON object or a JSON string (both accepted). Reversible ' + 'via page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), node: z .any() @@ -722,7 +723,7 @@ export class AiChatToolsService { 'object or a JSON string (both accepted). Omit content for a ' + 'title-only update. Reversible: the previous version is kept in page ' + 'history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to update.'), content: z .any() @@ -753,7 +754,7 @@ export class AiChatToolsService { description: 'Insert a row of plain-text cells into a table. Reversible via ' + 'page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), tableRef: z .string() @@ -772,7 +773,7 @@ export class AiChatToolsService { tableDeleteRow: tool({ description: 'Delete a table row at a 0-based index. Reversible via page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), tableRef: z .string() @@ -787,7 +788,7 @@ export class AiChatToolsService { description: 'Set the plain-text content of a table cell at [row, col] (0-based). ' + 'Reversible via page history.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page.'), tableRef: z .string() @@ -817,7 +818,7 @@ export class AiChatToolsService { 'Make a page PUBLICLY accessible and return its public URL. ' + 'Reversible via unsharePage. Only share when the user explicitly ' + 'asked, since this exposes the page to anyone with the link.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to share.'), searchIndexing: z .boolean() @@ -844,7 +845,7 @@ export class AiChatToolsService { "page's ProseMirror document for complex/scripted rewrites. dryRun " + '(default true) previews a diff WITHOUT writing; set dryRun:false to ' + 'apply. Reversible: applying creates a new page-history snapshot.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z.string().describe('The id of the page to transform.'), transformJs: z .string() diff --git a/apps/server/src/core/ai-chat/tools/model-friendly-input.spec.ts b/apps/server/src/core/ai-chat/tools/model-friendly-input.spec.ts new file mode 100644 index 00000000..de05cfaf --- /dev/null +++ b/apps/server/src/core/ai-chat/tools/model-friendly-input.spec.ts @@ -0,0 +1,89 @@ +import { z } from 'zod'; +import { modelFriendlyInput } from './model-friendly-input'; + +/** + * Unit tests for the model-friendly input wrapper (issue #190): validation + * failures must report a human-readable, parameter-naming message (not the raw + * zod text), successful validation must strip unknown keys (preserving the + * strip guardrails), and the JSON schema handed to the model must keep the + * required/optional contract and field descriptions intact. + */ +describe('modelFriendlyInput', () => { + // A representative shape: a required id + description, plus an optional field. + const shape = { + pageId: z.string().describe('The id of the page to comment on.'), + content: z.string(), + limit: z.number().int().optional(), + }; + + // The AI SDK `Schema` exposes a `validate` callback and a `jsonSchema` field; + // type them loosely for the test. + type SchemaLike = { + validate?: ( + v: unknown, + ) => + | { success: boolean; value?: Record; error?: Error } + | PromiseLike<{ + success: boolean; + value?: Record; + error?: Error; + }>; + jsonSchema: unknown; + }; + + it('reports a model-friendly error naming the missing REQUIRED param + retry hint', async () => { + const schema = modelFriendlyInput(shape) as unknown as SchemaLike; + // Drop the required `pageId` (the parallel-batch failure mode). + const result = await schema.validate!({ content: 'hi' }); + + expect(result.success).toBe(false); + const message = result.error?.message ?? ''; + // Names the offending parameter by name. + expect(message).toContain('pageId'); + // Carries the fixed actionable retry hint. + expect(message).toContain('Include every REQUIRED parameter and retry'); + expect(message).toContain('do not drop ids like "pageId"'); + // It must NOT be the bare raw zod text alone — our wrapper prefix is present. + expect(message).toContain('Invalid tool input'); + }); + + it('accepts valid input and STRIPS unknown keys (keeps declared ones)', async () => { + const schema = modelFriendlyInput(shape) as unknown as SchemaLike; + const result = await schema.validate!({ + pageId: 'p-1', + content: 'hello', + // An extra unknown key a (compromised) model might emit. + permanentlyDelete: true, + }); + + expect(result.success).toBe(true); + expect(result.value).toEqual({ pageId: 'p-1', content: 'hello' }); + expect(result.value).not.toHaveProperty('permanentlyDelete'); + }); + + it('produces a draft-07 JSON schema that preserves required + descriptions', async () => { + const schema = modelFriendlyInput(shape) as unknown as SchemaLike; + // jsonSchema may be a value or a promise; await either way. + const json = (await Promise.resolve(schema.jsonSchema)) as { + required?: string[]; + properties?: Record; + }; + + // Required contract preserved: pageId + content required, limit optional. + expect(json.required).toEqual(expect.arrayContaining(['pageId', 'content'])); + expect(json.required).not.toContain('limit'); + // Field description preserved. + expect(json.properties?.pageId?.description).toBe( + 'The id of the page to comment on.', + ); + }); + + it('handles a root-level type error with a "(root)" parameter name', async () => { + const schema = modelFriendlyInput(shape) as unknown as SchemaLike; + // Passing a non-object yields an issue with an empty path. + const result = await schema.validate!('not an object'); + + expect(result.success).toBe(false); + expect(result.error?.message).toContain('(root)'); + }); +}); diff --git a/apps/server/src/core/ai-chat/tools/model-friendly-input.ts b/apps/server/src/core/ai-chat/tools/model-friendly-input.ts new file mode 100644 index 00000000..61537ade --- /dev/null +++ b/apps/server/src/core/ai-chat/tools/model-friendly-input.ts @@ -0,0 +1,72 @@ +import { jsonSchema, type JSONSchema7, type Schema } from 'ai'; +import { z } from 'zod'; + +// Centralized input-schema wrapper for in-app AI tools. The JSON schema handed +// to the model is derived from the same zod shape (so `required`/`description`/ +// constraints are unchanged), but validation failures are reported with a +// human-readable message that NAMES the offending parameter(s) and asks the +// model to retry with every required field — instead of the raw zod text. This +// matters for parallel tool-call batches where the model tends to drop a +// repeated id like `pageId`. + +// Fixed, actionable hint appended to every validation error. Kept as a constant +// so the message stays deterministic and the spec can assert on it verbatim. +const RETRY_HINT = + 'Include every REQUIRED parameter and retry; when issuing parallel tool ' + + 'calls, do not drop ids like "pageId".'; + +/** + * Turn a zod validation error into a concise, model-friendly message that names + * each offending parameter (by its dotted path; the root object is "(root)"), + * gives a short reason, and ends with the fixed retry hint. Repeated parameter + * names are de-duplicated and the output is deterministic. + */ +export function formatIssues(error: z.ZodError): string { + const seen = new Set(); + const parts: string[] = []; + for (const issue of error.issues) { + const name = + Array.isArray(issue.path) && issue.path.length > 0 + ? issue.path.join('.') + : '(root)'; + if (seen.has(name)) continue; + seen.add(name); + // Prefer zod's own message (e.g. "Invalid input: expected string, received + // undefined"); fall back to a generic reason when it is missing. + const reason = issue.message ? issue.message : 'missing or invalid'; + parts.push(`parameter "${name}": ${reason}`); + } + const summary = parts.length > 0 ? parts.join('; ') : 'invalid tool input'; + return `Invalid tool input — ${summary}. ${RETRY_HINT}`; +} + +/** + * Build an AI SDK `Schema` from a zod raw shape. The JSON schema exposed to the + * model is derived from the zod object (preserving `required`, `description`, + * and field constraints), so the required/optional contract is UNCHANGED. On a + * validation failure we return a model-friendly error (see `formatIssues`); on + * success we return the PARSED data, which has unknown keys stripped by zod — + * this preserves the existing strip guardrails (e.g. deletePage never forwards + * permanentlyDelete/forceDelete; transformPage never forwards deleteComments). + */ +export function modelFriendlyInput( + shape: Shape, +): Schema>> { + const object = z.object(shape); + // draft-07 JSON schema for the model (keeps required/description/constraints). + const schema = z.toJSONSchema(object, { target: 'draft-7' }) as JSONSchema7; + return jsonSchema>(schema, { + validate: (value: unknown) => { + const result = object.safeParse(value); + if (result.success) { + // Return the PARSED (unknown-key-stripped) data so the SDK forwards a + // clean object to execute — preserves the existing strip guardrails. + return { success: true as const, value: result.data }; + } + return { + success: false as const, + error: new Error(formatIssues(result.error)), + }; + }, + }); +} diff --git a/apps/server/src/core/ai-chat/tools/public-share-chat-tools.service.ts b/apps/server/src/core/ai-chat/tools/public-share-chat-tools.service.ts index 5d3e8df0..2d2da79d 100644 --- a/apps/server/src/core/ai-chat/tools/public-share-chat-tools.service.ts +++ b/apps/server/src/core/ai-chat/tools/public-share-chat-tools.service.ts @@ -5,6 +5,7 @@ import { ShareService } from '../../share/share.service'; import { SearchService } from '../../search/search.service'; import { PageRepo } from '@docmost/db/repos/page/page.repo'; import { jsonToMarkdown } from '../../../collaboration/collaboration.util'; +import { modelFriendlyInput } from './model-friendly-input'; /** * Isolated, READ-ONLY toolset for the ANONYMOUS public-share assistant. @@ -52,7 +53,7 @@ export class PublicShareChatToolsService { '(key terms and entities), not a full sentence. If the first ' + 'results look weak, search again with different wording before ' + 'answering. Only pages inside this share are ever returned.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ query: z.string().describe('The search query.'), limit: z .number() @@ -87,7 +88,7 @@ export class PublicShareChatToolsService { 'Markdown, by its page id. Returns the page title and its Markdown ' + 'content. Only pages inside this share can be read; reading any ' + 'other page fails.', - inputSchema: z.object({ + inputSchema: modelFriendlyInput({ pageId: z .string() .describe('The id (or slugId) of a page within this share.'), @@ -142,7 +143,7 @@ export class PublicShareChatToolsService { 'List the pages (titles + ids) that make up THIS published ' + 'documentation share, so you can orient yourself before reading or ' + 'searching. Only pages inside this share are listed.', - inputSchema: z.object({}), + inputSchema: modelFriendlyInput({}), execute: async () => { // Reuse the same share-tree logic the public /shares/tree route uses: // it validates the share + workspace, excludes restricted subtrees,