stack/apps/orchestrator/src/api/agents/agents.controller.ts

import {
  Controller,
  Post,
  Get,
  Body,
  Param,
  BadRequestException,
  NotFoundException,
  Logger,
  UsePipes,
  ValidationPipe,
  HttpCode,
} from "@nestjs/common";
import { QueueService } from "../../queue/queue.service";
import { AgentSpawnerService } from "../../spawner/agent-spawner.service";
import { AgentLifecycleService } from "../../spawner/agent-lifecycle.service";
import { KillswitchService } from "../../killswitch/killswitch.service";
import { SpawnAgentDto, SpawnAgentResponseDto } from "./dto/spawn-agent.dto";

/**
 * Controller for agent management endpoints
 */
@Controller("agents")
export class AgentsController {
  private readonly logger = new Logger(AgentsController.name);

  constructor(
    private readonly queueService: QueueService,
    private readonly spawnerService: AgentSpawnerService,
    private readonly lifecycleService: AgentLifecycleService,
    private readonly killswitchService: KillswitchService
  ) {}

  /**
   * Spawn a new agent for the given task
   * @param dto Spawn agent request
   * @returns Agent spawn response with agentId and status
   */
  @Post("spawn")
  @UsePipes(new ValidationPipe({ transform: true, whitelist: true }))
  async spawn(@Body() dto: SpawnAgentDto): Promise<SpawnAgentResponseDto> {
    this.logger.log(`Received spawn request for task: ${dto.taskId}`);

    try {
      // Validate request manually (in addition to ValidationPipe)
      this.validateSpawnRequest(dto);

      // Spawn agent using spawner service
      const spawnResponse = this.spawnerService.spawnAgent({
        taskId: dto.taskId,
        agentType: dto.agentType,
        context: dto.context,
      });

      // Queue task in Valkey
      await this.queueService.addTask(dto.taskId, dto.context, {
        priority: 5, // Default priority
      });

      this.logger.log(`Agent spawned successfully: ${spawnResponse.agentId}`);

      // Return response
      return {
        agentId: spawnResponse.agentId,
        status: "spawning",
      };
    } catch (error) {
      this.logger.error(`Failed to spawn agent: ${String(error)}`);
      throw error;
    }
  }

  /**
   * List all agents
   * @returns Array of all agent sessions with their status
   */
  @Get()
  listAgents(): {
    agentId: string;
    taskId: string;
    status: string;
    agentType: string;
    spawnedAt: string;
    completedAt?: string;
    error?: string;
  }[] {
    this.logger.log("Received request to list all agents");

    try {
      // Get all sessions from spawner service
      const sessions = this.spawnerService.listAgentSessions();

      // Map to response format
      const agents = sessions.map((session) => ({
        agentId: session.agentId,
        taskId: session.taskId,
        status: session.state,
        agentType: session.agentType,
        spawnedAt: session.spawnedAt.toISOString(),
        completedAt: session.completedAt?.toISOString(),
        error: session.error,
      }));

      this.logger.log(`Found ${agents.length.toString()} agents`);

      return agents;
    } catch (error: unknown) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      this.logger.error(`Failed to list agents: ${errorMessage}`);
      throw new Error(`Failed to list agents: ${errorMessage}`);
    }
  }

  /**
   * Get agent status
   * @param agentId Agent ID to query
   * @returns Agent status details
   */
  @Get(":agentId/status")
  async getAgentStatus(@Param("agentId") agentId: string): Promise<{
    agentId: string;
    taskId: string;
    status: string;
    spawnedAt: string;
    startedAt?: string;
    completedAt?: string;
    error?: string;
  }> {
    this.logger.log(`Received status request for agent: ${agentId}`);

    try {
      // Try to get from lifecycle service (Valkey)
      const lifecycleState = await this.lifecycleService.getAgentLifecycleState(agentId);

      if (lifecycleState) {
        return {
          agentId: lifecycleState.agentId,
          taskId: lifecycleState.taskId,
          status: lifecycleState.status,
          spawnedAt: lifecycleState.startedAt ?? new Date().toISOString(),
          startedAt: lifecycleState.startedAt,
          completedAt: lifecycleState.completedAt,
          error: lifecycleState.error,
        };
      }

      // Fallback to spawner service (in-memory)
      const session = this.spawnerService.getAgentSession(agentId);

      if (session) {
        return {
          agentId: session.agentId,
          taskId: session.taskId,
          status: session.state,
          spawnedAt: session.spawnedAt.toISOString(),
          completedAt: session.completedAt?.toISOString(),
          error: session.error,
        };
      }

      throw new NotFoundException(`Agent ${agentId} not found`);
    } catch (error: unknown) {
      if (error instanceof NotFoundException) {
        throw error;
      }
      const errorMessage = error instanceof Error ? error.message : String(error);
      this.logger.error(`Failed to get agent status: ${errorMessage}`);
      throw new Error(`Failed to get agent status: ${errorMessage}`);
    }
  }

  /**
   * Kill a single agent immediately
   * @param agentId Agent ID to kill
   * @returns Success message
   */
  @Post(":agentId/kill")
  @HttpCode(200)
  async killAgent(@Param("agentId") agentId: string): Promise<{ message: string }> {
    this.logger.warn(`Received kill request for agent: ${agentId}`);

    try {
      await this.killswitchService.killAgent(agentId);

      this.logger.warn(`Agent ${agentId} killed successfully`);

      return {
        message: `Agent ${agentId} killed successfully`,
      };
    } catch (error) {
      this.logger.error(`Failed to kill agent ${agentId}: ${String(error)}`);
      throw error;
    }
  }

  /**
   * Kill all active agents
   * @returns Summary of kill operation
   */
  @Post("kill-all")
  @HttpCode(200)
  async killAllAgents(): Promise<{
    message: string;
    total: number;
    killed: number;
    failed: number;
    errors?: string[];
  }> {
    this.logger.warn("Received kill-all request");

    try {
      const result = await this.killswitchService.killAllAgents();

      this.logger.warn(
        `Kill all completed: ${result.killed.toString()} killed, ${result.failed.toString()} failed out of ${result.total.toString()}`
      );

      return {
        message: `Kill all completed: ${result.killed.toString()} killed, ${result.failed.toString()} failed`,
        ...result,
      };
    } catch (error) {
      this.logger.error(`Failed to kill all agents: ${String(error)}`);
      throw error;
    }
  }

  /**
   * Validate spawn request
   * @param dto Spawn request to validate
   * @throws BadRequestException if validation fails
   */
  private validateSpawnRequest(dto: SpawnAgentDto): void {
    if (!dto.taskId || dto.taskId.trim() === "") {
      throw new BadRequestException("taskId is required");
    }

    const validAgentTypes = ["worker", "reviewer", "tester"];
    if (!validAgentTypes.includes(dto.agentType)) {
      throw new BadRequestException(`agentType must be one of: ${validAgentTypes.join(", ")}`);
    }

    if (!dto.context.repository || dto.context.repository.trim() === "") {
      throw new BadRequestException("context.repository is required");
    }

    if (!dto.context.branch || dto.context.branch.trim() === "") {
      throw new BadRequestException("context.branch is required");
    }

    if (dto.context.workItems.length === 0) {
      throw new BadRequestException("context.workItems must not be empty");
    }
  }
}