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

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

import {
  GadgetComponent,
  GadgetLogLevel,
  IUser,
  IDroneRegistration,
  ChatTurnStatus,
  GadgetId,
  WorkspaceMode,
  IChatToolCall,
  IChatTurnBlock,
  IWorkOrderCompleteStats,
} from "@gadget/api";
import {
  GadgetSocket,
  SocketSession,
  SocketSessionType,
} from "./socket-session.js";
import { SocketService } from "../services/index.js";
import { ChatTurn } from "../models/chat-turn.js";
import ChatSession from "../models/chat-session.js";
import MessageQueue, { type QueuedMessage } from "./message-queue.js";

interface IStreamingBuffer {
  currentMode: 'thinking' | 'responding' | null;
  thinkingContent: string;
  respondingContent: string;
  lastBlockCreatedAt?: Date;
}

export class DroneSession extends SocketSession {
  protected type: SocketSessionType = SocketSessionType.Drone;
  registration: IDroneRegistration;
  chatSessionId: GadgetId | undefined;
  currentTurnId: GadgetId | undefined;
  workspaceMode: WorkspaceMode = WorkspaceMode.Idle;
  private streamingBuffers: Map<string, IStreamingBuffer> = new Map();
  private isDrainingQueue = false;

  constructor(socket: GadgetSocket, registration: IDroneRegistration) {
    super(socket, registration.user as IUser);
    this.registration = registration;
  }

