socket protocol completeness, fixes, and tests

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<socket.id, DroneSession> - Primary storage by socket ID
+2.  **`droneRegistrationIndex`**: Map<registration._id, DroneSession> - Lookup by drone registration
+3.  **`codeSessions`**: Map<socket.id, CodeSession> - Primary storage by socket ID
+4.  **`codeSessionUserIndex`**: Map<user._id, CodeSession> - Lookup by user ID
+5.  **`chatSessionIndex`**: Map<chatSessionId, CodeSession> - Reverse lookup from chat session to IDE
+
+All indexes are kept in sync during connection and disconnection.
 
 ---
 

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 <rob.colbert@openplatform.us>
+// 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<SeedResult> {
+  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);
+  });

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<void> {
+    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);
+    });
+  }
 }

gadget-code/src/services/socket.ts

@@ -33,6 +33,8 @@ class SocketService extends DtpService {
   private codeSessions: CodeSessionMap = new Map<string, CodeSession>();
   private droneSessions: DroneSessionMap = new Map<string, DroneSession>();
   private chatSessionIndex: ChatSessionCodeSessionMap = new Map<string, CodeSession>();
+  private droneRegistrationIndex: DroneSessionMap = new Map<string, DroneSession>();
+  private codeSessionUserIndex: CodeSessionMap = new Map<string, CodeSession>();
 
   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;

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

@@ -0,0 +1,175 @@
+// tests/drone-service.test.ts
+// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
+// 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();
+  });
+});

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);
+    });
+  });
 });

gadget-code/tests/fixtures/index.ts

@@ -0,0 +1,14 @@
+// tests/fixtures/index.ts
+// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
+// All Rights Reserved
+
+export {
+  createMockSocket,
+  createMockUser,
+  createMockDroneRegistration,
+  createMockChatSession,
+  createMockProject,
+  captureSocketEmits,
+  extractCallback,
+  type EmitCall,
+} from '../helpers/socket-test-helpers';

gadget-code/tests/helpers/socket-test-helpers.ts

@@ -0,0 +1,136 @@
+// tests/helpers/socket-test-helpers.ts
+// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
+// 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>): 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>
+): 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>
+): 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>
+): 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;
+}

gadget-code/tests/socket-service.test.ts

@@ -0,0 +1,325 @@
+// tests/socket-service.test.ts
+// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
+// 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);
+    });
+  });
+});

gadget-drone/.gitignore

@@ -1,4 +1,5 @@
 .env
+.gadget/*
 
 dist
 logs