From 80d3339c4c0689669c0bef38b6d4244ecf9e0f45 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Thu, 16 Apr 2026 13:52:04 +0900 Subject: [PATCH] feat(background-agent): add wait-for-task-session helper Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/features/background-agent/index.ts | 2 + .../wait-for-task-session.test.ts | 95 +++++++++++++++++++ .../background-agent/wait-for-task-session.ts | 68 +++++++++++++ 3 files changed, 165 insertions(+) create mode 100644 src/features/background-agent/wait-for-task-session.test.ts create mode 100644 src/features/background-agent/wait-for-task-session.ts diff --git a/src/features/background-agent/index.ts b/src/features/background-agent/index.ts index e1d1a9b73..d21b2faa8 100644 --- a/src/features/background-agent/index.ts +++ b/src/features/background-agent/index.ts @@ -1,2 +1,4 @@ export * from "./types" export { BackgroundManager, type SubagentSessionCreatedEvent, type OnSubagentSessionCreated } from "./manager" +export { waitForTaskSessionID } from "./wait-for-task-session" +export type { WaitForTaskSessionIDOptions } from "./wait-for-task-session" diff --git a/src/features/background-agent/wait-for-task-session.test.ts b/src/features/background-agent/wait-for-task-session.test.ts new file mode 100644 index 000000000..812d9f700 --- /dev/null +++ b/src/features/background-agent/wait-for-task-session.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, test } from "bun:test" + +import type { BackgroundTaskStatus } from "./types" +import { waitForTaskSessionID } from "./wait-for-task-session" + +interface TaskSnapshot { + sessionID?: string + status?: BackgroundTaskStatus +} + +function createManager(responses: TaskSnapshot[]) { + let index = 0 + + return { + getTask(_taskID: string): TaskSnapshot { + const response = responses[Math.min(index, responses.length - 1)] + index += 1 + return response + }, + } +} + +describe("waitForTaskSessionID", () => { + test("#given task already has a session id #when waiting #then it returns immediately", async () => { + // given + const manager = createManager([{ sessionID: "ses_ready_123", status: "running" }]) + + // when + const sessionID = await waitForTaskSessionID(manager, "bg_ready") + + // then + expect(sessionID).toBe("ses_ready_123") + }) + + test("#given session appears later #when waiting #then it polls until resolved", async () => { + // given + const manager = createManager([ + { status: "running" }, + { status: "running" }, + { sessionID: "ses_late_123", status: "running" }, + ]) + + // when + const sessionID = await waitForTaskSessionID(manager, "bg_late", { + intervalMs: 1, + timeoutMs: 20, + }) + + // then + expect(sessionID).toBe("ses_late_123") + }) + + test("#given aborted signal #when waiting #then it returns undefined", async () => { + // given + const controller = new AbortController() + controller.abort() + const manager = createManager([{ status: "running" }]) + + // when + const sessionID = await waitForTaskSessionID(manager, "bg_abort", { + signal: controller.signal, + }) + + // then + expect(sessionID).toBeUndefined() + }) + + test("#given task never resolves #when waiting past timeout #then it returns undefined", async () => { + // given + const manager = createManager([{ status: "running" }, { status: "running" }, { status: "running" }]) + + // when + const sessionID = await waitForTaskSessionID(manager, "bg_timeout", { + intervalMs: 1, + timeoutMs: 3, + }) + + // then + expect(sessionID).toBeUndefined() + }) + + test.each(["error", "cancelled", "interrupt"] satisfies BackgroundTaskStatus[])( + "#given %s task state #when waiting #then it returns undefined", + async (status: BackgroundTaskStatus) => { + // given + const manager = createManager([{ status }]) + + // when + const sessionID = await waitForTaskSessionID(manager, `bg_${status}`) + + // then + expect(sessionID).toBeUndefined() + } + ) +}) diff --git a/src/features/background-agent/wait-for-task-session.ts b/src/features/background-agent/wait-for-task-session.ts new file mode 100644 index 000000000..eb5fe49d8 --- /dev/null +++ b/src/features/background-agent/wait-for-task-session.ts @@ -0,0 +1,68 @@ +import { getTimingConfig } from "../../tools/delegate-task/timing" +import type { BackgroundTaskStatus } from "./types" + +type SessionWaitTerminalStatus = Extract +type AbortSignalLike = { aborted: boolean } + +interface TaskReader { + getTask(taskID: string): { sessionID?: string; status?: BackgroundTaskStatus } | undefined +} + +export interface WaitForTaskSessionIDOptions { + timeoutMs?: number + intervalMs?: number + signal?: AbortSignalLike +} + +function isTerminalStatus(status: BackgroundTaskStatus | undefined): status is SessionWaitTerminalStatus { + return status === "error" || status === "cancelled" || status === "interrupt" +} + +function waitForInterval(intervalMs: number): Promise { + return new Promise(resolve => { + const scheduler = globalThis as { setTimeout: (handler: () => void, timeout?: number) => unknown } + scheduler.setTimeout(resolve, intervalMs) + }) +} + +export async function waitForTaskSessionID( + manager: TaskReader, + taskID: string, + options: WaitForTaskSessionIDOptions = {} +): Promise { + const timing = getTimingConfig() + const timeoutMs = options.timeoutMs ?? timing.WAIT_FOR_SESSION_TIMEOUT_MS + const intervalMs = options.intervalMs ?? timing.WAIT_FOR_SESSION_INTERVAL_MS + + if (options.signal?.aborted) { + return undefined + } + + const initialTask = manager.getTask(taskID) + if (initialTask?.sessionID) { + return initialTask.sessionID + } + if (isTerminalStatus(initialTask?.status)) { + return undefined + } + + const deadline = Date.now() + timeoutMs + + while (Date.now() < deadline) { + if (options.signal?.aborted) { + return undefined + } + + await waitForInterval(intervalMs) + + const task = manager.getTask(taskID) + if (task?.sessionID) { + return task.sessionID + } + if (isTerminalStatus(task?.status)) { + return undefined + } + } + + return undefined +}