diff --git a/packages/codemode/src/codemode.ts b/packages/codemode/src/codemode.ts index 786b8b4f4e..0d557169d8 100644 --- a/packages/codemode/src/codemode.ts +++ b/packages/codemode/src/codemode.ts @@ -40,7 +40,7 @@ export type ExecuteOptions = {}> = { limits?: ExecutionLimits /** Observes decoded tool input immediately before tool execution. */ onToolCallStart?: (call: ToolRuntime.ToolCallStarted) => Effect.Effect> - /** Observes each admitted tool call as it settles, with outcome and duration. */ + /** Observes each admitted tool call as it succeeds, fails, or is interrupted. */ onToolCallEnd?: (call: ToolRuntime.ToolCallEnded) => Effect.Effect> } diff --git a/packages/codemode/src/tool-runtime.ts b/packages/codemode/src/tool-runtime.ts index e39c30e38a..f1e189f733 100644 --- a/packages/codemode/src/tool-runtime.ts +++ b/packages/codemode/src/tool-runtime.ts @@ -1,4 +1,4 @@ -import { Cause, Effect, Schema } from "effect" +import { Cause, Effect, Exit, Schema } from "effect" import { ToolError, toolError } from "./tool-error.js" import { decodeInput as decodeToolInput, @@ -52,7 +52,7 @@ export type ToolCallEnded = { readonly name: string readonly input: unknown readonly durationMs: number - readonly outcome: "success" | "failure" + readonly outcome: "success" | "failure" | "interrupted" readonly message?: string } @@ -495,22 +495,19 @@ export const make = ( const root = toolTrie(tools) const searchTool = makeSearchTool(searchIndex) - // End hooks observe settled success or failure; interruption emits neither outcome. const observeEnd = (effect: Effect.Effect, call: ToolCallStarted): Effect.Effect => { const onEnd = hooks?.onToolCallEnd if (onEnd === undefined) return effect const startedAt = Date.now() return effect.pipe( - Effect.tap(() => onEnd({ ...call, durationMs: Date.now() - startedAt, outcome: "success" })), - Effect.tapError((error) => { + Effect.onExit((exit) => { + const durationMs = Date.now() - startedAt + if (Exit.isSuccess(exit)) return onEnd({ ...call, durationMs, outcome: "success" }) + if (Cause.hasInterruptsOnly(exit.cause)) return onEnd({ ...call, durationMs, outcome: "interrupted" }) + const error = Cause.squash(exit.cause) const message = error instanceof ToolError || error instanceof ToolRuntimeError ? error.message : "Tool execution failed" - return onEnd({ - ...call, - durationMs: Date.now() - startedAt, - outcome: "failure", - message, - }) + return onEnd({ ...call, durationMs, outcome: "failure", message }) }), ) } @@ -528,12 +525,6 @@ export const make = ( calls.push(call) } - const recordAndObserve = (name: string, input: unknown) => - Effect.sync(() => { - recordCall({ name }) - return calls.length - 1 - }).pipe(Effect.tap((index) => hooks?.onToolCallStart?.({ index, name, input }) ?? Effect.void)) - const executeTool = (name: string, tool: Tool, externalArgs: Array) => Effect.gen(function* () { if (externalArgs.length !== 1) @@ -547,9 +538,14 @@ export const make = ( name === "search" ? [] : ["The signature may have changed. Use search to get the current signature."], ), }) - const index = yield* recordAndObserve(name, input) + const index = yield* Effect.sync(() => { + recordCall({ name }) + return calls.length - 1 + }) + const call = { index, name, input } return yield* observeEnd( Effect.gen(function* () { + if (hooks?.onToolCallStart !== undefined) yield* hooks.onToolCallStart(call) const raw = yield* runHost(Effect.suspend(() => tool.execute(input))) const result = yield* Effect.try({ try: () => decodeToolOutput(tool, raw), @@ -557,7 +553,7 @@ export const make = ( }) return yield* decodeOutput(result, name) }), - { index, name, input }, + call, ) }) diff --git a/packages/codemode/test/codemode.test.ts b/packages/codemode/test/codemode.test.ts index 86b1e076f5..cac1d46024 100644 --- a/packages/codemode/test/codemode.test.ts +++ b/packages/codemode/test/codemode.test.ts @@ -189,7 +189,12 @@ describe("CodeMode tool-call observation", () => { description: "Look up a value", input: Schema.Struct({ query: Schema.String }), output: Schema.String, - execute: ({ query }) => (query === "boom" ? Effect.fail(toolError("Lookup refused")) : Effect.succeed(query)), + execute: ({ query }) => + query === "boom" + ? Effect.fail(toolError("Lookup refused")) + : query === "defect" + ? Effect.die("broken") + : Effect.succeed(query), }) const runtime = CodeMode.make({ @@ -215,14 +220,98 @@ describe("CodeMode tool-call observation", () => { expect(success.ok).toBe(true) const failure = await Effect.runPromise(runtime.execute(`return await tools.context.lookup({ query: "boom" })`)) expect(failure.ok).toBe(false) + const defect = await Effect.runPromise(runtime.execute(`return await tools.context.lookup({ query: "defect" })`)) + expect(defect.ok).toBe(false) expect(events).toStrictEqual([ { phase: "start", index: 0, name: "context.lookup" }, { phase: "end", index: 0, name: "context.lookup", outcome: "success" }, { phase: "start", index: 0, name: "context.lookup" }, { phase: "end", index: 0, name: "context.lookup", outcome: "failure", message: "Lookup refused" }, + { phase: "start", index: 0, name: "context.lookup" }, + { phase: "end", index: 0, name: "context.lookup", outcome: "failure", message: "Tool execution failed" }, ]) }) + + test("observes interrupted calls", async () => { + const events: Array = [] + const call = Tool.make({ + description: "Interrupt", + input: Schema.Struct({}), + output: Schema.String, + execute: () => Effect.interrupt, + }) + const exit = await Effect.runPromiseExit( + CodeMode.make({ + tools: { host: { call } }, + onToolCallStart: () => Effect.sync(() => events.push("start")), + onToolCallEnd: (call) => Effect.sync(() => events.push(`end:${call.outcome}`)), + }).execute("return await tools.host.call({})"), + ) + + expect(exit._tag).toBe("Failure") + expect(events).toEqual(["start", "end:interrupted"]) + }) + + test("observes running calls interrupted during completion", async () => { + const events: Array = [] + const call = Tool.make({ + description: "Pending", + input: Schema.Struct({}), + output: Schema.String, + execute: () => Effect.never, + }) + const result = await Effect.runPromise( + CodeMode.make({ + tools: { host: { call } }, + onToolCallStart: () => Effect.sync(() => events.push("start")), + onToolCallEnd: (call) => Effect.sync(() => events.push(`end:${call.outcome}`)), + }).execute('tools.host.call({}); return "done"'), + ) + + expect(result).toMatchObject({ ok: true, value: "done" }) + expect(events).toEqual(["start", "end:interrupted"]) + }) + + test("ends calls interrupted during start observation", async () => { + const events: Array = [] + const call = Tool.make({ + description: "Unused", + input: Schema.Struct({}), + output: Schema.String, + execute: () => Effect.succeed("unused"), + }) + const exit = await Effect.runPromiseExit( + CodeMode.make({ + tools: { host: { call } }, + onToolCallStart: () => Effect.interrupt, + onToolCallEnd: (call) => Effect.sync(() => events.push(call.outcome)), + }).execute("return await tools.host.call({})"), + ) + + expect(exit._tag).toBe("Failure") + expect(events).toEqual(["interrupted"]) + }) + + test("observes calls interrupted by the execution timeout", async () => { + const outcomes: Array = [] + const call = Tool.make({ + description: "Pending", + input: Schema.Struct({}), + output: Schema.String, + execute: () => Effect.never, + }) + const result = await Effect.runPromise( + CodeMode.make({ + tools: { host: { call } }, + limits: { timeoutMs: 10 }, + onToolCallEnd: (call) => Effect.sync(() => outcomes.push(call.outcome)), + }).execute("return await tools.host.call({})"), + ) + + expect(result).toMatchObject({ ok: false, error: { kind: "TimeoutExceeded" } }) + expect(outcomes).toEqual(["interrupted"]) + }) }) describe("CodeMode console capture", () => {