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 } from '@gadget/api';
import SocketService from '../src/services/socket';
import { ChatTurn } from '../src/models/chat-turn';

// 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: Types.ObjectId;
  let mockTurnId: Types.ObjectId;

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

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

    mockUser = {
      _id: new Types.ObjectId(),
      email: 'test@example.com',
      displayName: 'Test User',
    } as IUser;

    mockRegistration = {
      _id: new Types.ObjectId(),
      user: mockUser,
      hostname: 'test-host',
      workspaceDir: '/test/workspace',
      status: 'available',
    } as IDroneRegistration;

    mockChatSessionId = new Types.ObjectId();
    mockTurnId = new Types.ObjectId();

    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 and update ChatTurn', async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(mockCodeSession as any);
      vi.mocked(ChatTurn.findByIdAndUpdate).mockResolvedValue({} as any);

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

      await droneSession.onThinking('thinking content');

      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(mockChatSessionId);
      expect(mockCodeSession.socket.emit).toHaveBeenCalledWith('thinking', 'thinking content');
      expect(ChatTurn.findByIdAndUpdate).toHaveBeenCalledWith(mockTurnId, {
        thinking: '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 and update ChatTurn', async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
      };
      vi.mocked(SocketService.getCodeSessionByChatSessionId).mockReturnValue(mockCodeSession as any);
      vi.mocked(ChatTurn.findByIdAndUpdate).mockResolvedValue({} as any);

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

      await droneSession.onResponse('response content');

      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(mockChatSessionId);
      expect(mockCodeSession.socket.emit).toHaveBeenCalledWith('response', 'response content');
      expect(ChatTurn.findByIdAndUpdate).toHaveBeenCalledWith(mockTurnId, {
        response: '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() },
      };
      const mockTurn = {
        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.socket.emit).toHaveBeenCalledWith('toolCall', '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() },
      };
      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.toHexString(), true);

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

    it('should update ChatTurn to Error status on failure', async () => {
      const mockCodeSession = {
        socket: { emit: vi.fn() },
      };
      const mockTurn = {
        status: ChatTurnStatus.Processing,
        response: '',
        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.toHexString(), false, 'Agent crashed');

      expect(mockTurn.status).toBe(ChatTurnStatus.Error);
      expect(mockTurn.response).toBe('Agent crashed');
      expect(mockTurn.save).toHaveBeenCalled();
    });

    it('should log warning if no chat session is active', async () => {
      await droneSession.onWorkOrderComplete(mockTurnId.toHexString(), 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);
    });
  });
});