life-echo/app-expo/src/features/conversation/realtime-session.ts

import type { QueryClient } from '@tanstack/react-query';

import {
  WsClient,
  type WsEventListener,
  type WsStateListener,
} from '@/core/ws/client';
import type { WsConnectionState, WsEvent } from '@/core/ws/types';

import { handleWsEvent } from './event-handlers';
import { lastSegmentPreview } from './message-split';
import { conversationKeys } from './query-keys';
import type { ConversationListItem, MessageItem } from './types';

export type StreamingTextCallback = (text: string, isComplete: boolean) => void;
export type ErrorCallback = (message: string, code?: string) => void;

/** WebSocket `tts_audio`：服务端可能只带 base64、只带 COS URL，或两者都有 */
export type TtsSegmentPayload = {
  audioBase64?: string;
  audioUrl?: string;
};

interface RealtimeSessionOptions {
  conversationId: string;
  queryClient: QueryClient;
  onStreamingText?: StreamingTextCallback;
  /** 收到 TTS 片段时入队播放（与「气泡上的手动朗读按钮」无关） */
  onTtsSegment?: (payload: TtsSegmentPayload) => void;
  onError?: ErrorCallback;
  onStateChange?: WsStateListener;
}

/**
 * Orchestrates a single WS conversation session.
 * Responsibilities:
 * - Connect/disconnect WS
 * - Route confirmed events to cache handlers
 * - Accumulate agent_response chunks, commit final message to cache
 * - Forward streaming text to UI callback
 * - Expose send methods (return success/failure)
 *
 * Does NOT hold message lists or UI state.
 */
export class RealtimeSession {
  private client: WsClient;
  private conversationId: string;
  private queryClient: QueryClient;
  private onStreamingText?: StreamingTextCallback;
  private onTtsSegment?: (payload: TtsSegmentPayload) => void;
  private onError?: ErrorCallback;
  private unsubEvent: (() => void) | null = null;
  private unsubState: (() => void) | null = null;

  private streamingBuffer = '';

  constructor(options: RealtimeSessionOptions) {
    this.client = new WsClient(options.conversationId);
    this.conversationId = options.conversationId;
    this.queryClient = options.queryClient;
    this.onStreamingText = options.onStreamingText;
    this.onTtsSegment = options.onTtsSegment;
    this.onError = options.onError;

    this.unsubEvent = this.client.onEvent(this.handleEvent);

    if (options.onStateChange) {
      this.unsubState = this.client.onStateChange(options.onStateChange);
    }
  }

  async connect(): Promise<void> {
    await this.client.connect();
  }

  disconnect(): void {
    this.client.disconnect();
  }

  dispose(): void {
    this.commitStreamingBuffer();
    this.unsubEvent?.();
    this.unsubState?.();
    this.client.dispose();
  }

  /** Returns true if the message was sent over the socket. */
  sendText(text: string): boolean {
    return this.client.sendText(text);
  }

  sendAudioSegment(
    audioBase64: string,
    segmentIndex: number,
    options?: {
      voiceSessionId?: string;
      clientSegmentId?: string;
      isLast?: boolean;
      duration?: number;
    },
  ): boolean {
    return this.client.send({
      type: 'audio_segment',
      data: {
        audio_base64: audioBase64,
        segment_index: segmentIndex,
        voice_session_id: options?.voiceSessionId,
        client_segment_id: options?.clientSegmentId,
        is_last: options?.isLast,
        duration: options?.duration,
      },
    });
  }

  sendEndConversation(): boolean {
    return this.client.sendEndConversation();
  }

  getConnectionState(): WsConnectionState {
    return this.client.getState();
  }

  // ─── Internal ───

  private handleEvent: WsEventListener = (event: WsEvent) => {
    if (event.kind === 'agent_response') {
      this.handleAgentChunk(event);
      return;
    }

    if (event.kind === 'tts_audio_received') {
      const b64 = event.audioBase64?.trim();
      const url = event.audioUrl?.trim();
      if (b64 || url) {
        this.onTtsSegment?.({
          audioBase64: b64 || undefined,
          audioUrl: url || undefined,
        });
      }
      return;
    }

    handleWsEvent(this.queryClient, event);

    if (event.kind === 'session_error') {
      this.onError?.(event.message, event.code);
    }
  };

  /**
   * Accumulates agent_response chunks into streamingBuffer.
   * Only commits the final aggregated message to Query cache
   * when the last chunk arrives (index >= total - 1).
   * Individual chunks are forwarded to onStreamingText for UI display.
   */
  private handleAgentChunk(
    event: Extract<WsEvent, { kind: 'agent_response' }>,
  ): void {
    this.streamingBuffer += event.text;

    const isComplete =
      event.index !== undefined &&
      event.total !== undefined &&
      event.index >= event.total - 1;

    this.onStreamingText?.(this.streamingBuffer, isComplete);

    if (isComplete) {
      this.commitStreamingBuffer();
    }
  }

  private commitStreamingBuffer(): void {
    if (!this.streamingBuffer) return;

    const fullText = this.streamingBuffer;
    this.streamingBuffer = '';

    const messagesKey = conversationKeys.messages(this.conversationId);
    this.queryClient.setQueryData<MessageItem[]>(messagesKey, (old) => {
      const message: MessageItem = {
        id: `${this.conversationId}_agent_${Date.now()}`,
        conversationId: this.conversationId,
        content: fullText,
        senderType: 'assistant',
        timestamp: Date.now(),
        messageType: 'text',
      };
      return [...(old ?? []), message];
    });

    this.queryClient.setQueryData<ConversationListItem[]>(
      conversationKeys.lists(),
      (old) => {
        if (!old) return old;
        return old.map((item) =>
          item.id === this.conversationId
            ? {
                ...item,
                latestMessagePreview: lastSegmentPreview(fullText, 50),
                latestMessageTime: Date.now(),
              }
            : item,
        );
      },
    );
  }
}