diff --git a/src/tools/delegate-task/sync-poll-timeout.test.ts b/src/tools/delegate-task/sync-poll-timeout.test.ts index fa94568de..d5381840c 100644 --- a/src/tools/delegate-task/sync-poll-timeout.test.ts +++ b/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) }) }) }) diff --git a/src/tools/delegate-task/sync-session-poller.test.ts b/src/tools/delegate-task/sync-session-poller.test.ts index 279116a17..28407bde8 100644 --- a/src/tools/delegate-task/sync-session-poller.test.ts +++ b/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) }) }) diff --git a/src/tools/delegate-task/sync-session-poller.ts b/src/tools/delegate-task/sync-session-poller.ts index c0ba40a4f..316039c73 100644 --- a/src/tools/delegate-task/sync-session-poller.ts +++ b/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 { + 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 } @@ -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