import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { AgentLifecycleService } from "./agent-lifecycle.service";
import { AgentSpawnerService } from "./agent-spawner.service";
import { ValkeyService } from "../valkey/valkey.service";
import type { AgentState } from "../valkey/types";

describe("AgentLifecycleService", () => {
  let service: AgentLifecycleService;
  let mockValkeyService: {
    getAgentState: ReturnType<typeof vi.fn>;
    setAgentState: ReturnType<typeof vi.fn>;
    updateAgentStatus: ReturnType<typeof vi.fn>;
    publishEvent: ReturnType<typeof vi.fn>;
    listAgents: ReturnType<typeof vi.fn>;
  };
  let mockSpawnerService: {
    scheduleSessionCleanup: ReturnType<typeof vi.fn>;
  };

  const mockAgentId = "test-agent-123";
  const mockTaskId = "test-task-456";

  beforeEach(() => {
    // Create mocks
    mockValkeyService = {
      getAgentState: vi.fn(),
      setAgentState: vi.fn(),
      updateAgentStatus: vi.fn(),
      publishEvent: vi.fn(),
      listAgents: vi.fn(),
    };

    mockSpawnerService = {
      scheduleSessionCleanup: vi.fn(),
    };

    // Create service with mocks
    service = new AgentLifecycleService(
      mockValkeyService as unknown as ValkeyService,
      mockSpawnerService as unknown as AgentSpawnerService
    );
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe("transitionToRunning", () => {
    it("should transition from spawning to running", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "running",
        startedAt: "2026-02-02T10:00:00Z",
      });

      const result = await service.transitionToRunning(mockAgentId);

      expect(result.status).toBe("running");
      expect(result.startedAt).toBeDefined();
      expect(mockValkeyService.updateAgentStatus).toHaveBeenCalledWith(
        mockAgentId,
        "running",
        undefined
      );
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.running",
          agentId: mockAgentId,
          taskId: mockTaskId,
        })
      );
    });

    it("should throw error if agent not found", async () => {
      mockValkeyService.getAgentState.mockResolvedValue(null);

      await expect(service.transitionToRunning(mockAgentId)).rejects.toThrow(
        `Agent ${mockAgentId} not found`
      );
    });

    it("should throw error for invalid transition from running", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      await expect(service.transitionToRunning(mockAgentId)).rejects.toThrow(
        "Invalid state transition from running to running"
      );
    });

    it("should throw error for invalid transition from completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "completed",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      await expect(service.transitionToRunning(mockAgentId)).rejects.toThrow(
        "Invalid state transition from completed to running"
      );
    });
  });

  describe("transitionToCompleted", () => {
    it("should transition from running to completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "completed",
        completedAt: expect.any(String),
      });

      const result = await service.transitionToCompleted(mockAgentId);

      expect(result.status).toBe("completed");
      expect(result.completedAt).toBeDefined();
      expect(mockValkeyService.updateAgentStatus).toHaveBeenCalledWith(
        mockAgentId,
        "completed",
        undefined
      );
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.completed",
          agentId: mockAgentId,
          taskId: mockTaskId,
        })
      );
    });

    it("should throw error if agent not found", async () => {
      mockValkeyService.getAgentState.mockResolvedValue(null);

      await expect(service.transitionToCompleted(mockAgentId)).rejects.toThrow(
        `Agent ${mockAgentId} not found`
      );
    });

    it("should throw error for invalid transition from spawning", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      await expect(service.transitionToCompleted(mockAgentId)).rejects.toThrow(
        "Invalid state transition from spawning to completed"
      );
    });
  });

  describe("transitionToFailed", () => {
    it("should transition from spawning to failed with error", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };
      const errorMessage = "Failed to spawn agent";

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "failed",
        error: errorMessage,
        completedAt: expect.any(String),
      });

      const result = await service.transitionToFailed(mockAgentId, errorMessage);

      expect(result.status).toBe("failed");
      expect(result.error).toBe(errorMessage);
      expect(result.completedAt).toBeDefined();
      expect(mockValkeyService.updateAgentStatus).toHaveBeenCalledWith(
        mockAgentId,
        "failed",
        errorMessage
      );
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.failed",
          agentId: mockAgentId,
          taskId: mockTaskId,
          error: errorMessage,
        })
      );
    });

    it("should transition from running to failed with error", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };
      const errorMessage = "Runtime error occurred";

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "failed",
        error: errorMessage,
        completedAt: expect.any(String),
      });

      const result = await service.transitionToFailed(mockAgentId, errorMessage);

      expect(result.status).toBe("failed");
      expect(result.error).toBe(errorMessage);
    });

    it("should throw error if agent not found", async () => {
      mockValkeyService.getAgentState.mockResolvedValue(null);

      await expect(service.transitionToFailed(mockAgentId, "Error")).rejects.toThrow(
        `Agent ${mockAgentId} not found`
      );
    });

    it("should throw error for invalid transition from completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "completed",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      await expect(service.transitionToFailed(mockAgentId, "Error")).rejects.toThrow(
        "Invalid state transition from completed to failed"
      );
    });
  });

  describe("transitionToKilled", () => {
    it("should transition from spawning to killed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "killed",
        completedAt: expect.any(String),
      });

      const result = await service.transitionToKilled(mockAgentId);

      expect(result.status).toBe("killed");
      expect(result.completedAt).toBeDefined();
      expect(mockValkeyService.updateAgentStatus).toHaveBeenCalledWith(
        mockAgentId,
        "killed",
        undefined
      );
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.killed",
          agentId: mockAgentId,
          taskId: mockTaskId,
        })
      );
    });

    it("should transition from running to killed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "killed",
        completedAt: expect.any(String),
      });

      const result = await service.transitionToKilled(mockAgentId);

      expect(result.status).toBe("killed");
    });

    it("should throw error if agent not found", async () => {
      mockValkeyService.getAgentState.mockResolvedValue(null);

      await expect(service.transitionToKilled(mockAgentId)).rejects.toThrow(
        `Agent ${mockAgentId} not found`
      );
    });

    it("should throw error for invalid transition from completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "completed",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      await expect(service.transitionToKilled(mockAgentId)).rejects.toThrow(
        "Invalid state transition from completed to killed"
      );
    });
  });

  describe("getAgentLifecycleState", () => {
    it("should return agent state from Valkey", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      const result = await service.getAgentLifecycleState(mockAgentId);

      expect(result).toEqual(mockState);
      expect(mockValkeyService.getAgentState).toHaveBeenCalledWith(mockAgentId);
    });

    it("should return null if agent not found", async () => {
      mockValkeyService.getAgentState.mockResolvedValue(null);

      const result = await service.getAgentLifecycleState(mockAgentId);

      expect(result).toBeNull();
    });
  });

  describe("listAgentLifecycleStates", () => {
    it("should return all agent states from Valkey", async () => {
      const mockStates: AgentState[] = [
        {
          agentId: "agent-1",
          status: "running",
          taskId: "task-1",
          startedAt: "2026-02-02T10:00:00Z",
        },
        {
          agentId: "agent-2",
          status: "completed",
          taskId: "task-2",
          startedAt: "2026-02-02T09:00:00Z",
          completedAt: "2026-02-02T10:00:00Z",
        },
      ];

      mockValkeyService.listAgents.mockResolvedValue(mockStates);

      const result = await service.listAgentLifecycleStates();

      expect(result).toEqual(mockStates);
      expect(mockValkeyService.listAgents).toHaveBeenCalled();
    });

    it("should return empty array if no agents", async () => {
      mockValkeyService.listAgents.mockResolvedValue([]);

      const result = await service.listAgentLifecycleStates();

      expect(result).toEqual([]);
    });
  });

  describe("state persistence", () => {
    it("should update completedAt timestamp on terminal states", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);

      let capturedState: AgentState | undefined;
      mockValkeyService.updateAgentStatus.mockImplementation(async (agentId, status, error) => {
        capturedState = {
          ...mockState,
          status,
          error,
          completedAt: new Date().toISOString(),
        };
        return capturedState;
      });

      await service.transitionToCompleted(mockAgentId);

      expect(capturedState?.completedAt).toBeDefined();
    });

    it("should preserve startedAt timestamp through transitions", async () => {
      const startedAt = "2026-02-02T10:00:00Z";
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "completed",
        completedAt: "2026-02-02T11:00:00Z",
      });

      const result = await service.transitionToCompleted(mockAgentId);

      expect(result.startedAt).toBe(startedAt);
    });

    it("should set startedAt if not already set when transitioning to running", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "running",
        // No startedAt in response
      });
      mockValkeyService.setAgentState.mockResolvedValue(undefined);

      await service.transitionToRunning(mockAgentId);

      expect(mockValkeyService.setAgentState).toHaveBeenCalledWith(
        expect.objectContaining({
          agentId: mockAgentId,
          status: "running",
          startedAt: expect.any(String),
        })
      );
    });

    it("should not set startedAt if already present in response", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "running",
        startedAt: "2026-02-02T10:00:00Z",
      });

      await service.transitionToRunning(mockAgentId);

      // Should not call setAgentState since startedAt is already present
      expect(mockValkeyService.setAgentState).not.toHaveBeenCalled();
    });

    it("should set completedAt if not already set when transitioning to completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "completed",
        // No completedAt in response
      });
      mockValkeyService.setAgentState.mockResolvedValue(undefined);

      await service.transitionToCompleted(mockAgentId);

      expect(mockValkeyService.setAgentState).toHaveBeenCalledWith(
        expect.objectContaining({
          agentId: mockAgentId,
          status: "completed",
          completedAt: expect.any(String),
        })
      );
    });

    it("should set completedAt if not already set when transitioning to failed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "failed",
        error: "Test error",
        // No completedAt in response
      });
      mockValkeyService.setAgentState.mockResolvedValue(undefined);

      await service.transitionToFailed(mockAgentId, "Test error");

      expect(mockValkeyService.setAgentState).toHaveBeenCalledWith(
        expect.objectContaining({
          agentId: mockAgentId,
          status: "failed",
          completedAt: expect.any(String),
        })
      );
    });

    it("should set completedAt if not already set when transitioning to killed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "killed",
        // No completedAt in response
      });
      mockValkeyService.setAgentState.mockResolvedValue(undefined);

      await service.transitionToKilled(mockAgentId);

      expect(mockValkeyService.setAgentState).toHaveBeenCalledWith(
        expect.objectContaining({
          agentId: mockAgentId,
          status: "killed",
          completedAt: expect.any(String),
        })
      );
    });
  });

  describe("event emission", () => {
    it("should emit events with correct structure", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "running",
        startedAt: "2026-02-02T10:00:00Z",
      });

      await service.transitionToRunning(mockAgentId);

      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.running",
          agentId: mockAgentId,
          taskId: mockTaskId,
          timestamp: expect.any(String),
        })
      );
    });

    it("should include error in failed event", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
      };
      const errorMessage = "Test error";

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "failed",
        error: errorMessage,
      });

      await service.transitionToFailed(mockAgentId, errorMessage);

      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.failed",
          agentId: mockAgentId,
          taskId: mockTaskId,
          error: errorMessage,
        })
      );
    });
  });

  describe("session cleanup on terminal states", () => {
    it("should schedule session cleanup when transitioning to completed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "completed",
        completedAt: "2026-02-02T11:00:00Z",
      });

      await service.transitionToCompleted(mockAgentId);

      expect(mockSpawnerService.scheduleSessionCleanup).toHaveBeenCalledWith(mockAgentId);
    });

    it("should schedule session cleanup when transitioning to failed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };
      const errorMessage = "Runtime error occurred";

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "failed",
        error: errorMessage,
        completedAt: "2026-02-02T11:00:00Z",
      });

      await service.transitionToFailed(mockAgentId, errorMessage);

      expect(mockSpawnerService.scheduleSessionCleanup).toHaveBeenCalledWith(mockAgentId);
    });

    it("should schedule session cleanup when transitioning to killed", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "killed",
        completedAt: "2026-02-02T11:00:00Z",
      });

      await service.transitionToKilled(mockAgentId);

      expect(mockSpawnerService.scheduleSessionCleanup).toHaveBeenCalledWith(mockAgentId);
    });

    it("should not schedule session cleanup when transitioning to running", async () => {
      const mockState: AgentState = {
        agentId: mockAgentId,
        status: "spawning",
        taskId: mockTaskId,
      };

      mockValkeyService.getAgentState.mockResolvedValue(mockState);
      mockValkeyService.updateAgentStatus.mockResolvedValue({
        ...mockState,
        status: "running",
        startedAt: "2026-02-02T10:00:00Z",
      });

      await service.transitionToRunning(mockAgentId);

      expect(mockSpawnerService.scheduleSessionCleanup).not.toHaveBeenCalled();
    });
  });

  describe("TOCTOU race prevention (CQ-ORCH-5)", () => {
    it("should serialize concurrent transitions to the same agent", async () => {
      const executionOrder: string[] = [];

      // Simulate state that changes after first transition completes
      let currentStatus: "spawning" | "running" | "completed" = "spawning";

      mockValkeyService.getAgentState.mockImplementation(async () => {
        return {
          agentId: mockAgentId,
          status: currentStatus,
          taskId: mockTaskId,
        } as AgentState;
      });

      mockValkeyService.updateAgentStatus.mockImplementation(
        async (_agentId: string, status: string) => {
          // Simulate delay to allow interleaving if lock is broken
          await new Promise<void>((resolve) => {
            setTimeout(resolve, 10);
          });
          currentStatus = status as "spawning" | "running" | "completed";
          executionOrder.push(`updated-to-${status}`);
          return {
            agentId: mockAgentId,
            status,
            taskId: mockTaskId,
            startedAt: "2026-02-02T10:00:00Z",
            ...(status === "completed" && { completedAt: "2026-02-02T11:00:00Z" }),
          } as AgentState;
        }
      );

      // Launch both transitions concurrently
      const [result1, result2] = await Promise.allSettled([
        service.transitionToRunning(mockAgentId),
        service.transitionToCompleted(mockAgentId),
      ]);

      // First should succeed (spawning -> running)
      expect(result1.status).toBe("fulfilled");

      // Second should also succeed (running -> completed) because the lock
      // serializes them: first one completes, updates state to running,
      // then second reads the updated state and transitions to completed
      expect(result2.status).toBe("fulfilled");

      // Verify they executed in order, not interleaved
      expect(executionOrder).toEqual(["updated-to-running", "updated-to-completed"]);
    });

    it("should reject second concurrent transition if first makes it invalid", async () => {
      let currentStatus: "running" | "completed" | "killed" = "running";

      mockValkeyService.getAgentState.mockImplementation(async () => {
        return {
          agentId: mockAgentId,
          status: currentStatus,
          taskId: mockTaskId,
          startedAt: "2026-02-02T10:00:00Z",
        } as AgentState;
      });

      mockValkeyService.updateAgentStatus.mockImplementation(
        async (_agentId: string, status: string) => {
          await new Promise<void>((resolve) => {
            setTimeout(resolve, 10);
          });
          currentStatus = status as "running" | "completed" | "killed";
          return {
            agentId: mockAgentId,
            status,
            taskId: mockTaskId,
            startedAt: "2026-02-02T10:00:00Z",
            completedAt: "2026-02-02T11:00:00Z",
          } as AgentState;
        }
      );

      // Both try to transition from running to a terminal state concurrently
      const [result1, result2] = await Promise.allSettled([
        service.transitionToCompleted(mockAgentId),
        service.transitionToKilled(mockAgentId),
      ]);

      // First should succeed (running -> completed)
      expect(result1.status).toBe("fulfilled");

      // Second should fail because after first completes,
      // agent is in "completed" state which cannot transition to "killed"
      expect(result2.status).toBe("rejected");
      if (result2.status === "rejected") {
        expect(result2.reason).toBeInstanceOf(Error);
        expect((result2.reason as Error).message).toContain("Invalid state transition");
      }
    });

    it("should allow concurrent transitions to different agents", async () => {
      const agent1Id = "agent-1";
      const agent2Id = "agent-2";
      const executionOrder: string[] = [];

      mockValkeyService.getAgentState.mockImplementation(async (agentId: string) => {
        return {
          agentId,
          status: "spawning",
          taskId: `task-for-${agentId}`,
        } as AgentState;
      });

      mockValkeyService.updateAgentStatus.mockImplementation(
        async (agentId: string, status: string) => {
          executionOrder.push(`${agentId}-start`);
          await new Promise<void>((resolve) => {
            setTimeout(resolve, 10);
          });
          executionOrder.push(`${agentId}-end`);
          return {
            agentId,
            status,
            taskId: `task-for-${agentId}`,
            startedAt: "2026-02-02T10:00:00Z",
          } as AgentState;
        }
      );

      // Both should run concurrently since they target different agents
      const [result1, result2] = await Promise.allSettled([
        service.transitionToRunning(agent1Id),
        service.transitionToRunning(agent2Id),
      ]);

      expect(result1.status).toBe("fulfilled");
      expect(result2.status).toBe("fulfilled");

      // Both should start before either finishes (concurrent, not serialized)
      // The execution order should show interleaving
      expect(executionOrder).toContain("agent-1-start");
      expect(executionOrder).toContain("agent-2-start");
    });

    it("should release lock even when transition throws an error", async () => {
      let callCount = 0;

      mockValkeyService.getAgentState.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          // First call: throw error
          return null;
        }
        // Second call: return valid state
        return {
          agentId: mockAgentId,
          status: "spawning",
          taskId: mockTaskId,
        } as AgentState;
      });

      mockValkeyService.updateAgentStatus.mockResolvedValue({
        agentId: mockAgentId,
        status: "running",
        taskId: mockTaskId,
        startedAt: "2026-02-02T10:00:00Z",
      });

      // First transition should fail (agent not found)
      await expect(service.transitionToRunning(mockAgentId)).rejects.toThrow(
        `Agent ${mockAgentId} not found`
      );

      // Second transition should succeed (lock was released despite error)
      const result = await service.transitionToRunning(mockAgentId);
      expect(result.status).toBe("running");
    });

    it("should handle three concurrent transitions sequentially for same agent", async () => {
      const executionOrder: string[] = [];
      let currentStatus: "spawning" | "running" | "completed" | "failed" = "spawning";

      mockValkeyService.getAgentState.mockImplementation(async () => {
        return {
          agentId: mockAgentId,
          status: currentStatus,
          taskId: mockTaskId,
          ...(currentStatus !== "spawning" && { startedAt: "2026-02-02T10:00:00Z" }),
        } as AgentState;
      });

      mockValkeyService.updateAgentStatus.mockImplementation(
        async (_agentId: string, status: string) => {
          executionOrder.push(`update-${status}`);
          await new Promise<void>((resolve) => {
            setTimeout(resolve, 5);
          });
          currentStatus = status as "spawning" | "running" | "completed" | "failed";
          return {
            agentId: mockAgentId,
            status,
            taskId: mockTaskId,
            startedAt: "2026-02-02T10:00:00Z",
            ...(["completed", "failed"].includes(status) && {
              completedAt: "2026-02-02T11:00:00Z",
            }),
          } as AgentState;
        }
      );

      // Launch three transitions at once: spawning->running->completed, plus a failed attempt
      const [r1, r2, r3] = await Promise.allSettled([
        service.transitionToRunning(mockAgentId),
        service.transitionToCompleted(mockAgentId),
        service.transitionToFailed(mockAgentId, "late error"),
      ]);

      // First: spawning -> running (succeeds)
      expect(r1.status).toBe("fulfilled");
      // Second: running -> completed (succeeds, serialized after first)
      expect(r2.status).toBe("fulfilled");
      // Third: completed -> failed (fails, completed is terminal)
      expect(r3.status).toBe("rejected");

      // Verify sequential execution
      expect(executionOrder[0]).toBe("update-running");
      expect(executionOrder[1]).toBe("update-completed");
      // Third never gets to update because validation fails
      expect(executionOrder).toHaveLength(2);
    });
  });
});