Prechádzať zdrojové kódy

fix(codemode): report interrupted tool calls (#38741)

Aiden Cline 3 týždňov pred
rodič
commit
ee5460a152

+ 1 - 1
packages/codemode/src/codemode.ts

@@ -40,7 +40,7 @@ export type ExecuteOptions<Provided extends Record<string, unknown> = {}> = {
   limits?: ExecutionLimits
   /** Observes decoded tool input immediately before tool execution. */
   onToolCallStart?: (call: ToolRuntime.ToolCallStarted) => Effect.Effect<void, never, Services<Provided>>
-  /** 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<void, never, Services<Provided>>
 }
 

+ 15 - 19
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 = <R>(
   const root = toolTrie(tools)
   const searchTool = makeSearchTool(searchIndex)
 
-  // End hooks observe settled success or failure; interruption emits neither outcome.
   const observeEnd = <A, E>(effect: Effect.Effect<A, E, R>, call: ToolCallStarted): Effect.Effect<A, E, R> => {
     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 = <R>(
     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<R>, externalArgs: Array<unknown>) =>
     Effect.gen(function* () {
       if (externalArgs.length !== 1)
@@ -547,9 +538,14 @@ export const make = <R>(
             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 = <R>(
           })
           return yield* decodeOutput(result, name)
         }),
-        { index, name, input },
+        call,
       )
     })
 

+ 90 - 1
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<string> = []
+    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<string> = []
+    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<string> = []
+    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<string> = []
+    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", () => {