Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 103 additions & 0 deletions src/telemetry/ClearcutLogger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions src/third_party/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down
66 changes: 65 additions & 1 deletion tests/telemetry/ClearcutLogger.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand Down Expand Up @@ -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/,
);
});
});
});
Loading