From edbeec8c2c82573adf6d9494e2b40d8fcc2bd986 Mon Sep 17 00:00:00 2001 From: James Long Date: Wed, 4 Nov 2026 23:17:37 +0000 Subject: [PATCH] feat(codemode): pass each tool call its idempotency key, at the tool and at both hooks --- packages/codemode/src/interpreter/driver.ts | 6 +-- packages/codemode/src/tool-runtime.ts | 47 ++++++++++++++------- packages/codemode/src/tool.ts | 22 +++++++--- packages/codemode/test/codemode.test.ts | 4 +- packages/codemode/test/suspend.test.ts | 26 ++++++++++++ 5 files changed, 78 insertions(+), 27 deletions(-) diff --git a/packages/codemode/src/interpreter/driver.ts b/packages/codemode/src/interpreter/driver.ts index 54df6df3443..17b632fc252 100644 --- a/packages/codemode/src/interpreter/driver.ts +++ b/packages/codemode/src/interpreter/driver.ts @@ -186,8 +186,8 @@ export class Driver { this.attach(current) }) } - if (request.kind === "tool") return this.fork(call.id, this.tools.execute(request.path, request.args, again)) - if (request.kind === "search") return this.fork(call.id, this.tools.search(request.args, again)) + if (request.kind === "tool") return this.fork(call.id, this.tools.execute(request.path, request.args, call.key, again)) + if (request.kind === "search") return this.fork(call.id, this.tools.search(request.args, call.key, again)) // An extension's work, made by the extension bridge with the program's services in its requirements. return this.fork(call.id, this.machine.ctx.extensionWork!(request.call) as Effect.Effect) } @@ -195,7 +195,7 @@ export class Driver { private attach(call: HostCall): void { const request = call.request if (request.kind !== "tool") return - this.fork(call.id, this.tools.attach(request.path, request.args, call.handle!)) + this.fork(call.id, this.tools.attach(request.path, request.args, call.key, call.handle!)) } // A durable call that may have started before the run was saved: reattach to its work if it did, else start it. diff --git a/packages/codemode/src/tool-runtime.ts b/packages/codemode/src/tool-runtime.ts index 0df1ca75713..041a7a7f5ff 100644 --- a/packages/codemode/src/tool-runtime.ts +++ b/packages/codemode/src/tool-runtime.ts @@ -23,7 +23,7 @@ type ServicesOf> = Depth["length"] exten ? never : T extends { readonly _tag: "CodeModeTool" - readonly execute: (input: unknown) => Effect.Effect + readonly execute: (input: unknown, context?: { readonly key: string }) => Effect.Effect } ? R : T extends object @@ -37,7 +37,12 @@ export type ToolCall = { } /** A tool call the program is making, with its decoded input. */ -export type ToolInvocation = { readonly name: string; readonly input: unknown } +export type ToolInvocation = { + readonly name: string + readonly input: unknown + /** The call's idempotency key, `:`: the same at every hook, and across a suspended run's resumes. */ + readonly key: string +} /** * A call the program is making to an extension global or module function, with its arguments. A module function is @@ -375,13 +380,21 @@ export class ToolRuntimeError extends Error { export type ToolRuntime = { readonly calls: Array readonly hooks: Hooks - /** Runs a tool call. `again` runs one admitted before, re-run on resume: it is not counted a second time. */ + /** + * Runs a tool call under its idempotency `key`. `again` runs one admitted before, re-run on resume: it is not counted + * a second time. + */ readonly execute: ( path: ReadonlyArray, args: Array, + key: string, + again?: boolean, + ) => Effect.Effect + readonly search: ( + args: Array, + key: string, again?: boolean, ) => Effect.Effect - readonly search: (args: Array, again?: boolean) => Effect.Effect readonly keys: (path: ReadonlyArray) => ReadonlyArray /** The tool at `path`, if there is one, so a caller can see whether it is durable or safe to retry. */ readonly tool: (path: ReadonlyArray) => Tool | undefined @@ -396,6 +409,7 @@ export type ToolRuntime = { readonly attach: ( path: ReadonlyArray, args: Array, + key: string, handle: Handle, ) => Effect.Effect /** The handle of a durable tool's work started under `key`, if it was. */ @@ -461,16 +475,16 @@ export const make = ( catch: (cause) => new ToolRuntimeError("InvalidToolOutput", `Invalid output from tool '${name}': ${cause}`), }) - const executeTool = (name: string, tool: Tool, args: Array, again = false) => + const executeTool = (name: string, tool: Tool, args: Array, key: string, again = false) => Effect.gen(function* () { const input = yield* decode(name, tool, args) if (!again) yield* Effect.sync(() => recordCall({ name })) return yield* hooked( - { name, input }, + { name, input, key }, hooks["tool.before"], hooks["tool.after"], Effect.gen(function* () { - const raw = yield* failure(Effect.suspend(() => tool.execute(input))) + const raw = yield* failure(Effect.suspend(() => tool.execute(input, { key }))) return yield* output(name, tool, raw) }), ) @@ -486,15 +500,15 @@ export const make = ( calls, hooks, keys: (path) => namespaceKeys(root, path), - search: (args, again) => - Effect.suspend(() => executeTool("search", makeSearchTool(prepared.searchIndex), args, again)), - execute: (path, args, again) => + search: (args, key, again) => + Effect.suspend(() => executeTool("search", makeSearchTool(prepared.searchIndex), args, key, again)), + execute: (path, args, key, again) => Effect.suspend(() => { const segments = canonicalSegments(path) // Models often write `tools.search(...)` for the bare `search(...)`; honor it unless a tool owns that path. if (segments.length === 1 && segments[0] === "search" && lookup(root, segments) === undefined) - return executeTool("search", makeSearchTool(prepared.searchIndex), args, again) - return executeTool(segments.join("."), resolve(root, path, prepared.searchIndex), args, again) + return executeTool("search", makeSearchTool(prepared.searchIndex), args, key, again) + return executeTool(segments.join("."), resolve(root, path, prepared.searchIndex), args, key, again) }), tool: (path) => lookup(root, canonicalSegments(path))?.tool, start: (path, args, key, again) => @@ -504,11 +518,11 @@ export const make = ( const input = yield* decode(found.name, found.tool, args) if (!again) yield* Effect.sync(() => recordCall({ name: found.name })) const before = hooks["tool.before"] - if (before !== undefined) yield* before({ name: found.name, input }) + if (before !== undefined) yield* before({ name: found.name, input, key }) return yield* failure(Effect.suspend(() => found.durable.start(input, { key }))) }) }), - attach: (path, args, handle) => + attach: (path, args, key, handle) => Effect.suspend(() => { const found = durableTool(path) const after = hooks["tool.after"] @@ -520,9 +534,10 @@ export const make = ( // Detaching is not an outcome: `after` sees how the work ended, once it has. return Effect.flatMap(decode(found.name, found.tool, args), (input) => Effect.onExit(run, (exit) => { - if (Exit.isSuccess(exit)) return after({ name: found.name, input }, { status: "success", value: exit.value }) + const call = { name: found.name, input, key } + if (Exit.isSuccess(exit)) return after(call, { status: "success", value: exit.value }) if (Cause.hasInterruptsOnly(exit.cause)) return Effect.void - return after({ name: found.name, input }, { status: "failure", error: Cause.squash(exit.cause) }) + return after(call, { status: "failure", error: Cause.squash(exit.cause) }) }), ) }), diff --git a/packages/codemode/src/tool.ts b/packages/codemode/src/tool.ts index 64b3c109377..9e03fdadcf2 100644 --- a/packages/codemode/src/tool.ts +++ b/packages/codemode/src/tool.ts @@ -68,7 +68,11 @@ export type Tool = { readonly description: string readonly input: SchemaType readonly output: SchemaType | undefined - readonly execute: (input: unknown) => Effect.Effect + /** + * Runs a call. A run passes the call's idempotency key, `:`, stable across a suspended run's + * resumes; called directly, a random key stands in. + */ + readonly execute: (input: unknown, context?: { readonly key: string }) => Effect.Effect readonly durable?: Durable | undefined readonly retry?: Retry | undefined } @@ -85,7 +89,10 @@ export type Options) => Effect.Effect, unknown, R> + readonly execute: ( + input: InputType, + context: { readonly key: string }, + ) => Effect.Effect, unknown, R> readonly durable?: Durable, ResultType, R> } | { readonly execute?: undefined; readonly durable: Durable, ResultType, R> } @@ -105,16 +112,19 @@ export const make = , ): Tool => { const durable = options.durable as Durable | undefined - const execute = options.execute as ((input: unknown) => Effect.Effect) | undefined + const execute = options.execute as + | ((input: unknown, context: { readonly key: string }) => Effect.Effect) + | undefined if (execute === undefined && durable === undefined) throw new TypeError("A tool needs `execute` or `durable`.") return { _tag: "CodeModeTool", description: options.description, input: options.input, output: options.output, - execute: - execute ?? - ((input) => Effect.flatMap(durable!.start(input, { key: crypto.randomUUID() }), (handle) => durable!.attach(handle))), + execute: (input, context = { key: crypto.randomUUID() }) => + execute !== undefined + ? execute(input, context) + : Effect.flatMap(durable!.start(input, context), (handle) => durable!.attach(handle)), ...(durable === undefined ? {} : { durable }), ...(options.retry === undefined ? {} : { retry: options.retry }), } diff --git a/packages/codemode/test/codemode.test.ts b/packages/codemode/test/codemode.test.ts index 225805e403f..942210549be 100644 --- a/packages/codemode/test/codemode.test.ts +++ b/packages/codemode/test/codemode.test.ts @@ -167,7 +167,7 @@ describe("CodeMode call hooks", () => { } // Every hook logs its event, so the order across tools and extensions is visible. const observe = (events: Array): CodeMode.Hooks => ({ - "tool.before": (call) => Effect.sync(() => void events.push(`tool.before ${JSON.stringify(call)}`)), + "tool.before": (call) => Effect.sync(() => void events.push(`tool.before ${JSON.stringify({ name: call.name, input: call.input })}`)), "tool.after": (call, result) => Effect.sync(() => void events.push(`tool.after ${call.name} ${ended(result)}`)), "extension.before": (call) => Effect.sync(() => void events.push(`extension.before ${JSON.stringify(call)}`)), "extension.after": (call, result) => @@ -928,7 +928,7 @@ describe("CodeMode public contract", () => { const result = await Effect.runPromise(limited.execute(`search({}); return search({})`)) expect(result.ok).toBe(false) if (!result.ok) expect(result.error.kind).toBe("ToolCallLimitExceeded") - expect(observed).toEqual([{ name: "search", input: {} }]) + expect(observed).toEqual([{ name: "search", input: {}, key: expect.any(String) }]) }) test("search is an opaque, shadowable global like other built-ins", async () => { diff --git a/packages/codemode/test/suspend.test.ts b/packages/codemode/test/suspend.test.ts index a5db5ec20b7..17cf480bde5 100644 --- a/packages/codemode/test/suspend.test.ts +++ b/packages/codemode/test/suspend.test.ts @@ -375,6 +375,32 @@ describe("suspend and resume", () => { expect(jobs.jobs.get("e:0")!.cancelled).toBe(true) }) + test("tools and hooks see each call's idempotency key, the same at both hooks of a durable call", async () => { + const jobs = durableJobs() + const keys: Array = [] + const echo = Tool.make({ + description: "Echoes its key", + input: Schema.Struct({}), + output: Schema.String, + execute: (_, { key }) => Effect.succeed(key), + }) + const runtime = CodeMode.make({ + tools: { echo, job: jobs.tool }, + hooks: { + "tool.before": (call) => Effect.sync(() => void keys.push(`before ${call.name} ${call.key}`)), + "tool.after": (call) => Effect.sync(() => void keys.push(`after ${call.name} ${call.key}`)), + }, + }) + const run = runtime.start(`const job = tools.job({ value: 1 }); const key = await tools.echo({}); return [key, (await job).value]`, undefined, { + runId: "k", + }) + const fiber = Effect.runFork(run.result) + await jobs.started() + jobs.finish("k:0", 7) + expect(await Effect.runPromise(Fiber.join(fiber))).toMatchObject({ ok: true, value: ["k:1", 7] }) + expect(keys.toSorted()).toEqual(["after echo k:1", "after job k:0", "before echo k:1", "before job k:0"]) + }) + test("the snapshot codec does not import Effect", async () => { for (const file of ["codec.ts", "snapshot.ts"]) { expect(await Bun.file(`${import.meta.dir}/../src/snapshot/${file}`).text()).not.toMatch(/from "effect"/)