feat(#71): implement graph data API
Implemented three new API endpoints for knowledge graph visualization: 1. GET /api/knowledge/graph - Full knowledge graph - Returns all entries and links with optional filtering - Supports filtering by tags, status, and node count limit - Includes orphan detection (entries with no links) 2. GET /api/knowledge/graph/stats - Graph statistics - Total entries and links counts - Orphan entries detection - Average links per entry - Top 10 most connected entries - Tag distribution across entries 3. GET /api/knowledge/graph/:slug - Entry-centered subgraph - Returns graph centered on specific entry - Supports depth parameter (1-5) for traversal distance - Includes all connected nodes up to specified depth New Files: - apps/api/src/knowledge/graph.controller.ts - apps/api/src/knowledge/graph.controller.spec.ts Modified Files: - apps/api/src/knowledge/dto/graph-query.dto.ts (added GraphFilterDto) - apps/api/src/knowledge/entities/graph.entity.ts (extended with new types) - apps/api/src/knowledge/services/graph.service.ts (added new methods) - apps/api/src/knowledge/services/graph.service.spec.ts (added tests) - apps/api/src/knowledge/knowledge.module.ts (registered controller) - apps/api/src/knowledge/dto/index.ts (exported new DTOs) - docs/scratchpads/71-graph-data-api.md (implementation notes) Test Coverage: 21 tests (all passing) - 14 service tests including orphan detection, filtering, statistics - 7 controller tests for all three endpoints Follows TDD principles with tests written before implementation. All code quality gates passed (lint, typecheck, tests). Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
Jason Woltje
615
apps/orchestrator/src/spawner/agent-lifecycle.service.spec.ts
Normal file
615
apps/orchestrator/src/spawner/agent-lifecycle.service.spec.ts
Normal file
@@ -0,0 +1,615 @@
|
||||
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
|
||||
import { AgentLifecycleService } from './agent-lifecycle.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>;
|
||||
};
|
||||
|
||||
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(),
|
||||
};
|
||||
|
||||
// Create service with mock
|
||||
service = new AgentLifecycleService(mockValkeyService as any);
|
||||
});
|
||||
|
||||
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,
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
232
apps/orchestrator/src/spawner/agent-lifecycle.service.ts
Normal file
232
apps/orchestrator/src/spawner/agent-lifecycle.service.ts
Normal file
@@ -0,0 +1,232 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { ValkeyService } from '../valkey/valkey.service';
|
||||
import type { AgentState, AgentStatus, AgentEvent } from '../valkey/types';
|
||||
import { isValidAgentTransition } from '../valkey/types/state.types';
|
||||
|
||||
/**
|
||||
* Service responsible for managing agent lifecycle state transitions
|
||||
*
|
||||
* Manages state transitions through the agent lifecycle:
|
||||
* spawning → running → completed/failed/killed
|
||||
*
|
||||
* - Enforces valid state transitions using state machine
|
||||
* - Persists agent state changes to Valkey
|
||||
* - Emits pub/sub events on state changes
|
||||
* - Tracks agent metadata (startedAt, completedAt, error)
|
||||
*/
|
||||
@Injectable()
|
||||
export class AgentLifecycleService {
|
||||
private readonly logger = new Logger(AgentLifecycleService.name);
|
||||
|
||||
constructor(private readonly valkeyService: ValkeyService) {
|
||||
this.logger.log('AgentLifecycleService initialized');
|
||||
}
|
||||
|
||||
/**
|
||||
* Transition agent from spawning to running state
|
||||
* @param agentId Unique agent identifier
|
||||
* @returns Updated agent state
|
||||
* @throws Error if agent not found or invalid transition
|
||||
*/
|
||||
async transitionToRunning(agentId: string): Promise<AgentState> {
|
||||
this.logger.log(`Transitioning agent ${agentId} to running`);
|
||||
|
||||
const currentState = await this.getAgentState(agentId);
|
||||
this.validateTransition(currentState.status, 'running');
|
||||
|
||||
// Set startedAt timestamp if not already set
|
||||
const startedAt = currentState.startedAt || new Date().toISOString();
|
||||
|
||||
// Update state in Valkey
|
||||
const updatedState = await this.valkeyService.updateAgentStatus(
|
||||
agentId,
|
||||
'running',
|
||||
undefined,
|
||||
);
|
||||
|
||||
// Ensure startedAt is set
|
||||
if (!updatedState.startedAt) {
|
||||
updatedState.startedAt = startedAt;
|
||||
await this.valkeyService.setAgentState(updatedState);
|
||||
}
|
||||
|
||||
// Emit event
|
||||
await this.publishStateChangeEvent('agent.running', updatedState);
|
||||
|
||||
this.logger.log(`Agent ${agentId} transitioned to running`);
|
||||
return updatedState;
|
||||
}
|
||||
|
||||
/**
|
||||
* Transition agent to completed state
|
||||
* @param agentId Unique agent identifier
|
||||
* @returns Updated agent state
|
||||
* @throws Error if agent not found or invalid transition
|
||||
*/
|
||||
async transitionToCompleted(agentId: string): Promise<AgentState> {
|
||||
this.logger.log(`Transitioning agent ${agentId} to completed`);
|
||||
|
||||
const currentState = await this.getAgentState(agentId);
|
||||
this.validateTransition(currentState.status, 'completed');
|
||||
|
||||
// Set completedAt timestamp
|
||||
const completedAt = new Date().toISOString();
|
||||
|
||||
// Update state in Valkey
|
||||
const updatedState = await this.valkeyService.updateAgentStatus(
|
||||
agentId,
|
||||
'completed',
|
||||
undefined,
|
||||
);
|
||||
|
||||
// Ensure completedAt is set
|
||||
if (!updatedState.completedAt) {
|
||||
updatedState.completedAt = completedAt;
|
||||
await this.valkeyService.setAgentState(updatedState);
|
||||
}
|
||||
|
||||
// Emit event
|
||||
await this.publishStateChangeEvent('agent.completed', updatedState);
|
||||
|
||||
this.logger.log(`Agent ${agentId} transitioned to completed`);
|
||||
return updatedState;
|
||||
}
|
||||
|
||||
/**
|
||||
* Transition agent to failed state with error
|
||||
* @param agentId Unique agent identifier
|
||||
* @param error Error message
|
||||
* @returns Updated agent state
|
||||
* @throws Error if agent not found or invalid transition
|
||||
*/
|
||||
async transitionToFailed(agentId: string, error: string): Promise<AgentState> {
|
||||
this.logger.log(`Transitioning agent ${agentId} to failed: ${error}`);
|
||||
|
||||
const currentState = await this.getAgentState(agentId);
|
||||
this.validateTransition(currentState.status, 'failed');
|
||||
|
||||
// Set completedAt timestamp
|
||||
const completedAt = new Date().toISOString();
|
||||
|
||||
// Update state in Valkey
|
||||
const updatedState = await this.valkeyService.updateAgentStatus(
|
||||
agentId,
|
||||
'failed',
|
||||
error,
|
||||
);
|
||||
|
||||
// Ensure completedAt is set
|
||||
if (!updatedState.completedAt) {
|
||||
updatedState.completedAt = completedAt;
|
||||
await this.valkeyService.setAgentState(updatedState);
|
||||
}
|
||||
|
||||
// Emit event
|
||||
await this.publishStateChangeEvent('agent.failed', updatedState, error);
|
||||
|
||||
this.logger.error(`Agent ${agentId} transitioned to failed: ${error}`);
|
||||
return updatedState;
|
||||
}
|
||||
|
||||
/**
|
||||
* Transition agent to killed state
|
||||
* @param agentId Unique agent identifier
|
||||
* @returns Updated agent state
|
||||
* @throws Error if agent not found or invalid transition
|
||||
*/
|
||||
async transitionToKilled(agentId: string): Promise<AgentState> {
|
||||
this.logger.log(`Transitioning agent ${agentId} to killed`);
|
||||
|
||||
const currentState = await this.getAgentState(agentId);
|
||||
this.validateTransition(currentState.status, 'killed');
|
||||
|
||||
// Set completedAt timestamp
|
||||
const completedAt = new Date().toISOString();
|
||||
|
||||
// Update state in Valkey
|
||||
const updatedState = await this.valkeyService.updateAgentStatus(
|
||||
agentId,
|
||||
'killed',
|
||||
undefined,
|
||||
);
|
||||
|
||||
// Ensure completedAt is set
|
||||
if (!updatedState.completedAt) {
|
||||
updatedState.completedAt = completedAt;
|
||||
await this.valkeyService.setAgentState(updatedState);
|
||||
}
|
||||
|
||||
// Emit event
|
||||
await this.publishStateChangeEvent('agent.killed', updatedState);
|
||||
|
||||
this.logger.warn(`Agent ${agentId} transitioned to killed`);
|
||||
return updatedState;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get current agent lifecycle state
|
||||
* @param agentId Unique agent identifier
|
||||
* @returns Agent state or null if not found
|
||||
*/
|
||||
async getAgentLifecycleState(agentId: string): Promise<AgentState | null> {
|
||||
return this.valkeyService.getAgentState(agentId);
|
||||
}
|
||||
|
||||
/**
|
||||
* List all agent lifecycle states
|
||||
* @returns Array of all agent states
|
||||
*/
|
||||
async listAgentLifecycleStates(): Promise<AgentState[]> {
|
||||
return this.valkeyService.listAgents();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get agent state and throw if not found
|
||||
* @param agentId Unique agent identifier
|
||||
* @returns Agent state
|
||||
* @throws Error if agent not found
|
||||
*/
|
||||
private async getAgentState(agentId: string): Promise<AgentState> {
|
||||
const state = await this.valkeyService.getAgentState(agentId);
|
||||
|
||||
if (!state) {
|
||||
throw new Error(`Agent ${agentId} not found`);
|
||||
}
|
||||
|
||||
return state;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate state transition is allowed
|
||||
* @param from Current state
|
||||
* @param to Target state
|
||||
* @throws Error if transition is invalid
|
||||
*/
|
||||
private validateTransition(from: AgentStatus, to: AgentStatus): void {
|
||||
if (!isValidAgentTransition(from, to)) {
|
||||
throw new Error(`Invalid state transition from ${from} to ${to}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Publish state change event
|
||||
* @param eventType Type of event
|
||||
* @param state Updated agent state
|
||||
* @param error Optional error message
|
||||
*/
|
||||
private async publishStateChangeEvent(
|
||||
eventType: 'agent.running' | 'agent.completed' | 'agent.failed' | 'agent.killed',
|
||||
state: AgentState,
|
||||
error?: string,
|
||||
): Promise<void> {
|
||||
const event: AgentEvent = {
|
||||
type: eventType,
|
||||
agentId: state.agentId,
|
||||
taskId: state.taskId,
|
||||
timestamp: new Date().toISOString(),
|
||||
error,
|
||||
};
|
||||
|
||||
await this.valkeyService.publishEvent(event);
|
||||
}
|
||||
}
|
||||
341
apps/orchestrator/src/spawner/docker-sandbox.service.spec.ts
Normal file
341
apps/orchestrator/src/spawner/docker-sandbox.service.spec.ts
Normal file
@@ -0,0 +1,341 @@
|
||||
import { ConfigService } from "@nestjs/config";
|
||||
import { describe, it, expect, beforeEach, vi } from "vitest";
|
||||
import { DockerSandboxService } from "./docker-sandbox.service";
|
||||
import Docker from "dockerode";
|
||||
|
||||
describe("DockerSandboxService", () => {
|
||||
let service: DockerSandboxService;
|
||||
let mockConfigService: ConfigService;
|
||||
let mockDocker: Docker;
|
||||
let mockContainer: Docker.Container;
|
||||
|
||||
beforeEach(() => {
|
||||
// Create mock Docker container
|
||||
mockContainer = {
|
||||
id: "container-123",
|
||||
start: vi.fn().mockResolvedValue(undefined),
|
||||
stop: vi.fn().mockResolvedValue(undefined),
|
||||
remove: vi.fn().mockResolvedValue(undefined),
|
||||
inspect: vi.fn().mockResolvedValue({
|
||||
State: { Status: "running" },
|
||||
}),
|
||||
} as unknown as Docker.Container;
|
||||
|
||||
// Create mock Docker instance
|
||||
mockDocker = {
|
||||
createContainer: vi.fn().mockResolvedValue(mockContainer),
|
||||
getContainer: vi.fn().mockReturnValue(mockContainer),
|
||||
} as unknown as Docker;
|
||||
|
||||
// Create mock ConfigService
|
||||
mockConfigService = {
|
||||
get: vi.fn((key: string, defaultValue?: unknown) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.docker.socketPath": "/var/run/docker.sock",
|
||||
"orchestrator.sandbox.enabled": true,
|
||||
"orchestrator.sandbox.defaultImage": "node:20-alpine",
|
||||
"orchestrator.sandbox.defaultMemoryMB": 512,
|
||||
"orchestrator.sandbox.defaultCpuLimit": 1.0,
|
||||
"orchestrator.sandbox.networkMode": "bridge",
|
||||
};
|
||||
return config[key] !== undefined ? config[key] : defaultValue;
|
||||
}),
|
||||
} as unknown as ConfigService;
|
||||
|
||||
// Create service with mock Docker instance
|
||||
service = new DockerSandboxService(mockConfigService, mockDocker);
|
||||
});
|
||||
|
||||
describe("constructor", () => {
|
||||
it("should be defined", () => {
|
||||
expect(service).toBeDefined();
|
||||
});
|
||||
|
||||
it("should use provided Docker instance", () => {
|
||||
expect(service).toBeDefined();
|
||||
// Service should use the mockDocker instance we provided
|
||||
});
|
||||
});
|
||||
|
||||
describe("createContainer", () => {
|
||||
it("should create a container with default configuration", async () => {
|
||||
const agentId = "agent-123";
|
||||
const taskId = "task-456";
|
||||
const workspacePath = "/workspace/agent-123";
|
||||
|
||||
const result = await service.createContainer(
|
||||
agentId,
|
||||
taskId,
|
||||
workspacePath
|
||||
);
|
||||
|
||||
expect(result.containerId).toBe("container-123");
|
||||
expect(result.agentId).toBe(agentId);
|
||||
expect(result.taskId).toBe(taskId);
|
||||
expect(result.createdAt).toBeInstanceOf(Date);
|
||||
expect(mockDocker.createContainer).toHaveBeenCalledWith({
|
||||
Image: "node:20-alpine",
|
||||
name: expect.stringContaining(`mosaic-agent-${agentId}`),
|
||||
User: "node:node",
|
||||
HostConfig: {
|
||||
Memory: 512 * 1024 * 1024, // 512MB in bytes
|
||||
NanoCpus: 1000000000, // 1.0 CPU
|
||||
NetworkMode: "bridge",
|
||||
Binds: [`${workspacePath}:/workspace`],
|
||||
AutoRemove: false,
|
||||
ReadonlyRootfs: false,
|
||||
},
|
||||
WorkingDir: "/workspace",
|
||||
Env: [`AGENT_ID=${agentId}`, `TASK_ID=${taskId}`],
|
||||
});
|
||||
});
|
||||
|
||||
it("should create a container with custom resource limits", async () => {
|
||||
const agentId = "agent-123";
|
||||
const taskId = "task-456";
|
||||
const workspacePath = "/workspace/agent-123";
|
||||
const options = {
|
||||
memoryMB: 1024,
|
||||
cpuLimit: 2.0,
|
||||
};
|
||||
|
||||
await service.createContainer(agentId, taskId, workspacePath, options);
|
||||
|
||||
expect(mockDocker.createContainer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
HostConfig: expect.objectContaining({
|
||||
Memory: 1024 * 1024 * 1024, // 1024MB in bytes
|
||||
NanoCpus: 2000000000, // 2.0 CPU
|
||||
}),
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should create a container with network isolation", async () => {
|
||||
const agentId = "agent-123";
|
||||
const taskId = "task-456";
|
||||
const workspacePath = "/workspace/agent-123";
|
||||
const options = {
|
||||
networkMode: "none" as const,
|
||||
};
|
||||
|
||||
await service.createContainer(agentId, taskId, workspacePath, options);
|
||||
|
||||
expect(mockDocker.createContainer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
HostConfig: expect.objectContaining({
|
||||
NetworkMode: "none",
|
||||
}),
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should create a container with custom environment variables", async () => {
|
||||
const agentId = "agent-123";
|
||||
const taskId = "task-456";
|
||||
const workspacePath = "/workspace/agent-123";
|
||||
const options = {
|
||||
env: {
|
||||
CUSTOM_VAR: "value123",
|
||||
ANOTHER_VAR: "value456",
|
||||
},
|
||||
};
|
||||
|
||||
await service.createContainer(agentId, taskId, workspacePath, options);
|
||||
|
||||
expect(mockDocker.createContainer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
Env: expect.arrayContaining([
|
||||
`AGENT_ID=${agentId}`,
|
||||
`TASK_ID=${taskId}`,
|
||||
"CUSTOM_VAR=value123",
|
||||
"ANOTHER_VAR=value456",
|
||||
]),
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should throw error if container creation fails", async () => {
|
||||
const agentId = "agent-123";
|
||||
const taskId = "task-456";
|
||||
const workspacePath = "/workspace/agent-123";
|
||||
|
||||
(mockDocker.createContainer as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Docker daemon not available")
|
||||
);
|
||||
|
||||
await expect(
|
||||
service.createContainer(agentId, taskId, workspacePath)
|
||||
).rejects.toThrow("Failed to create container for agent agent-123");
|
||||
});
|
||||
});
|
||||
|
||||
describe("startContainer", () => {
|
||||
it("should start a container by ID", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
await service.startContainer(containerId);
|
||||
|
||||
expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
|
||||
expect(mockContainer.start).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should throw error if container start fails", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.start as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container not found")
|
||||
);
|
||||
|
||||
await expect(service.startContainer(containerId)).rejects.toThrow(
|
||||
"Failed to start container container-123"
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("stopContainer", () => {
|
||||
it("should stop a container by ID", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
await service.stopContainer(containerId);
|
||||
|
||||
expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
|
||||
expect(mockContainer.stop).toHaveBeenCalledWith({ t: 10 });
|
||||
});
|
||||
|
||||
it("should stop a container with custom timeout", async () => {
|
||||
const containerId = "container-123";
|
||||
const timeout = 30;
|
||||
|
||||
await service.stopContainer(containerId, timeout);
|
||||
|
||||
expect(mockContainer.stop).toHaveBeenCalledWith({ t: timeout });
|
||||
});
|
||||
|
||||
it("should throw error if container stop fails", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container already stopped")
|
||||
);
|
||||
|
||||
await expect(service.stopContainer(containerId)).rejects.toThrow(
|
||||
"Failed to stop container container-123"
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("removeContainer", () => {
|
||||
it("should remove a container by ID", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
await service.removeContainer(containerId);
|
||||
|
||||
expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
|
||||
expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
|
||||
});
|
||||
|
||||
it("should throw error if container removal fails", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.remove as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container not found")
|
||||
);
|
||||
|
||||
await expect(service.removeContainer(containerId)).rejects.toThrow(
|
||||
"Failed to remove container container-123"
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getContainerStatus", () => {
|
||||
it("should return container status", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
const status = await service.getContainerStatus(containerId);
|
||||
|
||||
expect(status).toBe("running");
|
||||
expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
|
||||
expect(mockContainer.inspect).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should throw error if container inspect fails", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.inspect as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container not found")
|
||||
);
|
||||
|
||||
await expect(service.getContainerStatus(containerId)).rejects.toThrow(
|
||||
"Failed to get container status for container-123"
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("cleanup", () => {
|
||||
it("should stop and remove container", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
await service.cleanup(containerId);
|
||||
|
||||
expect(mockContainer.stop).toHaveBeenCalledWith({ t: 10 });
|
||||
expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
|
||||
});
|
||||
|
||||
it("should remove container even if stop fails", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container already stopped")
|
||||
);
|
||||
|
||||
await service.cleanup(containerId);
|
||||
|
||||
expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
|
||||
});
|
||||
|
||||
it("should throw error if both stop and remove fail", async () => {
|
||||
const containerId = "container-123";
|
||||
|
||||
(mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container not found")
|
||||
);
|
||||
(mockContainer.remove as ReturnType<typeof vi.fn>).mockRejectedValue(
|
||||
new Error("Container not found")
|
||||
);
|
||||
|
||||
await expect(service.cleanup(containerId)).rejects.toThrow(
|
||||
"Failed to cleanup container container-123"
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("isEnabled", () => {
|
||||
it("should return true if sandbox is enabled in config", () => {
|
||||
expect(service.isEnabled()).toBe(true);
|
||||
});
|
||||
|
||||
it("should return false if sandbox is disabled in config", () => {
|
||||
const disabledConfigService = {
|
||||
get: vi.fn((key: string, defaultValue?: unknown) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.docker.socketPath": "/var/run/docker.sock",
|
||||
"orchestrator.sandbox.enabled": false,
|
||||
"orchestrator.sandbox.defaultImage": "node:20-alpine",
|
||||
"orchestrator.sandbox.defaultMemoryMB": 512,
|
||||
"orchestrator.sandbox.defaultCpuLimit": 1.0,
|
||||
"orchestrator.sandbox.networkMode": "bridge",
|
||||
};
|
||||
return config[key] !== undefined ? config[key] : defaultValue;
|
||||
}),
|
||||
} as unknown as ConfigService;
|
||||
|
||||
const disabledService = new DockerSandboxService(
|
||||
disabledConfigService,
|
||||
mockDocker
|
||||
);
|
||||
|
||||
expect(disabledService.isEnabled()).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
254
apps/orchestrator/src/spawner/docker-sandbox.service.ts
Normal file
254
apps/orchestrator/src/spawner/docker-sandbox.service.ts
Normal file
@@ -0,0 +1,254 @@
|
||||
import { Injectable, Logger } from "@nestjs/common";
|
||||
import { ConfigService } from "@nestjs/config";
|
||||
import Docker from "dockerode";
|
||||
import {
|
||||
DockerSandboxOptions,
|
||||
ContainerCreateResult,
|
||||
} from "./types/docker-sandbox.types";
|
||||
|
||||
/**
|
||||
* Service for managing Docker container isolation for agents
|
||||
* Provides secure sandboxing with resource limits and cleanup
|
||||
*/
|
||||
@Injectable()
|
||||
export class DockerSandboxService {
|
||||
private readonly logger = new Logger(DockerSandboxService.name);
|
||||
private readonly docker: Docker;
|
||||
private readonly sandboxEnabled: boolean;
|
||||
private readonly defaultImage: string;
|
||||
private readonly defaultMemoryMB: number;
|
||||
private readonly defaultCpuLimit: number;
|
||||
private readonly defaultNetworkMode: string;
|
||||
|
||||
constructor(
|
||||
private readonly configService: ConfigService,
|
||||
docker?: Docker
|
||||
) {
|
||||
const socketPath = this.configService.get<string>(
|
||||
"orchestrator.docker.socketPath",
|
||||
"/var/run/docker.sock"
|
||||
);
|
||||
|
||||
this.docker = docker ?? new Docker({ socketPath });
|
||||
|
||||
this.sandboxEnabled = this.configService.get<boolean>(
|
||||
"orchestrator.sandbox.enabled",
|
||||
false
|
||||
);
|
||||
|
||||
this.defaultImage = this.configService.get<string>(
|
||||
"orchestrator.sandbox.defaultImage",
|
||||
"node:20-alpine"
|
||||
);
|
||||
|
||||
this.defaultMemoryMB = this.configService.get<number>(
|
||||
"orchestrator.sandbox.defaultMemoryMB",
|
||||
512
|
||||
);
|
||||
|
||||
this.defaultCpuLimit = this.configService.get<number>(
|
||||
"orchestrator.sandbox.defaultCpuLimit",
|
||||
1.0
|
||||
);
|
||||
|
||||
this.defaultNetworkMode = this.configService.get<string>(
|
||||
"orchestrator.sandbox.networkMode",
|
||||
"bridge"
|
||||
);
|
||||
|
||||
this.logger.log(
|
||||
`DockerSandboxService initialized (enabled: ${this.sandboxEnabled}, socket: ${socketPath})`
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a Docker container for agent isolation
|
||||
* @param agentId Unique agent identifier
|
||||
* @param taskId Task identifier
|
||||
* @param workspacePath Path to workspace directory to mount
|
||||
* @param options Optional container configuration
|
||||
* @returns Container creation result
|
||||
*/
|
||||
async createContainer(
|
||||
agentId: string,
|
||||
taskId: string,
|
||||
workspacePath: string,
|
||||
options?: DockerSandboxOptions
|
||||
): Promise<ContainerCreateResult> {
|
||||
try {
|
||||
const image = options?.image ?? this.defaultImage;
|
||||
const memoryMB = options?.memoryMB ?? this.defaultMemoryMB;
|
||||
const cpuLimit = options?.cpuLimit ?? this.defaultCpuLimit;
|
||||
const networkMode = options?.networkMode ?? this.defaultNetworkMode;
|
||||
|
||||
// Convert memory from MB to bytes
|
||||
const memoryBytes = memoryMB * 1024 * 1024;
|
||||
|
||||
// Convert CPU limit to NanoCPUs (1.0 = 1,000,000,000 nanocpus)
|
||||
const nanoCpus = Math.floor(cpuLimit * 1000000000);
|
||||
|
||||
// Build environment variables
|
||||
const env = [
|
||||
`AGENT_ID=${agentId}`,
|
||||
`TASK_ID=${taskId}`,
|
||||
];
|
||||
|
||||
if (options?.env) {
|
||||
Object.entries(options.env).forEach(([key, value]) => {
|
||||
env.push(`${key}=${value}`);
|
||||
});
|
||||
}
|
||||
|
||||
// Container name with timestamp to ensure uniqueness
|
||||
const containerName = `mosaic-agent-${agentId}-${Date.now()}`;
|
||||
|
||||
this.logger.log(
|
||||
`Creating container for agent ${agentId} (image: ${image}, memory: ${memoryMB}MB, cpu: ${cpuLimit})`
|
||||
);
|
||||
|
||||
const container = await this.docker.createContainer({
|
||||
Image: image,
|
||||
name: containerName,
|
||||
User: "node:node", // Non-root user for security
|
||||
HostConfig: {
|
||||
Memory: memoryBytes,
|
||||
NanoCpus: nanoCpus,
|
||||
NetworkMode: networkMode,
|
||||
Binds: [`${workspacePath}:/workspace`],
|
||||
AutoRemove: false, // Manual cleanup for audit trail
|
||||
ReadonlyRootfs: false, // Allow writes within container
|
||||
},
|
||||
WorkingDir: "/workspace",
|
||||
Env: env,
|
||||
});
|
||||
|
||||
const createdAt = new Date();
|
||||
|
||||
this.logger.log(
|
||||
`Container created successfully: ${container.id} for agent ${agentId}`
|
||||
);
|
||||
|
||||
return {
|
||||
containerId: container.id,
|
||||
agentId,
|
||||
taskId,
|
||||
createdAt,
|
||||
};
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to create container for agent ${agentId}: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to create container for agent ${agentId}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a Docker container
|
||||
* @param containerId Container ID to start
|
||||
*/
|
||||
async startContainer(containerId: string): Promise<void> {
|
||||
try {
|
||||
this.logger.log(`Starting container: ${containerId}`);
|
||||
const container = this.docker.getContainer(containerId);
|
||||
await container.start();
|
||||
this.logger.log(`Container started successfully: ${containerId}`);
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to start container ${containerId}: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to start container ${containerId}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Stop a Docker container
|
||||
* @param containerId Container ID to stop
|
||||
* @param timeout Timeout in seconds (default: 10)
|
||||
*/
|
||||
async stopContainer(containerId: string, timeout = 10): Promise<void> {
|
||||
try {
|
||||
this.logger.log(`Stopping container: ${containerId} (timeout: ${timeout}s)`);
|
||||
const container = this.docker.getContainer(containerId);
|
||||
await container.stop({ t: timeout });
|
||||
this.logger.log(`Container stopped successfully: ${containerId}`);
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to stop container ${containerId}: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to stop container ${containerId}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove a Docker container
|
||||
* @param containerId Container ID to remove
|
||||
*/
|
||||
async removeContainer(containerId: string): Promise<void> {
|
||||
try {
|
||||
this.logger.log(`Removing container: ${containerId}`);
|
||||
const container = this.docker.getContainer(containerId);
|
||||
await container.remove({ force: true });
|
||||
this.logger.log(`Container removed successfully: ${containerId}`);
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to remove container ${containerId}: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to remove container ${containerId}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get container status
|
||||
* @param containerId Container ID to inspect
|
||||
* @returns Container status string
|
||||
*/
|
||||
async getContainerStatus(containerId: string): Promise<string> {
|
||||
try {
|
||||
const container = this.docker.getContainer(containerId);
|
||||
const info = await container.inspect();
|
||||
return info.State.Status;
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to get container status for ${containerId}: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to get container status for ${containerId}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup container (stop and remove)
|
||||
* @param containerId Container ID to cleanup
|
||||
*/
|
||||
async cleanup(containerId: string): Promise<void> {
|
||||
this.logger.log(`Cleaning up container: ${containerId}`);
|
||||
|
||||
try {
|
||||
// Try to stop first
|
||||
await this.stopContainer(containerId);
|
||||
} catch (error) {
|
||||
this.logger.warn(
|
||||
`Failed to stop container ${containerId} during cleanup (may already be stopped): ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
// Always try to remove
|
||||
await this.removeContainer(containerId);
|
||||
} catch (error) {
|
||||
this.logger.error(
|
||||
`Failed to remove container ${containerId} during cleanup: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
throw new Error(`Failed to cleanup container ${containerId}`);
|
||||
}
|
||||
|
||||
this.logger.log(`Container cleanup completed: ${containerId}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if sandbox mode is enabled
|
||||
* @returns True if sandbox is enabled
|
||||
*/
|
||||
isEnabled(): boolean {
|
||||
return this.sandboxEnabled;
|
||||
}
|
||||
}
|
||||
@@ -2,5 +2,8 @@
|
||||
* Spawner module exports
|
||||
*/
|
||||
export { AgentSpawnerService } from "./agent-spawner.service";
|
||||
export { AgentLifecycleService } from "./agent-lifecycle.service";
|
||||
export { DockerSandboxService } from "./docker-sandbox.service";
|
||||
export { SpawnerModule } from "./spawner.module";
|
||||
export * from "./types/agent-spawner.types";
|
||||
export * from "./types/docker-sandbox.types";
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
import { Module } from "@nestjs/common";
|
||||
import { AgentSpawnerService } from "./agent-spawner.service";
|
||||
import { AgentLifecycleService } from "./agent-lifecycle.service";
|
||||
import { DockerSandboxService } from "./docker-sandbox.service";
|
||||
import { ValkeyModule } from "../valkey/valkey.module";
|
||||
|
||||
@Module({
|
||||
providers: [AgentSpawnerService],
|
||||
exports: [AgentSpawnerService],
|
||||
imports: [ValkeyModule],
|
||||
providers: [AgentSpawnerService, AgentLifecycleService, DockerSandboxService],
|
||||
exports: [AgentSpawnerService, AgentLifecycleService, DockerSandboxService],
|
||||
})
|
||||
export class SpawnerModule {}
|
||||
|
||||
@@ -82,4 +82,6 @@ export interface AgentSession {
|
||||
completedAt?: Date;
|
||||
/** Error if failed */
|
||||
error?: string;
|
||||
/** Docker container ID if sandbox is enabled */
|
||||
containerId?: string;
|
||||
}
|
||||
|
||||
46
apps/orchestrator/src/spawner/types/docker-sandbox.types.ts
Normal file
46
apps/orchestrator/src/spawner/types/docker-sandbox.types.ts
Normal file
@@ -0,0 +1,46 @@
|
||||
/**
|
||||
* Network mode options for Docker containers
|
||||
*/
|
||||
export type NetworkMode = "bridge" | "host" | "none";
|
||||
|
||||
/**
|
||||
* Options for creating a Docker sandbox container
|
||||
*/
|
||||
export interface DockerSandboxOptions {
|
||||
/** Memory limit in MB (default: 512) */
|
||||
memoryMB?: number;
|
||||
/** CPU limit (1.0 = 1 core, default: 1.0) */
|
||||
cpuLimit?: number;
|
||||
/** Network mode (default: bridge) */
|
||||
networkMode?: NetworkMode;
|
||||
/** Docker image to use (default: node:20-alpine) */
|
||||
image?: string;
|
||||
/** Additional environment variables */
|
||||
env?: Record<string, string>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Result of creating a Docker container
|
||||
*/
|
||||
export interface ContainerCreateResult {
|
||||
/** Docker container ID */
|
||||
containerId: string;
|
||||
/** Agent ID associated with this container */
|
||||
agentId: string;
|
||||
/** Task ID associated with this container */
|
||||
taskId: string;
|
||||
/** Timestamp when container was created */
|
||||
createdAt: Date;
|
||||
}
|
||||
|
||||
/**
|
||||
* Container status information
|
||||
*/
|
||||
export interface ContainerStatus {
|
||||
/** Container ID */
|
||||
containerId: string;
|
||||
/** Current status (running, stopped, etc.) */
|
||||
status: string;
|
||||
/** Additional state information */
|
||||
state?: Record<string, unknown>;
|
||||
}
|
||||
Reference in New Issue
Block a user