feat(background-agent): enforce launch depth and descendant limits
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -1637,6 +1637,25 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
|
||||
}
|
||||
}
|
||||
|
||||
function createMockClientWithSessionChain(
|
||||
sessions: Record<string, { directory: string; parentID?: string }>
|
||||
) {
|
||||
return {
|
||||
session: {
|
||||
create: async (_args?: any) => ({ data: { id: `ses_${crypto.randomUUID()}` } }),
|
||||
get: async ({ path }: { path: { id: string } }) => ({
|
||||
data: sessions[path.id] ?? { directory: "/test/dir" },
|
||||
}),
|
||||
prompt: async () => ({}),
|
||||
promptAsync: async () => ({}),
|
||||
messages: async () => ({ data: [] }),
|
||||
todo: async () => ({ data: [] }),
|
||||
status: async () => ({ data: {} }),
|
||||
abort: async () => ({}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
// given
|
||||
mockClient = createMockClient()
|
||||
@@ -1831,6 +1850,98 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
|
||||
expect(updatedTask.startedAt.getTime()).toBeGreaterThanOrEqual(queuedAt.getTime())
|
||||
}
|
||||
})
|
||||
|
||||
test("should track rootSessionID and spawnDepth from the parent chain", async () => {
|
||||
// given
|
||||
manager.shutdown()
|
||||
manager = new BackgroundManager(
|
||||
{
|
||||
client: createMockClientWithSessionChain({
|
||||
"session-depth-2": { directory: "/test/dir", parentID: "session-depth-1" },
|
||||
"session-depth-1": { directory: "/test/dir", parentID: "session-root" },
|
||||
"session-root": { directory: "/test/dir" },
|
||||
}),
|
||||
directory: tmpdir(),
|
||||
} as unknown as PluginInput,
|
||||
{ maxDepth: 3 },
|
||||
)
|
||||
|
||||
const input = {
|
||||
description: "Test task",
|
||||
prompt: "Do something",
|
||||
agent: "test-agent",
|
||||
parentSessionID: "session-depth-2",
|
||||
parentMessageID: "parent-message",
|
||||
}
|
||||
|
||||
// when
|
||||
const task = await manager.launch(input)
|
||||
|
||||
// then
|
||||
expect(task.rootSessionID).toBe("session-root")
|
||||
expect(task.spawnDepth).toBe(3)
|
||||
})
|
||||
|
||||
test("should block launches that exceed maxDepth", async () => {
|
||||
// given
|
||||
manager.shutdown()
|
||||
manager = new BackgroundManager(
|
||||
{
|
||||
client: createMockClientWithSessionChain({
|
||||
"session-depth-3": { directory: "/test/dir", parentID: "session-depth-2" },
|
||||
"session-depth-2": { directory: "/test/dir", parentID: "session-depth-1" },
|
||||
"session-depth-1": { directory: "/test/dir", parentID: "session-root" },
|
||||
"session-root": { directory: "/test/dir" },
|
||||
}),
|
||||
directory: tmpdir(),
|
||||
} as unknown as PluginInput,
|
||||
{ maxDepth: 3 },
|
||||
)
|
||||
|
||||
const input = {
|
||||
description: "Test task",
|
||||
prompt: "Do something",
|
||||
agent: "test-agent",
|
||||
parentSessionID: "session-depth-3",
|
||||
parentMessageID: "parent-message",
|
||||
}
|
||||
|
||||
// when
|
||||
const result = manager.launch(input)
|
||||
|
||||
// then
|
||||
await expect(result).rejects.toThrow("background_task.maxDepth=3")
|
||||
})
|
||||
|
||||
test("should block launches when maxDescendants is reached", async () => {
|
||||
// given
|
||||
manager.shutdown()
|
||||
manager = new BackgroundManager(
|
||||
{
|
||||
client: createMockClientWithSessionChain({
|
||||
"session-root": { directory: "/test/dir" },
|
||||
}),
|
||||
directory: tmpdir(),
|
||||
} as unknown as PluginInput,
|
||||
{ maxDescendants: 1 },
|
||||
)
|
||||
|
||||
const input = {
|
||||
description: "Test task",
|
||||
prompt: "Do something",
|
||||
agent: "test-agent",
|
||||
parentSessionID: "session-root",
|
||||
parentMessageID: "parent-message",
|
||||
}
|
||||
|
||||
await manager.launch(input)
|
||||
|
||||
// when
|
||||
const result = manager.launch(input)
|
||||
|
||||
// then
|
||||
await expect(result).rejects.toThrow("background_task.maxDescendants=1")
|
||||
})
|
||||
})
|
||||
|
||||
describe("pending task can be cancelled", () => {
|
||||
|
||||
@@ -47,6 +47,14 @@ import { MESSAGE_STORAGE } from "../hook-message-injector"
|
||||
import { join } from "node:path"
|
||||
import { pruneStaleTasksAndNotifications } from "./task-poller"
|
||||
import { checkAndInterruptStaleTasks } from "./task-poller"
|
||||
import {
|
||||
createSubagentDepthLimitError,
|
||||
createSubagentDescendantLimitError,
|
||||
getMaxRootDescendants,
|
||||
getMaxSubagentDepth,
|
||||
resolveSubagentSpawnContext,
|
||||
type SubagentSpawnContext,
|
||||
} from "./subagent-spawn-limits"
|
||||
|
||||
type OpencodeClient = PluginInput["client"]
|
||||
|
||||
@@ -111,6 +119,7 @@ export class BackgroundManager {
|
||||
private completionTimers: Map<string, ReturnType<typeof setTimeout>> = new Map()
|
||||
private idleDeferralTimers: Map<string, ReturnType<typeof setTimeout>> = new Map()
|
||||
private notificationQueueByParent: Map<string, Promise<void>> = new Map()
|
||||
private rootDescendantCounts: Map<string, number>
|
||||
private enableParentSessionNotifications: boolean
|
||||
readonly taskHistory = new TaskHistory()
|
||||
|
||||
@@ -135,10 +144,42 @@ export class BackgroundManager {
|
||||
this.tmuxEnabled = options?.tmuxConfig?.enabled ?? false
|
||||
this.onSubagentSessionCreated = options?.onSubagentSessionCreated
|
||||
this.onShutdown = options?.onShutdown
|
||||
this.rootDescendantCounts = new Map()
|
||||
this.enableParentSessionNotifications = options?.enableParentSessionNotifications ?? true
|
||||
this.registerProcessCleanup()
|
||||
}
|
||||
|
||||
async assertCanSpawn(parentSessionID: string): Promise<SubagentSpawnContext> {
|
||||
const spawnContext = await resolveSubagentSpawnContext(this.client, parentSessionID)
|
||||
const maxDepth = getMaxSubagentDepth(this.config)
|
||||
if (spawnContext.childDepth > maxDepth) {
|
||||
throw createSubagentDepthLimitError({
|
||||
childDepth: spawnContext.childDepth,
|
||||
maxDepth,
|
||||
parentSessionID,
|
||||
rootSessionID: spawnContext.rootSessionID,
|
||||
})
|
||||
}
|
||||
|
||||
const maxDescendants = getMaxRootDescendants(this.config)
|
||||
const descendantCount = this.rootDescendantCounts.get(spawnContext.rootSessionID) ?? 0
|
||||
if (descendantCount >= maxDescendants) {
|
||||
throw createSubagentDescendantLimitError({
|
||||
rootSessionID: spawnContext.rootSessionID,
|
||||
descendantCount,
|
||||
maxDescendants,
|
||||
})
|
||||
}
|
||||
|
||||
return spawnContext
|
||||
}
|
||||
|
||||
private registerRootDescendant(rootSessionID: string): number {
|
||||
const nextCount = (this.rootDescendantCounts.get(rootSessionID) ?? 0) + 1
|
||||
this.rootDescendantCounts.set(rootSessionID, nextCount)
|
||||
return nextCount
|
||||
}
|
||||
|
||||
async launch(input: LaunchInput): Promise<BackgroundTask> {
|
||||
log("[background-agent] launch() called with:", {
|
||||
agent: input.agent,
|
||||
@@ -151,16 +192,28 @@ export class BackgroundManager {
|
||||
throw new Error("Agent parameter is required")
|
||||
}
|
||||
|
||||
const spawnContext = await this.assertCanSpawn(input.parentSessionID)
|
||||
const descendantCount = this.registerRootDescendant(spawnContext.rootSessionID)
|
||||
|
||||
log("[background-agent] spawn guard passed", {
|
||||
parentSessionID: input.parentSessionID,
|
||||
rootSessionID: spawnContext.rootSessionID,
|
||||
childDepth: spawnContext.childDepth,
|
||||
descendantCount,
|
||||
})
|
||||
|
||||
// Create task immediately with status="pending"
|
||||
const task: BackgroundTask = {
|
||||
id: `bg_${crypto.randomUUID().slice(0, 8)}`,
|
||||
status: "pending",
|
||||
queuedAt: new Date(),
|
||||
rootSessionID: spawnContext.rootSessionID,
|
||||
// Do NOT set startedAt - will be set when running
|
||||
// Do NOT set sessionID - will be set when running
|
||||
description: input.description,
|
||||
prompt: input.prompt,
|
||||
agent: input.agent,
|
||||
spawnDepth: spawnContext.childDepth,
|
||||
parentSessionID: input.parentSessionID,
|
||||
parentMessageID: input.parentMessageID,
|
||||
parentModel: input.parentModel,
|
||||
@@ -205,7 +258,7 @@ export class BackgroundManager {
|
||||
// Trigger processing (fire-and-forget)
|
||||
this.processKey(key)
|
||||
|
||||
return task
|
||||
return { ...task }
|
||||
}
|
||||
|
||||
private async processKey(key: string): Promise<void> {
|
||||
@@ -875,6 +928,7 @@ export class BackgroundManager {
|
||||
}
|
||||
}
|
||||
|
||||
this.rootDescendantCounts.delete(sessionID)
|
||||
SessionCategoryRegistry.remove(sessionID)
|
||||
}
|
||||
|
||||
@@ -1609,6 +1663,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
|
||||
this.pendingNotifications.clear()
|
||||
this.pendingByParent.clear()
|
||||
this.notificationQueueByParent.clear()
|
||||
this.rootDescendantCounts.clear()
|
||||
this.queuesByKey.clear()
|
||||
this.processingKeys.clear()
|
||||
this.unregisterProcessCleanup()
|
||||
|
||||
Reference in New Issue
Block a user