luna/oh-my-opencode
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(concurrency): prevent background task races and leaks

Browse Source 
Ensure queue waiters settle once, centralize completion with status guards, and release slots before async work so shutdown and cancellations don’t leak concurrency. Internal hardening only.
This commit is contained in:
Jeremy Gollehon
2026-01-14 15:09:32 -08:00
parent 7ab7eaf2a7
commit 09e8546d3a
4 changed files with 482 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/features/background-agent/manager.ts
+158 -61
View File
@@ -52,6 +52,8 @@ export class BackgroundManager {
private directory: string
private pollingInterval?: ReturnType<typeof setInterval>
private concurrencyManager: ConcurrencyManager
private cleanupRegistered = false
private shutdownTriggered = false
constructor(ctx: PluginInput, config?: BackgroundTaskConfig) {
this.tasks = new Map()
@@ -60,6 +62,7 @@ export class BackgroundManager {
this.client = ctx.client
this.directory = ctx.directory
this.concurrencyManager = new ConcurrencyManager(config)
this.registerProcessCleanup()
}
async launch(input: LaunchInput): Promise<BackgroundTask> {
@@ -186,7 +189,7 @@ export class BackgroundManager {
existingTask.completedAt = new Date()
if (existingTask.concurrencyKey) {
this.concurrencyManager.release(existingTask.concurrencyKey)
existingTask.concurrencyKey = undefined // Prevent double-release
existingTask.concurrencyKey = undefined
}
this.markForNotification(existingTask)
this.notifyParentSession(existingTask).catch(err => {
@@ -238,14 +241,20 @@ export class BackgroundManager {
* Register an external task (e.g., from sisyphus_task) for notification tracking.
* This allows tasks created by external tools to receive the same toast/prompt notifications.
*/
registerExternalTask(input: {
async registerExternalTask(input: {
taskId: string
sessionID: string
parentSessionID: string
description: string
agent?: string
parentAgent?: string
}): BackgroundTask {
concurrencyKey?: string
}): Promise<BackgroundTask> {
// Acquire concurrency slot if a key is provided
if (input.concurrencyKey) {
await this.concurrencyManager.acquire(input.concurrencyKey)
}
const task: BackgroundTask = {
id: input.taskId,
sessionID: input.sessionID,
@@ -261,6 +270,7 @@ export class BackgroundManager {
lastUpdate: new Date(),
},
parentAgent: input.parentAgent,
concurrencyKey: input.concurrencyKey,
}
this.tasks.set(task.id, task)
@@ -283,6 +293,21 @@ export class BackgroundManager {
throw new Error(`Task not found for session: ${input.sessionId}`)
}
if (existingTask.status === "running") {
log("[background-agent] Resume skipped - task already running:", {
taskId: existingTask.id,
sessionID: existingTask.sessionID,
})
return existingTask
}
// Re-acquire concurrency using the agent name as the key (same as launch()).
// Note: existingTask.concurrencyKey is cleared when tasks complete, so we
// derive the key from task.agent which persists through completion.
const concurrencyKey = existingTask.agent
await this.concurrencyManager.acquire(concurrencyKey)
existingTask.concurrencyKey = concurrencyKey
existingTask.status = "running"
existingTask.completedAt = undefined
existingTask.error = undefined
@@ -344,11 +369,12 @@ export class BackgroundManager {
const errorMessage = error instanceof Error ? error.message : String(error)
existingTask.error = errorMessage
existingTask.completedAt = new Date()
// Release concurrency on resume error (matches launch error handler)
if (existingTask.concurrencyKey) {
this.concurrencyManager.release(existingTask.concurrencyKey)
existingTask.concurrencyKey = undefined // Prevent double-release
}
// Release concurrency on error to prevent slot leaks
if (existingTask.concurrencyKey) {
this.concurrencyManager.release(existingTask.concurrencyKey)
existingTask.concurrencyKey = undefined
}
this.markForNotification(existingTask)
this.notifyParentSession(existingTask).catch(err => {
log("[background-agent] Failed to notify on resume error:", err)
@@ -417,29 +443,31 @@ export class BackgroundManager {
// Edge guard: Verify session has actual assistant output before completing
this.validateSessionHasOutput(sessionID).then(async (hasValidOutput) => {
// Re-check status after async operation (could have been completed by polling)
if (task.status !== "running") {
log("[background-agent] Task status changed during validation, skipping:", { taskId: task.id, status: task.status })
return
}
if (!hasValidOutput) {
log("[background-agent] Session.idle but no valid output yet, waiting:", task.id)
return
}
const hasIncompleteTodos = await this.checkSessionTodos(sessionID)
// Re-check status after async operation again
if (task.status !== "running") {
log("[background-agent] Task status changed during todo check, skipping:", { taskId: task.id, status: task.status })
return
}
if (hasIncompleteTodos) {
log("[background-agent] Task has incomplete todos, waiting for todo-continuation:", task.id)
return
}
task.status = "completed"
task.completedAt = new Date()
// Release concurrency immediately on completion
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined // Prevent double-release
}
// Clean up pendingByParent to prevent stale entries
this.cleanupPendingByParent(task)
this.markForNotification(task)
await this.notifyParentSession(task)
log("[background-agent] Task completed via session.idle event:", task.id)
await this.tryCompleteTask(task, "session.idle event")
}).catch(err => {
log("[background-agent] Error in session.idle handler:", err)
})
@@ -459,10 +487,10 @@ export class BackgroundManager {
task.error = "Session deleted"
}
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined // Prevent double-release
}
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined
}
// Clean up pendingByParent to prevent stale entries
this.cleanupPendingByParent(task)
this.tasks.delete(task.id)
@@ -587,11 +615,25 @@ export class BackgroundManager {
}
}
cleanup(): void {
this.stopPolling()
this.tasks.clear()
this.notifications.clear()
this.pendingByParent.clear()
private registerProcessCleanup(): void {
if (this.cleanupRegistered) return
this.cleanupRegistered = true
const cleanup = () => {
try {
this.shutdown()
} catch (error) {
log("[background-agent] Error during shutdown cleanup:", error)
}
}
registerProcessSignal("SIGINT", cleanup)
registerProcessSignal("SIGTERM", cleanup)
if (process.platform === "win32") {
registerProcessSignal("SIGBREAK", cleanup)
}
process.on("beforeExit", cleanup)
process.on("exit", cleanup)
}
/**
@@ -608,12 +650,44 @@ cleanup(): void {
return Array.from(this.tasks.values()).filter(t => t.status !== "running")
}
private async notifyParentSession(task: BackgroundTask): Promise<void> {
/**
* Safely complete a task with race condition protection.
* Returns true if task was successfully completed, false if already completed by another path.
*/
private async tryCompleteTask(task: BackgroundTask, source: string): Promise<boolean> {
// Guard: Check if task is still running (could have been completed by another path)
if (task.status !== "running") {
log("[background-agent] Task already completed, skipping:", { taskId: task.id, status: task.status, source })
return false
}
// Atomically mark as completed to prevent race conditions
task.status = "completed"
task.completedAt = new Date()
// Release concurrency BEFORE any async operations to prevent slot leaks
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined
}
this.markForNotification(task)
try {
await this.notifyParentSession(task)
log(`[background-agent] Task completed via ${source}:`, task.id)
} catch (err) {
log("[background-agent] Error in notifyParentSession:", { taskId: task.id, error: err })
// Concurrency already released, notification failed but task is complete
}
return true
}
private async notifyParentSession(task: BackgroundTask): Promise<void> {
// Note: Callers must release concurrency before calling this method
// to ensure slots are freed even if notification fails
const duration = this.formatDuration(task.startedAt, task.completedAt)
log("[background-agent] notifyParentSession called for task:", task.id)
@@ -715,10 +789,12 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
const taskId = task.id
setTimeout(() => {
// Concurrency already released at completion - just cleanup notifications and task
this.clearNotificationsForTask(taskId)
this.tasks.delete(taskId)
log("[background-agent] Removed completed task from memory:", taskId)
// Guard: Only delete if task still exists (could have been deleted by session.deleted event)
if (this.tasks.has(taskId)) {
this.clearNotificationsForTask(taskId)
this.tasks.delete(taskId)
log("[background-agent] Removed completed task from memory:", taskId)
}
}, 5 * 60 * 1000)
}
@@ -755,7 +831,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
task.completedAt = new Date()
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined // Prevent double-release
task.concurrencyKey = undefined
}
// Clean up pendingByParent to prevent stale entries
this.cleanupPendingByParent(task)
@@ -791,7 +867,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
for (const task of this.tasks.values()) {
if (task.status !== "running") continue
try {
try {
const sessionStatus = allStatuses[task.sessionID]
// Don't skip if session not in status - fall through to message-based detection
@@ -803,24 +879,16 @@ try {
continue
}
// Re-check status after async operation
if (task.status !== "running") continue
const hasIncompleteTodos = await this.checkSessionTodos(task.sessionID)
if (hasIncompleteTodos) {
log("[background-agent] Task has incomplete todos via polling, waiting:", task.id)
continue
}
task.status = "completed"
task.completedAt = new Date()
// Release concurrency immediately on completion
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined // Prevent double-release
}
// Clean up pendingByParent to prevent stale entries
this.cleanupPendingByParent(task)
this.markForNotification(task)
await this.notifyParentSession(task)
log("[background-agent] Task completed via polling:", task.id)
await this.tryCompleteTask(task, "polling (idle status)")
continue
}
@@ -860,7 +928,7 @@ try {
task.progress.toolCalls = toolCalls
task.progress.lastTool = lastTool
task.progress.lastUpdate = new Date()
if (lastMessage) {
if (lastMessage) {
task.progress.lastMessage = lastMessage
task.progress.lastMessageAt = new Date()
}
@@ -880,20 +948,12 @@ if (lastMessage) {
continue
}
// Re-check status after async operation
if (task.status !== "running") continue
const hasIncompleteTodos = await this.checkSessionTodos(task.sessionID)
if (!hasIncompleteTodos) {
task.status = "completed"
task.completedAt = new Date()
// Release concurrency immediately on completion
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined // Prevent double-release
}
// Clean up pendingByParent to prevent stale entries
this.cleanupPendingByParent(task)
this.markForNotification(task)
await this.notifyParentSession(task)
log("[background-agent] Task completed via stability detection:", task.id)
await this.tryCompleteTask(task, "stability detection")
continue
}
}
@@ -912,6 +972,43 @@ if (lastMessage) {
this.stopPolling()
}
}
/**
* Shutdown the manager gracefully.
* Cancels all pending concurrency waiters and clears timers.
* Should be called when the plugin is unloaded.
*/
shutdown(): void {
if (this.shutdownTriggered) return
this.shutdownTriggered = true
log("[background-agent] Shutting down BackgroundManager")
this.stopPolling()
// Release concurrency for all running tasks first
for (const task of this.tasks.values()) {
if (task.concurrencyKey) {
this.concurrencyManager.release(task.concurrencyKey)
task.concurrencyKey = undefined
}
}
// Then clear all state (cancels any remaining waiters)
this.concurrencyManager.clear()
this.tasks.clear()
this.notifications.clear()
this.pendingByParent.clear()
log("[background-agent] Shutdown complete")
}
}
function registerProcessSignal(
signal: NodeJS.Signals,
handler: () => void
): void {
process.on(signal, () => {
handler()
process.exit(0)
})
}
function getMessageDir(sessionID: string): string | null {