  register() {
    super.register();

    this.socket.on("status", this.onStatus.bind(this));
    this.socket.on(
      "workspaceModeChanged",
      this.onWorkspaceModeChanged.bind(this),
    );

    this.socket.on("thinking", this.onThinking.bind(this));
    this.socket.on("response", this.onResponse.bind(this));
    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));

    this.socket.on("log", this.onLog.bind(this));

    this.socket.on("agent:thinking", this.onAgentThinking.bind(this));
    this.socket.on("agent:response", this.onAgentResponse.bind(this));
    this.socket.on("agent:tool-call", this.onAgentToolCall.bind(this));
    this.socket.on("agent:tool-result", this.onAgentToolResult.bind(this));
    this.socket.on("agent:complete", this.onAgentComplete.bind(this));
  }

  async onLog(
    timestamp: Date,
    component: GadgetComponent,
    level: GadgetLogLevel,
    message: string,
    metadata?: unknown,
  ): Promise<void> {
    // Route to chat session if one is active
    if (this.chatSessionId) {
      try {
        const codeSession = SocketService.getCodeSessionByChatSessionId(
          this.chatSessionId,
        );
        codeSession.onLog(timestamp, component, level, message, metadata);
      } catch (error) {
        await MessageQueue.enqueue(this.chatSessionId, {
          type: 'log',
          args: [timestamp, component, level, message, metadata],
          timestamp: Date.now(),
        });
        this.log.debug("queued log message", { chatSessionId: this.chatSessionId });
      }
    }

    // Broadcast to monitoring sessions (DroneManager, DroneInspector, etc.)
    SocketService.broadcastToMonitors(
      "drone:log",
      this.registration._id,
      {
        timestamp: timestamp instanceof Date ? timestamp.toISOString() : String(timestamp),
        level: String(level),
        component: typeof component === 'object' && component !== null ? (component as any).name ?? String(component) : String(component),
        message,
        metadata,
      },
    );
  }

  async onStatus(message: string): Promise<void> {
    // Route to chat session if one is active
    if (this.chatSessionId) {
      try {
        const codeSession = SocketService.getCodeSessionByChatSessionId(
          this.chatSessionId,
        );
        codeSession.socket.emit("status", message);
      } catch (error) {
        await MessageQueue.enqueue(this.chatSessionId, {
          type: 'status',
          args: [message],
          timestamp: Date.now(),
        });
        this.log.debug("queued status message", { chatSessionId: this.chatSessionId });
      }
    }

    // Broadcast to monitoring sessions
    SocketService.broadcastToMonitors(
      "drone:status",
      this.registration._id,
      {
        timestamp: new Date().toISOString(),
        message,
      },
    );
  }

  /**
   * Called when the drone emits thinking content from the agent.
   * Aggregates thinking tokens in memory and persists at mode changes.
   */
  async onThinking(content: string): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("thinking event received but no chat session is active");
      return;
    }

    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(
        this.chatSessionId,
      );
      codeSession.onThinking(content);

      if (this.currentTurnId) {
        const buffer = this.getOrCreateBuffer(this.currentTurnId);
        
        // Check for mode transition
        if (buffer.currentMode !== 'thinking') {
          // Flush previous mode if exists
          await this.flushBuffer(this.currentTurnId);
          buffer.currentMode = 'thinking';
          buffer.thinkingContent = '';
          buffer.lastBlockCreatedAt = new Date();
        }
        
        // Aggregate content
        buffer.thinkingContent += content;
      }
    } catch (error) {
      // Routing failed - queue to Redis
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'thinking',
        args: [content],
        timestamp: Date.now(),
      });
      this.log.debug("queued thinking message", { chatSessionId: this.chatSessionId });
    }
  }

  /**
   * Called when the drone emits response content from the agent.
   * Aggregates response tokens in memory and persists at mode changes.
   */
  async onResponse(content: string): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("response event received but no chat session is active");
      return;
    }

    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(
        this.chatSessionId,
      );
      codeSession.onResponse(content);

      if (this.currentTurnId) {
        const buffer = this.getOrCreateBuffer(this.currentTurnId);
        
        // Check for mode transition
        if (buffer.currentMode !== 'responding') {
          // Flush previous mode if exists
          await this.flushBuffer(this.currentTurnId);
          buffer.currentMode = 'responding';
          buffer.respondingContent = '';
          buffer.lastBlockCreatedAt = new Date();
        }
        
        // Aggregate content
        buffer.respondingContent += content;
      }
    } catch (error) {
      // Routing failed - queue to Redis
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'response',
        args: [content],
        timestamp: Date.now(),
      });
      this.log.debug("queued response message", { chatSessionId: this.chatSessionId });
    }
  }

  /**
   * Called when the drone emits a tool call event from the agent.
   * Flushes current buffer and adds tool block immediately.
   */
  async onToolCall(
    callId: string,
    name: string,
    params: string,
    response: string,
  ): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("toolCall event received but no chat session is active");
      return;
    }

    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(
        this.chatSessionId,
      );
      codeSession.onToolCall(callId, name, params, response);

      if (this.currentTurnId) {
        // Flush current buffer before adding tool block
        await this.flushBuffer(this.currentTurnId);
        
        // Add tool block immediately
        const turn = await ChatTurn.findById(this.currentTurnId);
        if (turn) {
          turn.blocks.push({
            mode: 'tool',
            createdAt: new Date(),
            content: {
              callId,
              name,
              parameters: params,
              response,
            },
          });
          turn.toolCalls.push({
            callId,
            name,
            parameters: params,
            response,
          });
          turn.stats.toolCallCount = turn.toolCalls.length;
          await turn.save();
        }
      }
    } catch (error) {
      // Routing failed - queue to Redis
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'toolCall',
        args: [callId, name, params, response],
        timestamp: Date.now(),
      });
      this.log.debug("queued toolCall message", { chatSessionId: this.chatSessionId });
    }
  }

  /**
   * Called when the drone completes a work order.
   * Flushes any remaining buffered content before finalizing.
   */
  async onWorkOrderComplete(
    turnId: string,
    success: boolean,
    message?: string,
    stats?: IWorkOrderCompleteStats,
  ): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("workOrderComplete event received but no chat session is active");
      return;
    }

    try {
      // Flush any remaining buffered content
      await this.flushBuffer(turnId);
      this.streamingBuffers.delete(turnId);

      const turn = await ChatTurn.findById(turnId);
      if (turn) {
        if (success && message === "aborted") {
          turn.status = ChatTurnStatus.Aborted;
          turn.errorMessage = "The turn was aborted by you.";
        } else {
          turn.status = success ? ChatTurnStatus.Finished : ChatTurnStatus.Error;
          if (!success && message) {
            turn.errorMessage = message;
          }
        }

        // UPDATE TURN STATS — master agent tokens only
        if (stats) {
          turn.stats.inputTokens = stats.masterInputTokens;
          turn.stats.thinkingTokenCount = stats.masterThinkingTokens;
          turn.stats.responseTokens = stats.masterOutputTokens;
          turn.stats.toolCallCount = stats.toolCallCount;
          turn.stats.durationMs = stats.durationMs;
          turn.stats.durationLabel = this.formatDurationLabel(stats.durationMs);
        }

        await turn.save();

        // COMPUTE AGGREGATE: walk the turn's subagent records for aggregate totals
        let totalInputTokens = stats?.masterInputTokens ?? 0;
        let totalOutputTokens = stats?.masterOutputTokens ?? 0;
        if (stats) {
          for (const tc of turn.toolCalls) {
            if ((tc as any).subagent?.stats) {
              totalInputTokens += (tc as any).subagent.stats.inputTokens;
              totalOutputTokens += (tc as any).subagent.stats.responseTokens + (tc as any).subagent.stats.thinkingTokenCount;
            }
          }
        }

        // INCREMENT SESSION STATS — aggregate totals (master + subagents)
        await ChatSession.updateOne(
          { _id: this.chatSessionId },
          { $inc: {
            "stats.inputTokens": totalInputTokens,
            "stats.outputTokens": totalOutputTokens,
            "stats.toolCallCount": stats?.toolCallCount ?? 0,
            // Master-only field for the fuel gauge
            "contextWindowUsage": stats?.masterInputTokens ?? 0,
          }},
        );
      }

      const codeSession = SocketService.getCodeSessionByChatSessionId(
        this.chatSessionId,
      );
      codeSession.onWorkOrderComplete(turnId, success, message, stats);

      this.currentTurnId = undefined;
    } catch (error) {
      // Routing failed - queue to Redis
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'workOrderComplete',
        args: [turnId, success, message, stats],
        timestamp: Date.now(),
      });
      this.log.debug("queued workOrderComplete message", { chatSessionId: this.chatSessionId });
    }
  }

  /**
   * Sets the active chat session ID for this drone session.
   */
  setChatSessionId(chatSessionId: GadgetId): void {
    this.chatSessionId = chatSessionId;
    // Clear buffer for this turn if exists
    this.streamingBuffers.clear();
  }

  private formatDurationLabel(durationMs: number): string {
    if (durationMs < 1000) return `${durationMs}ms`;
    const seconds = Math.floor(durationMs / 1000);
    const minutes = Math.floor(seconds / 60);
    const secs = seconds % 60;
    if (minutes > 0) {
      return `${minutes}m ${secs}s`;
    }
    return `${seconds}.${Math.floor((durationMs % 1000) / 100)}s`;
  }

  /**
   * Gets or creates a streaming buffer for a turn.
   */
  private getOrCreateBuffer(turnId: string): IStreamingBuffer {
    if (!this.streamingBuffers.has(turnId)) {
      this.streamingBuffers.set(turnId, {
        currentMode: null,
        thinkingContent: '',
        respondingContent: '',
      });
    }
    return this.streamingBuffers.get(turnId)!;
  }

  /**
   * Flushes the current buffer to the database.
   */
  private async flushBuffer(turnId: string): Promise<void> {
    const buffer = this.streamingBuffers.get(turnId);
    if (!buffer) return;

    const turn = await ChatTurn.findById(turnId);
    if (!turn) return;

    // Flush thinking content
    if (buffer.currentMode === 'thinking' && buffer.thinkingContent) {
      turn.blocks.push({
        mode: 'thinking',
        createdAt: buffer.lastBlockCreatedAt || new Date(),
        content: buffer.thinkingContent,
      });
      buffer.thinkingContent = '';
    }

    // Flush responding content
    if (buffer.currentMode === 'responding' && buffer.respondingContent) {
      turn.blocks.push({
        mode: 'responding',
        createdAt: buffer.lastBlockCreatedAt || new Date(),
        content: buffer.respondingContent,
      });
      buffer.respondingContent = '';
    }

    if (turn.blocks.length > 0) {
      await turn.save();
    }
  }

  /**
   * Sets the current turn ID being processed by this drone.
   */
  setCurrentTurnId(turnId: GadgetId): void {
    this.currentTurnId = turnId;
  }

  /**
   * Called when the drone emits a workspace mode change.
   */
  async onWorkspaceModeChanged(mode: WorkspaceMode): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("workspaceModeChanged event received but no chat session is active");
      return;
    }

    this.workspaceMode = mode;
    this.log.info("workspace mode changed", { mode });

    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(
        this.chatSessionId,
      );
      codeSession.onWorkspaceModeChanged(mode);
    } catch (error) {
      this.log.error("failed to route workspaceModeChanged event", { error });
    }
  }

  async onAgentThinking(data: { agentId: string; thinking: string }): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("agent:thinking event received but no chat session is active");
      return;
    }
    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);
      codeSession.socket.emit("agent:thinking", data);
    } catch (error) {
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'agent:thinking',
        args: [data],
        timestamp: Date.now(),
      });
      this.log.debug("queued agent:thinking message", { chatSessionId: this.chatSessionId });
    }
  }

  async onAgentResponse(data: { agentId: string; chunk: string }): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("agent:response event received but no chat session is active");
      return;
    }
    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);
      codeSession.socket.emit("agent:response", data);
    } catch (error) {
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'agent:response',
        args: [data],
        timestamp: Date.now(),
      });
      this.log.debug("queued agent:response message", { chatSessionId: this.chatSessionId });
    }
  }

  async onAgentToolCall(data: { agentId: string; tool: string; args: unknown }): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("agent:tool-call event received but no chat session is active");
      return;
    }
    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);
      codeSession.socket.emit("agent:tool-call", data);
    } catch (error) {
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'agent:tool-call',
        args: [data],
        timestamp: Date.now(),
      });
      this.log.debug("queued agent:tool-call message", { chatSessionId: this.chatSessionId });
    }
  }

  async onAgentToolResult(data: { agentId: string; tool: string; result: unknown }): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("agent:tool-result event received but no chat session is active");
      return;
    }
    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);
      codeSession.socket.emit("agent:tool-result", data);
    } catch (error) {
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'agent:tool-result',
        args: [data],
        timestamp: Date.now(),
      });
      this.log.debug("queued agent:tool-result message", { chatSessionId: this.chatSessionId });
    }
  }

  async onAgentComplete(data: { agentId: string; response?: string; subagent?: Record<string, unknown>; stats?: Record<string, unknown> }): Promise<void> {
    if (!this.chatSessionId) {
      this.log.warn("agent:complete event received but no chat session is active");
      return;
    }
    try {
      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);
      codeSession.socket.emit("agent:complete", data);
    } catch (error) {
      await MessageQueue.enqueue(this.chatSessionId, {
        type: 'agent:complete',
        args: [data],
        timestamp: Date.now(),
      });
      this.log.debug("queued agent:complete message", { chatSessionId: this.chatSessionId });
    }

    // Update the persisted tool call with the final response and subagent data
    if (this.currentTurnId && data.agentId) {
      try {
        const turn = await ChatTurn.findById(this.currentTurnId);
        if (turn) {
          const toolCall = turn.toolCalls.find((tc: IChatToolCall) => tc.callId === data.agentId);
          if (toolCall) {
            if (data.response) toolCall.response = data.response;
            if (data.subagent) (toolCall as any).subagent = data.subagent;
          }
          const block = turn.blocks.find((b: IChatTurnBlock) => b.mode === 'tool' && (b.content as IChatToolCall).callId === data.agentId);
          if (block && block.mode === 'tool') {
            if (data.response) (block.content as IChatToolCall).response = data.response;
            if (data.subagent) (block.content as any).subagent = data.subagent;
          }
          await turn.save();
        }
      } catch (error) {
        this.log.error("failed to update subagent tool call in DB", { error });
      }
    }
  }

  /**
   * Called when the drone requests crash recovery for an incomplete work order.
   */
  async onRequestCrashRecovery(data: {
    workspaceId: string;
    turnId: string;
    chatSessionId: string;
  }): Promise<void> {
    this.log.info("crash recovery request received", {
      workspaceId: data.workspaceId,
      turnId: data.turnId,
    });

    try {
      const turn = await ChatTurn.findById(data.turnId);

      if (!turn) {
        this.log.warn("crash recovery: turn not found", {
          turnId: data.turnId,
        });
        this.socket.emit("crashRecoveryResponse", {
          turnId: data.turnId,
          action: "discard",
        });
        return;
      }

      if (turn.status === ChatTurnStatus.Finished) {
        this.log.info("crash recovery: turn already finished", {
          turnId: data.turnId,
        });
        this.socket.emit("crashRecoveryResponse", {
          turnId: data.turnId,
          action: "discard",
        });
        return;
      }

      // Turn is still processing - mark for retry
      turn.status = ChatTurnStatus.Error;
      turn.errorMessage = "Drone crashed during processing - retrying";
      await turn.save();

      this.socket.emit("crashRecoveryResponse", {
        turnId: data.turnId,
        action: "retry",
        retryDelay: 5000,
      });

      this.log.info("crash recovery: scheduled retry", {
        turnId: data.turnId,
      });

      // Schedule retry (will route to same workspaceId)
      setTimeout(() => {
        this.retryWorkOrder(turn);
      }, 5000);
    } catch (error) {
      const err = error as Error;
      this.log.error("crash recovery failed", {
        error: err.message,
      });
      this.socket.emit("crashRecoveryResponse", {
        turnId: data.turnId,
        action: "discard",
      });
    }
  }

  /**
   * Retries a work order after crash recovery.
   */
  private async retryWorkOrder(turn: any): Promise<void> {
    // TODO: Re-emit processWorkOrder to this drone
    this.log.info("work order retry not yet implemented", {
      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,
    });

    this.socket.emit("requestTermination", (success: boolean) => {
      this.log.info("requestTermination forwarded to drone", { success });
      cb(success);
    });
  }

  /**
   * Drains queued messages from Redis and delivers to reconnected CodeSession.
   * Aggregates adjacent same-type streaming messages to reduce message count.
   */
  async drainMessageQueue(): Promise<void> {
    if (!this.chatSessionId || this.isDrainingQueue) return;
    this.isDrainingQueue = true;

    try {
      const messages = await MessageQueue.drain(this.chatSessionId);

      if (messages.length === 0) return;

      this.log.info("draining message queue", { count: messages.length });

      const codeSession = SocketService.getCodeSessionByChatSessionId(this.chatSessionId);

      // Aggregate adjacent same-type streaming messages
      const aggregated = this.aggregateMessages(messages);

      for (const msg of aggregated) {
        try {
          switch (msg.type) {
            case 'thinking':
              codeSession.onThinking(msg.args[0] as string);
              break;
            case 'response':
              codeSession.onResponse(msg.args[0] as string);
              break;
            case 'toolCall':
              codeSession.onToolCall(
                msg.args[0] as string,
                msg.args[1] as string,
                msg.args[2] as string,
                msg.args[3] as string,
              );
              break;
            case 'workOrderComplete':
              codeSession.onWorkOrderComplete(
                msg.args[0] as string,
                msg.args[1] as boolean,
                msg.args[2] as string | undefined,
                msg.args[3] as IWorkOrderCompleteStats | undefined,
              );
              break;
            case 'log':
              codeSession.onLog(
                msg.args[0] as Date,
                msg.args[1] as GadgetComponent,
                msg.args[2] as GadgetLogLevel,
                msg.args[3] as string,
                msg.args[4] as unknown,
              );
              break;
            case 'status':
              codeSession.socket.emit("status", msg.args[0] as string);
              break;
            case 'agent:thinking':
              codeSession.socket.emit("agent:thinking", msg.args[0] as { agentId: string; thinking: string });
              break;
            case 'agent:response':
              codeSession.socket.emit("agent:response", msg.args[0] as { agentId: string; chunk: string });
              break;
            case 'agent:tool-call':
              codeSession.socket.emit("agent:tool-call", msg.args[0] as { agentId: string; tool: string; args: unknown });
              break;
            case 'agent:tool-result':
              codeSession.socket.emit("agent:tool-result", msg.args[0] as { agentId: string; tool: string; result: unknown });
              break;
            case 'agent:complete':
              codeSession.socket.emit("agent:complete", msg.args[0] as { agentId: string; response?: string; subagent?: Record<string, unknown>; stats?: Record<string, unknown> });
              break;
            default:
              this.log.warn("unknown queued message type", { type: (msg as any).type });
          }
          // Small delay to avoid flooding
          await new Promise(resolve => setTimeout(resolve, 5));
        } catch (error) {
          this.log.error("failed to deliver queued message", {
            type: msg.type,
            error,
          });
        }
      }
    } catch (error) {
      this.log.error("failed to drain message queue", { error });
    } finally {
      this.isDrainingQueue = false;
    }
  }

  /**
   * Aggregates adjacent same-type streaming messages (thinking/response).
   * Preserves order and only aggregates during drain (not real-time).
   */
  private aggregateMessages(messages: QueuedMessage[]): QueuedMessage[] {
    const aggregated: QueuedMessage[] = [];
    let currentAggregate: QueuedMessage | null = null;

    for (const msg of messages) {
      // Only aggregate thinking and response messages
      if (msg.type === 'thinking' || msg.type === 'response') {
        if (currentAggregate && currentAggregate.type === msg.type) {
          // Continue aggregating same type
          const currentContent = (currentAggregate.args[0] as string) || '';
          const newContent = (msg.args[0] as string) || '';
          currentAggregate.args = [currentContent + newContent];
        } else {
          // Type changed - push current aggregate and start new
          if (currentAggregate) {
            aggregated.push(currentAggregate);
          }
          currentAggregate = { ...msg };
        }
      } else {
        // Non-aggregatable message type
        if (currentAggregate) {
          aggregated.push(currentAggregate);
          currentAggregate = null;
        }
        aggregated.push(msg);
      }
    }

    // Push final aggregate if exists
    if (currentAggregate) {
      aggregated.push(currentAggregate);
    }

    return aggregated;
  }
}