diff --git a/apps/orchestrator/src/config/orchestrator.config.spec.ts b/apps/orchestrator/src/config/orchestrator.config.spec.ts index 0c44f9d..c3f2263 100644 --- a/apps/orchestrator/src/config/orchestrator.config.spec.ts +++ b/apps/orchestrator/src/config/orchestrator.config.spec.ts @@ -83,4 +83,30 @@ describe("orchestratorConfig", () => { expect(config.valkey.url).toBe("redis://localhost:6379"); }); }); + + describe("spawner config", () => { + it("should use default maxConcurrentAgents of 20 when not set", () => { + delete process.env.MAX_CONCURRENT_AGENTS; + + const config = orchestratorConfig(); + + expect(config.spawner.maxConcurrentAgents).toBe(20); + }); + + it("should use provided maxConcurrentAgents when MAX_CONCURRENT_AGENTS is set", () => { + process.env.MAX_CONCURRENT_AGENTS = "50"; + + const config = orchestratorConfig(); + + expect(config.spawner.maxConcurrentAgents).toBe(50); + }); + + it("should handle MAX_CONCURRENT_AGENTS of 10", () => { + process.env.MAX_CONCURRENT_AGENTS = "10"; + + const config = orchestratorConfig(); + + expect(config.spawner.maxConcurrentAgents).toBe(10); + }); + }); }); diff --git a/apps/orchestrator/src/config/orchestrator.config.ts b/apps/orchestrator/src/config/orchestrator.config.ts index 2904344..ead5fa2 100644 --- a/apps/orchestrator/src/config/orchestrator.config.ts +++ b/apps/orchestrator/src/config/orchestrator.config.ts @@ -37,4 +37,7 @@ export const orchestratorConfig = registerAs("orchestrator", () => ({ yolo: { enabled: process.env.YOLO_MODE === "true", }, + spawner: { + maxConcurrentAgents: parseInt(process.env.MAX_CONCURRENT_AGENTS ?? "20", 10), + }, })); diff --git a/apps/orchestrator/src/spawner/agent-spawner.service.spec.ts b/apps/orchestrator/src/spawner/agent-spawner.service.spec.ts index 2a322d1..8eb2a42 100644 --- a/apps/orchestrator/src/spawner/agent-spawner.service.spec.ts +++ b/apps/orchestrator/src/spawner/agent-spawner.service.spec.ts @@ -1,4 +1,5 @@ import { ConfigService } from "@nestjs/config"; +import { HttpException, HttpStatus } from "@nestjs/common"; import { describe, it, expect, beforeEach, vi } from "vitest"; import { AgentSpawnerService } from "./agent-spawner.service"; import { SpawnAgentRequest } from "./types/agent-spawner.types"; @@ -14,6 +15,9 @@ describe("AgentSpawnerService", () => { if (key === "orchestrator.claude.apiKey") { return "test-api-key"; } + if (key === "orchestrator.spawner.maxConcurrentAgents") { + return 20; + } return undefined; }), } as unknown as ConfigService; @@ -252,4 +256,149 @@ describe("AgentSpawnerService", () => { expect(sessions[1].agentType).toBe("reviewer"); }); }); + + describe("max concurrent agents limit", () => { + const createValidRequest = (taskId: string): SpawnAgentRequest => ({ + taskId, + agentType: "worker", + context: { + repository: "https://github.com/test/repo.git", + branch: "main", + workItems: ["Implement feature X"], + }, + }); + + it("should allow spawning when under the limit", () => { + // Default limit is 20, spawn 5 agents + for (let i = 0; i < 5; i++) { + const response = service.spawnAgent(createValidRequest(`task-${i}`)); + expect(response.agentId).toBeDefined(); + } + + expect(service.listAgentSessions()).toHaveLength(5); + }); + + it("should reject spawn when at the limit", () => { + // Create service with low limit for testing + const limitedConfigService = { + get: vi.fn((key: string) => { + if (key === "orchestrator.claude.apiKey") { + return "test-api-key"; + } + if (key === "orchestrator.spawner.maxConcurrentAgents") { + return 3; + } + return undefined; + }), + } as unknown as ConfigService; + + const limitedService = new AgentSpawnerService(limitedConfigService); + + // Spawn up to the limit + limitedService.spawnAgent(createValidRequest("task-1")); + limitedService.spawnAgent(createValidRequest("task-2")); + limitedService.spawnAgent(createValidRequest("task-3")); + + // Next spawn should throw 429 Too Many Requests + expect(() => limitedService.spawnAgent(createValidRequest("task-4"))).toThrow(HttpException); + + try { + limitedService.spawnAgent(createValidRequest("task-5")); + } catch (error) { + expect(error).toBeInstanceOf(HttpException); + expect((error as HttpException).getStatus()).toBe(HttpStatus.TOO_MANY_REQUESTS); + expect((error as HttpException).message).toContain("Maximum concurrent agents limit"); + } + }); + + it("should provide appropriate error message when limit reached", () => { + const limitedConfigService = { + get: vi.fn((key: string) => { + if (key === "orchestrator.claude.apiKey") { + return "test-api-key"; + } + if (key === "orchestrator.spawner.maxConcurrentAgents") { + return 2; + } + return undefined; + }), + } as unknown as ConfigService; + + const limitedService = new AgentSpawnerService(limitedConfigService); + + // Spawn up to the limit + limitedService.spawnAgent(createValidRequest("task-1")); + limitedService.spawnAgent(createValidRequest("task-2")); + + // Next spawn should throw with appropriate message + try { + limitedService.spawnAgent(createValidRequest("task-3")); + expect.fail("Should have thrown"); + } catch (error) { + expect(error).toBeInstanceOf(HttpException); + const httpError = error as HttpException; + expect(httpError.getStatus()).toBe(HttpStatus.TOO_MANY_REQUESTS); + expect(httpError.message).toContain("2"); + } + }); + + it("should use default limit of 20 when not configured", () => { + const defaultConfigService = { + get: vi.fn((key: string) => { + if (key === "orchestrator.claude.apiKey") { + return "test-api-key"; + } + // Return undefined for maxConcurrentAgents to test default + return undefined; + }), + } as unknown as ConfigService; + + const defaultService = new AgentSpawnerService(defaultConfigService); + + // Should be able to spawn 20 agents + for (let i = 0; i < 20; i++) { + const response = defaultService.spawnAgent(createValidRequest(`task-${i}`)); + expect(response.agentId).toBeDefined(); + } + + // 21st should fail + expect(() => defaultService.spawnAgent(createValidRequest("task-21"))).toThrow(HttpException); + }); + + it("should return current and max count in error response", () => { + const limitedConfigService = { + get: vi.fn((key: string) => { + if (key === "orchestrator.claude.apiKey") { + return "test-api-key"; + } + if (key === "orchestrator.spawner.maxConcurrentAgents") { + return 5; + } + return undefined; + }), + } as unknown as ConfigService; + + const limitedService = new AgentSpawnerService(limitedConfigService); + + // Spawn 5 agents + for (let i = 0; i < 5; i++) { + limitedService.spawnAgent(createValidRequest(`task-${i}`)); + } + + try { + limitedService.spawnAgent(createValidRequest("task-6")); + expect.fail("Should have thrown"); + } catch (error) { + expect(error).toBeInstanceOf(HttpException); + const httpError = error as HttpException; + const response = httpError.getResponse() as { + message: string; + currentCount: number; + maxLimit: number; + }; + expect(response.currentCount).toBe(5); + expect(response.maxLimit).toBe(5); + } + }); + }); }); diff --git a/apps/orchestrator/src/spawner/agent-spawner.service.ts b/apps/orchestrator/src/spawner/agent-spawner.service.ts index eb23c77..fc6f0d4 100644 --- a/apps/orchestrator/src/spawner/agent-spawner.service.ts +++ b/apps/orchestrator/src/spawner/agent-spawner.service.ts @@ -1,4 +1,4 @@ -import { Injectable, Logger } from "@nestjs/common"; +import { Injectable, Logger, HttpException, HttpStatus } from "@nestjs/common"; import { ConfigService } from "@nestjs/config"; import Anthropic from "@anthropic-ai/sdk"; import { randomUUID } from "crypto"; @@ -17,6 +17,7 @@ export class AgentSpawnerService { private readonly logger = new Logger(AgentSpawnerService.name); private readonly anthropic: Anthropic; private readonly sessions = new Map(); + private readonly maxConcurrentAgents: number; constructor(private readonly configService: ConfigService) { const apiKey = this.configService.get("orchestrator.claude.apiKey"); @@ -29,7 +30,13 @@ export class AgentSpawnerService { apiKey, }); - this.logger.log("AgentSpawnerService initialized with Claude SDK"); + // Default to 20 if not configured + this.maxConcurrentAgents = + this.configService.get("orchestrator.spawner.maxConcurrentAgents") ?? 20; + + this.logger.log( + `AgentSpawnerService initialized with Claude SDK (max concurrent agents: ${String(this.maxConcurrentAgents)})` + ); } /** @@ -40,6 +47,9 @@ export class AgentSpawnerService { spawnAgent(request: SpawnAgentRequest): SpawnAgentResponse { this.logger.log(`Spawning agent for task: ${request.taskId}`); + // Check concurrent agent limit before proceeding + this.checkConcurrentAgentLimit(); + // Validate request this.validateSpawnRequest(request); @@ -90,6 +100,27 @@ export class AgentSpawnerService { return Array.from(this.sessions.values()); } + /** + * Check if the concurrent agent limit has been reached + * @throws HttpException with 429 Too Many Requests if limit reached + */ + private checkConcurrentAgentLimit(): void { + const currentCount = this.sessions.size; + if (currentCount >= this.maxConcurrentAgents) { + this.logger.warn( + `Maximum concurrent agents limit reached: ${String(currentCount)}/${String(this.maxConcurrentAgents)}` + ); + throw new HttpException( + { + message: `Maximum concurrent agents limit reached (${String(this.maxConcurrentAgents)}). Please wait for existing agents to complete.`, + currentCount, + maxLimit: this.maxConcurrentAgents, + }, + HttpStatus.TOO_MANY_REQUESTS + ); + } + } + /** * Validate spawn agent request * @param request Spawn request to validate