Agentic-Compaction-and-Pipleline-by-GLM-5/agent-system/core/subagent-spawner.ts

/**
 * Subagent Spawner Module
 * 
 * Creates and manages child agents (subagents) for parallel task execution.
 * Implements communication channels, result aggregation, and lifecycle management.
 */

import { randomUUID } from 'crypto';
import ZAI from 'z-ai-web-dev-sdk';
import { AgentOrchestrator, AgentConfig, Task, TaskPriority } from './orchestrator';

export type SubagentType = 
  | 'explorer'      // For code exploration
  | 'researcher'    // For information gathering
  | 'coder'         // For code generation
  | 'reviewer'      // For code review
  | 'planner'       // For task planning
  | 'executor'      // For task execution
  | 'custom';       // Custom subagent

export interface SubagentDefinition {
  type: SubagentType;
  name: string;
  description: string;
  systemPrompt: string;
  capabilities: string[];
  maxTasks?: number;
  timeout?: number;
}

export interface SubagentResult {
  subagentId: string;
  taskId: string;
  success: boolean;
  output: unknown;
  error?: string;
  tokens: {
    input: number;
    output: number;
  };
  duration: number;
}

export interface SpawnOptions {
  priority?: TaskPriority;
  timeout?: number;
  context?: string;
  dependencies?: string[];
  metadata?: Record<string, unknown>;
}

export interface SubagentPool {
  id: string;
  name: string;
  subagents: Map<string, SubagentInstance>;
  createdAt: Date;
}

/**
 * SubagentInstance - A running subagent
 */
export class SubagentInstance {
  id: string;
  definition: SubagentDefinition;
  orchestrator: AgentOrchestrator;
  private zai: Awaited<ReturnType<typeof ZAI.create>> | null = null;
  private initialized = false;

  constructor(
    definition: SubagentDefinition,
    orchestrator: AgentOrchestrator
  ) {
    this.id = `${definition.type}-${randomUUID().substring(0, 8)}`;
    this.definition = definition;
    this.orchestrator = orchestrator;
  }

  /**
   * Initialize the subagent
   */
  async initialize(): Promise<void> {
    if (this.initialized) return;

    this.zai = await ZAI.create();

    // Register with orchestrator
    const config: AgentConfig = {
      id: this.id,
      name: this.definition.name,
      type: this.definition.type,
      capabilities: this.definition.capabilities,
      maxConcurrentTasks: this.definition.maxTasks || 3,
      timeout: this.definition.timeout || 60000,
      metadata: {
        systemPrompt: this.definition.systemPrompt
      }
    };

    this.orchestrator.registerAgent(config);
    this.initialized = true;
  }

  /**
   * Execute a task
   */
  async execute(input: string, context?: string): Promise<SubagentResult> {
    const startTime = Date.now();
    
    if (!this.initialized || !this.zai) {
      await this.initialize();
    }

    const task = this.orchestrator.createTask(
      this.definition.type,
      `Execute ${this.definition.type} task`,
      { input, context },
      { assignedAgent: this.id }
    );

    try {
      const messages = [
        {
          role: 'assistant' as const,
          content: this.definition.systemPrompt
        },
        {
          role: 'user' as const,
          content: context 
            ? `Context: ${context}\n\nTask: ${input}`
            : input
        }
      ];

      const response = await this.zai!.chat.completions.create({
        messages,
        thinking: { type: 'disabled' }
      });

      const output = response.choices?.[0]?.message?.content || '';
      
      const result: SubagentResult = {
        subagentId: this.id,
        taskId: task.id,
        success: true,
        output,
        tokens: {
          input: 0, // Would need tokenizer to calculate
          output: 0
        },
        duration: Date.now() - startTime
      };

      return result;

    } catch (error) {
      return {
        subagentId: this.id,
        taskId: task.id,
        success: false,
        output: null,
        error: error instanceof Error ? error.message : String(error),
        tokens: { input: 0, output: 0 },
        duration: Date.now() - startTime
      };
    }
  }

