Skip to content

Commit 91107e2

Browse files
committed
fix(ptc): hash full tool definitions for cache invalidation
Previously the type cache only hashed tool names, so schema changes (e.g., updated MCP server, different workspace with same tool names) wouldn't invalidate cached types. Now hashes names + schemas + descriptions.
1 parent 4719592 commit 91107e2

File tree

2 files changed

+67
-8
lines changed

2 files changed

+67
-8
lines changed

src/node/services/ptc/typeGenerator.test.ts

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
import { describe, test, expect } from "bun:test";
1+
import { describe, test, expect, beforeEach } from "bun:test";
22
import { z } from "zod";
33
import type { Tool } from "ai";
4-
import { generateMuxTypes } from "./typeGenerator";
4+
import { generateMuxTypes, getCachedMuxTypes, clearTypeCache } from "./typeGenerator";
55

66
/**
77
* Create a mock tool with the given schema and optional execute function.
@@ -231,3 +231,53 @@ describe("generateMuxTypes", () => {
231231
expect(types).toContain("declare var console");
232232
});
233233
});
234+
235+
describe("getCachedMuxTypes", () => {
236+
beforeEach(() => {
237+
clearTypeCache();
238+
});
239+
240+
test("invalidates cache when tool schema changes", async () => {
241+
const toolV1 = createMockTool(z.object({ name: z.string() }));
242+
const toolV2 = createMockTool(z.object({ name: z.string(), age: z.number() }));
243+
244+
const types1 = await getCachedMuxTypes({ my_tool: toolV1 });
245+
expect(types1).toContain("name: string");
246+
expect(types1).not.toContain("age");
247+
248+
// Same tool name, different schema - should regenerate
249+
const types2 = await getCachedMuxTypes({ my_tool: toolV2 });
250+
expect(types2).toContain("name: string");
251+
expect(types2).toContain("age: number");
252+
});
253+
254+
test("invalidates cache when tool description changes", async () => {
255+
const tool1: Tool = {
256+
description: "Version 1",
257+
inputSchema: z.object({ x: z.string() }),
258+
execute: () => Promise.resolve({ success: true }),
259+
} as unknown as Tool;
260+
261+
const tool2: Tool = {
262+
description: "Version 2",
263+
inputSchema: z.object({ x: z.string() }),
264+
execute: () => Promise.resolve({ success: true }),
265+
} as unknown as Tool;
266+
267+
const types1 = await getCachedMuxTypes({ my_tool: tool1 });
268+
expect(types1).toContain("Version 1");
269+
270+
const types2 = await getCachedMuxTypes({ my_tool: tool2 });
271+
expect(types2).toContain("Version 2");
272+
});
273+
274+
test("returns cached types when tools are identical", async () => {
275+
const tool = createMockTool(z.object({ value: z.string() }));
276+
277+
const types1 = await getCachedMuxTypes({ my_tool: tool });
278+
const types2 = await getCachedMuxTypes({ my_tool: tool });
279+
280+
// Should be the exact same object reference (cached)
281+
expect(types1).toBe(types2);
282+
});
283+
});

src/node/services/ptc/typeGenerator.ts

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,18 +56,27 @@ export function clearTypeCache(): void {
5656
}
5757

5858
/**
59-
* Hash tool names to detect when tool set changes.
59+
* Hash tool definitions (names, schemas, descriptions) to detect when tools change.
60+
* This ensures cache invalidation when schemas are updated, not just when tool names change.
6061
*/
61-
function hashToolNames(tools: Record<string, Tool>): string {
62-
const names = Object.keys(tools).sort().join(",");
63-
return createHash("md5").update(names).digest("hex");
62+
function hashToolDefinitions(tools: Record<string, Tool>): string {
63+
const sortedNames = Object.keys(tools).sort();
64+
const toolData = sortedNames.map((name) => {
65+
const tool = tools[name];
66+
return {
67+
name,
68+
schema: getInputJsonSchema(tool),
69+
description: tool.description ?? "",
70+
};
71+
});
72+
return createHash("md5").update(JSON.stringify(toolData)).digest("hex");
6473
}
6574

6675
/**
67-
* Get cached mux types or generate new ones if tool set changed.
76+
* Get cached mux types or generate new ones if tool definitions changed.
6877
*/
6978
export async function getCachedMuxTypes(tools: Record<string, Tool>): Promise<string> {
70-
const hash = hashToolNames(tools);
79+
const hash = hashToolDefinitions(tools);
7180
const cached = cache.fullTypes.get(hash);
7281
if (cached) {
7382
return cached;

0 commit comments

Comments
 (0)