stack/apps/orchestrator/src/spawner/docker-sandbox.service.ts

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;
  }
}