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.js";
import {
  FileTreeEntry,
  FileTreeRequestArgs,
  GadgetComponent,
  GadgetLogLevel,
  IChatSession,
  IDroneRegistration,
  IProject,
  IUser,
  ChatTurnStatus,
  GadgetId,
  ChatTurnDocument,
  WorkspaceMode,
  SubmitPromptCallback,
  AbortWorkOrderCallback,
} from "@gadget/api";

import ChatSession from "../models/chat-session.ts";
import DroneRegistration from "../models/drone-registration.ts";
import { ChatTurn } from "../models/chat-turn.ts";

import { ChatSessionService, SocketService } from "../services/index.ts";
import TabLock from "./tab-lock.js";

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;
  protected workspaceMode: WorkspaceMode = WorkspaceMode.Idle;
  
  private chatSessionId: GadgetId | undefined;
  private isReconnecting = false;
  private tabLockAcquired = false;

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

  register() {
    super.register();

    this.socket.on("disconnect", this.onDisconnect.bind(this));
    this.socket.on("requestSessionLock", this.onRequestSessionLock.bind(this));
    this.socket.on(
      "requestWorkspaceMode",
      this.onRequestWorkspaceMode.bind(this),
    );
    this.socket.on("submitPrompt", this.onSubmitPrompt.bind(this));
    this.socket.on("abortWorkOrder", this.onAbortWorkOrder.bind(this));
    this.socket.on("releaseSessionLock", this.onReleaseSessionLock.bind(this));
    this.socket.on("sessionHeartbeat", this.onSessionHeartbeat.bind(this));
    this.socket.on("fileTreeRequest", this.onFileTreeRequest.bind(this));
    
    // Check for active session on connect
    this.checkAndReestablishActiveSession();
  }
  
  private async onDisconnect(): Promise<void> {
    // Release tab lock on disconnect
    if (this.chatSessionId && this.tabLockAcquired) {
      await TabLock.release(this.chatSessionId, this.socket.id);
      this.tabLockAcquired = false;
    }
  }

  get hasLock(): boolean {
    return this.selectedDrone !== undefined && this.chatSession !== undefined;
  }

  get selectedDroneId(): GadgetId | undefined {
    return this.selectedDrone?._id;
  }

  get activeChatSession(): IChatSession | undefined {
    return this.chatSession;
  }

  get activeProject(): IProject | undefined {
    return this.project;
  }

  private async checkAndReestablishActiveSession(): Promise<void> {
    if (this.isReconnecting) return;
    this.isReconnecting = true;

    try {
      // Get user's most recent chat session
      const recentSessions = await ChatSession.find({ user: this.user._id })
        .sort({ createdAt: -1 })
        .limit(5);

      for (const session of recentSessions) {
        // Check if this session has a processing turn
        const latestTurn = await ChatTurn.findOne({ session: session._id })
          .sort({ createdAt: -1 });

        if (latestTurn && latestTurn.status === ChatTurnStatus.Processing) {
          // Found active session - attempt to reestablish connection
          this.chatSessionId = session._id;

          // Get the drone that was processing this turn
          const droneReg = await DroneRegistration.findOne({
            chatSessionId: session._id,
          }).populate('user');

          if (droneReg && droneReg.user) {
            await this.autoRelock(droneReg, session);
          }
          break;
        }
      }
    } catch (error) {
      this.log.error("failed to check for active session", { error });
    } finally {
      this.isReconnecting = false;
    }
  }

  private async autoRelock(
    registration: IDroneRegistration,
    session: IChatSession
  ): Promise<void> {
    try {
      // Try to acquire tab lock
      const lockResult = await TabLock.acquire(
        session._id,
        this.user._id,
        this.socket.id,
      );

      if (!lockResult.success) {
        this.log.warn("tab lock denied - session open in another tab", {
          chatSessionId: session._id,
          lockedBy: lockResult.info?.socketId,
        });
        this.socket.emit("tabLockDenied", {
          message: "Chat session is open in another browser tab",
        });
        return;
      }

      this.tabLockAcquired = true;

      const droneSession = SocketService.getDroneSession(registration);

      // Re-establish the chat session index
      SocketService.registerChatSession(session._id, this);
      droneSession.setChatSessionId(session._id);

      // Update CodeSession state
      this.chatSession = session;
      this.selectedDrone = registration;

      // Drain any queued messages
      await droneSession.drainMessageQueue();

      this.log.info("auto-reestablished session connection", {
        chatSessionId: session._id,
        droneId: registration._id,
      });

      // Emit status to client
      this.socket.emit("status", "Reconnected to active session");
    } catch (error) {
      this.log.error("failed to auto-relock session", { error });
      this.socket.emit("status", "Failed to reconnect to session");
    }
  }

  /**
   * 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,
  ) {
    try {
      const droneSession = SocketService.getDroneSession(registration);
      droneSession.socket.emit(
        "requestSessionLock",
        registration,
        project,
        chatSession,
        async (success: boolean, chatSessionId: string): Promise<void> => {
          if (success) {
            this.selectedDrone = registration;
            this.chatSession = chatSession;
            this.project = project;
            this.chatSessionId = chatSession._id;

            SocketService.registerChatSession(chatSession._id, this);
            droneSession.setChatSessionId(chatSession._id);

            // Acquire tab lock
            const lockResult = await TabLock.acquire(
              chatSession._id,
              this.user._id,
              this.socket.id,
            );

            this.tabLockAcquired = lockResult.success;
          }
          cb(success, chatSessionId);
        },
      );
    } catch (error) {
      cb(false, "");
    }
  }

  /**
   * Called when the IDE sends a requestWorkspaceMode event to change the drone's
   * workspace mode.
   * @param registration the gadget-drone registration
   * @param project the project
   * @param chatSession the chat session
   * @param mode the requested workspace mode
   * @param cb response callback with success and current mode
   */
  onRequestWorkspaceMode(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    mode: WorkspaceMode,
    cb: (success: boolean, currentMode: WorkspaceMode) => void,
  ) {
    if (!this.selectedDrone) {
      this.log.warn("workspace mode request rejected: no drone selected");
      return cb(false, this.workspaceMode);
    }

    try {
      const droneSession = SocketService.getDroneSession(this.selectedDrone);
      droneSession.socket.emit(
        "requestWorkspaceMode",
        registration,
        project,
        chatSession,
        mode,
        (success: boolean, currentMode: WorkspaceMode) => {
          if (success) {
            this.workspaceMode = currentMode;
          }
          cb(success, currentMode);
        },
      );
    } catch (error) {
      cb(false, this.workspaceMode);
    }
  }

  /**
   * 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,
    cb: SubmitPromptCallback,
  ): Promise<void> {
    if (!this.selectedDrone) {
      this.log.warn("prompt rejected: no drone selected");
      return cb(false, { message: "No drone selected" });
    }
    if (!this.chatSession) {
      this.log.warn("prompt rejected: no chat session active");
      return cb(false, { message: "No chat session active" });
    }
    if (!this.project) {
      this.log.warn("prompt rejected: no project selected");
      return cb(false, { message: "No project selected" });
    }

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

      const latestSession = await ChatSessionService.getById(
        this.chatSession._id,
      );
      this.chatSession = latestSession;

      let turn: ChatTurnDocument = await ChatSessionService.createTurn(
        this.chatSession,
        content,
      );
      this.log.info("ChatTurn created", {
        turnId: turn._id,
        chatSessionId: this.chatSession._id,
      });

      this.currentTurnId = turn._id;
      droneSession.setCurrentTurnId(turn._id);

      /*
       * Increment the chat session turn count and replace our cached view of
       * it. Reject the prompt if the chat session has been removed.
       */
      const newSession = await ChatSession.findOneAndUpdate(
        { _id: this.chatSession._id },
        { $inc: { "stats.turnCount": 1 } },
        {
          new: true,
          populate: ChatSessionService.populateChatSession,
        },
      );
      if (!newSession) {
        // remove the turn
        await ChatSessionService.delete(this.chatSession._id);

        // reject the prompt
        const error = new Error("chat session has been removed");
        error.statusCode = 404;
        throw error;
      }
      this.chatSession = newSession;

      /*
       * Signal the IDE that the turn was created successfully. This moves the
       * IDE to the Processing state, and it will start expecting events to be
       * streamed in from the drone while processing the prompt.
       */
      cb(true, { turnId: turn._id, message: "turn created successfully" });

      /*
       * Forward to gadget-drone as a work order for processing.
       */
      droneSession.socket.emit(
        "processWorkOrder",
        this.selectedDrone,
        turn,
        async (success: boolean, message?: string) => {
          if (success) {
            this.log.info("work order accepted by drone", {
              turnId: turn._id,
              message,
            });

            /*
             * Auto-generate a session name from the first prompt. Only do this when
             * the name is still the default (user hasn't set a custom name) and
             * we're on the first turn (turnCount === 1 after the increment above).
             */
            if (
              this.chatSession &&
              this.chatSession.name === "New Chat Session" &&
              this.chatSession.stats.turnCount === 1
            ) {
              this.chatSession =
                await ChatSessionService.generateSessionNameFromPrompt(
                  this.chatSession,
                  content,
                );
              const update: Partial<IChatSession> = {
                name: this.chatSession.name,
              };
              this.socket.emit("sessionUpdated", update);
            }
          } else {
            this.log.error("work order rejected by drone", {
              turnId: turn._id,
              message,
            });
            turn.status = ChatTurnStatus.Error;
            turn.errorMessage = message || "Drone rejected work order";
            await turn.save();
          }
        },
      );
    } catch (error) {
      this.log.error("prompt rejected", { error });
      cb(false, {});
    }
  }

  /**
   * Called when the IDE sends an abortWorkOrder event to cancel the
   * currently running work order. Forwards to the drone.
   */
  onAbortWorkOrder(cb: AbortWorkOrderCallback): void {
    if (!this.selectedDrone) {
      return cb(false, "No drone selected");
    }
    try {
      const droneSession = SocketService.getDroneSession(this.selectedDrone);
      droneSession.socket.emit("abortWorkOrder", cb);
    } catch (error) {
      this.log.error("failed to forward abortWorkOrder to drone", { error });
      cb(false, "Failed to reach drone");
    }
  }

  /**
   * Called when the IDE sends a fileTreeRequest event to list directory contents.
   * Forwards to the drone and relays response back to IDE.
   */
  onFileTreeRequest(
    args: FileTreeRequestArgs,
    cb: (success: boolean, data: { entries?: FileTreeEntry[]; error?: string }) => void,
  ): void {
    if (!this.selectedDrone) {
      return cb(false, { error: "No drone selected" });
    }
    try {
      const droneSession = SocketService.getDroneSession(this.selectedDrone);
      droneSession.socket.emit("fileTreeRequest", args, (success: boolean, data: { entries?: FileTreeEntry[]; error?: string }) => {
        // Forward response to IDE
        if (success && data?.entries) {
          this.socket.emit("fileTreeResponse", args.path || "", data.entries);
        } else {
          this.socket.emit("fileTreeResponse", args.path || "", [], data?.error);
        }
        cb(success, data);
      });
    } catch (error) {
      this.log.error("failed to forward fileTreeRequest to drone", { error });
      cb(false, { error: "Failed to reach drone" });
    }
  }

  /**
   * Called when the IDE sends a releaseSessionLock event to release a
   * previously-acquired session lock on a gadget-drone instance.
   */
  onReleaseSessionLock(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    cb: (success: boolean) => void,
  ) {
    try {
      const droneSession = SocketService.getDroneSession(registration);
      droneSession.socket.emit(
        "releaseSessionLock",
        registration,
        project,
        chatSession,
        async (success: boolean) => {
          if (success) {
            SocketService.unregisterChatSession(chatSession._id);
            droneSession.chatSessionId = undefined;
            this.selectedDrone = undefined;
            this.chatSession = undefined;
            this.project = undefined;
            this.chatSessionId = undefined;

            // Release tab lock
            if (this.tabLockAcquired) {
              await TabLock.release(chatSession._id, this.socket.id);
              this.tabLockAcquired = false;
            }
          }
          cb(success);
        },
      );
    } catch (error) {
      cb(false);
    }
  }

  /**
   * Called when the IDE sends a sessionHeartbeat event to keep the session
   * lock alive. Forwards to the drone which maintains a timeout.
   */
  onSessionHeartbeat(cb: (ack: boolean) => void) {
    if (!this.selectedDrone) {
      return cb(false);
    }
    try {
      const droneSession = SocketService.getDroneSession(this.selectedDrone);
      droneSession.socket.emit("sessionHeartbeat", (ack: boolean) => {
        cb(ack);
      });
    } catch (error) {
      this.log.error("failed to forward session heartbeat to drone", {
        error,
      });
      cb(false);
    }
  }

  /**
   * Called by DroneSession when the drone emits a workspace mode change.
   * Updates local state and forwards to the IDE socket.
   */
  onWorkspaceModeChanged(mode: WorkspaceMode): void {
    this.workspaceMode = mode;
    this.socket.emit("workspaceModeChanged", mode);
  }

  onStatus(content: string): void {
    this.socket.emit("status", content);
  }

  onLog(
    timestamp: Date,
    component: GadgetComponent,
    level: GadgetLogLevel,
    message: string,
    metadata?: unknown,
  ): void {
    this.socket.emit("log", timestamp, component, level, message, metadata);
  }

  onThinking(content: string): void {
    this.socket.emit("thinking", content);
  }

  onResponse(content: string): void {
    this.socket.emit("response", content);
  }

  onToolCall(
    callId: string,
    name: string,
    params: string,
    response: string,
  ): void {
    this.socket.emit("toolCall", callId, name, params, response);
  }

  onWorkOrderComplete(
    turnId: string,
    success: boolean,
    message?: string,
  ): void {
    this.socket.emit("workOrderComplete", turnId, success, message);
  }
}