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:
Jason Woltje
2026-02-05 18:47:14 -06:00
parent 8d57191a91
commit a42f88d64c
4 changed files with 347 additions and 7 deletions

View File

@@ -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();
});
});
});