  /**
   * Terminate the subagent
   */
  terminate(): void {
    this.orchestrator.unregisterAgent(this.id);
    this.initialized = false;
  }
}

/**
 * SubagentSpawner - Factory for creating and managing subagents
 */
export class SubagentSpawner {
  private orchestrator: AgentOrchestrator;
  private subagents: Map<string, SubagentInstance> = new Map();
  private pools: Map<string, SubagentPool> = new Map();
  private definitions: Map<SubagentType, SubagentDefinition> = new Map();

  constructor(orchestrator?: AgentOrchestrator) {
    this.orchestrator = orchestrator || new AgentOrchestrator();
    this.registerDefaultDefinitions();
    this.orchestrator.start();
  }

  /**
   * Register default subagent definitions
   */
  private registerDefaultDefinitions(): void {
    const defaults: SubagentDefinition[] = [
      {
        type: 'explorer',
        name: 'Code Explorer',
        description: 'Explores codebases to find relevant files and code',
        systemPrompt: `You are a code explorer agent. Your job is to search through codebases to find relevant files, functions, and code patterns. Be thorough but concise in your findings.`,
        capabilities: ['explore', 'search', 'find']
      },
      {
        type: 'researcher',
        name: 'Research Agent',
        description: 'Gathers information and researches topics',
        systemPrompt: `You are a research agent. Your job is to gather comprehensive information on given topics. Focus on accuracy and completeness.`,
        capabilities: ['research', 'gather', 'analyze']
      },
      {
        type: 'coder',
        name: 'Code Generator',
        description: 'Generates code based on specifications',
        systemPrompt: `You are a code generation agent. Your job is to write clean, efficient, and well-documented code. Follow best practices and include appropriate error handling.`,
        capabilities: ['code', 'generate', 'implement']
      },
      {
        type: 'reviewer',
        name: 'Code Reviewer',
        description: 'Reviews code for quality, bugs, and improvements',
        systemPrompt: `You are a code review agent. Your job is to analyze code for bugs, security issues, performance problems, and best practice violations. Provide constructive feedback.`,
        capabilities: ['review', 'analyze', 'validate']
      },
      {
        type: 'planner',
        name: 'Task Planner',
        description: 'Plans and breaks down complex tasks',
        systemPrompt: `You are a planning agent. Your job is to break down complex tasks into smaller, manageable steps. Consider dependencies and optimal execution order.`,
        capabilities: ['plan', 'decompose', 'organize']
      },
      {
        type: 'executor',
        name: 'Task Executor',
        description: 'Executes specific tasks with precision',
        systemPrompt: `You are an execution agent. Your job is to carry out specific tasks accurately and efficiently. Report results clearly and flag any issues encountered.`,
        capabilities: ['execute', 'run', 'process']
      }
    ];

    for (const def of defaults) {
      this.definitions.set(def.type, def);
    }
  }

  /**
   * Register a custom subagent definition
   */
  registerDefinition(definition: SubagentDefinition): void {
    this.definitions.set(definition.type, definition);
  }

  /**
   * Spawn a single subagent
   */
  async spawn(type: SubagentType): Promise<SubagentInstance> {
    const definition = this.definitions.get(type);
    if (!definition) {
      throw new Error(`Unknown subagent type: ${type}`);
    }

    const subagent = new SubagentInstance(definition, this.orchestrator);
    await subagent.initialize();
    this.subagents.set(subagent.id, subagent);

    return subagent;
  }

  /**
   * Spawn multiple subagents of the same type
   */
  async spawnPool(
    type: SubagentType, 
    count: number,
    poolName?: string
  ): Promise<SubagentPool> {
    const pool: SubagentPool = {
      id: randomUUID(),
      name: poolName || `${type}-pool-${Date.now()}`,
      subagents: new Map(),
      createdAt: new Date()
    };

    for (let i = 0; i < count; i++) {
      const subagent = await this.spawn(type);
      pool.subagents.set(subagent.id, subagent);
    }

    this.pools.set(pool.id, pool);
    return pool;
  }

