stellar-cyber-mcp-agents
Version:
Model Context Protocol (MCP) server for Stellar Cyber security operations with specialized multi-agent analysis capabilities
578 lines • 22.9 kB
JavaScript
import { BaseAgent } from '../core/base-agent.js';
import { AgentState, AgentHealth, AgentEventType, RequestPriority } from '../types/agent.js';
export class HubAgent extends BaseAgent {
config;
workflows = new Map();
executions = new Map();
agentLoadBalancer = new Map();
requestQueues = new Map();
activeRequestsMap = new Map();
constructor(metadata, registry, channel, logger, metrics, config = {}) {
super(metadata, registry, channel, logger, metrics);
this.config = {
maxConcurrentRequests: 100,
requestTimeout: 30000,
heartbeatInterval: 10000,
healthCheckInterval: 30000,
loadBalancing: 'least-busy',
retryPolicy: {
maxRetries: 3,
backoffMultiplier: 2,
initialDelay: 1000
},
...config
};
}
async onInitialize() {
this.logger.info('Initializing Hub Agent');
// Subscribe to agent events
this.registry.subscribeToEvents(this.handleRegistryEvent.bind(this));
// Initialize default workflows
await this.initializeDefaultWorkflows();
this.logger.info('Hub Agent initialized successfully');
}
async onStart() {
this.logger.info('Starting Hub Agent');
// Start workflow execution monitoring
this.startWorkflowMonitoring();
this.logger.info('Hub Agent started successfully');
}
async onStop() {
this.logger.info('Stopping Hub Agent');
// Cancel all active workflows
await this.cancelAllActiveWorkflows();
this.logger.info('Hub Agent stopped successfully');
}
async onDestroy() {
this.logger.info('Destroying Hub Agent');
// Clear all state
this.workflows.clear();
this.executions.clear();
this.agentLoadBalancer.clear();
this.requestQueues.clear();
this.activeRequestsMap.clear();
this.logger.info('Hub Agent destroyed successfully');
}
async onHealthCheck() {
try {
// Check registry health
const allStatuses = await this.registry.getAllAgentStatuses();
const healthy = allStatuses.filter(s => s.health === AgentHealth.HEALTHY).length;
const unhealthy = allStatuses.length - healthy;
const registryHealth = { healthy, unhealthy, total: allStatuses.length };
// Check active workflows
const activeWorkflows = Array.from(this.executions.values())
.filter(exec => exec.status === 'running').length;
// Check queue sizes
const totalQueueSize = Array.from(this.requestQueues.values())
.reduce((sum, queue) => sum + queue.length, 0);
if (registryHealth.unhealthy > registryHealth.healthy * 0.5) {
return AgentHealth.DEGRADED;
}
if (activeWorkflows > 50 || totalQueueSize > 1000) {
return AgentHealth.DEGRADED;
}
return AgentHealth.HEALTHY;
}
catch (error) {
this.logger.error('Health check failed', { error });
return AgentHealth.CRITICAL;
}
}
async handleRequest(request, context) {
const { capability, payload } = request;
this.logger.debug('Handling Hub Agent request', {
capability,
requestId: request.id,
sourceAgent: request.sourceAgentId
});
switch (capability) {
case 'orchestrate_workflow':
return await this.orchestrateWorkflow(payload.workflowId, payload.input, context);
case 'execute_workflow':
return await this.executeWorkflow(payload.workflow, payload.input, context);
case 'get_workflow_status':
return await this.getWorkflowStatus(payload.executionId);
case 'cancel_workflow':
return await this.cancelWorkflow(payload.executionId);
case 'register_workflow':
return await this.registerWorkflow(payload.workflow);
case 'find_agents':
return await this.findAgents(payload.capability, payload.criteria);
case 'distribute_request':
return await this.distributeRequest(payload.targetCapability, payload.request, payload.strategy);
case 'get_agent_metrics':
return await this.getAgentMetrics(payload.agentId);
case 'get_system_status':
return await this.getSystemStatus();
default:
throw new Error(`Unknown capability: ${capability}`);
}
}
async orchestrateWorkflow(workflowId, input, context) {
const workflow = this.workflows.get(workflowId);
if (!workflow) {
throw new Error(`Workflow not found: ${workflowId}`);
}
const execution = {
id: crypto.randomUUID(),
workflowId,
status: 'pending',
currentStep: 0,
startTime: new Date(),
results: new Map(),
errors: new Map(),
context: { ...context, input }
};
this.executions.set(execution.id, execution);
// Start workflow execution
this.executeWorkflowSteps(execution);
this.logger.info('Workflow orchestration started', {
workflowId,
executionId: execution.id
});
return execution.id;
}
async executeWorkflow(workflow, input, context) {
// Register workflow if not exists
this.workflows.set(workflow.id, workflow);
return await this.orchestrateWorkflow(workflow.id, input, context);
}
async executeWorkflowSteps(execution) {
const workflow = this.workflows.get(execution.workflowId);
if (!workflow) {
throw new Error(`Workflow not found: ${execution.workflowId}`);
}
execution.status = 'running';
try {
for (let i = execution.currentStep; i < workflow.steps.length; i++) {
const step = workflow.steps[i];
execution.currentStep = i;
this.logger.debug('Executing workflow step', {
workflowId: execution.workflowId,
executionId: execution.id,
stepId: step.id,
stepName: step.name
});
// Find suitable agent for this step
const agents = await this.findAgents(step.capability, { type: step.agentType });
if (agents.length === 0) {
throw new Error(`No agents found for capability: ${step.capability}, type: ${step.agentType}`);
}
// Select agent using load balancing
const selectedAgent = await this.selectAgent(agents, this.config.loadBalancing);
// Prepare step input
const stepInput = this.prepareStepInput(step, execution);
// Execute step with retry logic
const result = await this.executeStepWithRetry(selectedAgent.id, step.capability, stepInput, step.retryPolicy || this.config.retryPolicy);
execution.results.set(step.id, result);
this.logger.debug('Workflow step completed', {
workflowId: execution.workflowId,
executionId: execution.id,
stepId: step.id,
result
});
// Check step condition if exists
if (step.condition && !this.evaluateCondition(step.condition, execution)) {
this.logger.info('Workflow step condition failed, skipping remaining steps', {
workflowId: execution.workflowId,
executionId: execution.id,
stepId: step.id,
condition: step.condition
});
break;
}
}
execution.status = 'completed';
execution.endTime = new Date();
this.logger.info('Workflow execution completed', {
workflowId: execution.workflowId,
executionId: execution.id,
duration: execution.endTime.getTime() - execution.startTime.getTime()
});
}
catch (error) {
execution.status = 'failed';
execution.endTime = new Date();
execution.errors.set('execution', error);
this.logger.error('Workflow execution failed', {
workflowId: execution.workflowId,
executionId: execution.id,
error
});
throw error;
}
}
async findAgents(capability, criteria = {}) {
let agents = await this.registry.findAgentsByCapability(capability);
// Apply additional filtering criteria
if (criteria.type) {
agents = agents.filter(agent => agent.id.type === criteria.type);
}
if (criteria.health) {
const healthyAgents = [];
for (const agent of agents) {
const status = await this.registry.getAgentStatus(agent.id);
if (status && status.health === criteria.health) {
healthyAgents.push(agent);
}
}
agents = healthyAgents;
}
if (criteria.state) {
const readyAgents = [];
for (const agent of agents) {
const status = await this.registry.getAgentStatus(agent.id);
if (status && status.state === criteria.state) {
readyAgents.push(agent);
}
}
agents = readyAgents;
}
return agents;
}
async selectAgent(agents, strategy) {
switch (strategy) {
case 'round-robin':
return this.selectAgentRoundRobin(agents);
case 'least-busy':
return await this.selectAgentLeastBusy(agents);
case 'random':
return agents[Math.floor(Math.random() * agents.length)];
default:
return agents[0];
}
}
selectAgentRoundRobin(agents) {
const agentKeys = agents.map(a => `${a.id.type}:${a.id.instance}`);
const key = agentKeys.join(',');
const currentIndex = this.agentLoadBalancer.get(key) || 0;
const nextIndex = (currentIndex + 1) % agents.length;
this.agentLoadBalancer.set(key, nextIndex);
return agents[currentIndex];
}
async selectAgentLeastBusy(agents) {
let leastBusyAgent = agents[0];
let minActiveRequests = Infinity;
for (const agent of agents) {
const status = await this.registry.getAgentStatus(agent.id);
if (status && status.activeRequests < minActiveRequests) {
minActiveRequests = status.activeRequests;
leastBusyAgent = agent;
}
}
return leastBusyAgent;
}
prepareStepInput(step, execution) {
let input = { ...step.input };
// Replace placeholders with results from previous steps
input = this.replacePlaceholders(input, execution);
// Add execution context
input._executionContext = {
workflowId: execution.workflowId,
executionId: execution.id,
stepId: step.id,
correlationId: execution.context.correlationId
};
return input;
}
replacePlaceholders(obj, execution) {
if (typeof obj === 'string') {
return obj.replace(/\${([^}]+)}/g, (match, key) => {
const [stepId, field] = key.split('.');
const result = execution.results.get(stepId);
return field ? result?.[field] : result;
});
}
if (Array.isArray(obj)) {
return obj.map(item => this.replacePlaceholders(item, execution));
}
if (typeof obj === 'object' && obj !== null) {
const result = {};
for (const [key, value] of Object.entries(obj)) {
result[key] = this.replacePlaceholders(value, execution);
}
return result;
}
return obj;
}
evaluateCondition(condition, execution) {
// Simple condition evaluation - in production, use a proper expression evaluator
try {
const context = {
results: Object.fromEntries(execution.results),
errors: Object.fromEntries(execution.errors)
};
// Replace variables in condition
const evalCondition = condition.replace(/\${([^}]+)}/g, (match, key) => {
const [stepId, field] = key.split('.');
const result = execution.results.get(stepId);
return field ? result?.[field] : result;
});
return new Function('context', `return ${evalCondition}`)(context);
}
catch (error) {
this.logger.error('Error evaluating condition', { condition, error });
return false;
}
}
async executeStepWithRetry(agentId, capability, input, retryPolicy) {
let lastError = null;
for (let attempt = 0; attempt <= retryPolicy.maxRetries; attempt++) {
try {
const response = await this.sendRequest(agentId, capability, input, RequestPriority.MEDIUM, this.config.requestTimeout);
if (response.success) {
return response.result;
}
else {
throw new Error(response.error?.message || 'Request failed');
}
}
catch (error) {
lastError = error instanceof Error ? error : new Error('Unknown error');
if (attempt < retryPolicy.maxRetries) {
const delay = retryPolicy.initialDelay * Math.pow(retryPolicy.backoffMultiplier, attempt);
await new Promise(resolve => setTimeout(resolve, delay));
}
}
}
throw lastError;
}
async getWorkflowStatus(executionId) {
return this.executions.get(executionId) || null;
}
async cancelWorkflow(executionId) {
const execution = this.executions.get(executionId);
if (!execution) {
return false;
}
execution.status = 'cancelled';
execution.endTime = new Date();
this.logger.info('Workflow cancelled', {
workflowId: execution.workflowId,
executionId
});
return true;
}
async registerWorkflow(workflow) {
this.workflows.set(workflow.id, workflow);
this.logger.info('Workflow registered', {
workflowId: workflow.id,
name: workflow.name
});
}
async cancelAllActiveWorkflows() {
const activeExecutions = Array.from(this.executions.values())
.filter(exec => exec.status === 'running' || exec.status === 'pending');
for (const execution of activeExecutions) {
await this.cancelWorkflow(execution.id);
}
}
async distributeRequest(targetCapability, request, strategy = 'least-busy') {
const agents = await this.findAgents(targetCapability, { state: AgentState.READY });
if (agents.length === 0) {
throw new Error(`No agents available for capability: ${targetCapability}`);
}
const selectedAgent = await this.selectAgent(agents, strategy);
const response = await this.sendRequest(selectedAgent.id, targetCapability, request, RequestPriority.MEDIUM);
return response.result;
}
async getAgentMetrics(agentId) {
if (agentId) {
return await this.registry.getAgentStatus(agentId);
}
else {
return await this.registry.getRegistryStats();
}
}
async getSystemStatus() {
const registryStats = await this.registry.getRegistryStats();
const activeWorkflows = Array.from(this.executions.values())
.filter(exec => exec.status === 'running').length;
return {
registry: registryStats,
workflows: {
total: this.workflows.size,
active: activeWorkflows,
completed: Array.from(this.executions.values())
.filter(exec => exec.status === 'completed').length,
failed: Array.from(this.executions.values())
.filter(exec => exec.status === 'failed').length
},
system: {
uptime: this.getUptime(),
health: this.getHealth(),
state: this.getState()
}
};
}
handleRegistryEvent(event) {
// Handle agent lifecycle events
switch (event.type) {
case AgentEventType.AGENT_STARTED:
this.logger.info('Agent started', { agentId: event.sourceAgentId });
break;
case AgentEventType.AGENT_STOPPED:
this.logger.info('Agent stopped', { agentId: event.sourceAgentId });
break;
case AgentEventType.AGENT_ERROR:
this.logger.warn('Agent error', { agentId: event.sourceAgentId, error: event.data });
break;
}
}
startWorkflowMonitoring() {
// Skip background timers in MCP mode to prevent EPIPE errors
if (process.env.MCP_MODE === 'true')
return;
setInterval(() => {
this.cleanupCompletedExecutions();
}, 60000); // Clean up every minute
}
cleanupCompletedExecutions() {
const cutoffTime = new Date(Date.now() - 24 * 60 * 60 * 1000); // 24 hours ago
for (const [executionId, execution] of this.executions) {
if (execution.status === 'completed' || execution.status === 'failed') {
if (execution.endTime && execution.endTime < cutoffTime) {
this.executions.delete(executionId);
}
}
}
}
async initializeDefaultWorkflows() {
// Case Investigation Workflow
const caseInvestigationWorkflow = {
id: 'case-investigation',
name: 'Case Investigation',
description: 'Comprehensive case investigation workflow',
steps: [
{
id: 'fetch-case',
name: 'Fetch Case Details',
agentType: 'investigation',
capability: 'get_case_details',
input: { caseId: '${input.caseId}' }
},
{
id: 'analyze-network',
name: 'Analyze Network Activity',
agentType: 'network',
capability: 'analyze_network_activity',
input: {
caseId: '${input.caseId}',
timeRange: '${fetch-case.timeRange}'
}
},
{
id: 'analyze-malware',
name: 'Analyze Malware',
agentType: 'malware',
capability: 'analyze_malware',
input: {
caseId: '${input.caseId}',
artifacts: '${fetch-case.artifacts}'
},
condition: '${fetch-case.hasArtifacts}'
},
{
id: 'correlate-findings',
name: 'Correlate Findings',
agentType: 'correlation',
capability: 'correlate_case_data',
input: {
caseId: '${input.caseId}',
networkFindings: '${analyze-network.findings}',
malwareFindings: '${analyze-malware.findings}'
}
}
],
triggers: [
{
type: 'manual',
condition: 'true',
priority: RequestPriority.HIGH
}
],
timeout: 300000 // 5 minutes
};
await this.registerWorkflow(caseInvestigationWorkflow);
}
}
export function createHubAgentMetadata() {
const capabilities = [
{
name: 'orchestrate_workflow',
description: 'Orchestrate workflow execution',
inputSchema: {
type: 'object',
properties: {
workflowId: { type: 'string' },
input: { type: 'object' }
},
required: ['workflowId', 'input']
},
outputSchema: {
type: 'object',
properties: {
executionId: { type: 'string' }
}
}
},
{
name: 'execute_workflow',
description: 'Execute workflow definition',
inputSchema: {
type: 'object',
properties: {
workflow: { type: 'object' },
input: { type: 'object' }
},
required: ['workflow', 'input']
},
outputSchema: {
type: 'object',
properties: {
executionId: { type: 'string' }
}
}
},
{
name: 'find_agents',
description: 'Find agents by capability and criteria',
inputSchema: {
type: 'object',
properties: {
capability: { type: 'string' },
criteria: { type: 'object' }
},
required: ['capability']
},
outputSchema: {
type: 'array',
items: { type: 'object' }
}
},
{
name: 'get_system_status',
description: 'Get system status and metrics',
inputSchema: { type: 'object' },
outputSchema: { type: 'object' }
}
];
return {
id: {
type: 'hub',
instance: 'primary',
uuid: crypto.randomUUID()
},
name: 'Hub Agent',
description: 'Central orchestration agent for multi-agent workflows',
version: '1.0.0',
capabilities,
dependencies: [],
resources: {
memory: 512,
cpu: 2
}
};
}
//# sourceMappingURL=hub-agent.js.map