rwadurian/backend/services/mpc-service/src/domain/entities/session-state.entity.ts

/**
 * SessionState Entity
 *
 * Tracks the state of an MPC session from this party's perspective.
 * Used for monitoring and recovery purposes.
 */

import { v4 as uuidv4 } from 'uuid';
import { SessionId, PartyId, PublicKey, MessageHash, Signature } from '../value-objects';
import { SessionType, SessionStatus, ParticipantStatus } from '../enums';
import {
  DomainEvent,
  PartyJoinedSessionEvent,
  KeygenCompletedEvent,
  SigningCompletedEvent,
  SessionFailedEvent,
  SessionTimeoutEvent,
} from '../events';

export interface Participant {
  partyId: string;
  partyIndex: number;
  status: ParticipantStatus;
}

export interface SessionStateCreateParams {
  sessionId: SessionId;
  partyId: PartyId;
  partyIndex: number;
  sessionType: SessionType;
  participants: Participant[];
  thresholdN: number;
  thresholdT: number;
  publicKey?: PublicKey;
  messageHash?: MessageHash;
}

export interface SessionStateReconstructParams extends SessionStateCreateParams {
  id: string;
  status: SessionStatus;
  currentRound: number;
  errorMessage?: string;
  signature?: Signature;
  startedAt: Date;
  completedAt?: Date;
}

export class SessionState {
  // ============================================================================
  // Private Fields
  // ============================================================================
  private readonly _id: string;
  private readonly _sessionId: SessionId;
  private readonly _partyId: PartyId;
  private readonly _partyIndex: number;
  private readonly _sessionType: SessionType;
  private readonly _participants: Participant[];
  private readonly _thresholdN: number;
  private readonly _thresholdT: number;
  private _status: SessionStatus;
  private _currentRound: number;
  private _errorMessage?: string;
  private _publicKey?: PublicKey;
  private _messageHash?: MessageHash;
  private _signature?: Signature;
  private readonly _startedAt: Date;
  private _completedAt?: Date;

  // Domain events collection
  private readonly _domainEvents: DomainEvent[] = [];

  // ============================================================================
  // Constructor (private - use factory methods)
  // ============================================================================
  private constructor(
    id: string,
    sessionId: SessionId,
    partyId: PartyId,
    partyIndex: number,
    sessionType: SessionType,
    participants: Participant[],
    thresholdN: number,
    thresholdT: number,
    status: SessionStatus,
    currentRound: number,
    startedAt: Date,
    errorMessage?: string,
    publicKey?: PublicKey,
    messageHash?: MessageHash,
    signature?: Signature,
    completedAt?: Date,
  ) {
    this._id = id;
    this._sessionId = sessionId;
    this._partyId = partyId;
    this._partyIndex = partyIndex;
    this._sessionType = sessionType;
    this._participants = participants;
    this._thresholdN = thresholdN;
    this._thresholdT = thresholdT;
    this._status = status;
    this._currentRound = currentRound;
    this._startedAt = startedAt;
    this._errorMessage = errorMessage;
    this._publicKey = publicKey;
    this._messageHash = messageHash;
    this._signature = signature;
    this._completedAt = completedAt;
  }

  // ============================================================================
  // Factory Methods
  // ============================================================================

  /**
   * Create a new session state when joining a session
   */
  static create(params: SessionStateCreateParams): SessionState {
    const session = new SessionState(
      uuidv4(),
      params.sessionId,
      params.partyId,
      params.partyIndex,
      params.sessionType,
      params.participants,
      params.thresholdN,
      params.thresholdT,
      SessionStatus.IN_PROGRESS,
      0,
      new Date(),
      undefined,
      params.publicKey,
      params.messageHash,
    );

    // Emit joined event
    session.addDomainEvent(new PartyJoinedSessionEvent(
      params.sessionId.value,
      params.partyId.value,
      params.partyIndex,
      params.sessionType,
    ));

    return session;
  }

  /**
   * Reconstruct from persistence
   */
  static reconstruct(params: SessionStateReconstructParams): SessionState {
    return new SessionState(
      params.id,
      params.sessionId,
      params.partyId,
      params.partyIndex,
      params.sessionType,
      params.participants,
      params.thresholdN,
      params.thresholdT,
      params.status,
      params.currentRound,
      params.startedAt,
      params.errorMessage,
      params.publicKey,
      params.messageHash,
      params.signature,
      params.completedAt,
    );
  }

  // ============================================================================
  // Getters
  // ============================================================================
  get id(): string {
    return this._id;
  }

  get sessionId(): SessionId {
    return this._sessionId;
  }

  get partyId(): PartyId {
    return this._partyId;
  }

  get partyIndex(): number {
    return this._partyIndex;
  }

  get sessionType(): SessionType {
    return this._sessionType;
  }

  get participants(): Participant[] {
    return [...this._participants];
  }

