diff --git a/src/services/agentInstance/__tests__/workerCleanup.test.ts b/src/services/agentInstance/__tests__/workerCleanup.test.ts new file mode 100644 index 00000000..a36e7f64 --- /dev/null +++ b/src/services/agentInstance/__tests__/workerCleanup.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from 'vitest'; +import { container } from '@services/container'; +import type { IAgentInstanceService } from '@services/agentInstance/interface'; +import serviceIdentifier from '@services/serviceIdentifier'; + +describe('AgentInstanceService worker cleanup', () => { + it('cleans reverse worker conversation mapping when a single worker conversation is cleaned up', () => { + const service = container.get(serviceIdentifier.AgentInstance) as unknown as { + workerConversationByAgentId: Map; + workerAgentIdByConversationId: Map; + workerConversationCleanupByAgentId: Map void>; + cleanupWorkerConversation: (agentId: string) => void; + }; + + const agentId = 'agent-cleanup-test'; + const conversationId = 'worker-conversation-cleanup-test'; + service.workerConversationByAgentId.set(agentId, conversationId); + service.workerAgentIdByConversationId.set(conversationId, agentId); + service.workerConversationCleanupByAgentId.set(agentId, () => undefined); + + service.cleanupWorkerConversation(agentId); + + expect(service.workerConversationCleanupByAgentId.has(agentId)).toBe(false); + expect(service.workerAgentIdByConversationId.has(conversationId)).toBe(false); + + service.workerConversationByAgentId.delete(agentId); + }); + + it('clears reverse worker conversation mappings when disposing the worker', async () => { + const service = container.get(serviceIdentifier.AgentInstance) as unknown as { + workerConversationByAgentId: Map; + workerAgentIdByConversationId: Map; + workerConversationCleanupByAgentId: Map void>; + memeLoopNativeWorker?: { terminate: () => Promise }; + memeLoopWorker?: unknown; + disposeMemeLoopWorker: () => Promise; + }; + + const agentId = 'agent-dispose-test'; + const conversationId = 'worker-conversation-dispose-test'; + service.workerConversationByAgentId.set(agentId, conversationId); + service.workerAgentIdByConversationId.set(conversationId, agentId); + service.workerConversationCleanupByAgentId.set(agentId, () => undefined); + service.memeLoopNativeWorker = { + terminate: async () => undefined, + }; + service.memeLoopWorker = {}; + + await service.disposeMemeLoopWorker(); + + expect(service.workerConversationByAgentId.size).toBe(0); + expect(service.workerAgentIdByConversationId.size).toBe(0); + expect(service.workerConversationCleanupByAgentId.size).toBe(0); + expect(service.memeLoopNativeWorker).toBeUndefined(); + expect(service.memeLoopWorker).toBeUndefined(); + }); +}); \ No newline at end of file diff --git a/src/services/agentInstance/index.ts b/src/services/agentInstance/index.ts index 05315dd6..278013e0 100644 --- a/src/services/agentInstance/index.ts +++ b/src/services/agentInstance/index.ts @@ -1513,6 +1513,7 @@ Result: ${JSON.stringify(approvalPrompt)} } private cleanupWorkerConversation(agentId: string): void { + const workerConversationId = this.workerConversationByAgentId.get(agentId); const cleanup = this.workerConversationCleanupByAgentId.get(agentId); if (cleanup) { try { @@ -1522,6 +1523,9 @@ Result: ${JSON.stringify(approvalPrompt)} } this.workerConversationCleanupByAgentId.delete(agentId); } + if (workerConversationId) { + this.workerAgentIdByConversationId.delete(workerConversationId); + } } private async cancelWorkerConversation(agentId: string): Promise { @@ -2055,6 +2059,7 @@ Result: ${JSON.stringify(approvalPrompt)} } this.workerConversationCleanupByAgentId.delete(agentId); } + this.workerAgentIdByConversationId.clear(); this.workerConversationByAgentId.clear(); if (this.memeLoopNativeWorker) {