gadget/gadget-code/src/services/drone.ts

// services/drone.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

import { PopulateOptions } from "mongoose";

import { IUser, DroneStatus, IDroneRegistration, GadgetId } from "@gadget/api";
import DroneRegistration from "../models/drone-registration.ts";

import { DtpService } from "../lib/service.ts";
import SocketService from "./socket.ts";

export interface IDroneDefinition {
  hostname: string;
  workspaceDir: string;
}

class DroneService extends DtpService {
  private populateDroneRegistration: PopulateOptions[];

  constructor() {
    super();
    this.populateDroneRegistration = [
      {
        path: "user",
        select: "-passwordSalt -password",
      },
    ];
  }

  get name(): string {
    return "DroneService";
  }

  get slug(): string {
    return "svc:drone";
  }

  async start(): Promise<void> {
    this.log.info("service started");
  }

  async stop(): Promise<void> {
    this.log.info("service stopped");
  }

  /**
   * The drone calls this service to register itself into the platform as
   * belonging to a specific User. Users can only ever access their own drones.
   * @param user The user for which a drone is being registered
   * @param definition The definition of the drone being registered
   * @returns the new drone registration
   */
  async register(
    user: IUser,
    definition: IDroneDefinition,
  ): Promise<IDroneRegistration> {
    const NOW = new Date();
    let registration = new DroneRegistration();

    registration.createdAt = NOW;
    registration.updatedAt = NOW;
    registration.user = user._id;
    registration.hostname = definition.hostname;
    registration.workspaceDir = definition.workspaceDir;
    registration.status = DroneStatus.Starting;

    this.log.info("registering drone", { hostname: registration.hostname });
    await registration.save();

    return registration.populate(this.populateDroneRegistration);
  }

  /**
   * The drone calls this service to unregister from the platform when shutting
   * down. This will set the status to offline and remove any associated
   * tasks.
   * @param registration The drone registration being unregistered (offline)
   * @returns the updated drone registration
   */
  async unregister(
    registration: IDroneRegistration,
  ): Promise<IDroneRegistration> {
    this.log.info("unregistering drone", { hostname: registration.hostname });
    return this.setStatus(registration, DroneStatus.Offline);
  }

  /**
   * Retrieve a drone registration by _id. Throws if the registration doesn't
   * exist.
   * @param registrationId The _id of the drone registration to be fetched
   * @returns the drone registration document
   */
  async getById(registrationId: GadgetId): Promise<IDroneRegistration> {
    const registration = await DroneRegistration.findById(
      registrationId,
    ).populate(this.populateDroneRegistration);
    if (!registration) {
      const error = new Error("drone registration not found");
      error.statusCode = 404;
      throw error;
    }
    return registration;
  }

  /**
   * Request a User's list of registered drones.
   * @param user The user for which a list of registered drones is being requested
   * @param statusFilter Optional array of statuses to exclude from the results
   * @returns The list of drone registrations for the user.
   */
  async getForUser(
    user: IUser,
    statusFilter?: DroneStatus[],
  ): Promise<IDroneRegistration[]> {
    const query: Record<string, unknown> = { user: user._id };

    if (statusFilter && statusFilter.length > 0) {
      query.status = { $nin: statusFilter };
    }

    const registrations = await DroneRegistration.find(query)
      .sort({ hostname: 1, workspaceDir: 1 })
      .populate(this.populateDroneRegistration);
    return registrations;
  }

  /**
   * Request a User's complete list of drone registrations for the Drone Manager.
   * Returns all drones (online and offline) without filtering.
   * @param user The user for which a list of drone registrations is being requested
   * @returns The complete list of drone registrations for the user.
   */
  async getForManager(user: IUser): Promise<IDroneRegistration[]> {
    const registrations = await DroneRegistration.find({ user: user._id })
      .sort({ status: 1, hostname: 1, workspaceDir: 1 })
      .populate(this.populateDroneRegistration);
    return registrations;
  }

  /**
   * Request a drone to terminate itself. Sets a 10-second timeout to wait for
   * the drone to mark itself as offline. If the timeout expires, forces the
   * status to offline.
   * @param registrationId The _id of the drone registration to terminate
   * @returns Promise that resolves when termination is complete (success or timeout)
   */
  async requestTermination(
    registrationId: GadgetId,
  ): Promise<{ success: boolean; message?: string }> {
    const registration = await this.getById(registrationId);

    if (registration.status === DroneStatus.Offline) {
      return { success: false, message: "Drone is already offline" };
    }

    let droneSession: any;
    try {
      droneSession = SocketService.getDroneSession(registration);
    } catch (error) {
      // Drone is not connected - mark as offline immediately
      this.log.warn("drone not connected, marking offline", {
        registrationId: registrationId,
        hostname: registration.hostname,
      });
      await this.setStatus(registration, DroneStatus.Offline);
      return {
        success: true,
        message: "Drone was not connected. Marked offline.",
      };
    }

    return new Promise((resolve) => {
      let resolved = false;
      let timeoutHandle: NodeJS.Timeout;

      const cleanup = () => {
        clearTimeout(timeoutHandle);
        resolved = true;
      };

      timeoutHandle = setTimeout(async () => {
        if (resolved) return;
        cleanup();

        this.log.warn("drone termination timed out, forcing offline", {
          registrationId: registrationId,
          hostname: registration.hostname,
        });

        await this.setStatus(registration, DroneStatus.Offline);
        resolve({
          success: true,
          message: "Drone termination timed out. Marked offline.",
        });
      }, 10000);

      droneSession.socket.emit(
        "requestTermination",
        async (success: boolean) => {
          if (resolved) return;

          if (!success) {
            cleanup();
            resolve({
              success: false,
              message: "Drone rejected termination request",
            });
            return;
          }

          this.log.info("drone accepted termination request", {
            registrationId: registrationId,
            hostname: registration.hostname,
          });

          const checkInterval = setInterval(async () => {
            try {
              const updated = await this.getById(registrationId);
              if (updated.status === DroneStatus.Offline) {
                clearInterval(checkInterval);
                cleanup();
                resolve({ success: true });
              }
            } catch (error) {
              clearInterval(checkInterval);
              cleanup();
              resolve({
                success: true,
                message: "Drone termination timed out. Marked offline.",
              });
            }
          }, 500);
        },
      );
    });
  }

  /**
   * The drone calls this to update it's status in the database.
   * @param registration The registration for which a status is being updated.
   * @param status The new status of the drone.
   * @returns the updated drone registration
   */
  async setStatus(
    registration: IDroneRegistration,
    status: DroneStatus,
  ): Promise<IDroneRegistration> {
    this.log.info("setting drone status", {
      hostname: registration.hostname,
      old: registration.status,
      new: status,
    });
    const newRegistration = await DroneRegistration.findOneAndUpdate(
      { _id: registration._id },
      { $set: { status } },
      { new: true, populate: this.populateDroneRegistration },
    );
    if (!newRegistration) {
      const error = new Error("drone registration has been removed");
      error.statusCode = 404;
      throw error;
    }
    return newRegistration;
  }
}

export default new DroneService();