chore: add function to sanitize params for tool calls

Author: yulunz

src/telemetry/ClearcutLogger.ts

@@ -8,6 +8,7 @@ import process from 'node:process';
 
 import {DAEMON_CLIENT_NAME} from '../daemon/utils.js';
 import {logger} from '../logger.js';
+import {zod, type ShapeOutput} from '../third_party/index.js';
 
 import type {LocalState, Persistence} from './persistence.js';
 import {FilePersistence} from './persistence.js';
@@ -20,6 +21,78 @@ import {
 import {WatchdogClient} from './WatchdogClient.js';
 
 const MS_PER_DAY = 24 * 60 * 60 * 1000;
+const PARAM_BLOCKLIST = ['uid'];
+
+function getZodType(zodType: any): string {
+  const def = zodType._def;
+  const typeName = def.typeName;
+
+  if (
+    typeName === 'ZodOptional' ||
+    typeName === 'ZodDefault' ||
+    typeName === 'ZodNullable'
+  ) {
+    return getZodType(def.innerType);
+  }
+  if (typeName === 'ZodEffects') {
+    return getZodType(def.schema);
+  }
+
+  let type: string;
+  switch (typeName) {
+    case 'ZodString':
+      return 'string';
+    case 'ZodNumber':
+      return 'number';
+    case 'ZodBoolean':
+      return 'boolean';
+    case 'ZodArray':
+      return 'array';
+    case 'ZodEnum':
+      return 'enum';
+    default:
+      throw new Error(`Unsupported zod type for tool parameter: ${typeName}`);
+  }
+}
+
+function filterBlockedParams(
+  params: ShapeOutput<zod.ZodRawShape>,
+): ShapeOutput<zod.ZodRawShape> {
+  const filteredParams = {...params};
+  for (const key of PARAM_BLOCKLIST) {
+    if (key in filteredParams) {
+      delete filteredParams[key];
+    }
+  }
+  return filteredParams;
+}
+
+function transformParams(
+  params: Record<string, any>,
+  schema: zod.ZodRawShape,
+): Record<string, any> {
+  const transformed: Record<string, any> = {};
+  for (const [name, value] of Object.entries(params)) {
+    const zodType = getZodType(schema[name]);
+
+    if (zodType === 'string') {
+      transformed[`${name}_length`] = (value as string).length;
+    } else if (zodType === 'array') {
+      transformed[`${name}_count`] = (value as any[]).length;
+    } else {
+      transformed[name] = value;
+    }
+  }
+  return transformed;
+}
+
+export function sanitizeParams(
+  params: ShapeOutput<zod.ZodRawShape>,
+  schema: zod.ZodRawShape,
+): Record<string, any> {
+  const filtered = filterBlockedParams(params);
+  return transformParams(filtered, schema);
+}
 
 function detectOsType(): OsType {
   switch (process.platform) {

src/third_party/index.ts

@@ -19,6 +19,7 @@ export {hideBin} from 'yargs/helpers';
 export {default as debug} from 'debug';
 export type {Debugger} from 'debug';
 export {McpServer} from '@modelcontextprotocol/sdk/server/mcp.js';
+export {type ShapeOutput} from '@modelcontextprotocol/sdk/server/zod-compat.js';
 export {StdioServerTransport} from '@modelcontextprotocol/sdk/server/stdio.js';
 export {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
 export {Client} from '@modelcontextprotocol/sdk/client/index.js';

tests/telemetry/ClearcutLogger.test.ts

@@ -10,11 +10,15 @@ import {describe, it, afterEach, beforeEach} from 'node:test';
 import sinon from 'sinon';
 
 import {DAEMON_CLIENT_NAME} from '../../src/daemon/utils.js';
-import {ClearcutLogger} from '../../src/telemetry/ClearcutLogger.js';
+import {
+  ClearcutLogger,
+  sanitizeParams,
+} from '../../src/telemetry/ClearcutLogger.js';
 import type {Persistence} from '../../src/telemetry/persistence.js';
 import {FilePersistence} from '../../src/telemetry/persistence.js';
 import {WatchdogMessageType} from '../../src/telemetry/types.js';
 import {WatchdogClient} from '../../src/telemetry/WatchdogClient.js';
+import {zod} from '../../src/third_party/index.js';
 
 describe('ClearcutLogger', () => {
   let mockPersistence: sinon.SinonStubbedInstance<Persistence>;
@@ -163,4 +167,50 @@ describe('ClearcutLogger', () => {
       assert(mockPersistence.saveState.called);
     });
   });
+
+  describe('sanitizeParams', () => {
+    it('filters out uid and transforms strings and arrays', () => {
+      const schema = {
+        uid: zod.string(),
+        myString: zod.string(),
+        myArray: zod.array(zod.string()),
+        myNumber: zod.number(),
+        myBool: zod.boolean(),
+        myEnum: zod.enum(['a', 'b']),
+      };
+
+      const params = {
+        uid: 'sensitive',
+        myString: 'hello',
+        myArray: ['one', 'two'],
+        myNumber: 42,
+        myBool: true,
+        myEnum: 'a' as const,
+      };
+
+      const sanitized = sanitizeParams(params, schema);
+
+      assert.deepStrictEqual(sanitized, {
+        myString_length: 5,
+        myArray_count: 2,
+        myNumber: 42,
+        myBool: true,
+        myEnum: 'a',
+      });
+    });
+
+    it('throws error for unsupported types', () => {
+      const schema = {
+        myObj: zod.object({foo: zod.string()}),
+      };
+      const params = {
+        myObj: {foo: 'bar'},
+      };
+
+      assert.throws(
+        () => sanitizeParams(params, schema),
+        /Unsupported zod type for tool parameter: ZodObject/,
+      );
+    });
+  });
 });