  get thresholdN(): number {
    return this._thresholdN;
  }

  get thresholdT(): number {
    return this._thresholdT;
  }

  get status(): SessionStatus {
    return this._status;
  }

  get currentRound(): number {
    return this._currentRound;
  }

  get errorMessage(): string | undefined {
    return this._errorMessage;
  }

  get publicKey(): PublicKey | undefined {
    return this._publicKey;
  }

  get messageHash(): MessageHash | undefined {
    return this._messageHash;
  }

  get signature(): Signature | undefined {
    return this._signature;
  }

  get startedAt(): Date {
    return this._startedAt;
  }

  get completedAt(): Date | undefined {
    return this._completedAt;
  }

  get domainEvents(): DomainEvent[] {
    return [...this._domainEvents];
  }

  // ============================================================================
  // Business Methods
  // ============================================================================

  /**
   * Update the current round number
   */
  advanceRound(roundNumber: number): void {
    this.ensureInProgress();
    if (roundNumber <= this._currentRound) {
      throw new Error(`Cannot go back to round ${roundNumber} from ${this._currentRound}`);
    }
    this._currentRound = roundNumber;
  }

  /**
   * Mark keygen as completed
   */
  completeKeygen(publicKey: PublicKey, shareId: string): void {
    this.ensureInProgress();
    if (this._sessionType !== SessionType.KEYGEN) {
      throw new Error('Cannot complete keygen for non-keygen session');
    }

    this._publicKey = publicKey;
    this._status = SessionStatus.COMPLETED;
    this._completedAt = new Date();

    this.addDomainEvent(new KeygenCompletedEvent(
      this._sessionId.value,
      this._partyId.value,
      publicKey.toHex(),
      shareId,
      `${this._thresholdT}-of-${this._thresholdN}`,
    ));
  }

  /**
   * Mark signing as completed
   */
  completeSigning(signature: Signature): void {
    this.ensureInProgress();
    if (this._sessionType !== SessionType.SIGN) {
      throw new Error('Cannot complete signing for non-signing session');
    }
    if (!this._messageHash) {
      throw new Error('Message hash not set for signing session');
    }
    if (!this._publicKey) {
      throw new Error('Public key not set for signing session');
    }

    this._signature = signature;
    this._status = SessionStatus.COMPLETED;
    this._completedAt = new Date();

    this.addDomainEvent(new SigningCompletedEvent(
      this._sessionId.value,
      this._partyId.value,
      this._messageHash.toHex(),
      signature.toHex(),
      this._publicKey.toHex(),
    ));
  }

  /**
   * Mark session as failed
   */
  fail(errorMessage: string, errorCode?: string): void {
    if (this._status === SessionStatus.COMPLETED) {
      throw new Error('Cannot fail a completed session');
    }

    this._status = SessionStatus.FAILED;
    this._errorMessage = errorMessage;
    this._completedAt = new Date();

    this.addDomainEvent(new SessionFailedEvent(
      this._sessionId.value,
      this._partyId.value,
      this._sessionType,
      errorMessage,
      errorCode,
    ));
  }

  /**
   * Mark session as timed out
   */
  timeout(): void {
    if (this._status === SessionStatus.COMPLETED) {
      throw new Error('Cannot timeout a completed session');
    }

    this._status = SessionStatus.TIMEOUT;
    this._completedAt = new Date();

    this.addDomainEvent(new SessionTimeoutEvent(
      this._sessionId.value,
      this._partyId.value,
      this._sessionType,
      this._currentRound,
    ));
  }

  /**
   * Update participant status
   */
  updateParticipantStatus(partyId: string, status: ParticipantStatus): void {
    const participant = this._participants.find(p => p.partyId === partyId);
    if (!participant) {
      throw new Error(`Participant ${partyId} not found`);
    }
    participant.status = status;
  }

  /**
   * Check if session is in progress
   */
  isInProgress(): boolean {
    return this._status === SessionStatus.IN_PROGRESS;
  }

  /**
   * Check if session is completed
   */
  isCompleted(): boolean {
    return this._status === SessionStatus.COMPLETED;
  }

  /**
   * Get the duration of the session in milliseconds
   */
  getDuration(): number | undefined {
    if (!this._completedAt) {
      return Date.now() - this._startedAt.getTime();
    }
    return this._completedAt.getTime() - this._startedAt.getTime();
  }

  // ============================================================================
  // Domain Event Management
  // ============================================================================

  addDomainEvent(event: DomainEvent): void {
    this._domainEvents.push(event);
  }

  clearDomainEvents(): void {
    this._domainEvents.length = 0;
  }

  // ============================================================================
  // Private Helper Methods
  // ============================================================================

  private ensureInProgress(): void {
    if (this._status !== SessionStatus.IN_PROGRESS) {
      throw new Error(`Cannot perform operation on ${this._status} session`);
    }
  }
}