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

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Types } from "mongoose";
import { CodeSession } from "../src/lib/code-session";
import {
  IChatSession,
  IProject,
  IUser,
  IDroneRegistration,
  ChatTurnStatus,
} from "@gadget/api";
import SocketService from "../src/services/socket";
import ChatSessionService from "../src/services/chat-session";
import { ChatTurn } from "../src/models/chat-turn";
import ChatSession from "../src/models/chat-session";
import { nanoid } from "nanoid";

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

describe("CodeSession", () => {
  let mockSocket: any;
  let mockUser: IUser;
  let mockDrone: IDroneRegistration;
  let mockChatSession: IChatSession;
  let mockProject: IProject;
  let codeSession: CodeSession;

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

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

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

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

    mockProject = {
      _id: nanoid(),
      slug: "test-project",
      name: "Test Project",
    } as IProject;

    mockChatSession = {
      _id: nanoid(),
      name: "New Chat Session",
      mode: "build",
      provider: nanoid(),
      selectedModel: "llama3.1",
      user: mockUser,
      project: mockProject,
      stats: {
        turnCount: 0,
        toolCallCount: 0,
        inputTokens: 0,
        outputTokens: 0,
      },
    } as IChatSession;

    codeSession = new CodeSession(mockSocket, mockUser);
  });

  describe("setSelectedDrone", () => {
    it("should set the selected drone", () => {
      codeSession.setSelectedDrone(mockDrone);
      expect(() => codeSession.setSelectedDrone(mockDrone)).not.toThrow();
    });
  });

  describe("setChatSession", () => {
    it("should set the chat session and project", () => {
      codeSession.setChatSession(mockChatSession, mockProject);
      expect(() =>
        codeSession.setChatSession(mockChatSession, mockProject),
      ).not.toThrow();
    });
  });

  describe("onSubmitPrompt", () => {
    let mockTurn: any;
    let mockDroneSession: any;
    let cb: any;

    beforeEach(() => {
      mockTurn = {
        _id: nanoid(),
        status: ChatTurnStatus.Processing,
        errorMessage: "",
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(ChatTurn).mockImplementation(function () {
        return mockTurn as any;
      });
      (vi.mocked(ChatTurn) as any).populate = vi
        .fn()
        .mockResolvedValue(mockTurn);

      mockDroneSession = {
        socket: {
          emit: vi.fn(),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(
        mockDroneSession as any,
      );

      vi.mocked(ChatSessionService.createTurn).mockResolvedValue(
        mockTurn as any,
      );

      cb = vi.fn();

      codeSession.setSelectedDrone(mockDrone);
      codeSession.setChatSession(mockChatSession, mockProject);
    });

    it("should return error if no drone is selected", async () => {
      codeSession = new CodeSession(mockSocket, mockUser);
      codeSession.setChatSession(mockChatSession, mockProject);

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, { message: "No drone selected" });
    });

    it("should return error if no chat session is active", async () => {
      codeSession = new CodeSession(mockSocket, mockUser);
      codeSession.setSelectedDrone(mockDrone);

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, {
        message: "No chat session active",
      });
    });

    it("should return error if no project is selected", async () => {
      codeSession.setChatSession(mockChatSession, undefined as any);
      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, { message: "No project selected" });
    });

    it("should create a ChatTurn, increment turnCount, and emit processWorkOrder", async () => {
      const updatedSession = {
        ...mockChatSession,
        stats: { ...mockChatSession.stats, turnCount: 1 },
      };
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(
        updatedSession as any,
      );
      vi.mocked(
        ChatSessionService.generateSessionNameFromPrompt,
      ).mockResolvedValue(updatedSession as any);

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(ChatSessionService.createTurn).toHaveBeenCalledWith(
        mockChatSession,
        "test prompt",
      );
      expect(ChatSession.findOneAndUpdate).toHaveBeenCalledWith(
        { _id: mockChatSession._id },
        { $inc: { "stats.turnCount": 1 } },
        expect.objectContaining({ new: true }),
      );
      expect(mockDroneSession.socket.emit).toHaveBeenCalledWith(
        "processWorkOrder",
        mockDrone,
        mockTurn,
        expect.any(Function),
      );
    });

    it("should update ChatTurn to Error status if drone rejects work order", async () => {
      const updatedSession = {
        ...mockChatSession,
        stats: { ...mockChatSession.stats, turnCount: 1 },
      };
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(
        updatedSession as any,
      );
      mockDroneSession.socket.emit = vi.fn((event: string, ...args: any[]) => {
        const callback = args[args.length - 1];
        callback(false, "Drone is busy");
      });

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(mockTurn.status).toBe(ChatTurnStatus.Error);
      expect(mockTurn.errorMessage).toBe("Drone is busy");
      expect(mockTurn.save).toHaveBeenCalled();
    });

    it("should reject prompt and delete turn if session was removed during increment", async () => {
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(null);
      vi.mocked(ChatSessionService.delete).mockResolvedValue(undefined);

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(ChatSessionService.delete).toHaveBeenCalledWith(mockChatSession._id);
      expect(cb).toHaveBeenCalledWith(false, {});
    });

    it("should auto-generate session name and emit sessionUpdated on first prompt with default name", async () => {
      const updatedSession = {
        ...mockChatSession,
        stats: { ...mockChatSession.stats, turnCount: 1 },
      };
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(
        updatedSession as any,
      );

      const namedSession = {
        ...updatedSession,
        name: "Build a REST API",
      };
      vi.mocked(
        ChatSessionService.generateSessionNameFromPrompt,
      ).mockResolvedValue(namedSession as any);

      await codeSession.onSubmitPrompt("build a rest api", cb);

      expect(
        ChatSessionService.generateSessionNameFromPrompt,
      ).toHaveBeenCalledWith(updatedSession, "build a rest api");
      expect(mockSocket.emit).toHaveBeenCalledWith("sessionUpdated", {
        name: "Build a REST API",
      });
    });

    it("should NOT auto-generate name if session already has a custom name", async () => {
      const customSession = {
        ...mockChatSession,
        name: "My Custom Session",
      };
      codeSession.setChatSession(customSession, mockProject);

      const updatedSession = {
        ...customSession,
        stats: { ...customSession.stats, turnCount: 1 },
      };
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(
        updatedSession as any,
      );

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(
        ChatSessionService.generateSessionNameFromPrompt,
      ).not.toHaveBeenCalled();
      expect(mockSocket.emit).not.toHaveBeenCalledWith(
        "sessionUpdated",
        expect.anything(),
      );
    });

    it("should NOT auto-generate name on subsequent prompts (turnCount > 1)", async () => {
      const multiTurnSession = {
        ...mockChatSession,
        stats: { ...mockChatSession.stats, turnCount: 5 },
      };
      codeSession.setChatSession(multiTurnSession, mockProject);

      const updatedSession = {
        ...multiTurnSession,
        stats: { ...multiTurnSession.stats, turnCount: 6 },
      };
      vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(
        updatedSession as any,
      );

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(
        ChatSessionService.generateSessionNameFromPrompt,
      ).not.toHaveBeenCalled();
      expect(mockSocket.emit).not.toHaveBeenCalledWith(
        "sessionUpdated",
        expect.anything(),
      );
    });
  });

  describe("onRequestSessionLock", () => {
    it("should set selected drone, chat session, and project on success", () => {
      const mockDroneSession = {
        socket: {
          emit: vi.fn((event: string, ...args: any[]) => {
            const callback = args[args.length - 1];
            callback(true, mockChatSession._id);
          }),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(
        mockDroneSession as any,
      );

      const callback = vi.fn();
      codeSession.onRequestSessionLock(
        mockDrone,
        mockProject,
        mockChatSession,
        callback,
      );

      expect(callback).toHaveBeenCalledWith(true, mockChatSession._id);
    });

    it("should not set session data on failure", () => {
      const mockDroneSession = {
        socket: {
          emit: vi.fn((event: string, ...args: any[]) => {
            const callback = args[args.length - 1];
            callback(false, "");
          }),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(
        mockDroneSession as any,
      );

      const callback = vi.fn();
      codeSession.onRequestSessionLock(
        mockDrone,
        mockProject,
        mockChatSession,
        callback,
      );

      expect(callback).toHaveBeenCalledWith(false, "");
    });
  });
});