624 lines
20 KiB
TypeScript
624 lines
20 KiB
TypeScript
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";
|
|
|
|
describe("AgentSpawnerService", () => {
|
|
let service: AgentSpawnerService;
|
|
let mockConfigService: ConfigService;
|
|
|
|
beforeEach(() => {
|
|
// Create mock ConfigService
|
|
mockConfigService = {
|
|
get: vi.fn((key: string) => {
|
|
if (key === "orchestrator.aiProvider") {
|
|
return "ollama";
|
|
}
|
|
if (key === "orchestrator.claude.apiKey") {
|
|
return "test-api-key";
|
|
}
|
|
if (key === "orchestrator.spawner.maxConcurrentAgents") {
|
|
return 20;
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
// Create service with mock
|
|
service = new AgentSpawnerService(mockConfigService);
|
|
});
|
|
|
|
describe("constructor", () => {
|
|
it("should be defined", () => {
|
|
expect(service).toBeDefined();
|
|
});
|
|
|
|
it("should initialize with default AI provider when API key is omitted", () => {
|
|
const noClaudeConfigService = {
|
|
get: vi.fn((key: string) => {
|
|
if (key === "orchestrator.aiProvider") {
|
|
return "ollama";
|
|
}
|
|
if (key === "orchestrator.spawner.maxConcurrentAgents") {
|
|
return 20;
|
|
}
|
|
if (key === "orchestrator.spawner.sessionCleanupDelayMs") {
|
|
return 30000;
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
const serviceNoKey = new AgentSpawnerService(noClaudeConfigService);
|
|
expect(serviceNoKey).toBeDefined();
|
|
});
|
|
|
|
it("should initialize with Claude provider when key is present", () => {
|
|
expect(mockConfigService.get).toHaveBeenCalledWith("orchestrator.claude.apiKey");
|
|
});
|
|
|
|
it("should initialize with CLAUDE provider when API key is present", () => {
|
|
const claudeConfigService = {
|
|
get: vi.fn((key: string) => {
|
|
if (key === "orchestrator.aiProvider") {
|
|
return "claude";
|
|
}
|
|
if (key === "orchestrator.claude.apiKey") {
|
|
return "test-api-key";
|
|
}
|
|
if (key === "orchestrator.spawner.maxConcurrentAgents") {
|
|
return 20;
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
const claudeService = new AgentSpawnerService(claudeConfigService);
|
|
expect(claudeService).toBeDefined();
|
|
});
|
|
|
|
it("should throw error if Claude API key is missing when provider is claude", () => {
|
|
const badConfigService = {
|
|
get: vi.fn((key: string) => {
|
|
if (key === "orchestrator.aiProvider") {
|
|
return "claude";
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
expect(() => new AgentSpawnerService(badConfigService)).toThrow(
|
|
"CLAUDE_API_KEY is required when AI_PROVIDER is set to 'claude'"
|
|
);
|
|
});
|
|
|
|
it("should still initialize when CLAUDE_API_KEY is missing for non-Claude provider", () => {
|
|
const nonClaudeConfigService = {
|
|
get: vi.fn((key: string) => {
|
|
if (key === "orchestrator.aiProvider") {
|
|
return "ollama";
|
|
}
|
|
if (key === "orchestrator.spawner.maxConcurrentAgents") {
|
|
return 20;
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
expect(() => new AgentSpawnerService(nonClaudeConfigService)).not.toThrow();
|
|
});
|
|
});
|
|
|
|
describe("spawnAgent", () => {
|
|
const validRequest: SpawnAgentRequest = {
|
|
taskId: "task-123",
|
|
agentType: "worker",
|
|
context: {
|
|
repository: "https://github.com/test/repo.git",
|
|
branch: "main",
|
|
workItems: ["Implement feature X"],
|
|
},
|
|
};
|
|
|
|
it("should spawn an agent and return agentId", () => {
|
|
const response = service.spawnAgent(validRequest);
|
|
|
|
expect(response).toBeDefined();
|
|
expect(response.agentId).toBeDefined();
|
|
expect(typeof response.agentId).toBe("string");
|
|
expect(response.state).toBe("spawning");
|
|
expect(response.spawnedAt).toBeInstanceOf(Date);
|
|
});
|
|
|
|
it("should generate unique agentId for each spawn", () => {
|
|
const response1 = service.spawnAgent(validRequest);
|
|
const response2 = service.spawnAgent(validRequest);
|
|
|
|
expect(response1.agentId).not.toBe(response2.agentId);
|
|
});
|
|
|
|
it("should track agent session", () => {
|
|
const response = service.spawnAgent(validRequest);
|
|
const session = service.getAgentSession(response.agentId);
|
|
|
|
expect(session).toBeDefined();
|
|
expect(session?.agentId).toBe(response.agentId);
|
|
expect(session?.taskId).toBe(validRequest.taskId);
|
|
expect(session?.agentType).toBe(validRequest.agentType);
|
|
expect(session?.state).toBe("spawning");
|
|
});
|
|
|
|
it("should validate taskId is provided", () => {
|
|
const invalidRequest = {
|
|
...validRequest,
|
|
taskId: "",
|
|
};
|
|
|
|
expect(() => service.spawnAgent(invalidRequest)).toThrow("taskId is required");
|
|
});
|
|
|
|
it("should validate agentType is valid", () => {
|
|
const invalidRequest = {
|
|
...validRequest,
|
|
agentType: "invalid" as unknown as "worker",
|
|
};
|
|
|
|
expect(() => service.spawnAgent(invalidRequest)).toThrow(
|
|
"agentType must be one of: worker, reviewer, tester"
|
|
);
|
|
});
|
|
|
|
it("should validate context.repository is provided", () => {
|
|
const invalidRequest = {
|
|
...validRequest,
|
|
context: {
|
|
...validRequest.context,
|
|
repository: "",
|
|
},
|
|
};
|
|
|
|
expect(() => service.spawnAgent(invalidRequest)).toThrow("context.repository is required");
|
|
});
|
|
|
|
it("should validate context.branch is provided", () => {
|
|
const invalidRequest = {
|
|
...validRequest,
|
|
context: {
|
|
...validRequest.context,
|
|
branch: "",
|
|
},
|
|
};
|
|
|
|
expect(() => service.spawnAgent(invalidRequest)).toThrow("context.branch is required");
|
|
});
|
|
|
|
it("should validate context.workItems is not empty", () => {
|
|
const invalidRequest = {
|
|
...validRequest,
|
|
context: {
|
|
...validRequest.context,
|
|
workItems: [],
|
|
},
|
|
};
|
|
|
|
expect(() => service.spawnAgent(invalidRequest)).toThrow(
|
|
"context.workItems must not be empty"
|
|
);
|
|
});
|
|
|
|
it("should accept optional skills in context", () => {
|
|
const requestWithSkills: SpawnAgentRequest = {
|
|
...validRequest,
|
|
context: {
|
|
...validRequest.context,
|
|
skills: ["typescript", "nestjs"],
|
|
},
|
|
};
|
|
|
|
const response = service.spawnAgent(requestWithSkills);
|
|
const session = service.getAgentSession(response.agentId);
|
|
|
|
expect(session?.context.skills).toEqual(["typescript", "nestjs"]);
|
|
});
|
|
|
|
it("should accept optional options", () => {
|
|
const requestWithOptions: SpawnAgentRequest = {
|
|
...validRequest,
|
|
options: {
|
|
sandbox: true,
|
|
timeout: 3600000,
|
|
maxRetries: 3,
|
|
},
|
|
};
|
|
|
|
const response = service.spawnAgent(requestWithOptions);
|
|
const session = service.getAgentSession(response.agentId);
|
|
|
|
expect(session?.options).toEqual({
|
|
sandbox: true,
|
|
timeout: 3600000,
|
|
maxRetries: 3,
|
|
});
|
|
});
|
|
|
|
it("should handle spawn errors gracefully", () => {
|
|
// Mock Claude SDK to throw error
|
|
const errorRequest = {
|
|
...validRequest,
|
|
context: {
|
|
...validRequest.context,
|
|
repository: "invalid-repo-that-will-fail",
|
|
},
|
|
};
|
|
|
|
// For now, this should not throw but handle gracefully
|
|
// We'll implement error handling in the service
|
|
const response = service.spawnAgent(errorRequest);
|
|
expect(response.agentId).toBeDefined();
|
|
});
|
|
});
|
|
|
|
describe("getAgentSession", () => {
|
|
it("should return undefined for non-existent agentId", () => {
|
|
const session = service.getAgentSession("non-existent-id");
|
|
expect(session).toBeUndefined();
|
|
});
|
|
|
|
it("should return session for existing agentId", () => {
|
|
const request: SpawnAgentRequest = {
|
|
taskId: "task-123",
|
|
agentType: "worker",
|
|
context: {
|
|
repository: "https://github.com/test/repo.git",
|
|
branch: "main",
|
|
workItems: ["Implement feature X"],
|
|
},
|
|
};
|
|
|
|
const response = service.spawnAgent(request);
|
|
const session = service.getAgentSession(response.agentId);
|
|
|
|
expect(session).toBeDefined();
|
|
expect(session?.agentId).toBe(response.agentId);
|
|
});
|
|
});
|
|
|
|
describe("listAgentSessions", () => {
|
|
it("should return empty array when no agents spawned", () => {
|
|
const sessions = service.listAgentSessions();
|
|
expect(sessions).toEqual([]);
|
|
});
|
|
|
|
it("should return all spawned agent sessions", () => {
|
|
const request1: SpawnAgentRequest = {
|
|
taskId: "task-1",
|
|
agentType: "worker",
|
|
context: {
|
|
repository: "https://github.com/test/repo1.git",
|
|
branch: "main",
|
|
workItems: ["Task 1"],
|
|
},
|
|
};
|
|
|
|
const request2: SpawnAgentRequest = {
|
|
taskId: "task-2",
|
|
agentType: "reviewer",
|
|
context: {
|
|
repository: "https://github.com/test/repo2.git",
|
|
branch: "develop",
|
|
workItems: ["Task 2"],
|
|
},
|
|
};
|
|
|
|
service.spawnAgent(request1);
|
|
service.spawnAgent(request2);
|
|
|
|
const sessions = service.listAgentSessions();
|
|
expect(sessions).toHaveLength(2);
|
|
expect(sessions[0].agentType).toBe("worker");
|
|
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);
|
|
}
|
|
});
|
|
});
|
|
|
|
describe("session cleanup", () => {
|
|
const createValidRequest = (taskId: string): SpawnAgentRequest => ({
|
|
taskId,
|
|
agentType: "worker",
|
|
context: {
|
|
repository: "https://github.com/test/repo.git",
|
|
branch: "main",
|
|
workItems: ["Implement feature X"],
|
|
},
|
|
});
|
|
|
|
it("should remove session immediately", () => {
|
|
const response = service.spawnAgent(createValidRequest("task-1"));
|
|
expect(service.getAgentSession(response.agentId)).toBeDefined();
|
|
|
|
const removed = service.removeSession(response.agentId);
|
|
|
|
expect(removed).toBe(true);
|
|
expect(service.getAgentSession(response.agentId)).toBeUndefined();
|
|
});
|
|
|
|
it("should return false when removing non-existent session", () => {
|
|
const removed = service.removeSession("non-existent-id");
|
|
expect(removed).toBe(false);
|
|
});
|
|
|
|
it("should schedule session cleanup with delay", async () => {
|
|
vi.useFakeTimers();
|
|
|
|
const response = service.spawnAgent(createValidRequest("task-1"));
|
|
expect(service.getAgentSession(response.agentId)).toBeDefined();
|
|
|
|
// Schedule cleanup with short delay
|
|
service.scheduleSessionCleanup(response.agentId, 100);
|
|
|
|
// Session should still exist before delay
|
|
expect(service.getAgentSession(response.agentId)).toBeDefined();
|
|
expect(service.getPendingCleanupCount()).toBe(1);
|
|
|
|
// Advance timer past the delay
|
|
vi.advanceTimersByTime(150);
|
|
|
|
// Session should be cleaned up
|
|
expect(service.getAgentSession(response.agentId)).toBeUndefined();
|
|
expect(service.getPendingCleanupCount()).toBe(0);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("should replace existing cleanup timer when rescheduled", async () => {
|
|
vi.useFakeTimers();
|
|
|
|
const response = service.spawnAgent(createValidRequest("task-1"));
|
|
|
|
// Schedule cleanup with 100ms delay
|
|
service.scheduleSessionCleanup(response.agentId, 100);
|
|
expect(service.getPendingCleanupCount()).toBe(1);
|
|
|
|
// Advance by 50ms (halfway)
|
|
vi.advanceTimersByTime(50);
|
|
expect(service.getAgentSession(response.agentId)).toBeDefined();
|
|
|
|
// Reschedule with 100ms delay (should reset the timer)
|
|
service.scheduleSessionCleanup(response.agentId, 100);
|
|
expect(service.getPendingCleanupCount()).toBe(1);
|
|
|
|
// Advance by 75ms (past original but not new)
|
|
vi.advanceTimersByTime(75);
|
|
expect(service.getAgentSession(response.agentId)).toBeDefined();
|
|
|
|
// Advance by remaining 25ms
|
|
vi.advanceTimersByTime(50);
|
|
expect(service.getAgentSession(response.agentId)).toBeUndefined();
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("should clear cleanup timer when session is removed directly", () => {
|
|
vi.useFakeTimers();
|
|
|
|
const response = service.spawnAgent(createValidRequest("task-1"));
|
|
|
|
// Schedule cleanup
|
|
service.scheduleSessionCleanup(response.agentId, 1000);
|
|
expect(service.getPendingCleanupCount()).toBe(1);
|
|
|
|
// Remove session directly
|
|
service.removeSession(response.agentId);
|
|
|
|
// Timer should be cleared
|
|
expect(service.getPendingCleanupCount()).toBe(0);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("should decrease session count after cleanup", async () => {
|
|
vi.useFakeTimers();
|
|
|
|
// 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 2;
|
|
}
|
|
return undefined;
|
|
}),
|
|
} as unknown as ConfigService;
|
|
|
|
const limitedService = new AgentSpawnerService(limitedConfigService);
|
|
|
|
// Spawn up to the limit
|
|
const response1 = limitedService.spawnAgent(createValidRequest("task-1"));
|
|
limitedService.spawnAgent(createValidRequest("task-2"));
|
|
|
|
// Should be at limit
|
|
expect(limitedService.listAgentSessions()).toHaveLength(2);
|
|
expect(() => limitedService.spawnAgent(createValidRequest("task-3"))).toThrow(HttpException);
|
|
|
|
// Schedule cleanup for first agent
|
|
limitedService.scheduleSessionCleanup(response1.agentId, 100);
|
|
vi.advanceTimersByTime(150);
|
|
|
|
// Should have freed a slot
|
|
expect(limitedService.listAgentSessions()).toHaveLength(1);
|
|
|
|
// Should be able to spawn another agent now
|
|
const response3 = limitedService.spawnAgent(createValidRequest("task-3"));
|
|
expect(response3.agentId).toBeDefined();
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("should clear all timers on module destroy", () => {
|
|
vi.useFakeTimers();
|
|
|
|
const response1 = service.spawnAgent(createValidRequest("task-1"));
|
|
const response2 = service.spawnAgent(createValidRequest("task-2"));
|
|
|
|
service.scheduleSessionCleanup(response1.agentId, 1000);
|
|
service.scheduleSessionCleanup(response2.agentId, 1000);
|
|
|
|
expect(service.getPendingCleanupCount()).toBe(2);
|
|
|
|
// Call module destroy
|
|
service.onModuleDestroy();
|
|
|
|
expect(service.getPendingCleanupCount()).toBe(0);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
});
|
|
});
|