gadget/gadget-code/src/lib/code-session.ts

// src/lib/code-session.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

import {
  GadgetSocket,
  SocketSession,
  SocketSessionType,
} from "./socket-session";
import {
  IChatSession,
  IDroneRegistration,
  IProject,
  IUser,
  ChatTurnStatus,
  GadgetId,
} from "@gadget/api";

import SocketService from "../services/socket.ts";
import { ChatTurn } from "../models/chat-turn.ts";

export class CodeSession extends SocketSession {
  protected type: SocketSessionType = SocketSessionType.Code;

  protected project: IProject | undefined;
  protected chatSession: IChatSession | undefined;
  protected selectedDrone: IDroneRegistration | undefined;
  protected currentTurnId: GadgetId | undefined;

  constructor(socket: GadgetSocket, user: IUser) {
    super(socket, user);
  }

  register() {
    super.register();

    this.socket.on("requestSessionLock", this.onRequestSessionLock.bind(this));
    this.socket.on("submitPrompt", this.onSubmitPrompt.bind(this));
  }

  /**
   * Sets the selected drone for this code session.
   */
  setSelectedDrone(registration: IDroneRegistration): void {
    this.selectedDrone = registration;
  }

  /**
   * Sets the active chat session and project for this code session.
   */
  setChatSession(chatSession: IChatSession, project: IProject): void {
    this.chatSession = chatSession;
    this.project = project;
  }

  /**
   * Called when the IDE sends a requestSessionLock event to lock a gadget-drone
   * instance to this code session.
   * @param registration the gadget-drone registration to which the request will
   * be sent.
   * @param project the project we're locking the drone to
   * @param chatSession the chat session we're locking the drone to
   * @param cb response callback to call with the result of the request
   */
  onRequestSessionLock(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    cb: (success: boolean, chatSessionId: string) => void,
  ) {
    const droneSession = SocketService.getDroneSession(registration);
    droneSession.socket.emit(
      "requestSessionLock",
      registration,
      project,
      chatSession,
      (success: boolean, chatSessionId: string): void => {
        if (success) {
          this.selectedDrone = registration;
          this.chatSession = chatSession;
          this.project = project;
          SocketService.registerChatSession(chatSession._id, this);
          droneSession.setChatSessionId(chatSession._id);
        }
        cb(success, chatSessionId);
      },
    );
  }

  /**
   * Called when the IDE submits a prompt to be processed by the agent.
   * Creates a ChatTurn document and sends a work order to the selected drone.
   */
  async onSubmitPrompt(content: string): Promise<void> {
    if (!this.selectedDrone) {
      this.log.warn("prompt rejected: no drone selected");
      throw new Error("No drone selected");
    }
    if (!this.chatSession) {
      this.log.warn("prompt rejected: no chat session active");
      throw new Error("No chat session active");
    }
    if (!this.project) {
      this.log.warn("prompt rejected: no project selected");
      throw new Error("No project selected");
    }

    const droneSession = SocketService.getDroneSession(this.selectedDrone);

    const turn = new ChatTurn({
      createdAt: new Date(),
      user: this.user._id,
      project: this.project._id,
      session: this.chatSession._id,
      provider: this.chatSession.provider,
      llm: this.chatSession.selectedModel,
      mode: this.chatSession.mode,
      status: ChatTurnStatus.Processing,
      prompts: {
        user: content,
        system: undefined,
      },
      toolCalls: [],
      subagents: [],
      stats: {
        toolCallCount: 0,
        inputTokens: 0,
        thinkingTokenCount: 0,
        responseTokens: 0,
        durationMs: 0,
        durationLabel: "pending",
      },
    });
    await turn.save();
    this.currentTurnId = turn._id;

    this.log.info("ChatTurn created", {
      turnId: turn._id,
      chatSessionId: this.chatSession._id,
    });

    droneSession.setCurrentTurnId(turn._id);

    droneSession.socket.emit(
      "processWorkOrder",
      this.selectedDrone,
      this.project,
      this.chatSession,
      turn,
      (success: boolean, message?: string) => {
        if (success) {
          this.log.info("work order accepted by drone", {
            turnId: turn._id,
            message,
          });
        } else {
          this.log.error("work order rejected by drone", {
            turnId: turn._id,
            message,
          });
          turn.status = ChatTurnStatus.Error;
          turn.response = message || "Drone rejected work order";
          turn.save();
        }
      },
    );
  }
}