feat: add metadata filtering and typed upsert guardrails

cursoragent · cursoragent · commit 38c2e269578c · 2026-03-20T18:58:43.000Z
diff --git a/src/test-utils/mock-pinecone.ts b/src/test-utils/mock-pinecone.ts
@@ -1,13 +1,13 @@
 import {vi} from 'vitest';
 
-export function createMockNamespace() {
+export function createMockNamespace(): any {
   return {
     searchRecords: vi.fn(),
     upsertRecords: vi.fn(),
   };
 }
 
-export function createMockIndex() {
+export function createMockIndex(): any {
   const mockNamespace = createMockNamespace();
   return {
     describeIndexStats: vi.fn(),
@@ -16,7 +16,7 @@ export function createMockIndex() {
   };
 }
 
-export function createMockPinecone() {
+export function createMockPinecone(): any {
   const mockIndex = createMockIndex();
   return {
     listIndexes: vi.fn(),
diff --git a/src/tools/database/search-records.ts b/src/tools/database/search-records.ts
@@ -57,25 +57,60 @@ const SCHEMA = {
   namespace: z.string().describe('The namespace to search.'),
   query: SEARCH_QUERY_SCHEMA,
   rerank: RERANK_SCHEMA,
+  selectedMetadataKeys: z
+    .array(z.string())
+    .optional()
+    .describe(
+      'Optional: List of metadata keys to return. If omitted, all metadata is returned (Potential PII risk).',
+    ),
 };
 
 type SearchArgs = {
   name: string;
   namespace: string;
   query: SearchQuery;
   rerank?: SearchRerank;
+  selectedMetadataKeys?: string[];
 };
+type SearchRecord = {
+  metadata?: Record<string, unknown>;
+} & Record<string, unknown>;
+type SearchResults = {
+  records?: SearchRecord[];
+  matches?: SearchRecord[];
+} & Record<string, unknown>;
 
 export function addSearchRecordsTool(server: McpServer) {
   registerDatabaseTool(
     server,
     'search-records',
     {description: INSTRUCTIONS, inputSchema: SCHEMA},
     async (args, pc) => {
-      const {name, namespace, query, rerank} = args as SearchArgs;
+      const {name, namespace, query, rerank, selectedMetadataKeys} = args as SearchArgs;
       try {
         const ns = pc.index(name).namespace(namespace);
-        const results = await ns.searchRecords({query, rerank});
+        const results = (await ns.searchRecords({query, rerank})) as SearchResults;
+
+        // --- SECURITY PATCH: Metadata Filtering
+        if (selectedMetadataKeys) {
+          const filterRecordMetadata = (recordArray: SearchRecord[]) =>
+            recordArray.map((record) => {
+              if (!record.metadata) return record;
+              const filteredMetadata: Record<string, unknown> = {};
+              selectedMetadataKeys.forEach((key) => {
+                if (record.metadata[key] !== undefined)
+                  filteredMetadata[key] = record.metadata[key];
+              });
+              return {...record, metadata: filteredMetadata};
+            });
+
+          if (Array.isArray(results.records)) {
+            results.records = filterRecordMetadata(results.records);
+          }
+          if (Array.isArray(results.matches)) {
+            results.matches = filterRecordMetadata(results.matches);
+          }
+        }
         return {
           content: [
             {
diff --git a/src/tools/database/security.test.ts b/src/tools/database/security.test.ts
@@ -0,0 +1,109 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+// 1. HOISTED MOCK: This must be at the very top to intercept the tools' imports
+vi.mock('./common/pinecone-client.js', () => ({
+  getPineconeClient: vi.fn(),
+}));
+
+// Now import the tools and utilities
+import { createMockServer } from '../../test-utils/mock-server.js';
+import { createMockPinecone } from '../../test-utils/mock-pinecone.js';
+import { getPineconeClient } from './common/pinecone-client.js';
+import { addSearchRecordsTool } from './search-records.js';
+import { addUpsertRecordsTool } from './upsert-records.js';
+
+describe('Pinecone MCP Security Hardening', () => {
+  let mockServer: any;
+  let mockPc: any;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockServer = createMockServer();
+    mockPc = createMockPinecone();
+
+    // 2. Force the client factory to return our mock for every call
+    vi.mocked(getPineconeClient).mockReturnValue(mockPc);
+
+    addSearchRecordsTool(mockServer);
+    addUpsertRecordsTool(mockServer);
+    
+    // Satisfy the initial constant check
+    vi.stubEnv('PINECONE_API_KEY', 'test-key');
+  });
+
+  describe('Metadata Firewall (search-records)', () => {
+    it('should strip unauthorized metadata keys from search results', async () => {
+      const mockRecords = {
+        records: [{
+          id: 'vec-1',
+          metadata: { text: 'public', ssn: 'secret-123' }
+        }]
+      };
+
+      mockPc._mockIndex._mockNamespace.searchRecords.mockResolvedValue(mockRecords);
+
+      const tool = mockServer.getRegisteredTool('search-records');
+      const result = await tool.handler({
+        name: 'idx', namespace: 'ns',
+        query: { inputs: { text: 'test' }, topK: 1 },
+        selectedMetadataKeys: ['text']
+      });
+
+      // Helpful debugging: if it fails, show why
+      if (result.isError) {
+        throw new Error(`Tool returned error: ${result.content[0].text}`);
+      }
+
+      const output = JSON.parse(result.content[0].text);
+      expect(output.records[0].metadata.ssn).toBeUndefined();
+      expect(output.records[0].metadata.text).toBe('public');
+    });
+
+    it('should return all metadata if selectedMetadataKeys is omitted', async () => {
+      mockPc._mockIndex._mockNamespace.searchRecords.mockResolvedValue({
+        records: [{ id: '1', metadata: { secret: 'data' } }]
+      });
+
+      const tool = mockServer.getRegisteredTool('search-records');
+      const result = await tool.handler({
+        name: 'idx', namespace: 'ns',
+        query: { inputs: { text: 'test' }, topK: 1 }
+      });
+
+      if (result.isError) {
+        throw new Error(`Tool returned error: ${result.content[0].text}`);
+      }
+
+      const output = JSON.parse(result.content[0].text);
+      expect(output.records[0].metadata.secret).toBe('data');
+    });
+  });
+
+  describe('Upsert Guardrail (upsert-records)', () => {
+    it('should block upsert if confirmOverwrite flag is missing', async () => {
+      const tool = mockServer.getRegisteredTool('upsert-records');
+      const result = await tool.handler({
+        name: 'idx', namespace: 'ns',
+        records: [{ id: '1', text: 'val' }]
+      });
+
+      expect(result.isError).toBe(true);
+      expect(result.content[0].text).toContain('confirmOverwrite must be set to true');
+    });
+
+    it('should execute upsert when confirmOverwrite is true', async () => {
+      mockPc._mockIndex._mockNamespace.upsertRecords.mockResolvedValue(undefined);
+      const tool = mockServer.getRegisteredTool('upsert-records');
+
+      const result = await tool.handler({
+        name: 'idx', namespace: 'ns',
+        records: [{ id: '1', text: 'val' }],
+        confirmOverwrite: true
+      });
+
+      // Final check: ensure no error was returned
+      expect(result.isError).toBeFalsy();
+      expect(result.content[0].text).toBe('Data upserted successfully');
+    });
+  });
+});
diff --git a/src/tools/database/upsert-records.handler.test.ts b/src/tools/database/upsert-records.handler.test.ts
@@ -42,6 +42,7 @@ describe('upsert-records tool handler', () => {
       name: 'test-index',
       namespace: 'test-ns',
       records: [{id: '1', content: 'test content'}],
+      confirmOverwrite: true,
     });
 
     expect(mockPc.index).toHaveBeenCalledWith('test-index');
@@ -63,6 +64,7 @@ describe('upsert-records tool handler', () => {
       name: 'test-index',
       namespace: 'test-ns',
       records: [{id: '1', content: 'test content'}],
+      confirmOverwrite: true,
     });
 
     expect(result).toEqual({
@@ -80,11 +82,32 @@ describe('upsert-records tool handler', () => {
       name: 'test-index',
       namespace: 'test-ns',
       records: [{_id: 'alt-1', content: 'test content'}],
+      confirmOverwrite: true,
     })) as {content: Array<{text: string}>};
 
     expect(mockPc._mockIndex._mockNamespace.upsertRecords).toHaveBeenCalledWith([
       {_id: 'alt-1', content: 'test content'},
     ]);
     expect(result.content[0].text).toContain('successfully');
   });
-});
+  
+  it('fails if confirmOverwrite is missing or false', async () => {
+    // 🛡️ BULLETPROOF SCOPING: Register the tool explicitly for this test
+    // This bypasses any nested 'beforeEach' boundary issues
+    addUpsertRecordsTool(mockServer);
+    
+    // Now we are 100% guaranteed the tool exists
+    const tool = mockServer.getRegisteredTool('upsert-records')!;
+
+    const result = (await tool.handler({
+      name: 'test-index',
+      namespace: 'test-ns',
+      records: [{id: '1', content: 'test content'}]
+    })) as any;
+
+    expect(result.isError).toBe(true);
+    expect(result.content[0].text).toContain('confirmOverwrite must be set to true');
+    expect(mockPc._mockIndex._mockNamespace.upsertRecords).not.toHaveBeenCalled();
+  });
+
+});
diff --git a/src/tools/database/upsert-records.ts b/src/tools/database/upsert-records.ts
@@ -35,12 +35,18 @@ const SCHEMA = {
   name: z.string().describe('The index to upsert into.'),
   namespace: z.string().describe('The namespace to upsert into.'),
   records: RECORD_SET_SCHEMA,
+  confirmOverwrite: z
+    .boolean()
+    .describe(
+      'CRITICAL: Set to true to acknowledge that this will overwrite any existing records with the same IDs.',
+    ),
 };
 
 type UpsertArgs = {
   name: string;
   namespace: string;
   records: IntegratedRecord[];
+  confirmOverwrite: boolean;
 };
 
 export function addUpsertRecordsTool(server: McpServer) {
@@ -49,7 +55,20 @@ export function addUpsertRecordsTool(server: McpServer) {
     'upsert-records',
     {description: INSTRUCTIONS, inputSchema: SCHEMA},
     async (args, pc) => {
-      const {name, namespace, records} = args as UpsertArgs;
+      const {name, namespace, records, confirmOverwrite} = args as UpsertArgs;
+
+      // --- SECURITY PATCH: Destructive Action Guardrail ---
+      if (confirmOverwrite !== true) {
+        return {
+          isError: true,
+          content: [
+            {
+              type: 'text' as const,
+              text: 'Error: confirmOverwrite must be set to true. Upserting records can overwrite existing data. Please ask the user for confirmation.',
+            },
+          ],
+        };
+      }
       try {
         const ns = pc.index(name).namespace(namespace);
         await ns.upsertRecords(records);
