diff --git a/src/memory/__tests__/search-nodes-schema.test.ts b/src/memory/__tests__/search-nodes-schema.test.ts new file mode 100644 index 0000000000..c03f384ebb --- /dev/null +++ b/src/memory/__tests__/search-nodes-schema.test.ts @@ -0,0 +1,25 @@ +import { describe, it, expect } from 'vitest'; +import { SearchNodesQuerySchema, SEARCH_QUERY_MAX_LENGTH } from '../index.js'; + +describe('search_nodes input schema', () => { + it('should accept a normal query', () => { + expect(SearchNodesQuerySchema.safeParse('Alice').success).toBe(true); + expect(SearchNodesQuerySchema.safeParse('works at Acme Corp').success).toBe(true); + }); + + it('should accept a query at exactly the max length', () => { + const atLimit = 'a'.repeat(SEARCH_QUERY_MAX_LENGTH); + expect(SearchNodesQuerySchema.safeParse(atLimit).success).toBe(true); + }); + + it('should reject a query longer than the max length', () => { + const oversized = 'a'.repeat(SEARCH_QUERY_MAX_LENGTH + 1); + const result = SearchNodesQuerySchema.safeParse(oversized); + expect(result.success).toBe(false); + }); + + it('should still reject non-string input', () => { + expect(SearchNodesQuerySchema.safeParse(42).success).toBe(false); + expect(SearchNodesQuerySchema.safeParse(null).success).toBe(false); + }); +}); diff --git a/src/memory/index.ts b/src/memory/index.ts index 9865c5318e..cea1cc35dd 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -486,6 +486,13 @@ server.registerTool( } ); +export const SEARCH_QUERY_MAX_LENGTH = 2048; + +export const SearchNodesQuerySchema = z + .string() + .max(SEARCH_QUERY_MAX_LENGTH) + .describe("The search query to match against entity names, types, and observation content"); + // Register search_nodes tool server.registerTool( "search_nodes", @@ -493,7 +500,7 @@ server.registerTool( title: "Search Nodes", description: "Search for nodes in the knowledge graph based on a query", inputSchema: { - query: z.string().describe("The search query to match against entity names, types, and observation content") + query: SearchNodesQuerySchema }, outputSchema: { entities: z.array(EntitySchema),