Merge pull request #2409 from ualtinok/fix/bgpollfix

fix(delegate-task): abort sync sessions on timeout and parent abort

src/tools/delegate-task/sync-poll-timeout.test.ts

@@ -13,9 +13,12 @@ function createMockCtx(aborted = false) {
   }
 }
 
-function createNeverCompleteClient(sessionID: string) {
+function createNeverCompleteClient(sessionID: string, onAbort?: () => void) {
   return {
     session: {
+      abort: async () => {
+        onAbort?.()
+      },
       messages: async () => ({
         data: [{ info: { id: "msg_001", role: "user", time: { created: 1000 } } }],
       }),
@@ -59,7 +62,10 @@ describe("syncPollTimeoutMs threading", () => {
     describe("#when custom timeout is provided", () => {
       test("#then custom timeout value is used", async () => {
         const { pollSyncSession } = require("./sync-session-poller")
-        const mockClient = createNeverCompleteClient("ses_custom")
+        let abortCount = 0
+        const mockClient = createNeverCompleteClient("ses_custom", () => {
+          abortCount++
+        })
 
         await withMockedDateNow(60_000, async () => {
           const result = await pollSyncSession(createMockCtx(), mockClient, {
@@ -70,6 +76,7 @@ describe("syncPollTimeoutMs threading", () => {
           }, 120_000)
 
           expect(result).toBe("Poll timeout reached after 120000ms for session ses_custom")
+          expect(abortCount).toBe(1)
         })
       })
     })

src/tools/delegate-task/sync-session-poller.test.ts

@@ -223,8 +223,12 @@ describe("pollSyncSession", () => {
     test("returns abort message when signal is aborted", async () => {
       //#given
       const { pollSyncSession } = require("./sync-session-poller")
+      let abortCount = 0
       const mockClient = {
         session: {
+          abort: async () => {
+            abortCount++
+          },
           messages: async () => ({ data: [] }),
           status: async () => ({ data: {} }),
         },
@@ -241,6 +245,7 @@ describe("pollSyncSession", () => {
       //#then
       expect(result).toContain("Task aborted")
       expect(result).toContain("ses_abort")
+      expect(abortCount).toBe(1)
     })
   })
 
@@ -256,8 +261,12 @@ describe("pollSyncSession", () => {
         MAX_POLL_TIME_MS: 0,
       })
 
+      let abortCount = 0
       const mockClient = {
         session: {
+          abort: async () => {
+            abortCount++
+          },
           messages: async () => ({
             data: [
               { info: { id: "msg_001", role: "user", time: { created: 1000 } } },
@@ -277,6 +286,7 @@ describe("pollSyncSession", () => {
 
       //#then - timeout returns error string
       expect(result).toBe("Poll timeout reached after 50ms for session ses_timeout")
+      expect(abortCount).toBe(1)
     })
   })
 

src/tools/delegate-task/sync-session-poller.ts

@@ -6,6 +6,22 @@ import { normalizeSDKResponse } from "../../shared"
 
 const NON_TERMINAL_FINISH_REASONS = new Set(["tool-calls", "unknown"])
 
+function wait(milliseconds: number): Promise<void> {
+  const sharedBuffer = new SharedArrayBuffer(Int32Array.BYTES_PER_ELEMENT)
+  const typedArray = new Int32Array(sharedBuffer)
+  const result = Atomics.waitAsync(typedArray, 0, 0, milliseconds)
+  return result.async ? result.value.then(() => undefined) : Promise.resolve()
+}
+
+function abortSyncSession(client: OpencodeClient, sessionID: string, reason: string): void {
+  log("[task] Aborting sync session", { sessionID, reason })
+  void client.session.abort({
+    path: { id: sessionID },
+  }).catch((error: unknown) => {
+    log("[task] Failed to abort sync session", { sessionID, reason, error: String(error) })
+  })
+}
+
 export function isSessionComplete(messages: SessionMessage[]): boolean {
   let lastUser: SessionMessage | undefined
   let lastAssistant: SessionMessage | undefined
@@ -46,11 +62,12 @@ export async function pollSyncSession(
   while (Date.now() - pollStart < maxPollTimeMs) {
     if (ctx.abort?.aborted) {
       log("[task] Aborted by user", { sessionID: input.sessionID })
+      abortSyncSession(client, input.sessionID, "parent_abort")
       if (input.toastManager && input.taskId) input.toastManager.removeTask(input.taskId)
       return `Task aborted.\n\nSession ID: ${input.sessionID}`
     }
 
-    await new Promise(resolve => setTimeout(resolve, syncTiming.POLL_INTERVAL_MS))
+    await wait(syncTiming.POLL_INTERVAL_MS)
     pollCount++
 
     let statusResult: { data?: Record<string, { type: string }> }
@@ -118,6 +135,7 @@ export async function pollSyncSession(
   if (Date.now() - pollStart >= maxPollTimeMs) {
     timedOut = true
     log("[task] Poll timeout reached", { sessionID: input.sessionID, pollCount })
+    abortSyncSession(client, input.sessionID, "poll_timeout")
   }
 
   return timedOut ? `Poll timeout reached after ${maxPollTimeMs}ms for session ${input.sessionID}` : null