From 15da8fc444abbcb8e69f6af7aec28603dbb8ad98 Mon Sep 17 00:00:00 2001 From: Rob Colbert Date: Thu, 30 Apr 2026 16:51:33 -0400 Subject: [PATCH] socket protocol completeness, fixes, and tests --- docs/socket-protocol.md | 49 +++ gadget-code/scripts/seed-socket-test-data.ts | 181 ++++++++++ gadget-code/src/lib/drone-session.ts | 17 + gadget-code/src/services/socket.ts | 19 +- gadget-code/tests/drone-service.test.ts | 175 ++++++++++ gadget-code/tests/drone-session.test.ts | 39 +++ gadget-code/tests/fixtures/index.ts | 14 + .../tests/helpers/socket-test-helpers.ts | 136 ++++++++ gadget-code/tests/socket-service.test.ts | 325 ++++++++++++++++++ gadget-drone/.gitignore | 1 + 10 files changed, 954 insertions(+), 2 deletions(-) create mode 100644 gadget-code/scripts/seed-socket-test-data.ts create mode 100644 gadget-code/tests/drone-service.test.ts create mode 100644 gadget-code/tests/fixtures/index.ts create mode 100644 gadget-code/tests/helpers/socket-test-helpers.ts create mode 100644 gadget-code/tests/socket-service.test.ts diff --git a/docs/socket-protocol.md b/docs/socket-protocol.md index 9322cd9..a8459df 100644 --- a/docs/socket-protocol.md +++ b/docs/socket-protocol.md @@ -27,10 +27,12 @@ Defined in `packages/api/src/messages/socket.ts`. * `toolCall`: Emit a specific tool execution event with result. * `workOrderComplete`: Signal that a prompt processing turn is finished. * `requestCrashRecovery`: Inbound from drone on restart if it finds a stalled work order. +* `requestTermination`: Acknowledgment from drone that termination request was received. ### Web -> Drone (Server to Client) * `processWorkOrder`: Command to start processing a specific prompt/turn. * `crashRecoveryResponse`: Command to `discard` or `retry` a stalled work order. +* `requestTermination`: Command to immediately terminate the drone process. --- @@ -62,6 +64,29 @@ Defined in `packages/api/src/messages/socket.ts`. * Forwards event to **IDE**. * Clears `currentTurnId` from the drone session. +### 3.4 Drone Termination Flow +1. **User** clicks "Terminate" button in Drone Manager UI. +2. **IDE** calls `POST /api/v1/drone/registration/:id/terminate`. +3. **Web (`DroneService.ts`)**: + * Checks if drone is already offline → returns error if so. + * Looks up `DroneSession` via `SocketService.getDroneSession()`. + * If drone not connected → marks as offline immediately, returns success. + * Emits `requestTermination` to drone socket with callback. + * Starts 10-second timeout. +4. **Web (`DroneSession.ts`)**: + * Receives `requestTermination` event. + * Logs the termination request. + * Forwards `requestTermination` to drone socket (passthrough). +5. **Drone (`gadget-drone.ts`)**: + * Receives `requestTermination` from platform. + * Calls callback with `success: true`. + * Sends `SIGINT` to self, triggering graceful shutdown. + * Updates status to `Offline` during shutdown. +6. **Web (`DroneService.ts`)**: + * Drone accepts termination → polls DB every 500ms waiting for `Offline` status. + * Drone goes offline → resolves with success. + * Timeout expires (10s) → forces status to `Offline`, resolves with success. + --- ## 4. Message Signatures (TS Reference) @@ -87,6 +112,9 @@ type ProcessWorkOrderMessage = ( turn: IChatTurn, cb: (success: boolean, message?: string) => void ) => void; +type RequestTerminationMessage = ( + cb: (success: boolean) => void +) => void; ``` ### Drone -> Web (Streaming) @@ -99,6 +127,14 @@ type ToolCallMessage = ( params: string, // JSON.stringify response: string // JSON.stringify ) => void; +type WorkOrderCompleteMessage = ( + workOrderId: string, + success: boolean, + message?: string +) => void; +type RequestTerminationMessage = ( + cb: (success: boolean) => void +) => void; ``` --- @@ -116,6 +152,19 @@ Manages the IDE socket. Manages the Drone socket. * **Logic**: Maps Drone Registration ID -> Socket ID. * **Routing**: When a drone streams, `DroneSession` looks up the `chatSessionId` in the `SocketService` index to find the return path to the IDE. +* **Session Lookup**: `SocketService` maintains a `droneRegistrationIndex` Map that maps `registration._id` → `DroneSession` for efficient lookup by registration ID. + +### Session Indexing Architecture + +The `SocketService` maintains multiple indexes for efficient session lookup: + +1. **`droneSessions`**: Map - Primary storage by socket ID +2. **`droneRegistrationIndex`**: Map - Lookup by drone registration +3. **`codeSessions`**: Map - Primary storage by socket ID +4. **`codeSessionUserIndex`**: Map - Lookup by user ID +5. **`chatSessionIndex`**: Map - Reverse lookup from chat session to IDE + +All indexes are kept in sync during connection and disconnection. --- diff --git a/gadget-code/scripts/seed-socket-test-data.ts b/gadget-code/scripts/seed-socket-test-data.ts new file mode 100644 index 0000000..1f984a1 --- /dev/null +++ b/gadget-code/scripts/seed-socket-test-data.ts @@ -0,0 +1,181 @@ +#!/usr/bin/env node +// scripts/seed-socket-test-data.ts +// Copyright (C) 2026 Robert Colbert +// All Rights Reserved + +/** + * Seeds the database with test data for socket messaging tests. + * + * This script creates: + * - Test user (if not exists) + * - Test project (unique per run) + * - Test chat session (unique per run) + * - Test drone registrations (unique per run) + * + * IMPORTANT: This script does NOT delete existing data. It only creates new, + * unique records. Tests are responsible for cleaning up their own data. + * + * Usage: + * npx tsx scripts/seed-socket-test-data.ts + * + * Output: + * JSON object with created resource IDs for test cleanup + */ + +import mongoose from 'mongoose'; +import { config } from 'dotenv'; +import { IUser } from '@gadget/api'; +import User from '../src/models/user'; +import Project from '../src/models/project'; +import ChatSession from '../src/models/chat-session'; +import DroneRegistration from '../src/models/drone-registration'; +import { AiProvider } from '../src/models/ai-provider'; + +config(); + +interface SeedResult { + userId: string; + projectId: string; + chatSessionId: string; + droneIds: string[]; + providerId?: string; + createdAt: string; + note: string; +} + +async function seedSocketTestData(): Promise { + const NOW = new Date(); + const timestamp = NOW.toISOString().replace(/[:.]/g, '-'); + + try { + // Connect to database + await mongoose.connect(process.env.MONGODB_URI || 'mongodb://localhost:27017/gadget-code'); + console.log('Connected to MongoDB'); + + // Find or create test user + let user = await User.findOne({ email: 'test-socket@gadget-code.test' }); + if (!user) { + user = new User({ + email: 'test-socket@gadget-code.test', + displayName: 'Socket Test User', + passwordSalt: 'test-salt', + passwordHash: 'test-hash', + banned: false, + admin: false, + createdAt: NOW, + updatedAt: NOW, + }); + await user.save(); + console.log('Created test user:', user._id); + } else { + console.log('Found existing test user:', user._id); + } + + // Find or create a test AI provider + let provider = await AiProvider.findOne({ name: 'Test Socket Provider' }); + if (!provider) { + provider = new AiProvider({ + name: 'Test Socket Provider', + apiType: 'ollama', + baseUrl: 'http://localhost:11434', + enabled: true, + models: [{ + id: 'llama3.2', + name: 'Llama 3.2', + capabilities: { + canCallTools: true, + hasVision: false, + hasEmbedding: false, + hasThinking: false, + isInstructTuned: true, + }, + }], + lastModelRefresh: NOW, + }); + await provider.save(); + console.log('Created test provider:', provider._id); + } else { + console.log('Found existing test provider:', provider._id); + } + + // Create test project (unique per run) + const project = new Project({ + user: user._id, + slug: `socket-test-${timestamp}`, + name: `Socket Test Project ${timestamp}`, + gitUrl: 'https://github.com/test/socket-test.git', + createdAt: NOW, + updatedAt: NOW, + }); + await project.save(); + console.log('Created test project:', project._id); + + // Create test chat session (unique per run) + const chatSession = new ChatSession({ + user: user._id, + project: project._id, + provider: provider._id, + selectedModel: 'llama3.2', + mode: 'build', + name: `Socket Test Session ${timestamp}`, + stats: { + toolCallCount: 0, + fileOpCount: 0, + subagentCount: 0, + }, + createdAt: NOW, + updatedAt: NOW, + }); + await chatSession.save(); + console.log('Created test chat session:', chatSession._id); + + // Create test drone registrations (unique per run) + const droneCount = 3; + const droneIds: string[] = []; + + for (let i = 0; i < droneCount; i++) { + const drone = new DroneRegistration({ + user: user._id, + hostname: `test-drone-${timestamp}-${i}`, + workspaceDir: `/tmp/socket-test-${timestamp}-${i}`, + status: 'available', + createdAt: NOW, + updatedAt: NOW, + }); + await drone.save(); + droneIds.push(drone._id.toHexString()); + console.log(`Created test drone ${i + 1}/${droneCount}:`, drone._id); + } + + const result: SeedResult = { + userId: user._id.toHexString(), + projectId: project._id.toHexString(), + chatSessionId: chatSession._id.toHexString(), + droneIds, + providerId: provider._id.toHexString(), + createdAt: NOW.toISOString(), + note: 'TEST DATA - Safe to delete. Created for socket messaging tests.', + }; + + console.log('\n✅ Seed complete! Save this output for test cleanup:'); + console.log(JSON.stringify(result, null, 2)); + + return result; + } catch (error) { + console.error('❌ Seed failed:', error); + process.exit(1); + } finally { + await mongoose.disconnect(); + console.log('\nDisconnected from MongoDB'); + } +} + +// Run the seed +seedSocketTestData() + .then(() => { + process.exit(0); + }) + .catch((error) => { + console.error('Unhandled error:', error); + process.exit(1); + }); diff --git a/gadget-code/src/lib/drone-session.ts b/gadget-code/src/lib/drone-session.ts index 6d04cc7..f76198b 100644 --- a/gadget-code/src/lib/drone-session.ts +++ b/gadget-code/src/lib/drone-session.ts @@ -31,6 +31,7 @@ export class DroneSession extends SocketSession { this.socket.on("toolCall", this.onToolCall.bind(this)); this.socket.on("workOrderComplete", this.onWorkOrderComplete.bind(this)); this.socket.on("requestCrashRecovery", this.onRequestCrashRecovery.bind(this)); + this.socket.on("requestTermination", this.onRequestTermination.bind(this)); } /** @@ -247,4 +248,20 @@ export class DroneSession extends SocketSession { turnId: turn._id, }); } + + /** + * Called when the platform requests termination of this drone. + * Forwards the termination request to the drone socket with logging. + * @param cb Callback to invoke with termination result + */ + async onRequestTermination(cb: (success: boolean) => void): Promise { + this.log.info("requestTermination received, forwarding to drone", { + registrationId: this.registration._id.toHexString(), + }); + + this.socket.emit("requestTermination", (success: boolean) => { + this.log.info("requestTermination forwarded to drone", { success }); + cb(success); + }); + } } diff --git a/gadget-code/src/services/socket.ts b/gadget-code/src/services/socket.ts index e66808c..6b38a2b 100644 --- a/gadget-code/src/services/socket.ts +++ b/gadget-code/src/services/socket.ts @@ -33,6 +33,8 @@ class SocketService extends DtpService { private codeSessions: CodeSessionMap = new Map(); private droneSessions: DroneSessionMap = new Map(); private chatSessionIndex: ChatSessionCodeSessionMap = new Map(); + private droneRegistrationIndex: DroneSessionMap = new Map(); + private codeSessionUserIndex: CodeSessionMap = new Map(); private io?: Server< ClientToServerEvents, @@ -96,6 +98,7 @@ class SocketService extends DtpService { const session: CodeSession = new CodeSession(socket, user); this.codeSessions.set(socket.id, session); + this.codeSessionUserIndex.set(user._id.toHexString(), session); socket.data = { sessionType: SocketSessionType.Code }; socket.on("disconnect", (reason: DisconnectReason, extra?: unknown) => { @@ -123,6 +126,7 @@ class SocketService extends DtpService { const droneSession: DroneSession = new DroneSession(socket, registration); this.droneSessions.set(socket.id, droneSession); + this.droneRegistrationIndex.set(registration._id.toHexString(), droneSession); socket.data = { sessionType: SocketSessionType.Drone }; socket.on("disconnect", (reason: DisconnectReason, extra?: unknown) => { @@ -186,10 +190,21 @@ class SocketService extends DtpService { case SocketSessionType.Code: this.log.info("closing code socket session", { id: socket.id }); this.codeSessions.delete(socket.id); + const codeUserIndex = (this as any).codeSessionUserIndex; + if (codeUserIndex) { + const session = this.codeSessions.get(socket.id); + if (session) { + codeUserIndex.delete(session.user._id.toHexString()); + } + } return; case SocketSessionType.Drone: this.log.info("closing drone socket session", { id: socket.id }); + const droneSession = this.droneSessions.get(socket.id); + if (droneSession) { + this.droneRegistrationIndex.delete(droneSession.registration._id.toHexString()); + } this.droneSessions.delete(socket.id); return; @@ -204,7 +219,7 @@ class SocketService extends DtpService { } getCodeSession(ideSession: IIdeSession): CodeSession { - const session = this.codeSessions.get(ideSession._id.toHexString()); + const session = this.codeSessionUserIndex.get(ideSession._id.toHexString()); if (!session) { const error = new Error("code session not found"); error.statusCode = 404; @@ -214,7 +229,7 @@ class SocketService extends DtpService { } getDroneSession(registration: IDroneRegistration): DroneSession { - const session = this.droneSessions.get(registration._id.toHexString()); + const session = this.droneRegistrationIndex.get(registration._id.toHexString()); if (!session) { const error = new Error("drone session not found"); error.statusCode = 404; diff --git a/gadget-code/tests/drone-service.test.ts b/gadget-code/tests/drone-service.test.ts new file mode 100644 index 0000000..7b66dd8 --- /dev/null +++ b/gadget-code/tests/drone-service.test.ts @@ -0,0 +1,175 @@ +// tests/drone-service.test.ts +// Copyright (C) 2026 Robert Colbert +// All Rights Reserved + +import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { Types } from 'mongoose'; +import { DroneStatus, IDroneRegistration } from '@gadget/api'; +import { createMockUser, createMockDroneRegistration } from './fixtures'; + +// Mock SocketService +vi.mock('../src/services/socket'); + +import DroneService from '../src/services/drone'; +import SocketService from '../src/services/socket'; + +describe('DroneService.requestTermination', () => { + let mockUser: any; + let mockDrone: any; + let mockDroneSession: any; + let mockSocket: any; + + beforeEach(() => { + vi.clearAllMocks(); + mockUser = createMockUser(); + mockDrone = createMockDroneRegistration(mockUser); + mockSocket = { + emit: vi.fn(), + }; + mockDroneSession = { + socket: mockSocket, + registration: mockDrone, + }; + + // Setup mocks for DroneService methods + vi.spyOn(DroneService, 'getById').mockResolvedValue(mockDrone); + vi.spyOn(DroneService, 'setStatus').mockResolvedValue(mockDrone); + vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession); + }); + + it('should return error if drone is already offline', async () => { + const offlineDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Offline, + }); + vi.spyOn(DroneService, 'getById').mockResolvedValue(offlineDrone); + + const result = await DroneService.requestTermination(offlineDrone._id); + + expect(result.success).toBe(false); + expect(result.message).toBe('Drone is already offline'); + expect(SocketService.getDroneSession).not.toHaveBeenCalled(); + }); + + it('should mark drone offline if not connected and return success', async () => { + const connectedDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Available, + }); + vi.spyOn(DroneService, 'getById').mockResolvedValue(connectedDrone); + vi.spyOn(SocketService, 'getDroneSession').mockImplementation(() => { + throw new Error('drone session not found'); + }); + + const result = await DroneService.requestTermination(connectedDrone._id); + + expect(result.success).toBe(true); + expect(result.message).toContain('not connected'); + expect(DroneService.setStatus).toHaveBeenCalledWith(connectedDrone, DroneStatus.Offline); + }); + + it('should emit requestTermination to drone socket', async () => { + const availableDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Available, + }); + vi.spyOn(DroneService, 'getById') + .mockResolvedValueOnce(availableDrone) + .mockResolvedValueOnce(availableDrone) + .mockResolvedValueOnce(createMockDroneRegistration(mockUser, { status: DroneStatus.Offline })); + vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession); + + // Mock socket emit to call the callback with success + mockSocket.emit.mockImplementation((event: string, cb: Function) => { + if (event === 'requestTermination') { + cb(true); + } + }); + + const result = await DroneService.requestTermination(availableDrone._id); + + expect(SocketService.getDroneSession).toHaveBeenCalledWith(availableDrone); + expect(mockSocket.emit).toHaveBeenCalledWith('requestTermination', expect.any(Function)); + expect(result.success).toBe(true); + }); + + it('should handle drone rejection of termination', async () => { + const availableDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Available, + }); + vi.spyOn(DroneService, 'getById').mockResolvedValue(availableDrone); + vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession); + + // Mock socket emit to call the callback with failure + mockSocket.emit.mockImplementation((event: string, cb: Function) => { + if (event === 'requestTermination') { + cb(false); + } + }); + + const result = await DroneService.requestTermination(availableDrone._id); + + expect(result.success).toBe(false); + expect(result.message).toBe('Drone rejected termination request'); + }); + + it('should timeout and force offline after 10 seconds', async () => { + vi.useFakeTimers(); + + const availableDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Available, + }); + vi.spyOn(DroneService, 'getById').mockResolvedValue(availableDrone); + vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession); + + // Mock socket emit but never call the callback (simulating no response) + mockSocket.emit.mockImplementation(() => { + // Never call callback + }); + + const terminationPromise = DroneService.requestTermination(availableDrone._id); + + // Advance time past 10 second timeout + await vi.advanceTimersByTimeAsync(10000); + + const result = await terminationPromise; + + expect(result.success).toBe(true); + expect(result.message).toContain('timed out'); + expect(DroneService.setStatus).toHaveBeenCalledWith(availableDrone, DroneStatus.Offline); + + vi.useRealTimers(); + }); + + it('should poll for offline status after drone accepts termination', async () => { + vi.useFakeTimers(); + + const availableDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Available, + }); + const offlineDrone = createMockDroneRegistration(mockUser, { + status: DroneStatus.Offline, + }); + + vi.spyOn(DroneService, 'getById') + .mockResolvedValueOnce(availableDrone) // Initial check + .mockResolvedValueOnce(offlineDrone); // Poll after accept + vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession); + + // Mock socket emit to call callback with success immediately + mockSocket.emit.mockImplementation((event: string, cb: Function) => { + if (event === 'requestTermination') { + cb(true); + } + }); + + const terminationPromise = DroneService.requestTermination(availableDrone._id); + + // Advance time to allow callback and one poll cycle (500ms interval) + await vi.advanceTimersByTimeAsync(600); + + const result = await terminationPromise; + + expect(result.success).toBe(true); + expect(result.message).toBeUndefined(); + + vi.useRealTimers(); + }); +}); diff --git a/gadget-code/tests/drone-session.test.ts b/gadget-code/tests/drone-session.test.ts index b0f5ab1..c28b617 100644 --- a/gadget-code/tests/drone-session.test.ts +++ b/gadget-code/tests/drone-session.test.ts @@ -54,6 +54,8 @@ describe('DroneSession', () => { 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)); }); }); @@ -214,4 +216,41 @@ describe('DroneSession', () => { 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); + }); + }); }); diff --git a/gadget-code/tests/fixtures/index.ts b/gadget-code/tests/fixtures/index.ts new file mode 100644 index 0000000..cfffe0f --- /dev/null +++ b/gadget-code/tests/fixtures/index.ts @@ -0,0 +1,14 @@ +// tests/fixtures/index.ts +// Copyright (C) 2026 Robert Colbert +// All Rights Reserved + +export { + createMockSocket, + createMockUser, + createMockDroneRegistration, + createMockChatSession, + createMockProject, + captureSocketEmits, + extractCallback, + type EmitCall, +} from '../helpers/socket-test-helpers'; diff --git a/gadget-code/tests/helpers/socket-test-helpers.ts b/gadget-code/tests/helpers/socket-test-helpers.ts new file mode 100644 index 0000000..f41b60d --- /dev/null +++ b/gadget-code/tests/helpers/socket-test-helpers.ts @@ -0,0 +1,136 @@ +// tests/helpers/socket-test-helpers.ts +// Copyright (C) 2026 Robert Colbert +// All Rights Reserved + +import { vi } from 'vitest'; +import { Types } from 'mongoose'; +import { + IUser, + IDroneRegistration, + IChatSession, + IProject, + DroneStatus, + ChatSessionMode, +} from '@gadget/api'; + +/** + * Creates a mock socket object with common methods stubbed. + */ +export function createMockSocket(id?: string) { + return { + id: id || `socket-${new Types.ObjectId().toHexString()}`, + on: vi.fn(), + emit: vi.fn(), + disconnect: vi.fn(), + data: {}, + handshake: { + auth: { token: '' }, + }, + }; +} + +/** + * Creates a mock user object for testing. + */ +export function createMockUser(overrides?: Partial): IUser { + return { + _id: new Types.ObjectId(), + email: `user-${new Types.ObjectId().toHexString()}@example.com`, + displayName: 'Test User', + banned: false, + admin: false, + createdAt: new Date(), + updatedAt: new Date(), + ...overrides, + } as IUser; +} + +/** + * Creates a mock drone registration for testing. + */ +export function createMockDroneRegistration( + user?: IUser, + overrides?: Partial +): IDroneRegistration { + const testUser = user || createMockUser(); + return { + _id: new Types.ObjectId(), + createdAt: new Date(), + updatedAt: new Date(), + user: testUser, + hostname: `drone-${new Types.ObjectId().toHexString()}`, + workspaceDir: '/test/workspace', + status: DroneStatus.Available, + ...overrides, + } as IDroneRegistration; +} + +/** + * Creates a mock chat session for testing. + */ +export function createMockChatSession( + user?: IUser, + project?: IProject, + overrides?: Partial +): IChatSession { + return { + _id: new Types.ObjectId(), + createdAt: new Date(), + updatedAt: new Date(), + user: user?._id || new Types.ObjectId(), + project: project?._id || new Types.ObjectId(), + name: 'Test Chat Session', + mode: ChatSessionMode.Build, + provider: new Types.ObjectId(), + selectedModel: 'llama3.2', + stats: { + toolCallCount: 0, + fileOpCount: 0, + subagentCount: 0, + }, + ...overrides, + } as IChatSession; +} + +/** + * Creates a mock project for testing. + */ +export function createMockProject( + user?: IUser, + overrides?: Partial +): IProject { + return { + _id: new Types.ObjectId(), + createdAt: new Date(), + updatedAt: new Date(), + user: user?._id || new Types.ObjectId(), + slug: `project-${new Types.ObjectId().toHexString()}`, + name: 'Test Project', + gitUrl: 'https://github.com/test/test.git', + ...overrides, + } as IProject; +} + +/** + * Captures all emit calls for a mock socket. + */ +export interface EmitCall { + event: string; + args: any[]; +} + +export function captureSocketEmits(socket: any): EmitCall[] { + const calls: EmitCall[] = []; + socket.emit.mockImplementation((event: string, ...args: any[]) => { + calls.push({ event, args }); + }); + return calls; +} + +/** + * Extracts callback from emit arguments for testing. + */ +export function extractCallback(emitCall: EmitCall): Function | null { + const lastArg = emitCall.args[emitCall.args.length - 1]; + return typeof lastArg === 'function' ? lastArg : null; +} diff --git a/gadget-code/tests/socket-service.test.ts b/gadget-code/tests/socket-service.test.ts new file mode 100644 index 0000000..4fc1d08 --- /dev/null +++ b/gadget-code/tests/socket-service.test.ts @@ -0,0 +1,325 @@ +// tests/socket-service.test.ts +// Copyright (C) 2026 Robert Colbert +// All Rights Reserved + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { Types } from 'mongoose'; +import { DroneStatus } from '@gadget/api'; +import { + createMockSocket, + createMockUser, + createMockDroneRegistration, + createMockChatSession, + createMockProject, +} from './fixtures'; + +// Mock the entire socket service module +vi.mock('../src/services/socket', async () => { + const chatSessionIndex = new Map(); + return { + default: { + codeSessions: new Map(), + droneSessions: new Map(), + chatSessionIndex: chatSessionIndex, + droneRegistrationIndex: new Map(), + codeSessionUserIndex: new Map(), + getDroneSession: vi.fn(), + getCodeSession: vi.fn(), + getCodeSessionByChatSessionId: vi.fn(), + registerChatSession: vi.fn((chatSessionId: string, codeSession: any) => { + chatSessionIndex.set(chatSessionId, codeSession); + }), + unregisterChatSession: vi.fn((chatSessionId: string) => { + chatSessionIndex.delete(chatSessionId); + }), + }, + }; +}); + +describe('SocketService Session Indexing', () => { + let SocketService: any; + let mockUser: any; + let mockDrone: any; + let mockChatSession: any; + let mockProject: any; + + beforeEach(async () => { + vi.clearAllMocks(); + + const mod = await import('../src/services/socket'); + SocketService = mod.default; + + // Clear all session maps + SocketService.codeSessions.clear(); + SocketService.droneSessions.clear(); + SocketService.chatSessionIndex.clear(); + SocketService.droneRegistrationIndex.clear(); + SocketService.codeSessionUserIndex.clear(); + + mockUser = createMockUser(); + mockDrone = createMockDroneRegistration(mockUser); + mockChatSession = createMockChatSession(mockUser); + mockProject = createMockProject(mockUser); + }); + + describe('Drone Session Indexing', () => { + it('should store drone session by socket.id and registration._id', () => { + const mockSocket = createMockSocket('drone-socket-123'); + const mockDroneSession = { + socket: mockSocket, + registration: mockDrone, + type: 'drone', + }; + + // Store in both indexes + SocketService.droneSessions.set(mockSocket.id, mockDroneSession); + SocketService.droneRegistrationIndex.set( + mockDrone._id.toHexString(), + mockDroneSession + ); + + // Verify both lookups work + expect(SocketService.droneSessions.get(mockSocket.id)).toBe(mockDroneSession); + expect(SocketService.droneRegistrationIndex.get(mockDrone._id.toHexString())).toBe(mockDroneSession); + }); + + it('should find drone session by registration._id', () => { + const mockSocket = createMockSocket('drone-socket-456'); + const mockDroneSession = { + socket: mockSocket, + registration: mockDrone, + type: 'drone', + }; + + // Store in both indexes + SocketService.droneSessions.set(mockSocket.id, mockDroneSession); + SocketService.droneRegistrationIndex.set( + mockDrone._id.toHexString(), + mockDroneSession + ); + + // Mock getDroneSession to use the registration index + SocketService.getDroneSession.mockImplementation((registration: any) => { + const session = SocketService.droneRegistrationIndex.get(registration._id.toHexString()); + if (!session) { + const error = new Error('drone session not found'); + (error as any).statusCode = 404; + throw error; + } + return session; + }); + + const found = SocketService.getDroneSession(mockDrone); + expect(found).toBe(mockDroneSession); + }); + + it('should throw 404 when drone session not found', () => { + const nonExistentDrone = createMockDroneRegistration(mockUser); + + SocketService.getDroneSession.mockImplementation(() => { + const error = new Error('drone session not found'); + (error as any).statusCode = 404; + throw error; + }); + + expect(() => SocketService.getDroneSession(nonExistentDrone)).toThrow('drone session not found'); + expect(() => SocketService.getDroneSession(nonExistentDrone)).toThrowError(expect.objectContaining({ + statusCode: 404, + })); + }); + + it('should remove drone session from all indexes on disconnect', () => { + const mockSocket = createMockSocket('drone-socket-789'); + const mockDroneSession = { + socket: mockSocket, + registration: mockDrone, + type: 'drone', + }; + + // Store in indexes + SocketService.droneSessions.set(mockSocket.id, mockDroneSession); + SocketService.droneRegistrationIndex.set( + mockDrone._id.toHexString(), + mockDroneSession + ); + + // Simulate disconnect + SocketService.droneSessions.delete(mockSocket.id); + SocketService.droneRegistrationIndex.delete(mockDrone._id.toHexString()); + + // Verify removal from all indexes + expect(SocketService.droneSessions.get(mockSocket.id)).toBeUndefined(); + expect(SocketService.droneRegistrationIndex.get(mockDrone._id.toHexString())).toBeUndefined(); + }); + }); + + describe('Code Session Indexing', () => { + it('should store code session by socket.id and user._id', () => { + const mockSocket = createMockSocket('code-socket-123'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + // Store in both indexes + SocketService.codeSessions.set(mockSocket.id, mockCodeSession); + SocketService.codeSessionUserIndex.set( + mockUser._id.toHexString(), + mockCodeSession + ); + + // Verify both lookups work + expect(SocketService.codeSessions.get(mockSocket.id)).toBe(mockCodeSession); + expect(SocketService.codeSessionUserIndex.get(mockUser._id.toHexString())).toBe(mockCodeSession); + }); + + it('should find code session by user._id', () => { + const mockSocket = createMockSocket('code-socket-456'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + // Store in both indexes + SocketService.codeSessions.set(mockSocket.id, mockCodeSession); + SocketService.codeSessionUserIndex.set( + mockUser._id.toHexString(), + mockCodeSession + ); + + // Mock getCodeSession to use the user index + SocketService.getCodeSession.mockImplementation((user: any) => { + const session = SocketService.codeSessionUserIndex.get(user._id.toHexString()); + if (!session) { + const error = new Error('code session not found'); + (error as any).statusCode = 404; + throw error; + } + return session; + }); + + const found = SocketService.getCodeSession(mockUser); + expect(found).toBe(mockCodeSession); + }); + + it('should throw 404 when code session not found', () => { + const nonExistentUser = createMockUser(); + + SocketService.getCodeSession.mockImplementation(() => { + const error = new Error('code session not found'); + (error as any).statusCode = 404; + throw error; + }); + + expect(() => SocketService.getCodeSession(nonExistentUser)).toThrow('code session not found'); + expect(() => SocketService.getCodeSession(nonExistentUser)).toThrowError(expect.objectContaining({ + statusCode: 404, + })); + }); + + it('should remove code session from all indexes on disconnect', () => { + const mockSocket = createMockSocket('code-socket-789'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + // Store in indexes + SocketService.codeSessions.set(mockSocket.id, mockCodeSession); + SocketService.codeSessionUserIndex.set( + mockUser._id.toHexString(), + mockCodeSession + ); + + // Simulate disconnect + SocketService.codeSessions.delete(mockSocket.id); + SocketService.codeSessionUserIndex.delete(mockUser._id.toHexString()); + + // Verify removal from all indexes + expect(SocketService.codeSessions.get(mockSocket.id)).toBeUndefined(); + expect(SocketService.codeSessionUserIndex.get(mockUser._id.toHexString())).toBeUndefined(); + }); + }); + + describe('Chat Session Index', () => { + it('should register and retrieve code session by chatSessionId', () => { + const mockSocket = createMockSocket('code-socket-chat'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + const chatSessionId = mockChatSession._id.toHexString(); + + // Mock the retrieval to return our session after registration + SocketService.registerChatSession(chatSessionId, mockCodeSession); + SocketService.getCodeSessionByChatSessionId.mockReturnValue(mockCodeSession); + + const found = SocketService.getCodeSessionByChatSessionId(mockChatSession._id); + expect(found).toBe(mockCodeSession); + expect(SocketService.registerChatSession).toHaveBeenCalledWith(chatSessionId, mockCodeSession); + expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(mockChatSession._id); + }); + + it('should handle chatSessionId as string or ObjectId', () => { + const mockSocket = createMockSocket('code-socket-chat2'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + const chatSessionId = mockChatSession._id.toHexString(); + SocketService.registerChatSession(chatSessionId, mockCodeSession); + SocketService.getCodeSessionByChatSessionId.mockReturnValue(mockCodeSession); + + // Test with string + const found1 = SocketService.getCodeSessionByChatSessionId(chatSessionId); + expect(found1).toBe(mockCodeSession); + + // Test with ObjectId + const found2 = SocketService.getCodeSessionByChatSessionId(mockChatSession._id); + expect(found2).toBe(mockCodeSession); + }); + + it('should throw 404 when chat session not found', () => { + const nonExistentChatSessionId = new Types.ObjectId(); + + SocketService.getCodeSessionByChatSessionId.mockImplementation(() => { + const error = new Error('code session not found for chat session'); + (error as any).statusCode = 404; + throw error; + }); + + expect(() => SocketService.getCodeSessionByChatSessionId(nonExistentChatSessionId)).toThrow('code session not found for chat session'); + expect(() => SocketService.getCodeSessionByChatSessionId(nonExistentChatSessionId)).toThrowError(expect.objectContaining({ + statusCode: 404, + })); + }); + + it('should unregister chat session', () => { + const mockSocket = createMockSocket('code-socket-chat3'); + const mockCodeSession = { + socket: mockSocket, + user: mockUser, + type: 'code', + }; + + const chatSessionId = mockChatSession._id.toHexString(); + SocketService.registerChatSession(chatSessionId, mockCodeSession); + SocketService.getCodeSessionByChatSessionId.mockReturnValue(mockCodeSession); + + // Verify it's registered + const found = SocketService.getCodeSessionByChatSessionId(chatSessionId); + expect(found).toBe(mockCodeSession); + + // Unregister + SocketService.unregisterChatSession(chatSessionId); + expect(SocketService.unregisterChatSession).toHaveBeenCalledWith(chatSessionId); + }); + }); +}); diff --git a/gadget-drone/.gitignore b/gadget-drone/.gitignore index 1bf742c..e07f84d 100644 --- a/gadget-drone/.gitignore +++ b/gadget-drone/.gitignore @@ -1,4 +1,5 @@ .env +.gadget/* dist logs