chore: add function to sanitize params for tool calls (#1250)

This adds functions to sanitize the tool call parameters. They are not
called as of now since we don't have server side changes landed yet to
support these.

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 type {zod, ShapeOutput} from '../third_party/index.js';
 
 import type {LocalState, Persistence} from './persistence.js';
 import {FilePersistence} from './persistence.js';
@@ -20,6 +21,108 @@ import {
 import {WatchdogClient} from './WatchdogClient.js';
 
 const MS_PER_DAY = 24 * 60 * 60 * 1000;
+const PARAM_BLOCKLIST = new Set(['uid']);
+
+const SUPPORTED_ZOD_TYPES = [
+  'ZodString',
+  'ZodNumber',
+  'ZodBoolean',
+  'ZodArray',
+  'ZodEnum',
+] as const;
+type ZodType = (typeof SUPPORTED_ZOD_TYPES)[number];
+
+function isZodType(type: string): type is ZodType {
+  return SUPPORTED_ZOD_TYPES.includes(type as ZodType);
+}
+
+function getZodType(zodType: zod.ZodTypeAny): ZodType {
+  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);
+  }
+
+  if (isZodType(typeName)) {
+    return typeName;
+  }
+  throw new Error(`Unsupported zod type for tool parameter: ${typeName}`);
+}
+
+type LoggedToolCallArgValue = string | number | boolean;
+
+function transformName(zodType: ZodType, name: string): string {
+  if (zodType === 'ZodString') {
+    return `${name}_length`;
+  } else if (zodType === 'ZodArray') {
+    return `${name}_count`;
+  } else {
+    return name;
+  }
+}
+
+function transformValue(
+  zodType: ZodType,
+  value: unknown,
+): LoggedToolCallArgValue {
+  if (zodType === 'ZodString') {
+    return (value as string).length;
+  } else if (zodType === 'ZodArray') {
+    return (value as unknown[]).length;
+  } else {
+    return value as LoggedToolCallArgValue;
+  }
+}
+
+function hasEquivalentType(zodType: ZodType, value: unknown): boolean {
+  if (zodType === 'ZodString') {
+    return typeof value === 'string';
+  } else if (zodType === 'ZodArray') {
+    return Array.isArray(value);
+  } else if (zodType === 'ZodNumber') {
+    return typeof value === 'number';
+  } else if (zodType === 'ZodBoolean') {
+    return typeof value === 'boolean';
+  } else if (zodType === 'ZodEnum') {
+    return (
+      typeof value === 'string' ||
+      typeof value === 'number' ||
+      typeof value === 'boolean'
+    );
+  } else {
+    return false;
+  }
+}
+
+export function sanitizeParams(
+  params: ShapeOutput<zod.ZodRawShape>,
+  schema: zod.ZodRawShape,
+): ShapeOutput<zod.ZodRawShape> {
+  const transformed: ShapeOutput<zod.ZodRawShape> = {};
+  for (const [name, value] of Object.entries(params)) {
+    if (PARAM_BLOCKLIST.has(name)) {
+      continue;
+    }
+    const zodType = getZodType(schema[name]);
+    if (!hasEquivalentType(zodType, value)) {
+      throw new Error(
+        `parameter ${name} has type ${zodType} but value ${value} is not of equivalent type`,
+      );
+    }
+    const transformedName = transformName(zodType, name);
+    const transformedValue = transformValue(zodType, value);
+    transformed[transformedName] = transformedValue;
+  }
+  return transformed;
+}
 
 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,64 @@ 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/,
+      );
+    });
+
+    it('throws error when value is not of equivalent type', () => {
+      const schema = {
+        myString: zod.string(),
+      };
+      const params = {
+        myString: 123,
+      };
+
+      assert.throws(
+        () => sanitizeParams(params, schema),
+        /parameter myString has type ZodString but value 123 is not of equivalent type/,
+      );
+    });
+  });
 });