Selaa lähdekoodia

fix(core): stream Code Mode tool progress (#38718)

Aiden Cline 3 viikkoa sitten
vanhempi
sitoutus
b31747124b

+ 25 - 24
packages/core/src/tool/execute.ts

@@ -3,7 +3,7 @@ export type { Registration } from "./tool"
 
 import { CodeMode, Tool, toolError } from "@opencode-ai/codemode"
 import type { ToolContent } from "@opencode-ai/ai"
-import { Effect, Ref, Schema } from "effect"
+import { Effect, Ref, Schema, Semaphore } from "effect"
 import { execute, make, toLLMDefinition, type Content, type Metadata, type Registration } from "./tool"
 
 const ExecuteFile = Schema.Struct({
@@ -50,12 +50,11 @@ export const create = (registrations: ReadonlyMap<string, Registration>) => {
         const callIndex = yield* Ref.make(0)
         const files = yield* Ref.make<Array<CollectedFiles>>([])
         const calls = yield* Ref.make<Array<ExecuteCall>>([])
-        // TODO: Publish live call-list updates once V2 has a generic tool progress API.
-        const finalCalls = Ref.get(calls).pipe(
-          Effect.map((items) =>
-            items.map((call) => (call.status === "running" ? { ...call, status: "error" as const } : call)),
-          ),
-        )
+        const lock = Semaphore.makeUnsafe(1)
+        const updateCalls = (update: (items: Array<ExecuteCall>) => Array<ExecuteCall>) =>
+          lock.withPermit(
+            Ref.updateAndGet(calls, update).pipe(Effect.flatMap((toolCalls) => context.progress({ toolCalls }))),
+          )
         const result = yield* runtime(
           registrations,
           (name, registration, input) =>
@@ -66,7 +65,7 @@ export const create = (registrations: ReadonlyMap<string, Registration>) => {
                 agent: context.agent,
                 messageID: context.messageID,
                 callID: context.callID,
-                progress: context.progress,
+                progress: () => Effect.void,
               }).pipe(Effect.mapError((failure) => toolError(failure.message, failure)))
               const outputFileParts = outputFiles(executed.content)
               if (outputFileParts.length > 0)
@@ -74,26 +73,28 @@ export const create = (registrations: ReadonlyMap<string, Registration>) => {
               return executed.output
             }),
           {
-            onToolCallStart: ({ index, name, input }) =>
-              Effect.gen(function* () {
-                const shown = displayInput(input)
-                yield* Ref.update(calls, (items) => {
-                  const next = [...items]
-                  next[index] = { tool: name, status: "running", ...(shown ? { input: shown } : {}) }
-                  return next
-                })
-              }),
-            onToolCallEnd: ({ index, outcome }) =>
-              Ref.update(calls, (items) => {
-                const current = items[index]
-                if (!current) return items
+            onToolCallStart: ({ index, name, input }) => {
+              const shown = displayInput(input)
+              return updateCalls((items) => {
+                const next = [...items]
+                next[index] = { tool: name, status: "running", ...(shown ? { input: shown } : {}) }
+                return next
+              })
+            },
+            onToolCallEnd: ({ index, name, input, outcome }) => {
+              const shown = displayInput(input)
+              return updateCalls((items) => {
                 const next = [...items]
-                next[index] = { ...current, status: outcome === "success" ? "completed" : "error" }
+                next[index] = {
+                  ...(items[index] ?? { tool: name, ...(shown ? { input: shown } : {}) }),
+                  status: outcome === "success" ? "completed" : "error",
+                }
                 return next
-              }),
+              })
+            },
           },
         ).execute(code)
-        const toolCalls = yield* finalCalls
+        const toolCalls = yield* Ref.get(calls)
         const collected = (yield* Ref.get(files))
           .toSorted((left, right) => left.index - right.index)
           .flatMap((item) => item.files)

+ 12 - 3
packages/core/test/session-runner-tool-registry.test.ts

@@ -515,7 +515,7 @@ describe("ToolRegistry", () => {
     }),
   )
 
-  it.effect("executes codemode tools advertised in a model request", () =>
+  it.effect("executes and reports progress for codemode tools advertised in a model request", () =>
     Effect.gen(function* () {
       const service = yield* ToolRegistry.Service
       const executed: string[] = []
@@ -526,8 +526,11 @@ describe("ToolRegistry", () => {
             description: "Echo text",
             input: Schema.Struct({ text: Schema.String }),
             output: Schema.Struct({ text: Schema.String }),
-            execute: ({ text }) =>
-              Effect.sync(() => executed.push(`old:${text}`)).pipe(Effect.as({ output: { text } })),
+            execute: ({ text }, context) =>
+              Effect.sync(() => executed.push(`old:${text}`)).pipe(
+                Effect.andThen(context.progress({ stage: "old" })),
+                Effect.as({ output: { text } }),
+              ),
           }),
         })
         .pipe(Scope.provide(scope))
@@ -546,6 +549,7 @@ describe("ToolRegistry", () => {
         }),
       })
 
+      const progress: ToolRegistry.Progress[] = []
       const execution = yield* toolSet.execute({
         ...call("execute"),
         call: {
@@ -554,10 +558,15 @@ describe("ToolRegistry", () => {
           name: "execute",
           input: { code: 'return await tools.echo({ text: "request" })' },
         },
+        progress: (update) => Effect.sync(() => progress.push(update)),
       })
 
       expect(execution).toMatchObject({ status: "completed", content: [{ type: "text" }] })
       expect(executed).toEqual(["old:request"])
+      expect(progress).toEqual([
+        { toolCalls: [{ tool: "echo", status: "running", input: { text: "request" } }] },
+        { toolCalls: [{ tool: "echo", status: "completed", input: { text: "request" } }] },
+      ])
     }),
   )
 })

