diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index a478f6fc28f..cdc122205e2 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1139,6 +1139,19 @@ describe('loadCliConfig', () => { expect(config.getAgentsSettings().maxParallelAgents).toBe(2); }); + it('passes agents.maxParallelAgentsByModel from settings to core config', async () => { + process.argv = ['node', 'script.js']; + const argv = await parseArguments(); + const config = await loadCliConfig( + { agents: { maxParallelAgentsByModel: { 'weak-model': 1 } } }, + argv, + ); + + expect(config.getAgentsSettings().maxParallelAgentsByModel).toEqual({ + 'weak-model': 1, + }); + }); + it('passes tools.shell.defaultTimeoutMs from settings to core config', async () => { process.argv = ['node', 'script.js']; const argv = await parseArguments(); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 559cf7e5d18..2fbadccc3d2 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -2255,6 +2255,7 @@ export async function loadCliConfig( } : undefined, maxParallelAgents: settings.agents.maxParallelAgents, + maxParallelAgentsByModel: settings.agents.maxParallelAgentsByModel, displayMode: settings.agents.displayMode, arena: settings.agents.arena ? { diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index e5026bdbaef..f22f63c13ee 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -2822,13 +2822,31 @@ const SETTINGS_SCHEMA = { default: undefined as number | undefined, minimum: 1, description: - 'Global maximum number of background sub-agents that can run concurrently. Additional background agents wait in a queue until a slot is available. Per-model limits are not supported yet.', + 'Global maximum number of background sub-agents that can run concurrently. Additional background agents wait in a queue until a slot is available. Use maxParallelAgentsByModel to cap a specific model below this global limit.', showInDialog: false, jsonSchemaOverride: { type: 'integer', minimum: 1, }, }, + maxParallelAgentsByModel: { + type: 'object', + label: 'Max Parallel Agents Per Model', + category: 'Advanced', + requiresRestart: true, + default: undefined as Record | undefined, + description: + 'Per-model maximum number of background sub-agents that can run concurrently, keyed by model ID (e.g. { "qwen3-max": 2 }). Useful when a model has a lower concurrency capacity. Takes precedence over the global maxParallelAgents for the matched model; models not listed here fall back to the global limit.', + showInDialog: false, + mergeStrategy: MergeStrategy.SHALLOW_MERGE, + jsonSchemaOverride: { + type: 'object', + additionalProperties: { + type: 'integer', + minimum: 1, + }, + }, + }, displayMode: { type: 'enum', label: 'Display Mode', diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 7274b265c65..750c48bcb60 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -2970,7 +2970,7 @@ export const AppContainer = (props: AppContainerProps) => { stickyTodos !== null && !dialogsVisible && !isFeedbackDialogOpen && - streamingState !== StreamingState.WaitingForConfirmation; + streamingState === StreamingState.Responding; const stickyTodoWidth = Math.min(mainAreaWidth, 64); const stickyTodoMaxVisibleItems = getStickyTodoMaxVisibleItems(terminalHeight); diff --git a/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx b/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx index 355cc45e9de..322b28d54f0 100644 --- a/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx +++ b/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx @@ -84,7 +84,7 @@ const baseUIState: Partial = { terminalHeight: 24, staticExtraHeight: 0, constrainHeight: true, - streamingState: StreamingState.Idle, + streamingState: StreamingState.Responding, historyManager: { addItem: vi.fn(), history: [], @@ -212,6 +212,22 @@ describe('DefaultAppLayout', () => { expect(output).toContain('Composer'); }); + it('does not render sticky todo list when agent is idle', () => { + mockedUseAgentViewState.mockReturnValue({ + activeView: 'main', + agents: new Map(), + }); + + const { lastFrame } = renderLayout({ + ...baseUIState, + streamingState: StreamingState.Idle, + }); + + const output = lastFrame() ?? ''; + expect(output).not.toContain('StickyTodoList'); + expect(output).toContain('Composer'); + }); + it('does not render sticky todo list when feedback dialog is open', () => { mockedUseAgentViewState.mockReturnValue({ activeView: 'main', diff --git a/packages/cli/src/ui/layouts/DefaultAppLayout.tsx b/packages/cli/src/ui/layouts/DefaultAppLayout.tsx index eadc1ae9355..7684b5b96b4 100644 --- a/packages/cli/src/ui/layouts/DefaultAppLayout.tsx +++ b/packages/cli/src/ui/layouts/DefaultAppLayout.tsx @@ -45,7 +45,7 @@ export const DefaultAppLayout: React.FC = () => { uiState.stickyTodos !== null && !uiState.dialogsVisible && !uiState.isFeedbackDialogOpen && - uiState.streamingState !== StreamingState.WaitingForConfirmation; + uiState.streamingState === StreamingState.Responding; // Clear terminal on view switch so previous view's output // is removed. refreshStatic clears the terminal and bumps the diff --git a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.test.tsx b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.test.tsx index 5fec67bd7fb..dbbfb8e2796 100644 --- a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.test.tsx +++ b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.test.tsx @@ -60,7 +60,7 @@ const baseUIState: Partial = { terminalHeight: 24, staticExtraHeight: 0, constrainHeight: true, - streamingState: StreamingState.Idle, + streamingState: StreamingState.Responding, historyManager: { addItem: vi.fn(), history: [], @@ -164,6 +164,17 @@ describe('ScreenReaderAppLayout', () => { expect(output).toContain('Composer'); }); + it('does not render sticky todo list when agent is idle', () => { + const { lastFrame } = renderLayout({ + ...baseUIState, + streamingState: StreamingState.Idle, + }); + + const output = lastFrame() ?? ''; + expect(output).not.toContain('StickyTodoList'); + expect(output).toContain('Composer'); + }); + it('does not render sticky todo list when feedback dialog is open', () => { const { lastFrame } = renderLayout({ ...baseUIState, diff --git a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx index 1e5205fe7f5..5c8da8631f7 100644 --- a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx +++ b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx @@ -34,7 +34,7 @@ export const ScreenReaderAppLayout: React.FC = () => { uiState.stickyTodos !== null && !uiState.dialogsVisible && !uiState.isFeedbackDialogOpen && - uiState.streamingState !== StreamingState.WaitingForConfirmation; + uiState.streamingState === StreamingState.Responding; return ( diff --git a/packages/core/src/agents/agent-transcript.ts b/packages/core/src/agents/agent-transcript.ts index b4081015691..689bc4cc97c 100644 --- a/packages/core/src/agents/agent-transcript.ts +++ b/packages/core/src/agents/agent-transcript.ts @@ -127,6 +127,11 @@ export interface AgentMeta { * via {@link normalizeResumedAgentDepth} — never trust the raw value. */ depth?: number; + /** + * Concrete model ID this agent runs with. Persisted so a process-restart + * recovery can enforce per-model concurrency caps on the revive path. + */ + model?: string; /** Last terminal error, if any. */ lastError?: string; } diff --git a/packages/core/src/agents/background-agent-resume.test.ts b/packages/core/src/agents/background-agent-resume.test.ts index fb157718412..a58c6d5527c 100644 --- a/packages/core/src/agents/background-agent-resume.test.ts +++ b/packages/core/src/agents/background-agent-resume.test.ts @@ -264,6 +264,45 @@ describe('BackgroundAgentResumeService', () => { expect(subagentManager.loadSubagent).not.toHaveBeenCalled(); }); + it('restores the model from the meta sidecar for per-model cap accounting', async () => { + const sessionId = 'session-model-resume'; + const agentId = 'agent-model-resume'; + const metaPath = getAgentMetaPath(tempDir, sessionId, agentId); + + writeAgentMeta(metaPath, { + agentId, + agentType: 'researcher', + description: 'Model-capped background task', + parentSessionId: sessionId, + parentAgentId: null, + createdAt: '2026-04-20T00:00:00.000Z', + status: 'running', + subagentName: 'researcher', + resolvedApprovalMode: 'default', + model: 'gemini-2.5-pro', + }); + fs.writeFileSync( + getAgentJsonlPath(tempDir, sessionId, agentId), + JSON.stringify({ + uuid: 'u1', + parentUuid: null, + sessionId, + timestamp: '2026-04-20T00:00:00.000Z', + type: 'user', + message: { + role: 'user', + parts: [{ text: 'Model-capped background task' }], + }, + }) + '\n', + 'utf8', + ); + + const { service } = createService(); + await service.loadPausedBackgroundAgents(sessionId); + + expect(registry.get(agentId)?.model).toBe('gemini-2.5-pro'); + }); + it('keeps missing subagents visible so they can be abandoned later', async () => { const sessionId = 'session-missing'; const agentId = 'agent-missing'; diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index 9edd0e08c46..0cf813baa0b 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -444,6 +444,7 @@ export class BackgroundAgentResumeService { // UI falls back to its generic orphan annotation. parentAgentId: meta.parentAgentId, depth: meta.depth, + model: meta.model, }; const entry = registry.register(registration); recovered.push(entry); @@ -532,7 +533,7 @@ export class BackgroundAgentResumeService { // entry back to paused, so an at-capacity revive fails cleanly instead of // stranding the entry as paused. try { - registry.assertCanStartBackgroundAgent(); + registry.assertCanStartBackgroundAgent(entry.model); } catch (error) { debugLogger.warn( `[BackgroundAgentResume] Cannot revive "${agentId}": ` + @@ -579,7 +580,7 @@ export class BackgroundAgentResumeService { ); } try { - registry.assertCanStartBackgroundAgent(); + registry.assertCanStartBackgroundAgent(entry.model); } catch (error) { debugLogger.warn( `[BackgroundAgentResume] Cannot revive "${agentId}": ` + diff --git a/packages/core/src/agents/background-tasks.test.ts b/packages/core/src/agents/background-tasks.test.ts index d06a7868fcc..061934fb2aa 100644 --- a/packages/core/src/agents/background-tasks.test.ts +++ b/packages/core/src/agents/background-tasks.test.ts @@ -756,6 +756,173 @@ describe('BackgroundTaskRegistry', () => { }); }); + describe('per-model background concurrency limit', () => { + it('caps a single model while leaving room for others', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 10, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 1 }, + }); + + registry.register(makeRegistration('bg-1', { model: 'weak-model' })); + + // The capped model is full... + expect(() => + registry.register(makeRegistration('bg-2', { model: 'weak-model' })), + ).toThrow( + 'Cannot start background agent: maximum concurrent background agents ' + + 'for model "weak-model" (1) reached. Stop an existing agent on that ' + + 'model first.', + ); + expect(registry.get('bg-2')).toBeUndefined(); + + // ...but a different model is unaffected. + registry.register(makeRegistration('bg-3', { model: 'strong-model' })); + expect(registry.get('bg-3')?.status).toBe('running'); + }); + + it('lets a model without a per-model cap use the global limit', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 2, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 1 }, + }); + + registry.register(makeRegistration('bg-1', { model: 'uncapped-model' })); + registry.register(makeRegistration('bg-2', { model: 'uncapped-model' })); + + // The global cap still bounds uncapped models. + expect(() => + registry.register( + makeRegistration('bg-3', { model: 'uncapped-model' }), + ), + ).toThrow('maximum concurrent background agents (2) reached'); + }); + + it('enforces the global cap even when the per-model cap has room', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 1, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 5 }, + }); + + registry.register(makeRegistration('bg-1', { model: 'other-model' })); + + expect(() => + registry.register(makeRegistration('bg-2', { model: 'weak-model' })), + ).toThrow('maximum concurrent background agents (1) reached'); + }); + + it('counts reservations against the per-model cap', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 10, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 1 }, + }); + + const reservation = registry.tryReserveBackgroundSlot('weak-model'); + expect(reservation).toBeDefined(); + expect(reservation?.model).toBe('weak-model'); + + // A second reservation for the same model is refused while the first + // is outstanding. + expect(registry.tryReserveBackgroundSlot('weak-model')).toBeUndefined(); + // A reservation for a different model is still granted. + expect(registry.tryReserveBackgroundSlot('strong-model')).toBeDefined(); + + // Releasing frees the per-model slot. + registry.releaseBackgroundSlot(reservation!); + expect(registry.tryReserveBackgroundSlot('weak-model')).toBeDefined(); + }); + + it('frees the per-model cap when an agent on that model completes', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 10, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 1 }, + }); + + registry.register(makeRegistration('bg-1', { model: 'weak-model' })); + expect(registry.tryReserveBackgroundSlot('weak-model')).toBeUndefined(); + + registry.complete('bg-1', 'done'); + registry.register(makeRegistration('bg-2', { model: 'weak-model' })); + expect(registry.get('bg-2')?.status).toBe('running'); + }); + + it('drains a different-model waiter while a capped-model waiter stays queued', async () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 2, + maxConcurrentBackgroundAgentsByModel: { 'weak-model': 1 }, + }); + // Fill both the weak-model cap (1) and the global cap (2) so neither + // waiter below can reserve a slot immediately. + registry.register(makeRegistration('bg-weak', { model: 'weak-model' })); + registry.register(makeRegistration('bg-other', { model: 'other-model' })); + + // Both queue because the global cap is full. + const weakWaiter = registry.waitForBackgroundSlot( + new AbortController().signal, + 'weak-model', + ); + const strongWaiter = registry.waitForBackgroundSlot( + new AbortController().signal, + 'strong-model', + ); + expect(registry.getQueuedCount()).toBe(2); + + // Freeing one global slot (but NOT the weak-model slot) lets the + // strong-model waiter through while the weak-model waiter stays queued. + registry.complete('bg-other', 'done'); + + const strongReservation = await strongWaiter; + expect(strongReservation.model).toBe('strong-model'); + expect(registry.getQueuedCount()).toBe(1); + + // The weak-model waiter is released only once a weak-model slot frees. + registry.complete('bg-weak', 'done'); + const weakReservation = await weakWaiter; + expect(weakReservation.model).toBe('weak-model'); + expect(registry.getQueuedCount()).toBe(0); + }); + + it('ignores malformed per-model cap values', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 10, + maxConcurrentBackgroundAgentsByModel: { + 'bad-zero': 0, + 'bad-negative': -3, + 'bad-float': 1.5, + good: 1, + } as Record, + }); + + // Malformed entries are dropped, so those models fall back to the + // global cap and can each start. + registry.register(makeRegistration('bg-zero', { model: 'bad-zero' })); + registry.register( + makeRegistration('bg-negative', { model: 'bad-negative' }), + ); + registry.register(makeRegistration('bg-float', { model: 'bad-float' })); + expect(registry.get('bg-zero')?.status).toBe('running'); + expect(registry.get('bg-negative')?.status).toBe('running'); + expect(registry.get('bg-float')?.status).toBe('running'); + + // The one valid entry is still enforced. + registry.register(makeRegistration('bg-good', { model: 'good' })); + expect(() => + registry.register(makeRegistration('bg-good-2', { model: 'good' })), + ).toThrow('for model "good" (1) reached'); + }); + + it('accepts a ReadonlyMap for the per-model caps', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 10, + maxConcurrentBackgroundAgentsByModel: new Map([['weak-model', 1]]), + }); + + registry.register(makeRegistration('bg-1', { model: 'weak-model' })); + expect(() => + registry.register(makeRegistration('bg-2', { model: 'weak-model' })), + ).toThrow('for model "weak-model" (1) reached'); + }); + }); + it('aborts all running agents and emits fallback notifications', () => { const callback = vi.fn(); registry.setNotificationCallback(callback); diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index beb43ac3c52..cca7d28f005 100644 --- a/packages/core/src/agents/background-tasks.ts +++ b/packages/core/src/agents/background-tasks.ts @@ -89,6 +89,38 @@ export function resolveMaxConcurrentBackgroundAgents( export const MAX_CONCURRENT_BACKGROUND_AGENTS = resolveMaxConcurrentBackgroundAgents(); +/** + * Normalize the `agents.maxParallelAgentsByModel` setting into a clean + * model-ID → cap map. Drops entries whose key is blank or whose value is not + * a positive integer (mirrors the validation the global cap goes through) so + * a malformed settings file degrades to "no per-model cap" rather than + * throwing at construction. + */ +function normalizePerModelConcurrency( + raw: ReadonlyMap | Record | undefined, +): Map { + const result = new Map(); + if (!raw) { + return result; + } + const entries = raw instanceof Map ? raw.entries() : Object.entries(raw); + for (const [model, value] of entries) { + const key = model?.trim(); + if (!key) { + continue; + } + if (!Number.isInteger(value) || value < 1) { + debugLogger.warn( + `Invalid maxParallelAgentsByModel[${JSON.stringify(model)}]=` + + `${JSON.stringify(value)}; ignoring (must be a positive integer).`, + ); + continue; + } + result.set(key, value); + } + return result; +} + /** * Cap on how many fully-finalized terminal entries (those that have * already emitted their terminal `task-notification`) the registry @@ -234,6 +266,13 @@ export interface AgentTask extends TaskBase { */ agentId: string; subagentType?: string; + /** + * Concrete model ID this agent runs with (resolved from the subagent's + * model selector at launch time). Used to enforce per-model concurrency + * caps (`agents.maxParallelAgentsByModel`); undefined when the model + * could not be resolved, in which case only the global cap applies. + */ + model?: string; /** * AgentId of the sub-agent that spawned this one; null when launched * from the top-level session. Drives the nested-agent tree display in @@ -394,14 +433,32 @@ type MessageWaiter = () => void; export interface BackgroundTaskRegistryOptions { maxConcurrentBackgroundAgents?: number; + /** + * Per-model concurrency caps keyed by concrete model ID. Each value is the + * maximum number of background sub-agents that may run concurrently on that + * model. A model not present here is bounded only by the global + * `maxConcurrentBackgroundAgents` cap. Useful when a model has a lower + * concurrency capacity than the rest of the fleet. + */ + maxConcurrentBackgroundAgentsByModel?: + | ReadonlyMap + | Record; } export interface BackgroundSlotReservation { readonly id: symbol; + /** + * Concrete model ID the slot was reserved for; undefined when the launch + * path could not resolve a model. Carried so the per-model cap can be + * checked consistently across reserve → consume → release. + */ + readonly model?: string; } interface BackgroundSlotWaiter { readonly signal?: AbortSignal; + /** Concrete model ID the waiter needs a slot for (per-model cap check). */ + readonly model?: string; readonly resolve: (reservation: BackgroundSlotReservation) => void; readonly reject: (error: Error) => void; readonly onAbort: () => void; @@ -414,8 +471,19 @@ export class BackgroundTaskRegistry { private readonly agents = new Map(); private readonly messageWaiters = new Map>(); private readonly waitQueue: BackgroundSlotWaiter[] = []; - private readonly reservedBackgroundSlots = new Set(); + // Maps each outstanding slot reservation to the concrete model ID it was + // reserved for (undefined when unresolved). A Map rather than a Set so the + // per-model cap can count reservations against the same model the running + // agents are tallied under. + private readonly reservedBackgroundSlots = new Map< + symbol, + string | undefined + >(); private readonly maxConcurrentBackgroundAgents: number; + // Per-model concurrency caps keyed by concrete model ID. Empty when no + // `agents.maxParallelAgentsByModel` is configured, in which case only the + // global cap is enforced. + private readonly maxConcurrentBackgroundAgentsByModel: Map; private notificationCallback?: BackgroundNotificationCallback; private registerCallback?: BackgroundRegisterCallback; private statusChangeCallback?: BackgroundStatusChangeCallback; @@ -429,15 +497,33 @@ export class BackgroundTaskRegistry { Number.isInteger(configured) && configured >= 1 ? configured : MAX_CONCURRENT_BACKGROUND_AGENTS; + this.maxConcurrentBackgroundAgentsByModel = normalizePerModelConcurrency( + options.maxConcurrentBackgroundAgentsByModel, + ); } - canStartBackgroundAgent(): boolean { - return ( - this.getClaimedBackgroundSlotCount() < this.maxConcurrentBackgroundAgents - ); + /** + * Whether a new background agent may start. Always bounded by the global + * cap; when `model` is given and a per-model cap is configured for it, the + * per-model cap must also have room. + */ + canStartBackgroundAgent(model?: string): boolean { + if ( + this.getClaimedBackgroundSlotCount() >= this.maxConcurrentBackgroundAgents + ) { + return false; + } + const perModelCap = this.resolvePerModelCap(model); + if ( + perModelCap !== undefined && + this.getClaimedBackgroundSlotCount(model) >= perModelCap + ) { + return false; + } + return true; } - assertCanStartBackgroundAgent(): void { + assertCanStartBackgroundAgent(model?: string): void { const claimed = this.getClaimedBackgroundSlotCount(); if (claimed >= this.maxConcurrentBackgroundAgents) { debugLogger.warn( @@ -451,15 +537,40 @@ export class BackgroundTaskRegistry { `agent first.`, ); } + const perModelCap = this.resolvePerModelCap(model); + if (perModelCap !== undefined) { + const claimedForModel = this.getClaimedBackgroundSlotCount(model); + if (claimedForModel >= perModelCap) { + debugLogger.warn( + `Background agent per-model concurrency cap reached for ` + + `${JSON.stringify(model)}: ${claimedForModel}/${perModelCap}. ` + + `Refusing new background agent.`, + ); + throw new Error( + `Cannot start background agent: maximum concurrent background agents ` + + `for model "${model}" (${perModelCap}) reached. Stop an existing ` + + `agent on that model first.`, + ); + } + } + } + + /** Configured per-model cap for `model`, or undefined when none applies. */ + private resolvePerModelCap(model?: string): number | undefined { + if (model === undefined) { + return undefined; + } + return this.maxConcurrentBackgroundAgentsByModel.get(model); } async waitForBackgroundSlot( signal?: AbortSignal, + model?: string, ): Promise { if (signal?.aborted) { throw new Error(BACKGROUND_SLOT_WAIT_CANCELLED); } - const reservation = this.tryReserveBackgroundSlot(); + const reservation = this.tryReserveBackgroundSlot(model); if (reservation) { return reservation; } @@ -474,6 +585,7 @@ export class BackgroundTaskRegistry { }; const waiter: BackgroundSlotWaiter = { signal, + model, resolve, reject, onAbort, @@ -483,11 +595,13 @@ export class BackgroundTaskRegistry { }); } - tryReserveBackgroundSlot(): BackgroundSlotReservation | undefined { - if (!this.canStartBackgroundAgent()) { + tryReserveBackgroundSlot( + model?: string, + ): BackgroundSlotReservation | undefined { + if (!this.canStartBackgroundAgent(model)) { return undefined; } - return this.reserveBackgroundSlot(); + return this.reserveBackgroundSlot(model); } getQueuedCount(): number { @@ -513,7 +627,7 @@ export class BackgroundTaskRegistry { if (options.slotReservation) { this.consumeBackgroundSlot(options.slotReservation); } else { - this.assertCanStartBackgroundAgent(); + this.assertCanStartBackgroundAgent(registration.model); } } } @@ -945,23 +1059,50 @@ export class BackgroundTaskRegistry { return Array.from(this.agents.values()); } - private getRunningBackgroundCount(): number { - return Array.from(this.agents.values()).filter( - (entry) => + // Counts backgrounded agents that still occupy a slot: running, or + // cancelled-but-not-yet-finalized. When `model` is given, only agents on + // that model are counted (per-model cap); otherwise all of them (global). + private getRunningBackgroundCount(model?: string): number { + let count = 0; + for (const entry of this.agents.values()) { + const occupiesSlot = entry.isBackgrounded && (entry.status === 'running' || - (entry.status === 'cancelled' && !entry.notified)), - ).length; + (entry.status === 'cancelled' && !entry.notified)); + if (!occupiesSlot) { + continue; + } + if (model === undefined || entry.model === model) { + count++; + } + } + return count; } - private getClaimedBackgroundSlotCount(): number { - return this.getRunningBackgroundCount() + this.reservedBackgroundSlots.size; + private getReservedBackgroundSlotCount(model?: string): number { + if (model === undefined) { + return this.reservedBackgroundSlots.size; + } + let count = 0; + for (const slotModel of this.reservedBackgroundSlots.values()) { + if (slotModel === model) { + count++; + } + } + return count; } - private reserveBackgroundSlot(): BackgroundSlotReservation { - const reservation = { id: Symbol('background-slot') }; - this.reservedBackgroundSlots.add(reservation.id); - return reservation; + private getClaimedBackgroundSlotCount(model?: string): number { + return ( + this.getRunningBackgroundCount(model) + + this.getReservedBackgroundSlotCount(model) + ); + } + + private reserveBackgroundSlot(model?: string): BackgroundSlotReservation { + const id = Symbol('background-slot'); + this.reservedBackgroundSlots.set(id, model); + return { id, model }; } private consumeBackgroundSlot(reservation: BackgroundSlotReservation): void { @@ -973,14 +1114,29 @@ export class BackgroundTaskRegistry { } private drainWaitQueue(): void { - while (this.waitQueue.length > 0 && this.canStartBackgroundAgent()) { - const waiter = this.waitQueue.shift()!; + for (let i = 0; i < this.waitQueue.length; ) { + // Once the global cap is hit no remaining waiter can be served, + // regardless of model — bail out instead of scanning the rest. + if ( + this.getClaimedBackgroundSlotCount() >= + this.maxConcurrentBackgroundAgents + ) { + break; + } + const waiter = this.waitQueue[i]!; + // A waiter whose model is at its per-model cap stays queued even while + // a different model's waiter behind it can still be served. + if (!this.canStartBackgroundAgent(waiter.model)) { + i++; + continue; + } + this.waitQueue.splice(i, 1); waiter.signal?.removeEventListener('abort', waiter.onAbort); if (waiter.signal?.aborted) { waiter.reject(new Error(BACKGROUND_SLOT_WAIT_CANCELLED)); continue; } - waiter.resolve(this.reserveBackgroundSlot()); + waiter.resolve(this.reserveBackgroundSlot(waiter.model)); } } diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 9f3ae1dae28..f441743e959 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -809,6 +809,42 @@ describe('Server Config (config.ts)', () => { }); }); + describe('agents.maxParallelAgentsByModel', () => { + it('configures a per-model background task concurrency cap', () => { + const config = new Config({ + ...baseParams, + agents: { + maxParallelAgentsByModel: { 'weak-model': 1 }, + }, + }); + const registry = config.getBackgroundTaskRegistry(); + + registry.register({ + agentId: 'bg-1', + description: 'one', + model: 'weak-model', + isBackgrounded: true, + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/bg-1.jsonl', + }); + + expect(() => + registry.register({ + agentId: 'bg-2', + description: 'two', + model: 'weak-model', + isBackgrounded: true, + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/bg-2.jsonl', + }), + ).toThrow('for model "weak-model" (1) reached'); + }); + }); + describe('getTeamMemoryEnabled', () => { const prevEnv = process.env['QWEN_CODE_MEMORY_TEAM']; afterEach(() => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 4d4fd887d17..99382c1362c 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -832,6 +832,13 @@ export interface AgentsCollabSettings { * When the cap is reached, additional launches wait for a slot. */ maxParallelAgents?: number; + /** + * Per-model maximum number of background sub-agents running concurrently, + * keyed by concrete model ID. Overrides the global `maxParallelAgents` for + * the matched model; models not listed here fall back to the global limit. + * Useful when a model has a lower concurrency capacity than the rest. + */ + maxParallelAgentsByModel?: Record; /** Display mode for multi-agent sessions ('in-process' | 'tmux' | 'iterm2') */ displayMode?: string; /** Arena-specific settings */ @@ -2132,14 +2139,20 @@ export class Config { this.eventEmitter = params.eventEmitter; this.arenaAgentClient = ArenaAgentClient.create(); this.agentsSettings = params.agents ?? {}; - this.backgroundTaskRegistry = new BackgroundTaskRegistry( - this.agentsSettings.maxParallelAgents === undefined - ? undefined - : { + this.backgroundTaskRegistry = new BackgroundTaskRegistry({ + ...(this.agentsSettings.maxParallelAgents !== undefined + ? { maxConcurrentBackgroundAgents: this.agentsSettings.maxParallelAgents, - }, - ); + } + : {}), + ...(this.agentsSettings.maxParallelAgentsByModel !== undefined + ? { + maxConcurrentBackgroundAgentsByModel: + this.agentsSettings.maxParallelAgentsByModel, + } + : {}), + }); this.worktreeSettings = params.worktree ?? {}; if (params.contextFileName) { setGeminiMdFilename(params.contextFileName); diff --git a/packages/core/src/tools/agent/agent.test.ts b/packages/core/src/tools/agent/agent.test.ts index 150d3debbbb..17b27eef34a 100644 --- a/packages/core/src/tools/agent/agent.test.ts +++ b/packages/core/src/tools/agent/agent.test.ts @@ -2679,6 +2679,74 @@ describe('AgentTool', () => { { notify: false }, ); }); + + it('reserves a background slot with the resolved parent model when fork runs in background', async () => { + // Removing the `!isFork` guard from the slot-reservation condition + // silently subjected fork agents to background slot reservation and + // per-model concurrency caps. This test pins the contract: a fork + // launched with run_in_background: true resolves its concrete model + // from the parent config (FORK_AGENT has no model selector, so it + // inherits) and passes that model to tryReserveBackgroundSlot and + // the registry register call. + (mockAgent as unknown as Record)['getCore'] = vi + .fn() + .mockReturnValue({ + getEventEmitter: () => ({ on: vi.fn(), off: vi.fn() }), + }); + (mockAgent as unknown as Record)[ + 'setExternalMessageProvider' + ] = vi.fn(); + (mockAgent as unknown as Record)[ + 'setExternalMessageWaiter' + ] = vi.fn(); + (mockAgent as unknown as Record)[ + 'setExternalMessageWaitPredicate' + ] = vi.fn(); + + const stubRegistry = ( + config as unknown as { + getBackgroundTaskRegistry: () => { + tryReserveBackgroundSlot: ReturnType; + register: ReturnType; + }; + } + ).getBackgroundTaskRegistry(); + + const params: AgentParams = { + description: 'fork task', + prompt: 'do the thing', + subagent_type: 'fork', + run_in_background: true, + }; + + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation(params); + await invocation.execute(); + + // createForkSubagent runs unconditionally before the background + // branch, so AgentHeadless.create fires once for the foreground + // probe and again for the background agent body. The load-bearing + // assertions are on the registry calls below. + expect(AgentHeadless.create).toHaveBeenCalled(); + // Fork inherits the parent model (FORK_AGENT has no model selector), + // so resolveModelId returns the parent's current model. + expect(stubRegistry.tryReserveBackgroundSlot).toHaveBeenCalledWith( + 'parent-model', + ); + expect(stubRegistry.register).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'parent-model', + isBackgrounded: true, + subagentType: 'fork', + }), + expect.objectContaining({ + slotReservation: expect.objectContaining({ + id: expect.any(Symbol), + }), + }), + ); + }); }); describe('SubagentStart hook integration', () => { diff --git a/packages/core/src/tools/agent/agent.ts b/packages/core/src/tools/agent/agent.ts index 4a28455b308..54e99a2ae87 100644 --- a/packages/core/src/tools/agent/agent.ts +++ b/packages/core/src/tools/agent/agent.ts @@ -112,6 +112,7 @@ import { } from '../../agents/agent-transcript.js'; import type { BackgroundSlotReservation } from '../../agents/background-tasks.js'; import { getGitBranch } from '../../utils/gitUtils.js'; +import { buildModelIdContext, resolveModelId } from '../../utils/modelId.js'; // Memoize git branch per cwd for the agent-launch path. `getGitBranch` // shells out to `git rev-parse` synchronously; caching avoids the per-launch @@ -2232,6 +2233,10 @@ class AgentToolInvocation extends BaseToolInvocation { let restoreParentPM: () => void = () => {}; let backgroundSlotReservation: BackgroundSlotReservation | undefined; let backgroundSlotReservationConsumed = false; + // Concrete model ID the sub-agent will run with, resolved from its model + // selector once subagentConfig is loaded. Used to enforce per-model + // background-agent concurrency caps (agents.maxParallelAgentsByModel). + let subagentModelId: string | undefined; const releaseBackgroundSlotReservation = () => { if (backgroundSlotReservation && !backgroundSlotReservationConsumed) { this.config @@ -2348,9 +2353,20 @@ class AgentToolInvocation extends BaseToolInvocation { ); } - if (!isFork && shouldRunInBackground) { + if (shouldRunInBackground) { + // Resolve the concrete model the sub-agent (or fork) will run with so the + // registry can apply a per-model cap. `subagentConfig.model` is a + // selector (omitted/"inherit"/"fast"/modelId/authType:modelId); + // resolveModelId maps it to the actual model ID, falling back to the + // parent's current model when the sub-agent inherits (forks always + // inherit, since FORK_AGENT has no model selector). + subagentModelId = resolveModelId( + subagentConfig.model, + buildModelIdContext(this.config), + )?.modelId; const registry = this.config.getBackgroundTaskRegistry(); - backgroundSlotReservation = registry.tryReserveBackgroundSlot(); + backgroundSlotReservation = + registry.tryReserveBackgroundSlot(subagentModelId); if (!backgroundSlotReservation) { const queuedCount = registry.getQueuedCount(); const queueText = @@ -2366,8 +2382,10 @@ class AgentToolInvocation extends BaseToolInvocation { }, updateOutput, ); - backgroundSlotReservation = - await registry.waitForBackgroundSlot(signal); + backgroundSlotReservation = await registry.waitForBackgroundSlot( + signal, + subagentModelId, + ); } this.updateDisplay( { @@ -2794,6 +2812,9 @@ class AgentToolInvocation extends BaseToolInvocation { agentId: hookOpts.agentId, description: this.params.description, subagentType: subagentConfig.name, + // Concrete model ID for per-model concurrency accounting; the + // slot reservation above was taken against this same model. + model: subagentModelId, isBackgrounded: true, status: 'running', startTime: Date.now(), @@ -2911,6 +2932,7 @@ class AgentToolInvocation extends BaseToolInvocation { // Persisted so resume restores the original nesting level; see // childLaunchDepth() for the rationale. depth: childLaunchDepth(), + model: subagentModelId, }); // Subscribe to the subagent's tool-call event stream so the diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index 6324e45de2c..8817e1793b7 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -1325,7 +1325,15 @@ "maxParallelAgents": { "type": "integer", "minimum": 1, - "description": "Global maximum number of background sub-agents that can run concurrently. Additional background agents wait in a queue until a slot is available. Per-model limits are not supported yet." + "description": "Global maximum number of background sub-agents that can run concurrently. Additional background agents wait in a queue until a slot is available. Use maxParallelAgentsByModel to cap a specific model below this global limit." + }, + "maxParallelAgentsByModel": { + "type": "object", + "additionalProperties": { + "type": "integer", + "minimum": 1 + }, + "description": "Per-model maximum number of background sub-agents that can run concurrently, keyed by model ID (e.g. { \"qwen3-max\": 2 }). Useful when a model has a lower concurrency capacity. Takes precedence over the global maxParallelAgents for the matched model; models not listed here fall back to the global limit." }, "displayMode": { "description": "Display mode for multi-agent sessions. Currently only \"in-process\" is supported. Options: in-process",