From 50a2264d759f68494ec6792e040eb8ba3239936f Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 8 Mar 2026 02:22:22 +0900 Subject: [PATCH] feat(background-agent): enforce launch depth and descendant limits Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus --- src/features/background-agent/manager.test.ts | 111 ++++++++++++++++++ src/features/background-agent/manager.ts | 57 ++++++++- 2 files changed, 167 insertions(+), 1 deletion(-) diff --git a/src/features/background-agent/manager.test.ts b/src/features/background-agent/manager.test.ts index 8a6ddc358..e12b99375 100644 --- a/src/features/background-agent/manager.test.ts +++ b/src/features/background-agent/manager.test.ts @@ -1637,6 +1637,25 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => { } } + function createMockClientWithSessionChain( + sessions: Record + ) { + 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", () => { diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index 247e9f948..eea43424e 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -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> = new Map() private idleDeferralTimers: Map> = new Map() private notificationQueueByParent: Map> = new Map() + private rootDescendantCounts: Map 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 { + 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 { 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 { @@ -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()