capsule-ai-cli/dist/services/sub-agent-manager.js

The AI Model Orchestrator - Intelligent multi-model workflows with device-locked licensing

import { v4 as uuidv4 } from 'uuid';
import { EventEmitter } from 'events';
import { ChatService } from './chat.js';
import { contextManager as globalContextManager } from './context.js';
import { providerRegistry } from '../providers/base.js';
import { stateService } from './state.js';
import { updateTodoByDescription } from '../tools/builtin/todo-list.js';
export class SubAgentManager extends EventEmitter {
    tasks = new Map();
    activeTasks = new Set();
    constructor() {
        super();
    }
    emit(event, ...args) {
        return super.emit(event, ...args);
    }
    on(event, listener) {
        return super.on(event, listener);
    }
    async spawnTask(config) {
        const task = {
            taskId: uuidv4().substring(0, 8),
            name: config.name,
            description: config.description,
            status: 'pending',
            provider: config.provider,
            model: config.model,
            toolCalls: [],
            tokensUsed: 0,
            cost: 0,
            startTime: new Date()
        };
        this.tasks.set(task.taskId, task);
        this.emit('task:created', task);
        if (config.parallel !== false) {
            this.executeSubAgent(task).catch(error => {
                task.status = 'failed';
                task.error = error.message;
                this.emit('task:failed', task);
            });
        }
        else {
            await this.executeSubAgent(task);
        }
        return task;
    }
    async executeSubAgent(task) {
        const originalContextId = globalContextManager.getCurrentContext().id;
        try {
            task.status = 'running';
            this.activeTasks.add(task.taskId);
            this.emit('task:started', task);
            const context = globalContextManager.createNewContext();
            globalContextManager.setCurrentContext(context.id);
            const chatService = new ChatService();
            const abortController = new AbortController();
            const contextManager = globalContextManager;
            task.agent = {
                context,
                contextManager,
                chatService,
                abortController
            };
            this.setupToolCallTracking(task);
            this.emit('task:progress', task.taskId, 'Starting task execution...');
            const contextPath = process.cwd();
            updateTodoByDescription(contextPath, task.name, { status: 'in_progress' });
            const agentPrompt = `You are a sub-agent tasked with the following:

${task.description}

Please complete this task using any available tools. Be thorough and provide detailed results.`;
            if (task.provider) {
                stateService.setProvider(task.provider);
            }
            if (task.model) {
                stateService.setModel(task.model);
            }
            const response = await chatService.chat(agentPrompt, {
                mode: 'agent',
                provider: task.provider || stateService.getProvider(),
                model: task.model || stateService.getModel(),
                stream: false
            });
            const toolCallMatch = response.content.match(/\[Used tools: ([^\]]+)\]/);
            ;
            if (toolCallMatch) {
                const toolNames = toolCallMatch[1].split(', ');
                for (const toolName of toolNames) {
                    if (!task.toolCalls.find(tc => tc.name === toolName)) {
                        task.toolCalls.push({
                            id: uuidv4(),
                            name: toolName,
                            arguments: {}
                        });
                    }
                }
            }
            task.result = response.content;
            task.tokensUsed = response.usage.totalTokens;
            const providerName = task.provider || 'openai';
            const model = task.model || 'gpt-4o';
            const provider = providerRegistry.get(providerName);
            if (provider) {
                const cost = provider.calculateCost(response.usage, model);
                task.cost = cost.amount;
            }
            else {
                task.cost = 0;
            }
            task.status = 'completed';
            task.endTime = new Date();
            const completionPath = process.cwd();
            updateTodoByDescription(completionPath, task.name, {
                status: 'completed',
                metadata: {
                    actualTime: task.endTime.getTime() - task.startTime.getTime(),
                    completionNotes: `Completed by sub-agent ${task.taskId}`
                }
            });
            this.emit('task:completed', task);
        }
        catch (error) {
            if (error.name === 'AbortError') {
                task.status = 'cancelled';
                const cancelPath = process.cwd();
                updateTodoByDescription(cancelPath, task.name, { status: 'cancelled' });
                this.emit('task:cancelled', task);
            }
            else {
                task.status = 'failed';
                task.error = error.message || 'Unknown error';
                task.endTime = new Date();
                const blockPath = process.cwd();
                updateTodoByDescription(blockPath, task.name, {
                    status: 'blocked',
                    metadata: { blockedReason: error.message || 'Task failed' }
                });
                this.emit('task:failed', task);
            }
        }
        finally {
            globalContextManager.setCurrentContext(originalContextId);
            this.activeTasks.delete(task.taskId);
        }
    }
    setupToolCallTracking(task) {
        const toolCallHandler = (toolCall) => {
            if (toolCall && toolCall.name) {
                task.toolCalls.push({
                    id: toolCall.id || uuidv4(),
                    name: toolCall.name,
                    arguments: toolCall.arguments || {}
                });
                this.emit('task:tool_call', task.taskId, toolCall);
            }
        };
        task.toolCallHandler = toolCallHandler;
    }
    getTask(taskId) {
        return this.tasks.get(taskId);
    }
    getAllTasks() {
        return Array.from(this.tasks.values());
    }
    getActiveTasks() {
        return Array.from(this.activeTasks).map(id => this.tasks.get(id));
    }
    async cancelTask(taskId) {
        const task = this.tasks.get(taskId);
        if (!task || task.status !== 'running') {
            return false;
        }
        task.agent?.abortController.abort();
        task.status = 'cancelled';
        task.endTime = new Date();
        this.activeTasks.delete(taskId);
        this.emit('task:cancelled', task);
        return true;
    }
    async waitForTask(taskId, timeoutMs) {
        const task = this.tasks.get(taskId);
        if (!task) {
            throw new Error(`Task ${taskId} not found`);
        }
        if (task.status !== 'pending' && task.status !== 'running') {
            return task;
        }
        return new Promise((resolve, reject) => {
            const timeout = timeoutMs ? setTimeout(() => {
                reject(new Error(`Timeout waiting for task ${taskId}`));
            }, timeoutMs) : null;
            const checkStatus = () => {
                const currentTask = this.tasks.get(taskId);
                if (!currentTask) {
                    if (timeout)
                        clearTimeout(timeout);
                    reject(new Error('Task disappeared while waiting'));
                    return;
                }
                if (currentTask.status === 'completed' ||
                    currentTask.status === 'failed' ||
                    currentTask.status === 'cancelled') {
                    if (timeout)
                        clearTimeout(timeout);
                    resolve(currentTask);
                }
            };
            this.on('task:completed', (completedTask) => {
                if (completedTask.taskId === taskId)
                    checkStatus();
            });
            this.on('task:failed', (failedTask) => {
                if (failedTask.taskId === taskId)
                    checkStatus();
            });
            this.on('task:cancelled', (cancelledTask) => {
                if (cancelledTask.taskId === taskId)
                    checkStatus();
            });
        });
    }
    getStats() {
        const tasks = this.getAllTasks();
        return {
            total: tasks.length,
            pending: tasks.filter(t => t.status === 'pending').length,
            running: tasks.filter(t => t.status === 'running').length,
            completed: tasks.filter(t => t.status === 'completed').length,
            failed: tasks.filter(t => t.status === 'failed').length,
            cancelled: tasks.filter(t => t.status === 'cancelled').length,
            totalTokensUsed: tasks.reduce((sum, t) => sum + t.tokensUsed, 0),
            totalCost: tasks.reduce((sum, t) => sum + t.cost, 0)
        };
    }
}
export const subAgentManager = new SubAgentManager();
//# sourceMappingURL=sub-agent-manager.js.map