2b356f6ca2
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Add per-agent mutex using promise chaining to serialize state transitions for the same agent. This prevents the Time-of-Check-Time-of-Use race condition where two concurrent requests could both read the current state, both validate it as valid for transition, and both write, causing one to overwrite the other's transition. The mutex uses a Map<string, Promise<void>> with promise chaining so that: - Concurrent transitions to the same agent are queued and executed sequentially - Different agents can still transition concurrently without contention - The lock is always released even if the transition throws an error Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
939 lines
30 KiB
TypeScript
939 lines
30 KiB
TypeScript
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);
|
|
});
|
|
});
|
|
});
|