feat(codemode): pass each tool call its idempotency key, at the tool and at both hooks

packages/codemode/src/interpreter/driver.ts

@@ -186,8 +186,8 @@ export class Driver<R> {
         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<unknown, unknown, R>)
   }
@@ -195,7 +195,7 @@ export class Driver<R> {
   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.

packages/codemode/src/tool-runtime.ts

@@ -23,7 +23,7 @@ type ServicesOf<T, Depth extends ReadonlyArray<unknown>> = Depth["length"] exten
   ? never
   : T extends {
         readonly _tag: "CodeModeTool"
-        readonly execute: (input: unknown) => Effect.Effect<unknown, unknown, infer R>
+        readonly execute: (input: unknown, context?: { readonly key: string }) => Effect.Effect<unknown, unknown, infer R>
       }
     ? 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, `<run id>:<call id>`: 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<R = never> = {
   readonly calls: Array<ToolCall>
   readonly hooks: Hooks<R>
-  /** 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<string>,
     args: Array<Json | undefined>,
+    key: string,
+    again?: boolean,
+  ) => Effect.Effect<Json | undefined, unknown, R>
+  readonly search: (
+    args: Array<Json | undefined>,
+    key: string,
     again?: boolean,
   ) => Effect.Effect<Json | undefined, unknown, R>
-  readonly search: (args: Array<Json | undefined>, again?: boolean) => Effect.Effect<Json | undefined, unknown, R>
   readonly keys: (path: ReadonlyArray<string>) => ReadonlyArray<string>
   /** 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<string>) => Tool<R> | undefined
@@ -396,6 +409,7 @@ export type ToolRuntime<R = never> = {
   readonly attach: (
     path: ReadonlyArray<string>,
     args: Array<Json | undefined>,
+    key: string,
     handle: Handle,
   ) => Effect.Effect<Json | undefined, unknown, R>
   /** The handle of a durable tool's work started under `key`, if it was. */
@@ -461,16 +475,16 @@ export const make = <R>(
       catch: (cause) => new ToolRuntimeError("InvalidToolOutput", `Invalid output from tool '${name}': ${cause}`),
     })
 
-  const executeTool = (name: string, tool: Tool<R>, args: Array<Json | undefined>, again = false) =>
+  const executeTool = (name: string, tool: Tool<R>, args: Array<Json | undefined>, 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 = <R>(
     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 = <R>(
           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 = <R>(
         // 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) })
           }),
         )
       }),

packages/codemode/src/tool.ts

@@ -68,7 +68,11 @@ export type Tool<R = never> = {
   readonly description: string
   readonly input: SchemaType
   readonly output: SchemaType | undefined
-  readonly execute: (input: unknown) => Effect.Effect<unknown, unknown, R>
+  /**
+   * Runs a call. A run passes the call's idempotency key, `<run id>:<call id>`, stable across a suspended run's
+   * resumes; called directly, a random key stands in.
+   */
+  readonly execute: (input: unknown, context?: { readonly key: string }) => Effect.Effect<unknown, unknown, R>
   readonly durable?: Durable<unknown, unknown, R> | undefined
   readonly retry?: Retry | undefined
 }
@@ -85,7 +89,10 @@ export type Options<I extends SchemaType, O extends SchemaType | undefined, R =
   readonly retry?: Retry
 } & (
   | {
-      readonly execute: (input: InputType<I>) => Effect.Effect<ResultType<O>, unknown, R>
+      readonly execute: (
+        input: InputType<I>,
+        context: { readonly key: string },
+      ) => Effect.Effect<ResultType<O>, unknown, R>
       readonly durable?: Durable<InputType<I>, ResultType<O>, R>
     }
   | { readonly execute?: undefined; readonly durable: Durable<InputType<I>, ResultType<O>, R> }
@@ -105,16 +112,19 @@ export const make = <I extends SchemaType, const O extends SchemaType | undefine
   options: Options<I, O, R>,
 ): Tool<R> => {
   const durable = options.durable as Durable<unknown, unknown, R> | undefined
-  const execute = options.execute as ((input: unknown) => Effect.Effect<unknown, unknown, R>) | undefined
+  const execute = options.execute as
+    | ((input: unknown, context: { readonly key: string }) => Effect.Effect<unknown, unknown, R>)
+    | 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 }),
   }

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<string>): 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 () => {

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<string> = []
+    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"/)