Skip to content

Commit 3025863

Browse files
vincentkocvincentkoc
committed
fix(agents): snapshot SDK custom tools safely
1 parent 286c8e3 commit 3025863

4 files changed

Lines changed: 183 additions & 4 deletions

File tree

‎src/agents/sessions/agent-session.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,10 @@ import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-promp
9393
import type { BashOperations } from "./tools/bash-operations.js";
9494
import { createLocalBashOperations } from "./tools/bash.js";
9595
import { createAllToolDefinitions } from "./tools/index.js";
96-
import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js";
96+
import {
97+
createToolDefinitionFromAgentTool,
98+
snapshotToolDefinitions,
99+
} from "./tools/tool-definition-wrapper.js";
97100

98101
function unwrapCoreResult<T>(result: { ok: true; value: T } | { ok: false; error: Error }): T {
99102
if (result.ok) {
@@ -383,7 +386,7 @@ export class AgentSession {
383386
this.settingsManager = config.settingsManager;
384387
this.scopedModelEntries = config.scopedModels ?? [];
385388
this.sessionResourceLoader = config.resourceLoader;
386-
this.customTools = config.customTools ?? [];
389+
this.customTools = snapshotToolDefinitions(config.customTools);
387390
this.cwd = config.cwd;
388391
this.sessionModelRegistry = config.modelRegistry;
389392
this.extensionRunnerRef = config.extensionRunnerRef;

‎src/agents/sessions/sdk.test.ts‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,115 @@ describe("createAgentSession tool defaults", () => {
117117
expect(session.getActiveToolNames()).toEqual(["custom_lookup"]);
118118
});
119119

120+
it("drops unreadable SDK custom tool descriptors before session setup", async () => {
121+
const healthyTool: ToolDefinition = {
122+
name: "healthy_lookup",
123+
label: "Healthy Lookup",
124+
description: "Looks up a test value.",
125+
promptSnippet: "Lookup test values",
126+
promptGuidelines: ["Use healthy_lookup for test values."],
127+
parameters: Type.Object({}),
128+
execute: async () => ({
129+
content: [{ type: "text", text: "ok" }],
130+
details: {},
131+
}),
132+
};
133+
const unreadableNameTool = {
134+
get name(): string {
135+
throw new Error("boom");
136+
},
137+
label: "Broken",
138+
description: "Breaks when read.",
139+
parameters: Type.Object({}),
140+
execute: async () => ({
141+
content: [{ type: "text", text: "bad" }],
142+
details: {},
143+
}),
144+
} as unknown as ToolDefinition;
145+
const unreadableGuidelinesTool = {
146+
name: "bad_guidelines",
147+
label: "Broken Guidelines",
148+
description: "Breaks when prompt metadata is read.",
149+
get promptGuidelines(): string[] {
150+
throw new Error("boom");
151+
},
152+
parameters: Type.Object({}),
153+
execute: async () => ({
154+
content: [{ type: "text", text: "bad" }],
155+
details: {},
156+
}),
157+
} as unknown as ToolDefinition;
158+
const unreadableParametersTool = {
159+
name: "bad_parameters",
160+
label: "Broken Parameters",
161+
description: "Breaks when schema is read.",
162+
get parameters(): ToolDefinition["parameters"] {
163+
throw new Error("boom");
164+
},
165+
execute: async () => ({
166+
content: [{ type: "text", text: "bad" }],
167+
details: {},
168+
}),
169+
} as unknown as ToolDefinition;
170+
171+
const { session } = await createAgentSession({
172+
model: testModel,
173+
noTools: "builtin",
174+
customTools: [
175+
unreadableNameTool,
176+
unreadableGuidelinesTool,
177+
healthyTool,
178+
unreadableParametersTool,
179+
],
180+
resourceLoader: createEmptyResourceLoader(),
181+
sessionManager: SessionManager.inMemory(),
182+
settingsManager: SettingsManager.inMemory(),
183+
modelRegistry: ModelRegistry.inMemory(AuthStorage.inMemory()),
184+
});
185+
186+
expect(session.getActiveToolNames()).toEqual(["healthy_lookup"]);
187+
expect(session.getAllTools().map((tool) => tool.name)).toEqual(["healthy_lookup"]);
188+
});
189+
190+
it("preserves the original custom tool receiver after snapshotting", async () => {
191+
const statefulTool = {
192+
name: "stateful_lookup",
193+
label: "Stateful Lookup",
194+
description: "Uses state from the tool object.",
195+
parameters: Type.Object({}),
196+
calls: 0,
197+
async execute() {
198+
this.calls += 1;
199+
return {
200+
content: [{ type: "text" as const, text: String(this.calls) }],
201+
details: {},
202+
};
203+
},
204+
} as ToolDefinition & { calls: number };
205+
206+
const { session } = await createAgentSession({
207+
model: testModel,
208+
noTools: "builtin",
209+
customTools: [statefulTool],
210+
resourceLoader: createEmptyResourceLoader(),
211+
sessionManager: SessionManager.inMemory(),
212+
settingsManager: SettingsManager.inMemory(),
213+
modelRegistry: ModelRegistry.inMemory(AuthStorage.inMemory()),
214+
});
215+
216+
const definition = session.getToolDefinition("stateful_lookup");
217+
const result = await definition?.execute(
218+
"call-1",
219+
{},
220+
undefined,
221+
undefined,
222+
undefined as unknown as Parameters<ToolDefinition["execute"]>[4],
223+
);
224+
225+
expect(result?.content).toEqual([{ type: "text", text: "1" }]);
226+
expect(statefulTool.calls).toBe(1);
227+
});
228+
120229
it("preserves an exact base system prompt when active tools change", async () => {
121230
const customTool: ToolDefinition = {
122231
name: "custom_lookup",

‎src/agents/sessions/sdk.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import {
3535
type ToolName,
3636
withFileMutationQueue,
3737
} from "./tools/index.js";
38+
import { snapshotToolDefinitions } from "./tools/tool-definition-wrapper.js";
3839

3940
export interface CreateAgentSessionOptions {
4041
/** Working directory for project-local discovery. Default: process.cwd() */
@@ -283,7 +284,8 @@ export async function createAgentSession(
283284
}
284285

285286
const defaultActiveToolNames: ToolName[] = ["read", "bash", "edit", "write"];
286-
const customToolNames = options.customTools?.map((tool) => tool.name) ?? [];
287+
const customTools = snapshotToolDefinitions(options.customTools);
288+
const customToolNames = customTools.map((tool) => tool.name);
287289
const allowedToolNames = options.tools ?? (options.noTools === "all" ? [] : undefined);
288290
const disableBuiltInTools = !options.tools && options.noTools === "builtin";
289291
const initialActiveToolNames: string[] = options.tools
@@ -426,7 +428,7 @@ export async function createAgentSession(
426428
cwd,
427429
scopedModels: options.scopedModels,
428430
resourceLoader,
429-
customTools: options.customTools,
431+
customTools,
430432
modelRegistry,
431433
initialActiveToolNames,
432434
allowedToolNames,

‎src/agents/sessions/tools/tool-definition-wrapper.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,71 @@ import type { TSchema } from "typebox";
22
import type { AgentTool } from "../../runtime/index.js";
33
import type { ExtensionContext, ToolDefinition } from "../extensions/types.js";
44

5+
/** Snapshot caller-owned tool definitions before registry code reads their fields. */
6+
export function snapshotToolDefinitions(
7+
definitions: readonly ToolDefinition[] | undefined,
8+
): ToolDefinition[] {
9+
const snapshots: ToolDefinition[] = [];
10+
for (const definition of definitions ?? []) {
11+
const snapshot = snapshotToolDefinition(definition);
12+
if (snapshot) {
13+
snapshots.push(snapshot);
14+
}
15+
}
16+
return snapshots;
17+
}
18+
19+
function snapshotToolDefinition(definition: ToolDefinition): ToolDefinition | undefined {
20+
try {
21+
if (!definition || typeof definition !== "object") {
22+
return undefined;
23+
}
24+
const name = Reflect.get(definition, "name");
25+
const execute = Reflect.get(definition, "execute");
26+
if (typeof name !== "string" || name.length === 0 || typeof execute !== "function") {
27+
return undefined;
28+
}
29+
const promptGuidelines = Reflect.get(definition, "promptGuidelines");
30+
const prepareArguments = Reflect.get(definition, "prepareArguments");
31+
const renderCall = Reflect.get(definition, "renderCall");
32+
const renderResult = Reflect.get(definition, "renderResult");
33+
const executeWithReceiver = ((...args: Parameters<ToolDefinition["execute"]>) =>
34+
Reflect.apply(execute, definition, args)) as ToolDefinition["execute"];
35+
return {
36+
name,
37+
label: Reflect.get(definition, "label"),
38+
description: Reflect.get(definition, "description"),
39+
promptSnippet: Reflect.get(definition, "promptSnippet"),
40+
promptGuidelines: Array.isArray(promptGuidelines) ? [...promptGuidelines] : promptGuidelines,
41+
parameters: Reflect.get(definition, "parameters"),
42+
renderShell: Reflect.get(definition, "renderShell"),
43+
prepareArguments:
44+
typeof prepareArguments === "function"
45+
? (((...args: Parameters<NonNullable<ToolDefinition["prepareArguments"]>>) =>
46+
Reflect.apply(
47+
prepareArguments,
48+
definition,
49+
args,
50+
)) as ToolDefinition["prepareArguments"])
51+
: prepareArguments,
52+
executionMode: Reflect.get(definition, "executionMode"),
53+
execute: executeWithReceiver,
54+
renderCall:
55+
typeof renderCall === "function"
56+
? (((...args: Parameters<NonNullable<ToolDefinition["renderCall"]>>) =>
57+
Reflect.apply(renderCall, definition, args)) as ToolDefinition["renderCall"])
58+
: renderCall,
59+
renderResult:
60+
typeof renderResult === "function"
61+
? (((...args: Parameters<NonNullable<ToolDefinition["renderResult"]>>) =>
62+
Reflect.apply(renderResult, definition, args)) as ToolDefinition["renderResult"])
63+
: renderResult,
64+
};
65+
} catch {
66+
return undefined;
67+
}
68+
}
69+
570
/** Wrap a ToolDefinition into an AgentTool for the core runtime. */
671
export function wrapToolDefinition<
772
TParams extends TSchema = TSchema,

0 commit comments

Comments
 (0)