fix(#338): Add session cleanup on terminal states
- Add removeSession and scheduleSessionCleanup methods to AgentSpawnerService - Schedule session cleanup after completed/failed/killed transitions - Default 30 second delay before cleanup to allow status queries - Implement OnModuleDestroy to clean up pending timers - Add forwardRef injection to avoid circular dependency - Add comprehensive tests for cleanup functionality Refs #338
This commit is contained in:
@@ -401,4 +401,159 @@ describe("AgentSpawnerService", () => {
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user