From dadd00543675905d114854f9e670fbfb352c59a8 Mon Sep 17 00:00:00 2001 From: Yulun Zeng <11618243+yulunz@users.noreply.github.com> Date: Fri, 27 Mar 2026 00:21:49 +0000 Subject: [PATCH] chore: add function to sanitize params for tool calls --- src/telemetry/ClearcutLogger.ts | 103 +++++++++++++++++++++++++ src/third_party/index.ts | 1 + tests/telemetry/ClearcutLogger.test.ts | 66 +++++++++++++++- 3 files changed, 169 insertions(+), 1 deletion(-) diff --git a/src/telemetry/ClearcutLogger.ts b/src/telemetry/ClearcutLogger.ts index 06cfe8df1..a53345746 100644 --- a/src/telemetry/ClearcutLogger.ts +++ b/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, + schema: zod.ZodRawShape, +): ShapeOutput { + const transformed: ShapeOutput = {}; + 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) { diff --git a/src/third_party/index.ts b/src/third_party/index.ts index f719f4df8..fc0fce620 100644 --- a/src/third_party/index.ts +++ b/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'; diff --git a/tests/telemetry/ClearcutLogger.test.ts b/tests/telemetry/ClearcutLogger.test.ts index 36e21bc67..792202856 100644 --- a/tests/telemetry/ClearcutLogger.test.ts +++ b/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; @@ -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/, + ); + }); + }); });