  /**
   * Execute task with a spawned subagent
   */
  async executeWithSubagent(
    type: SubagentType,
    task: string,
    context?: string,
    options?: SpawnOptions
  ): Promise<SubagentResult> {
    const subagent = await this.spawn(type);
    
    try {
      const result = await subagent.execute(task, context);
      return result;
    } finally {
      // Auto-terminate after execution
      subagent.terminate();
      this.subagents.delete(subagent.id);
    }
  }

  /**
   * Execute multiple tasks in parallel
   */
  async executeParallel(
    tasks: Array<{
      type: SubagentType;
      input: string;
      context?: string;
    }>,
    options?: { maxConcurrent?: number }
  ): Promise<SubagentResult[]> {
    const maxConcurrent = options?.maxConcurrent || 5;
    const results: SubagentResult[] = [];

    // Process in batches
    for (let i = 0; i < tasks.length; i += maxConcurrent) {
      const batch = tasks.slice(i, i + maxConcurrent);
      const batchPromises = batch.map(t => 
        this.executeWithSubagent(t.type, t.input, t.context)
      );
      
      const batchResults = await Promise.all(batchPromises);
      results.push(...batchResults);
    }

    return results;
  }

  /**
   * Execute tasks in a pipeline (sequential with context passing)
   */
  async executePipeline(
    steps: Array<{
      type: SubagentType;
      input: string | ((prevResult: unknown) => string);
    }>,
    initialContext?: string
  ): Promise<{ results: SubagentResult[]; finalOutput: unknown }> {
    const results: SubagentResult[] = [];
    let currentContext = initialContext;
    let currentOutput: unknown = null;

    for (const step of steps) {
      const input = typeof step.input === 'function' 
        ? step.input(currentOutput)
        : step.input;

      const result = await this.executeWithSubagent(
        step.type, 
        input, 
        currentContext
      );

      results.push(result);
      
      if (result.success) {
        currentOutput = result.output;
        currentContext = typeof result.output === 'string' 
          ? result.output 
          : JSON.stringify(result.output);
      } else {
        // Stop pipeline on failure
        break;
      }
    }

    return { results, finalOutput: currentOutput };
  }

  /**
   * Terminate a specific subagent
   */
  terminate(subagentId: string): boolean {
    const subagent = this.subagents.get(subagentId);
    if (subagent) {
      subagent.terminate();
      this.subagents.delete(subagentId);
      return true;
    }
    return false;
  }

  /**
   * Terminate all subagents in a pool
   */
  terminatePool(poolId: string): boolean {
    const pool = this.pools.get(poolId);
    if (!pool) return false;

    for (const subagent of pool.subagents.values()) {
      subagent.terminate();
      this.subagents.delete(subagent.id);
    }

    this.pools.delete(poolId);
    return true;
  }

  /**
   * Terminate all subagents
   */
  terminateAll(): void {
    for (const subagent of this.subagents.values()) {
      subagent.terminate();
    }
    this.subagents.clear();
    this.pools.clear();
  }

  /**
   * Get active subagents
   */
  getActiveSubagents(): SubagentInstance[] {
    return Array.from(this.subagents.values());
  }

  /**
   * Get orchestrator stats
   */
  getStats() {
    return {
      activeSubagents: this.subagents.size,
      pools: this.pools.size,
      orchestrator: this.orchestrator.getStats()
    };
  }
}

/**
 * Quick spawn function for simple use cases
 */
export async function spawnAndExecute(
  type: SubagentType,
  task: string,
  context?: string
): Promise<SubagentResult> {
  const spawner = new SubagentSpawner();
  return spawner.executeWithSubagent(type, task, context);
}

// Default spawner instance
export const defaultSpawner = new SubagentSpawner();