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

Count sync subagent spawns against descendant limits

Browse Source
This commit is contained in:
YeonGyu-Kim
2026-03-11 18:44:20 +09:00
parent c4d8dedb94
commit d2526878b6
8 changed files with 356 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/features/background-agent/manager.test.ts
+22
View File
@@ -2036,6 +2036,28 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
// then
await expect(result).rejects.toThrow("background_task.maxDescendants=1")
})
test("should consume descendant quota for reserved sync spawns", async () => {
// given
manager.shutdown()
manager = new BackgroundManager(
{
client: createMockClientWithSessionChain({
"session-root": { directory: "/test/dir" },
}),
directory: tmpdir(),
} as unknown as PluginInput,
{ maxDescendants: 1 },
)
await manager.reserveSubagentSpawn("session-root")
// when
const result = manager.assertCanSpawn("session-root")
// then
await expect(result).rejects.toThrow("background_task.maxDescendants=1")
})
})
describe("pending task can be cancelled", () => {
src/features/background-agent/manager.ts
+105 -72
View File
@@ -178,12 +178,47 @@ export class BackgroundManager {
return spawnContext
}
async reserveSubagentSpawn(parentSessionID: string): Promise<{
spawnContext: SubagentSpawnContext
descendantCount: number
commit: () => number
rollback: () => void
}> {
const spawnContext = await this.assertCanSpawn(parentSessionID)
const descendantCount = this.registerRootDescendant(spawnContext.rootSessionID)
let settled = false
return {
spawnContext,
descendantCount,
commit: () => {
settled = true
return descendantCount
},
rollback: () => {
if (settled) return
settled = true
this.unregisterRootDescendant(spawnContext.rootSessionID)
},
}
}
private registerRootDescendant(rootSessionID: string): number {
const nextCount = (this.rootDescendantCounts.get(rootSessionID) ?? 0) + 1
this.rootDescendantCounts.set(rootSessionID, nextCount)
return nextCount
}
private unregisterRootDescendant(rootSessionID: string): void {
const currentCount = this.rootDescendantCounts.get(rootSessionID) ?? 0
if (currentCount <= 1) {
this.rootDescendantCounts.delete(rootSessionID)
return
}
this.rootDescendantCounts.set(rootSessionID, currentCount - 1)
}
async launch(input: LaunchInput): Promise<BackgroundTask> {
log("[background-agent] launch() called with:", {
agent: input.agent,
@@ -196,73 +231,79 @@ export class BackgroundManager {
throw new Error("Agent parameter is required")
}
const spawnContext = await this.assertCanSpawn(input.parentSessionID)
const descendantCount = this.registerRootDescendant(spawnContext.rootSessionID)
const spawnReservation = await this.reserveSubagentSpawn(input.parentSessionID)
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,
parentAgent: input.parentAgent,
parentTools: input.parentTools,
model: input.model,
fallbackChain: input.fallbackChain,
attemptCount: 0,
category: input.category,
}
this.tasks.set(task.id, task)
this.taskHistory.record(input.parentSessionID, { id: task.id, agent: input.agent, description: input.description, status: "pending", category: input.category })
// Track for batched notifications immediately (pending state)
if (input.parentSessionID) {
const pending = this.pendingByParent.get(input.parentSessionID) ?? new Set()
pending.add(task.id)
this.pendingByParent.set(input.parentSessionID, pending)
}
// Add to queue
const key = this.getConcurrencyKeyFromInput(input)
const queue = this.queuesByKey.get(key) ?? []
queue.push({ task, input })
this.queuesByKey.set(key, queue)
log("[background-agent] Task queued:", { taskId: task.id, key, queueLength: queue.length })
const toastManager = getTaskToastManager()
if (toastManager) {
toastManager.addTask({
id: task.id,
description: input.description,
agent: input.agent,
isBackground: true,
status: "queued",
skills: input.skills,
try {
log("[background-agent] spawn guard passed", {
parentSessionID: input.parentSessionID,
rootSessionID: spawnReservation.spawnContext.rootSessionID,
childDepth: spawnReservation.spawnContext.childDepth,
descendantCount: spawnReservation.descendantCount,
})
// Create task immediately with status="pending"
const task: BackgroundTask = {
id: `bg_${crypto.randomUUID().slice(0, 8)}`,
status: "pending",
queuedAt: new Date(),
rootSessionID: spawnReservation.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: spawnReservation.spawnContext.childDepth,
parentSessionID: input.parentSessionID,
parentMessageID: input.parentMessageID,
parentModel: input.parentModel,
parentAgent: input.parentAgent,
parentTools: input.parentTools,
model: input.model,
fallbackChain: input.fallbackChain,
attemptCount: 0,
category: input.category,
}
this.tasks.set(task.id, task)
this.taskHistory.record(input.parentSessionID, { id: task.id, agent: input.agent, description: input.description, status: "pending", category: input.category })
// Track for batched notifications immediately (pending state)
if (input.parentSessionID) {
const pending = this.pendingByParent.get(input.parentSessionID) ?? new Set()
pending.add(task.id)
this.pendingByParent.set(input.parentSessionID, pending)
}
// Add to queue
const key = this.getConcurrencyKeyFromInput(input)
const queue = this.queuesByKey.get(key) ?? []
queue.push({ task, input })
this.queuesByKey.set(key, queue)
log("[background-agent] Task queued:", { taskId: task.id, key, queueLength: queue.length })
const toastManager = getTaskToastManager()
if (toastManager) {
toastManager.addTask({
id: task.id,
description: input.description,
agent: input.agent,
isBackground: true,
status: "queued",
skills: input.skills,
})
}
spawnReservation.commit()
// Trigger processing (fire-and-forget)
this.processKey(key)
return { ...task }
} catch (error) {
spawnReservation.rollback()
throw error
}
// Trigger processing (fire-and-forget)
this.processKey(key)
return { ...task }
}
private async processKey(key: string): Promise<void> {
@@ -1476,14 +1517,6 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
}
}
private formatDuration(start: Date, end?: Date): string {
return formatDuration(start, end)
}
private isAbortedSessionError(error: unknown): boolean {
return isAbortedSessionError(error)
}
private hasRunningTasks(): boolean {
for (const task of this.tasks.values()) {
if (task.status === "running") return true
src/tools/call-omo-agent/sync-executor.test.ts
+48
View File
@@ -249,4 +249,52 @@ describe("executeSync", () => {
expect(deps.waitForCompletion).not.toHaveBeenCalled()
expect(deps.processMessages).not.toHaveBeenCalled()
})
test("commits reserved descendant quota after creating a new sync session", async () => {
//#given
const { executeSync } = require("./sync-executor")
const deps = {
createOrGetSession: mock(async () => ({ sessionID: "ses-test-789", isNew: true })),
waitForCompletion: mock(async () => {}),
processMessages: mock(async () => "agent response"),
setSessionFallbackChain: mock(() => {}),
}
const spawnReservation = {
commit: mock(() => 1),
rollback: mock(() => {}),
}
const args = {
subagent_type: "explore",
description: "test task",
prompt: "find something",
}
const toolContext = {
sessionID: "parent-session",
messageID: "msg-4",
agent: "sisyphus",
abort: new AbortController().signal,
metadata: mock(async () => {}),
}
const ctx = {
client: {
session: {
promptAsync: mock(async () => ({ data: {} })),
},
},
}
//#when
await executeSync(args, toolContext, ctx as any, deps, undefined, spawnReservation)
//#then
expect(spawnReservation.commit).toHaveBeenCalledTimes(1)
expect(spawnReservation.rollback).toHaveBeenCalledTimes(0)
})
})
export {}
src/tools/call-omo-agent/sync-executor.ts
+59 -41
View File
@@ -19,6 +19,11 @@ type ExecuteSyncDeps = {
setSessionFallbackChain: typeof setSessionFallbackChain
}
type SpawnReservation = {
commit: () => number
rollback: () => void
}
const defaultDeps: ExecuteSyncDeps = {
createOrGetSession,
waitForCompletion,
@@ -33,54 +38,67 @@ export async function executeSync(
messageID: string
agent: string
abort: AbortSignal
metadata?: (input: { title?: string; metadata?: Record<string, unknown> }) => void
metadata?: (input: { title?: string; metadata?: Record<string, unknown> }) => void | Promise<void>
},
ctx: PluginInput,
deps: ExecuteSyncDeps = defaultDeps,
fallbackChain?: FallbackEntry[],
spawnReservation?: SpawnReservation,
): Promise<string> {
const { sessionID } = await deps.createOrGetSession(args, toolContext, ctx)
if (fallbackChain && fallbackChain.length > 0) {
deps.setSessionFallbackChain(sessionID, fallbackChain)
}
await toolContext.metadata?.({
title: args.description,
metadata: { sessionId: sessionID },
})
log(`[call_omo_agent] Sending prompt to session ${sessionID}`)
log(`[call_omo_agent] Prompt text:`, args.prompt.substring(0, 100))
let sessionID: string | undefined
try {
await (ctx.client.session as unknown as SessionWithPromptAsync).promptAsync({
path: { id: sessionID },
body: {
agent: args.subagent_type,
tools: {
...getAgentToolRestrictions(args.subagent_type),
task: false,
question: false,
},
parts: [{ type: "text", text: args.prompt }],
},
})
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error)
log(`[call_omo_agent] Prompt error:`, errorMessage)
if (errorMessage.includes("agent.name") || errorMessage.includes("undefined")) {
return `Error: Agent "${args.subagent_type}" not found. Make sure the agent is registered in your opencode.json or provided by a plugin.\n\n<task_metadata>\nsession_id: ${sessionID}\n</task_metadata>`
const session = await deps.createOrGetSession(args, toolContext, ctx)
sessionID = session.sessionID
if (session.isNew) {
spawnReservation?.commit()
}
return `Error: Failed to send prompt: ${errorMessage}\n\n<task_metadata>\nsession_id: ${sessionID}\n</task_metadata>`
if (fallbackChain && fallbackChain.length > 0) {
deps.setSessionFallbackChain(sessionID, fallbackChain)
}
await toolContext.metadata?.({
title: args.description,
metadata: { sessionId: sessionID },
})
log(`[call_omo_agent] Sending prompt to session ${sessionID}`)
log(`[call_omo_agent] Prompt text:`, args.prompt.substring(0, 100))
try {
await (ctx.client.session as unknown as SessionWithPromptAsync).promptAsync({
path: { id: sessionID },
body: {
agent: args.subagent_type,
tools: {
...getAgentToolRestrictions(args.subagent_type),
task: false,
question: false,
},
parts: [{ type: "text", text: args.prompt }],
},
})
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error)
log(`[call_omo_agent] Prompt error:`, errorMessage)
if (errorMessage.includes("agent.name") || errorMessage.includes("undefined")) {
return `Error: Agent "${args.subagent_type}" not found. Make sure the agent is registered in your opencode.json or provided by a plugin.\n\n<task_metadata>\nsession_id: ${sessionID}\n</task_metadata>`
}
return `Error: Failed to send prompt: ${errorMessage}\n\n<task_metadata>\nsession_id: ${sessionID}\n</task_metadata>`
}
await deps.waitForCompletion(sessionID, toolContext, ctx)
const responseText = await deps.processMessages(sessionID, ctx)
const output =
responseText + "\n\n" + ["<task_metadata>", `session_id: ${sessionID}`, "</task_metadata>"].join("\n")
return output
} catch (error) {
spawnReservation?.rollback()
throw error
}
await deps.waitForCompletion(sessionID, toolContext, ctx)
const responseText = await deps.processMessages(sessionID, ctx)
const output =
responseText + "\n\n" + ["<task_metadata>", `session_id: ${sessionID}`, "</task_metadata>"].join("\n")
return output
}
src/tools/call-omo-agent/tools.test.ts
+25 -10
View File
@@ -1,18 +1,24 @@
import { describe, test, expect, mock } from "bun:test"
import type { PluginInput } from "@opencode-ai/plugin"
import type { BackgroundManager } from "../../features/background-agent"
import type { FallbackEntry } from "../../shared/model-requirements"
import { createCallOmoAgent } from "./tools"
const { beforeEach, describe, test, expect, mock } = require("bun:test")
const { createCallOmoAgent } = require("./tools")
describe("createCallOmoAgent", () => {
const assertCanSpawnMock = mock(() => Promise.resolve(undefined))
const reserveCommitMock = mock(() => 1)
const reserveRollbackMock = mock(() => {})
const reserveSubagentSpawnMock = mock(() => Promise.resolve({
spawnContext: { rootSessionID: "root-session", parentDepth: 0, childDepth: 1 },
descendantCount: 1,
commit: reserveCommitMock,
rollback: reserveRollbackMock,
}))
const mockCtx = {
client: {},
directory: "/test",
} as unknown as PluginInput
}
const mockBackgroundManager = {
assertCanSpawn: assertCanSpawnMock,
reserveSubagentSpawn: reserveSubagentSpawnMock,
launch: mock(() => Promise.resolve({
id: "test-task-id",
sessionID: null,
@@ -20,7 +26,14 @@ describe("createCallOmoAgent", () => {
agent: "test-agent",
status: "pending",
})),
} as unknown as BackgroundManager
}
beforeEach(() => {
assertCanSpawnMock.mockClear()
reserveSubagentSpawnMock.mockClear()
reserveCommitMock.mockClear()
reserveRollbackMock.mockClear()
})
test("should reject agent in disabled_agents list", async () => {
//#given
@@ -105,7 +118,7 @@ describe("createCallOmoAgent", () => {
test("uses agent override fallback_models when launching background subagent", async () => {
//#given
const launch = mock((_input: { fallbackChain?: FallbackEntry[] }) => Promise.resolve({
const launch = mock((_input: { fallbackChain?: Array<{ providers: string[]; model: string; variant?: string }> }) => Promise.resolve({
id: "task-fallback",
sessionID: "sub-session",
description: "Test task",
@@ -115,7 +128,7 @@ describe("createCallOmoAgent", () => {
const managerWithLaunch = {
launch,
getTask: mock(() => undefined),
} as unknown as BackgroundManager
}
const toolDef = createCallOmoAgent(
mockCtx,
managerWithLaunch,
@@ -154,7 +167,7 @@ describe("createCallOmoAgent", () => {
test("should return a tool error when sync spawn depth validation fails", async () => {
//#given
assertCanSpawnMock.mockRejectedValueOnce(new Error("Subagent spawn blocked: child depth 4 exceeds background_task.maxDepth=3."))
reserveSubagentSpawnMock.mockRejectedValueOnce(new Error("Subagent spawn blocked: child depth 4 exceeds background_task.maxDepth=3."))
const toolDef = createCallOmoAgent(mockCtx, mockBackgroundManager, [])
const executeFunc = toolDef.execute as Function
@@ -173,3 +186,5 @@ describe("createCallOmoAgent", () => {
expect(result).toContain("background_task.maxDepth=3")
})
})
export {}
src/tools/call-omo-agent/tools.ts
+4 -1
View File
@@ -96,9 +96,12 @@ export function createCallOmoAgent(
}
if (!args.session_id) {
let spawnReservation: Awaited<ReturnType<BackgroundManager["reserveSubagentSpawn"]>> | undefined
try {
await backgroundManager.assertCanSpawn(toolCtx.sessionID)
spawnReservation = await backgroundManager.reserveSubagentSpawn(toolCtx.sessionID)
return await executeSync(args, toolCtx, ctx, undefined, fallbackChain, spawnReservation)
} catch (error) {
spawnReservation?.rollback()
return `Error: ${error instanceof Error ? error.message : String(error)}`
}
}
src/tools/delegate-task/sync-task.test.ts
+75
View File
@@ -110,6 +110,65 @@ describe("executeSyncTask - cleanup on error paths", () => {
expect(deleteCalls[0]).toBe("ses_test_12345678")
})
test("rolls back reserved descendant quota when sync session creation fails", async () => {
const mockClient = {
session: {
create: async () => ({ data: { id: "ses_test_12345678" } }),
},
}
const { executeSyncTask } = require("./sync-task")
const commit = mock(() => 1)
const rollback = mock(() => {})
const reserveSubagentSpawn = mock(async () => ({
spawnContext: { rootSessionID: "parent-session", parentDepth: 0, childDepth: 1 },
descendantCount: 1,
commit,
rollback,
}))
const deps = {
createSyncSession: async () => ({ ok: false as const, error: "Failed to create session" }),
sendSyncPrompt: async () => null,
pollSyncSession: async () => null,
fetchSyncResult: async () => ({ ok: true as const, textContent: "Result" }),
}
const mockCtx = {
sessionID: "parent-session",
callID: "call-123",
metadata: () => {},
}
const mockExecutorCtx = {
manager: { reserveSubagentSpawn },
client: mockClient,
directory: "/tmp",
onSyncSessionCreated: null,
}
const args = {
prompt: "test prompt",
description: "test task",
category: "test",
load_skills: [],
run_in_background: false,
command: null,
}
//#when
const result = await executeSyncTask(args, mockCtx, mockExecutorCtx, {
sessionID: "parent-session",
}, "test-agent", undefined, undefined, undefined, undefined, deps)
//#then
expect(result).toBe("Failed to create session")
expect(reserveSubagentSpawn).toHaveBeenCalledWith("parent-session")
expect(commit).toHaveBeenCalledTimes(0)
expect(rollback).toHaveBeenCalledTimes(1)
})
test("cleans up toast and subagentSessions when pollSyncSession returns error", async () => {
const mockClient = {
session: {
@@ -182,7 +241,18 @@ describe("executeSyncTask - cleanup on error paths", () => {
metadata: () => {},
}
const commit = mock(() => 1)
const rollback = mock(() => {})
const mockExecutorCtx = {
manager: {
reserveSubagentSpawn: mock(async () => ({
spawnContext: { rootSessionID: "parent-session", parentDepth: 0, childDepth: 1 },
descendantCount: 1,
commit,
rollback,
})),
},
client: mockClient,
directory: "/tmp",
onSyncSessionCreated: null,
@@ -204,9 +274,14 @@ describe("executeSyncTask - cleanup on error paths", () => {
//#then - should complete and cleanup resources
expect(result).toContain("Task completed")
expect(mockExecutorCtx.manager.reserveSubagentSpawn).toHaveBeenCalledWith("parent-session")
expect(commit).toHaveBeenCalledTimes(1)
expect(rollback).toHaveBeenCalledTimes(0)
expect(removeTaskCalls.length).toBe(1)
expect(removeTaskCalls[0]).toBe("sync_ses_test")
expect(deleteCalls.length).toBe(1)
expect(deleteCalls[0]).toBe("ses_test_12345678")
})
})
export {}
src/tools/delegate-task/sync-task.ts
+18 -7
View File
@@ -27,15 +27,23 @@ export async function executeSyncTask(
const toastManager = getTaskToastManager()
let taskId: string | undefined
let syncSessionID: string | undefined
let spawnReservation:
| Awaited<ReturnType<ExecutorContext["manager"]["reserveSubagentSpawn"]>>
| undefined
try {
const spawnContext = typeof manager?.assertCanSpawn === "function"
? await manager.assertCanSpawn(parentContext.sessionID)
: {
rootSessionID: parentContext.sessionID,
parentDepth: 0,
childDepth: 1,
}
if (typeof manager?.reserveSubagentSpawn === "function") {
spawnReservation = await manager.reserveSubagentSpawn(parentContext.sessionID)
}
const spawnContext = spawnReservation?.spawnContext
?? (typeof manager?.assertCanSpawn === "function"
? await manager.assertCanSpawn(parentContext.sessionID)
: {
rootSessionID: parentContext.sessionID,
parentDepth: 0,
childDepth: 1,
})
const createSessionResult = await deps.createSyncSession(client, {
parentSessionID: parentContext.sessionID,
@@ -45,10 +53,12 @@ export async function executeSyncTask(
})
if (!createSessionResult.ok) {
spawnReservation?.rollback()
return createSessionResult.error
}
const sessionID = createSessionResult.sessionID
spawnReservation?.commit()
syncSessionID = sessionID
subagentSessions.add(sessionID)
syncSubagentSessions.add(sessionID)
@@ -156,6 +166,7 @@ session_id: ${sessionID}
}
}
} catch (error) {
spawnReservation?.rollback()
return formatDetailedError(error, {
operation: "Execute task",
args,