Codewalkers/apps/server/agent/mutex-completion.test.ts

/**
 * Focused test for completion handler mutex functionality.
 * Tests the race condition fix without complex mocking.
 */

import { describe, it, beforeEach, expect } from 'vitest';
import { OutputHandler } from './output-handler.js';
import type { AgentRepository } from '../db/repositories/agent-repository.js';

describe('OutputHandler completion mutex', () => {
  let outputHandler: OutputHandler;
  let completionCallCount: number;
  let callOrder: string[];

  // Default agent for update return value
  const defaultAgent = {
    id: 'test-agent',
    name: 'test-agent',
    taskId: null,
    provider: 'claude',
    mode: 'execute' as const,
    status: 'idle' as const,
    worktreeId: 'test-worktree',
    outputFilePath: null,
    sessionId: null,
    result: null,
    pendingQuestions: null,
    initiativeId: null,
    accountId: null,
    userDismissedAt: null,
    pid: null,
    exitCode: null,
    createdAt: new Date(),
    updatedAt: new Date(),
  };

  // Simple mock that tracks completion attempts
  const mockRepository: AgentRepository = {
    async findById() {
      return null; // Return null to cause early exit after mutex check
    },
    async update(_id: string, data: any) { return { ...defaultAgent, ...data }; },
    async create() { throw new Error('Not implemented'); },
    async findAll() { throw new Error('Not implemented'); },
    async findByStatus() { throw new Error('Not implemented'); },
    async findByTaskId() { throw new Error('Not implemented'); },
    async findByName() { throw new Error('Not implemented'); },
    async findBySessionId() { throw new Error('Not implemented'); },
    async delete() { throw new Error('Not implemented'); },
    async findWaitingWithContext() { throw new Error('Not implemented'); }
  };

  beforeEach(() => {
    outputHandler = new OutputHandler(mockRepository);
    completionCallCount = 0;
    callOrder = [];
  });

  it('should prevent concurrent completion handling with mutex', async () => {
    const agentId = 'test-agent';

    // Mock the findById method to track calls and simulate processing time
    let firstCallCompleted = false;
    (mockRepository as any).findById = async (id: string) => {
      completionCallCount++;
      const callIndex = completionCallCount;
      callOrder.push(`call-${callIndex}-start`);

      if (callIndex === 1) {
        // First call - simulate some processing time
        await new Promise(resolve => setTimeout(resolve, 50));
        firstCallCompleted = true;
      }

      callOrder.push(`call-${callIndex}-end`);
      return null; // Return null to exit early
    };

    // Start two concurrent completion handlers
    const getAgentWorkdir = () => '/test/workdir';
    const completion1Promise = outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir);
    const completion2Promise = outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir);

    await Promise.all([completion1Promise, completion2Promise]);

    // Verify only one completion handler executed
    expect(completionCallCount, 'Should only execute one completion handler').toBe(1);
    expect(firstCallCompleted, 'First handler should have completed').toBe(true);
    expect(callOrder).toEqual(['call-1-start', 'call-1-end']);
  });

  it('should allow sequential completion handling after first completes', async () => {
    const agentId = 'test-agent';

    // Mock findById to track calls
    (mockRepository as any).findById = async (id: string) => {
      completionCallCount++;
      callOrder.push(`call-${completionCallCount}`);
      return null; // Return null to exit early
    };

    const getAgentWorkdir = () => '/test/workdir';

    // First completion
    await outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir);

    // Second completion (after first is done)
    await outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir);

    // Both should execute sequentially
    expect(completionCallCount, 'Should execute both handlers sequentially').toBe(2);
    expect(callOrder).toEqual(['call-1', 'call-2']);
  });

  it('should clean up mutex lock even when exception is thrown', async () => {
    const agentId = 'test-agent';

    let firstCallMadeThrowCall = false;
    let secondCallCompleted = false;

    // First call throws an error
    (mockRepository as any).findById = async (id: string) => {
      if (!firstCallMadeThrowCall) {
        firstCallMadeThrowCall = true;
        throw new Error('Database error');
      } else {
        secondCallCompleted = true;
        return null;
      }
    };

    const getAgentWorkdir = () => '/test/workdir';

    // First call should throw but clean up mutex
    await expect(outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir))
      .rejects.toThrow('Database error');

    expect(firstCallMadeThrowCall, 'First call should have thrown').toBe(true);

    // Second call should succeed (proving mutex was cleaned up)
    await outputHandler.handleCompletion(agentId, undefined, getAgentWorkdir);
    expect(secondCallCompleted, 'Second call should have completed').toBe(true);
  });

  it('should use agent ID as mutex key', async () => {
    const agentId1 = 'agent-1';
    const agentId2 = 'agent-2';

    // Both agents can process concurrently since they have different IDs
    let agent1Started = false;
    let agent2Started = false;

    (mockRepository as any).findById = async (id: string) => {
      if (id === agentId1) {
        agent1Started = true;
        await new Promise(resolve => setTimeout(resolve, 30));
      } else if (id === agentId2) {
        agent2Started = true;
        await new Promise(resolve => setTimeout(resolve, 30));
      }
      return null;
    };

    const getAgentWorkdir = () => '/test/workdir';

    // Start both agents concurrently - they should NOT block each other
    const agent1Promise = outputHandler.handleCompletion(agentId1, undefined, getAgentWorkdir);
    const agent2Promise = outputHandler.handleCompletion(agentId2, undefined, getAgentWorkdir);

    await Promise.all([agent1Promise, agent2Promise]);

    expect(agent1Started, 'Agent 1 should have started').toBe(true);
    expect(agent2Started, 'Agent 2 should have started').toBe(true);
  });
});