refactor: fix tool call state handling and clean up imports (#21709)

5 files changed, 45 insertions, 6 deletions

diff --git a/packages/opencode/src/server/routes/session.ts b/packages/opencode/src/server/routes/session.ts
index 68be1a6a5..b57ed9d47 100644
--- a/packages/opencode/src/server/routes/session.ts
+++ b/packages/opencode/src/server/routes/session.ts
@@ -121,7 +121,6 @@ export const SessionRoutes = lazy(() =>
       ),
       async (c) => {
         const sessionID = c.req.valid("param").sessionID
-        log.info("SEARCH", { url: c.req.url })
         const session = await Session.get(sessionID)
         return c.json(session)
       },
diff --git a/packages/opencode/src/session/index.ts b/packages/opencode/src/session/index.ts
index 65032de96..8e1ed9dcc 100644
--- a/packages/opencode/src/session/index.ts
+++ b/packages/opencode/src/session/index.ts
@@ -12,7 +12,7 @@ import { Installation } from "../installation"
 import { Database, NotFoundError, eq, and, gte, isNull, desc, like, inArray, lt } from "../storage/db"
 import { SyncEvent } from "../sync"
 import type { SQL } from "../storage/db"
-import { SessionTable } from "./session.sql"
+import { PartTable, SessionTable } from "./session.sql"
 import { ProjectTable } from "../project/project.sql"
 import { Storage } from "@/storage/storage"
 import { Log } from "../util/log"
@@ -345,6 +345,11 @@ export namespace Session {
       messageID: MessageID
       partID: PartID
     }) => Effect.Effect<PartID>
+    readonly getPart: (input: {
+      sessionID: SessionID
+      messageID: MessageID
+      partID: PartID
+    }) => Effect.Effect<MessageV2.Part | undefined>
     readonly updatePart: <T extends MessageV2.Part>(part: T) => Effect.Effect<T>
     readonly updatePartDelta: (input: {
       sessionID: SessionID
@@ -492,6 +497,29 @@ export namespace Session {
           return part
         }).pipe(Effect.withSpan("Session.updatePart"))
 
+      const getPart: Interface["getPart"] = Effect.fn("Session.getPart")(function* (input) {
+        const row = Database.use((db) =>
+          db
+            .select()
+            .from(PartTable)
+            .where(
+              and(
+                eq(PartTable.session_id, input.sessionID),
+                eq(PartTable.message_id, input.messageID),
+                eq(PartTable.id, input.partID),
+              ),
+            )
+            .get(),
+        )
+        if (!row) return
+        return {
+          ...row.data,
+          id: row.id,
+          sessionID: row.session_id,
+          messageID: row.message_id,
+        } as MessageV2.Part
+      })
+
       const create = Effect.fn("Session.create")(function* (input?: {
         parentID?: SessionID
         title?: string
@@ -675,6 +703,7 @@ export namespace Session {
         removeMessage,
         removePart,
         updatePart,
+        getPart,
         updatePartDelta,
         initialize,
       })
diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts
index 497747a5d..225961aef 100644
--- a/packages/opencode/src/session/processor.ts
+++ b/packages/opencode/src/session/processor.ts
@@ -176,12 +176,22 @@ export namespace SessionProcessor {
               if (ctx.assistantMessage.summary) {
                 throw new Error(`Tool call not allowed while generating summary: ${value.toolName}`)
               }
-              const match = ctx.toolcalls[value.toolCallId]
-              if (!match) return
+              const pointer = ctx.toolcalls[value.toolCallId]
+              const match = yield* session.getPart({
+                partID: pointer.id,
+                messageID: pointer.messageID,
+                sessionID: pointer.sessionID,
+              })
+              if (!match || match.type !== "tool") return
               ctx.toolcalls[value.toolCallId] = yield* session.updatePart({
                 ...match,
                 tool: value.toolName,
-                state: { status: "running", input: value.input, time: { start: Date.now() } },
+                state: {
+                  ...match.state,
+                  status: "running",
+                  input: value.input,
+                  time: { start: Date.now() },
+                },
                 metadata: match.metadata?.providerExecuted
                   ? { ...value.providerMetadata, providerExecuted: true }
                   : value.providerMetadata,
@@ -237,6 +247,7 @@ export namespace SessionProcessor {
             case "tool-error": {
               const match = ctx.toolcalls[value.toolCallId]
               if (!match || match.state.status !== "running") return
+
               yield* session.updatePart({
                 ...match,
                 state: {
diff --git a/packages/opencode/src/tool/task.ts b/packages/opencode/src/tool/task.ts
index 73b55a2fb..b97b53bb9 100644
--- a/packages/opencode/src/tool/task.ts
+++ b/packages/opencode/src/tool/task.ts
@@ -9,6 +9,7 @@ import { SessionPrompt } from "../session/prompt"
 import { Config } from "../config/config"
 import { Permission } from "@/permission"
 import { Effect } from "effect"
+import { Log } from "@/util/log"
 
 const id = "task"
 
diff --git a/packages/sdk/js/src/v2/index.ts b/packages/sdk/js/src/v2/index.ts
index d514784bc..9615eacc7 100644
--- a/packages/sdk/js/src/v2/index.ts
+++ b/packages/sdk/js/src/v2/index.ts
@@ -6,7 +6,6 @@ import { createOpencodeServer } from "./server.js"
 import type { ServerOptions } from "./server.js"
 
 export * as data from "./data.js"
-import * as data from "./data.js"
 
 export async function createOpencode(options?: ServerOptions) {
   const server = await createOpencodeServer({