gadget/gadget-code/tests/drone-session.test.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Types } from "mongoose";
import { DroneSession } from "../src/lib/drone-session";
import {
  IDroneRegistration,
  IUser,
  ChatTurnStatus,
  GadgetId,
} from "@gadget/api";
import SocketService from "../src/services/socket";
import { ChatTurn } from "../src/models/chat-turn";
import { nanoid } from "nanoid";

// Mock dependencies
vi.mock("../src/services/socket");
vi.mock("../src/models/chat-turn");

describe("DroneSession", () => {
  let mockSocket: any;
  let mockRegistration: IDroneRegistration;
  let mockUser: IUser;
  let droneSession: DroneSession;
  let mockChatSessionId: GadgetId;
  let mockTurnId: GadgetId;

  beforeEach(() => {
    vi.clearAllMocks();

    mockSocket = {
      id: "test-drone-socket",
      on: vi.fn(),
      emit: vi.fn(),
    };

    mockUser = {
      _id: nanoid(),
      email: "test@example.com",
      displayName: "Test User",
    } as IUser;

    mockRegistration = {
      _id: nanoid(),
      user: mockUser,
      hostname: "test-host",
      workspaceDir: "/test/workspace",
      status: "available",
    } as IDroneRegistration;

    mockChatSessionId = nanoid();
    mockTurnId = nanoid();

    droneSession = new DroneSession(mockSocket, mockRegistration);
  });

  describe("register", () => {
    it("should register event handlers for drone events", () => {
      droneSession.register();

      expect(mockSocket.on).toHaveBeenCalledWith(
        "thinking",
        expect.any(Function),
      );
      expect(mockSocket.on).toHaveBeenCalledWith(
        "response",
        expect.any(Function),
      );
      expect(mockSocket.on).toHaveBeenCalledWith(
        "toolCall",
        expect.any(Function),
      );
      expect(mockSocket.on).toHaveBeenCalledWith(
        "workOrderComplete",
        expect.any(Function),
      );
      expect(mockSocket.on).toHaveBeenCalledWith(
        "requestCrashRecovery",
        expect.any(Function),
      );
      expect(mockSocket.on).toHaveBeenCalledWith(
        "requestTermination",
        expect.any(Function),
      );
    });
  });

  describe("setChatSessionId", () => {
    it("should set the chat session ID", () => {
      droneSession.setChatSessionId(mockChatSessionId);
      expect(() =>
        droneSession.setChatSessionId(mockChatSessionId),
      ).not.toThrow();
    });
  });

  describe("setCurrentTurnId", () => {
    it("should set the current turn ID", () => {
      droneSession.setCurrentTurnId(mockTurnId);
      expect(() => droneSession.setCurrentTurnId(mockTurnId)).not.toThrow();
    });
  });

  describe("onThinking", () => {
    it("should route thinking event to code session", async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
        onThinking: vi.fn(),
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(
        mockCodeSession as any,
      );

      droneSession.setChatSessionId(mockChatSessionId);

      await droneSession.onThinking("thinking content");

      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(
        mockChatSessionId,
      );
      expect(mockCodeSession.onThinking).toHaveBeenCalledWith(
        "thinking content",
      );
    });

    it("should log warning if no chat session is active", async () => {
      await droneSession.onThinking("thinking content");
      // No exception thrown, warning logged internally
      expect(mockSocket.emit).not.toHaveBeenCalled();
    });
  });

  describe("onResponse", () => {
    it("should route response event to code session", async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
        onResponse: vi.fn(),
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(
        mockCodeSession as any,
      );

      droneSession.setChatSessionId(mockChatSessionId);

      await droneSession.onResponse("response content");

      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(
        mockChatSessionId,
      );
      expect(mockCodeSession.onResponse).toHaveBeenCalledWith(
        "response content",
      );
    });

    it("should log warning if no chat session is active", async () => {
      await droneSession.onResponse("response content");
      // No exception thrown, warning logged internally
      expect(mockSocket.emit).not.toHaveBeenCalled();
    });
  });

  describe("onToolCall", () => {
    it("should route toolCall event to code session and update ChatTurn", async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
        onToolCall: vi.fn(),
      };
      const mockTurn = {
        blocks: [],
        toolCalls: [],
        stats: { toolCallCount: 0 },
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(
        mockCodeSession as any,
      );
      vi.mocked(ChatTurn.findById).mockResolvedValue(mockTurn as any);

      droneSession.setChatSessionId(mockChatSessionId);
      droneSession.setCurrentTurnId(mockTurnId);

      await droneSession.onToolCall(
        "call-123",
        "readFile",
        '{"path":"test.ts"}',
        "file contents",
      );

      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(
        mockChatSessionId,
      );
      expect(mockCodeSession.onToolCall).toHaveBeenCalledWith(
        "call-123",
        "readFile",
        '{"path":"test.ts"}',
        "file contents",
      );
      expect(ChatTurn.findById).toHaveBeenCalledWith(mockTurnId);
      expect(mockTurn.toolCalls).toHaveLength(1);
      expect(mockTurn.toolCalls[0]).toEqual({
        callId: "call-123",
        name: "readFile",
        parameters: '{"path":"test.ts"}',
        response: "file contents",
      });
      expect(mockTurn.stats.toolCallCount).toBe(1);
      expect(mockTurn.save).toHaveBeenCalled();
    });

    it("should log warning if no chat session is active", async () => {
      await droneSession.onToolCall("call-1", "test", "{}", "result");
      // No exception thrown, warning logged internally
      expect(mockSocket.emit).not.toHaveBeenCalled();
    });
  });

  describe("onWorkOrderComplete", () => {
    it("should update ChatTurn status and emit to code session on success", async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
        onWorkOrderComplete: vi.fn(),
      };
      const mockTurn = {
        status: ChatTurnStatus.Processing,
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(
        mockCodeSession as any,
      );
      vi.mocked(ChatTurn.findById).mockResolvedValue(mockTurn as any);

      droneSession.setChatSessionId(mockChatSessionId);

      await droneSession.onWorkOrderComplete(mockTurnId, true);

      expect(ChatTurn.findById).toHaveBeenCalledWith(mockTurnId);
      expect(mockTurn.status).toBe(ChatTurnStatus.Finished);
      expect(mockTurn.save).toHaveBeenCalled();
      expect(mockCodeSession.onWorkOrderComplete).toHaveBeenCalledWith(
        mockTurnId,
        true,
        undefined,
      );
      expect(droneSession.currentTurnId).toBeUndefined();
    });

    it("should update ChatTurn to Error status on failure", async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
        onWorkOrderComplete: vi.fn(),
      };
      const mockTurn = {
        status: ChatTurnStatus.Processing,
        errorMessage: "",
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(
        mockCodeSession as any,
      );
      vi.mocked(ChatTurn.findById).mockResolvedValue(mockTurn as any);

      droneSession.setChatSessionId(mockChatSessionId);

      await droneSession.onWorkOrderComplete(
        mockTurnId,
        false,
        "Agent crashed",
      );

      expect(mockTurn.status).toBe(ChatTurnStatus.Error);
      expect(mockTurn.errorMessage).toBe("Agent crashed");
      expect(mockTurn.save).toHaveBeenCalled();
    });

    it("should log warning if no chat session is active", async () => {
      await droneSession.onWorkOrderComplete(mockTurnId, true);
      // No exception thrown, warning logged internally
      expect(mockSocket.emit).not.toHaveBeenCalled();
    });
  });

  describe("onRequestTermination", () => {
    it("should forward requestTermination to drone socket with logging", async () => {
      const mockCallback = vi.fn();
      const mockDroneCallback = vi.fn();
      mockSocket.emit.mockImplementation((event: string, ...args: any[]) => {
        if (event === "requestTermination" && args.length > 0) {
          const cb = args[args.length - 1];
          if (typeof cb === "function") {
            cb(true);
          }
        }
      });

      await droneSession.onRequestTermination(mockCallback);

      expect(mockSocket.emit).toHaveBeenCalledWith(
        "requestTermination",
        expect.any(Function),
      );
      expect(mockCallback).toHaveBeenCalledWith(true);
    });

    it("should pass through failure response from drone", async () => {
      const mockCallback = vi.fn();
      mockSocket.emit.mockImplementation((event: string, ...args: any[]) => {
        if (event === "requestTermination" && args.length > 0) {
          const cb = args[args.length - 1];
          if (typeof cb === "function") {
            cb(false);
          }
        }
      });

      await droneSession.onRequestTermination(mockCallback);

      expect(mockSocket.emit).toHaveBeenCalledWith(
        "requestTermination",
        expect.any(Function),
      );
      expect(mockCallback).toHaveBeenCalledWith(false);
    });
  });
});