feat(atlas): add background task session tracking with retry scheduling
- Add background-launch-session-tracking.ts to persist delegated sessions
- Add task-context.ts for task context resolution utilities
- Modify idle-event.ts to schedule retries when background tasks are running
- Update tool-execute-after.ts to integrate session tracking
- Add comprehensive tests for background task retry behavior
🤖 Generated with assistance of OhMyOpenCode
This commit is contained in:
YeonGyu-Kim
@@ -0,0 +1,225 @@
|
||||
import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"
|
||||
import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs"
|
||||
import { tmpdir } from "node:os"
|
||||
import { join } from "node:path"
|
||||
import { randomUUID } from "node:crypto"
|
||||
import type { PluginInput } from "@opencode-ai/plugin"
|
||||
import { createAtlasHook } from "./atlas-hook"
|
||||
import { clearBoulderState, writeBoulderState } from "../../features/boulder-state"
|
||||
import { _resetForTesting, registerAgentName } from "../../features/claude-code-session-state"
|
||||
|
||||
type LongTimerCallback = (...args: unknown[]) => void | Promise<void>
|
||||
|
||||
describe("atlas background task retry", () => {
|
||||
let testDir: string
|
||||
const sessionID = "main-session-123"
|
||||
const capturedTimers = new Map<number, { callback: () => Promise<void> | void; cleared: boolean }>()
|
||||
let nextFakeTimerId = 1000
|
||||
const originalSetTimeout = globalThis.setTimeout
|
||||
const originalClearTimeout = globalThis.clearTimeout
|
||||
|
||||
async function flushMicrotasks(): Promise<void> {
|
||||
await Promise.resolve()
|
||||
await Promise.resolve()
|
||||
}
|
||||
|
||||
async function firePendingTimers(): Promise<void> {
|
||||
const entries = [...capturedTimers.entries()]
|
||||
for (const [id, entry] of entries) {
|
||||
if (entry.cleared) {
|
||||
continue
|
||||
}
|
||||
|
||||
capturedTimers.delete(id)
|
||||
await entry.callback()
|
||||
}
|
||||
await flushMicrotasks()
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
_resetForTesting()
|
||||
registerAgentName("atlas")
|
||||
registerAgentName("sisyphus")
|
||||
|
||||
testDir = join(tmpdir(), `atlas-background-retry-${randomUUID()}`)
|
||||
mkdirSync(testDir, { recursive: true })
|
||||
|
||||
capturedTimers.clear()
|
||||
nextFakeTimerId = 1000
|
||||
|
||||
globalThis.setTimeout = ((callback: Parameters<typeof setTimeout>[0], delay?: number, ...args: unknown[]) => {
|
||||
const normalizedDelay = typeof delay === "number" ? delay : 0
|
||||
if (typeof callback !== "function") {
|
||||
return originalSetTimeout(callback, delay, ...args)
|
||||
}
|
||||
|
||||
if (normalizedDelay >= 5000) {
|
||||
const id = nextFakeTimerId++
|
||||
capturedTimers.set(id, {
|
||||
callback: () => (callback as LongTimerCallback)(...args),
|
||||
cleared: false,
|
||||
})
|
||||
return id as unknown as ReturnType<typeof setTimeout>
|
||||
}
|
||||
|
||||
return originalSetTimeout(callback, delay, ...args)
|
||||
}) as typeof setTimeout
|
||||
|
||||
globalThis.clearTimeout = ((id?: number | ReturnType<typeof setTimeout>) => {
|
||||
if (typeof id === "number" && capturedTimers.has(id)) {
|
||||
capturedTimers.get(id)!.cleared = true
|
||||
capturedTimers.delete(id)
|
||||
return
|
||||
}
|
||||
|
||||
originalClearTimeout(id as Parameters<typeof clearTimeout>[0])
|
||||
}) as typeof clearTimeout
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
globalThis.setTimeout = originalSetTimeout
|
||||
globalThis.clearTimeout = originalClearTimeout
|
||||
_resetForTesting()
|
||||
clearBoulderState(testDir)
|
||||
if (existsSync(testDir)) {
|
||||
rmSync(testDir, { recursive: true, force: true })
|
||||
}
|
||||
})
|
||||
|
||||
test("#given background tasks are still running #when retry fires before they finish #then atlas keeps retrying until continuation can resume", async () => {
|
||||
// given
|
||||
const planPath = join(testDir, "test-plan.md")
|
||||
writeFileSync(planPath, "# Plan\n- [ ] Task 1\n- [ ] Task 2")
|
||||
writeBoulderState(testDir, {
|
||||
active_plan: planPath,
|
||||
started_at: "2026-01-02T10:00:00Z",
|
||||
session_ids: [sessionID],
|
||||
plan_name: "test-plan",
|
||||
agent: "atlas",
|
||||
})
|
||||
|
||||
let backgroundRunning = true
|
||||
const promptMock = mock(async () => ({}))
|
||||
const hook = createAtlasHook({
|
||||
directory: testDir,
|
||||
client: {
|
||||
session: {
|
||||
promptAsync: promptMock,
|
||||
messages: async () => ({ data: [] }),
|
||||
},
|
||||
},
|
||||
} as unknown as PluginInput, {
|
||||
directory: testDir,
|
||||
backgroundManager: {
|
||||
getTasksByParentSession: () => backgroundRunning ? [{ status: "running" }] : [],
|
||||
} as unknown as NonNullable<Parameters<typeof createAtlasHook>[1]>["backgroundManager"] & {
|
||||
getTasksByParentSession: (sessionID: string) => Array<{ status: string }>
|
||||
},
|
||||
})
|
||||
|
||||
// when
|
||||
await hook.handler({ event: { type: "session.idle", properties: { sessionID } } })
|
||||
await firePendingTimers()
|
||||
backgroundRunning = false
|
||||
await firePendingTimers()
|
||||
|
||||
// then
|
||||
expect(promptMock).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
test("#given multiple idle events arrive while background retry is already pending #when tasks are still running #then atlas keeps only one retry timer active", async () => {
|
||||
// given
|
||||
const planPath = join(testDir, "test-plan.md")
|
||||
writeFileSync(planPath, "# Plan\n- [ ] Task 1\n- [ ] Task 2")
|
||||
writeBoulderState(testDir, {
|
||||
active_plan: planPath,
|
||||
started_at: "2026-01-02T10:00:00Z",
|
||||
session_ids: [sessionID],
|
||||
plan_name: "test-plan",
|
||||
agent: "atlas",
|
||||
})
|
||||
|
||||
let backgroundRunning = true
|
||||
const promptMock = mock(async () => ({}))
|
||||
const hook = createAtlasHook({
|
||||
directory: testDir,
|
||||
client: {
|
||||
session: {
|
||||
promptAsync: promptMock,
|
||||
messages: async () => ({ data: [] }),
|
||||
},
|
||||
},
|
||||
} as unknown as PluginInput, {
|
||||
directory: testDir,
|
||||
backgroundManager: {
|
||||
getTasksByParentSession: () => backgroundRunning ? [{ status: "running" }] : [],
|
||||
} as unknown as NonNullable<Parameters<typeof createAtlasHook>[1]>["backgroundManager"] & {
|
||||
getTasksByParentSession: (sessionID: string) => Array<{ status: string }>
|
||||
},
|
||||
})
|
||||
|
||||
// when
|
||||
await hook.handler({ event: { type: "session.idle", properties: { sessionID } } })
|
||||
await hook.handler({ event: { type: "session.idle", properties: { sessionID } } })
|
||||
await hook.handler({ event: { type: "session.idle", properties: { sessionID } } })
|
||||
|
||||
// then
|
||||
expect(capturedTimers.size).toBe(1)
|
||||
backgroundRunning = false
|
||||
await firePendingTimers()
|
||||
expect(promptMock).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
test("#given background tasks keep running across multiple retries #when they finally finish on a later retry #then atlas resumes exactly once", async () => {
|
||||
// given
|
||||
const planPath = join(testDir, "test-plan.md")
|
||||
writeFileSync(planPath, "# Plan\n- [ ] Task 1\n- [ ] Task 2")
|
||||
writeBoulderState(testDir, {
|
||||
active_plan: planPath,
|
||||
started_at: "2026-01-02T10:00:00Z",
|
||||
session_ids: [sessionID],
|
||||
plan_name: "test-plan",
|
||||
agent: "atlas",
|
||||
})
|
||||
|
||||
let remainingRunningRetries = 2
|
||||
const promptMock = mock(async () => ({}))
|
||||
const hook = createAtlasHook({
|
||||
directory: testDir,
|
||||
client: {
|
||||
session: {
|
||||
promptAsync: promptMock,
|
||||
messages: async () => ({ data: [] }),
|
||||
},
|
||||
},
|
||||
} as unknown as PluginInput, {
|
||||
directory: testDir,
|
||||
backgroundManager: {
|
||||
getTasksByParentSession: () => {
|
||||
if (remainingRunningRetries > 0) {
|
||||
remainingRunningRetries -= 1
|
||||
return [{ status: "running" }]
|
||||
}
|
||||
|
||||
return []
|
||||
},
|
||||
} as unknown as NonNullable<Parameters<typeof createAtlasHook>[1]>["backgroundManager"] & {
|
||||
getTasksByParentSession: (sessionID: string) => Array<{ status: string }>
|
||||
},
|
||||
})
|
||||
|
||||
// when
|
||||
await hook.handler({ event: { type: "session.idle", properties: { sessionID } } })
|
||||
expect(capturedTimers.size).toBe(1)
|
||||
|
||||
await firePendingTimers()
|
||||
expect(promptMock).toHaveBeenCalledTimes(0)
|
||||
expect(capturedTimers.size).toBe(1)
|
||||
|
||||
await firePendingTimers()
|
||||
|
||||
// then
|
||||
expect(promptMock).toHaveBeenCalledTimes(1)
|
||||
expect(capturedTimers.size).toBe(0)
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user