+ 44 - 1
packages/core/test/tool-execute.test.ts

@@ -4,7 +4,7 @@ import { Tool } from "@opencode-ai/core/tool/tool"
 import { Agent } from "@opencode-ai/schema/agent"
 import { Session } from "@opencode-ai/schema/session"
 import { SessionMessage } from "@opencode-ai/schema/session-message"
-import { Effect, Schema } from "effect"
+import { Deferred, Effect, Fiber, Schema } from "effect"
 
 const context = {
   sessionID: Session.ID.make("ses_execute"),
@@ -131,3 +131,46 @@ test("execute supports callable namespace tools", async () => {
   })
   expect(result.content).toEqual([{ type: "text", text: '[\n  "admin",\n  "created"\n]' }])
 })
+
+test("execute marks every admitted child call failed when interrupted", async () => {
+  const child = Tool.make({
+    description: "Wait forever",
+    input: Schema.Struct({ id: Schema.Number }),
+    output: Schema.String,
+    execute: () => Effect.never,
+  })
+  const execute = ExecuteTool.create(new Map([["wait", { tool: child, name: "wait", permission: "wait" }]]))
+  const updates: Tool.Metadata[] = []
+
+  await Effect.runPromise(
+    Effect.gen(function* () {
+      const started = yield* Deferred.make<void>()
+      const fiber = yield* Tool.execute(
+        execute,
+        { code: "return await Promise.all([tools.wait({ id: 1 }), tools.wait({ id: 2 })])" },
+        {
+          ...context,
+          progress: (update) =>
+            Effect.gen(function* () {
+              updates.push(update)
+              if (updates.length > 1) return
+              yield* Deferred.succeed(started, undefined)
+              yield* Effect.never
+            }),
+        },
+      ).pipe(Effect.forkChild)
+      yield* Deferred.await(started)
+      yield* Effect.yieldNow
+      yield* Effect.yieldNow
+      yield* Fiber.interrupt(fiber)
+    }),
+  )
+
+  expect(updates[0]).toEqual({ toolCalls: [{ tool: "wait", status: "running", input: { id: 1 } }] })
+  expect(updates.at(-1)).toEqual({
+    toolCalls: [
+      { tool: "wait", status: "error", input: { id: 1 } },
+      { tool: "wait", status: "error", input: { id: 2 } },
+    ],
+  })
+})

+ 1 - 1
packages/tui/src/routes/session/index.tsx

@@ -2872,7 +2872,7 @@ function Execute(props: ToolProps) {
   const isLoading = createMemo(() => props.part.state.status === "streaming" || props.part.state.status === "running")
   const calls = createMemo(() => executeCalls(props.metadata.toolCalls))
   const output = createMemo(() => stripAnsi(props.output?.trim() ?? ""))
-  const hasRuntimeError = createMemo(() => props.metadata.error === true)
+  const hasRuntimeError = createMemo(() => props.metadata.error === true || props.part.state.status === "error")
   const outputPreview = createMemo(() => collapseToolOutput(output(), 4, 4 * Math.max(20, ctx.width - 6)).output)
   const showOutput = createMemo(() => output() && hasRuntimeError())
   const content = createMemo(() => {