life-echo/app-expo/tests/features/conversation/realtime-session-sync-order.test.ts

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

import { RealtimeSession } from '@/features/conversation/realtime-session';
import { conversationKeys } from '@/features/conversation/query-keys';
import type { MessageItem } from '@/features/conversation/types';

jest.mock('@/core/auth/token-manager', () => ({
  tokenManager: {
    getAccessToken: jest.fn().mockResolvedValue('test-token'),
  },
}));

jest.mock('@/core/config', () => ({
  config: {
    wsBaseUrl: 'ws://localhost:8000/',
    ws: {
      reconnectMaxRetries: 3,
      reconnectBaseDelayMs: 10,
      reconnectMaxDelayMs: 100,
      heartbeatIntervalMs: 600000,
    },
  },
}));

class MockWebSocket {
  static OPEN = 1;
  static CLOSED = 3;
  static instances: MockWebSocket[] = [];

  readyState = MockWebSocket.OPEN;
  onopen: (() => void) | null = null;
  onmessage: ((event: { data: string }) => void) | null = null;
  onclose: (() => void) | null = null;
  onerror: (() => void) | null = null;

  constructor(public url: string) {
    MockWebSocket.instances.push(this);
    queueMicrotask(() => this.onopen?.());
  }

  send(): void {}

  close(): void {
    this.readyState = MockWebSocket.CLOSED;
  }

  simulateMessage(data: Record<string, unknown>): void {
    this.onmessage?.({ data: JSON.stringify(data) });
  }
}

(global as Record<string, unknown>).WebSocket = MockWebSocket;

function msgs(qc: QueryClient, cid: string): MessageItem[] {
  return qc.getQueryData<MessageItem[]>(conversationKeys.messages(cid)) ?? [];
}

describe('RealtimeSession sync TTS / agent ordering', () => {
  let qc: QueryClient;

  beforeEach(() => {
    jest.clearAllMocks();
    MockWebSocket.instances = [];
    qc = new QueryClient();
    qc.setQueryData(conversationKeys.messages('conv-x'), []);
  });

  afterEach(async () => {
    await new Promise((r) => setImmediate(r));
  });

  it('defers assistant commit when agent_response arrives before tts_audio (single segment)', async () => {
    const aid = 'aa11aa11-aaaa-aaaa-aaaa-aaaaaaaaaaaa';
    const onTts = jest.fn(() => {
      expect(msgs(qc, 'conv-x').some((m) => m.id === aid)).toBe(true);
    });
    const onStream = jest.fn();
    const session = new RealtimeSession({
      conversationId: 'conv-x',
      queryClient: qc,
      onStreamingText: onStream,
      onTtsSegment: onTts,
    });

    await session.connect();
    await new Promise((r) => setImmediate(r));

    const ws = MockWebSocket.instances[0]!;
    expect(session.sendText('hi', { ttsThisTurn: true })).toBe(true);

    ws.simulateMessage({
      type: 'agent_response',
      conversation_id: 'conv-x',
      data: {
        text: 'Hello segment',
        index: 0,
        total: 1,
        assistant_message_id: aid,
      },
      timestamp: new Date().toISOString(),
    });

    const afterAgentOnly = msgs(qc, 'conv-x').filter(
      (m) => m.senderType === 'assistant',
    );
    expect(afterAgentOnly).toHaveLength(0);

    ws.simulateMessage({
      type: 'tts_audio',
      conversation_id: 'conv-x',
      data: {
        audio_url: 'https://example.com/tts-a.mp3',
        index: 0,
        total: 1,
        assistant_message_id: aid,
      },
      timestamp: new Date().toISOString(),
    });

    expect(onTts).toHaveBeenCalledTimes(1);
    const committed = msgs(qc, 'conv-x').filter(
      (m) => m.senderType === 'assistant',
    );
    expect(committed).toHaveLength(1);
    expect(committed[0]!.content).toContain('Hello segment');

    session.dispose();
  });

  it('multi-segment sync clears pending UI without streaming footer text', async () => {
    const aid = 'bb22bb22-bbbb-bbbb-bbbb-bbbbbbbbbbbb';
    const onTts = jest.fn(() => {
      expect(
        msgs(qc, 'conv-x').some((m) => m.id === `${aid}_seg_0`),
      ).toBe(true);
    });
    const onStream = jest.fn();
    const session = new RealtimeSession({
      conversationId: 'conv-x',
      queryClient: qc,
      onStreamingText: onStream,
      onTtsSegment: onTts,
    });

    await session.connect();
    await new Promise((r) => setImmediate(r));

    const ws = MockWebSocket.instances[0]!;
    session.sendText('hi', { ttsThisTurn: true });

    ws.simulateMessage({
      type: 'tts_audio',
      conversation_id: 'conv-x',
      data: {
        audio_url: 'https://example.com/tts-b.mp3',
        index: 0,
        total: 2,
        assistant_message_id: aid,
      },
      timestamp: new Date().toISOString(),
    });

    ws.simulateMessage({
      type: 'agent_response',
      conversation_id: 'conv-x',
      data: {
        text: 'Part A',
        index: 0,
        total: 2,
        assistant_message_id: aid,
      },
      timestamp: new Date().toISOString(),
    });

    expect(onStream).toHaveBeenCalledWith('', true);
    expect(onStream).not.toHaveBeenCalledWith('Part A', true);
    expect(onTts).toHaveBeenCalled();
    session.dispose();
  });

  it('keeps active screen TTS callback when stale offscreen attach runs later', async () => {
    const aid = 'cc33cc33-cccc-cccc-cccc-cccccccccccc';
    const screenOnTts = jest.fn();
    const offscreenOnTts = jest.fn();
    const session = new RealtimeSession({
      conversationId: 'conv-x',
      queryClient: qc,
    });
    const owner = Symbol('screen-owner');

    session.attachUiCallbacks({ onTtsSegment: screenOnTts }, owner);
    session.attachUiCallbacks({ onTtsSegment: offscreenOnTts });

    await session.connect();
    await new Promise((r) => setImmediate(r));

    const ws = MockWebSocket.instances[0]!;
    ws.simulateMessage({
      type: 'tts_audio',
      conversation_id: 'conv-x',
      data: {
        audio_base64: 'ZmFrZS1tcDM=',
        audio_url: 'https://example.com/tts-c.mp3',
        index: 0,
        total: 1,
        assistant_message_id: aid,
        manual: true,
      },
      timestamp: new Date().toISOString(),
    });

    expect(screenOnTts).toHaveBeenCalledTimes(1);
    expect(offscreenOnTts).not.toHaveBeenCalled();
    session.dispose();
  });

  it('keeps active screen TTS callback when a stale screen owner attaches later', async () => {
    const aid = 'dd44dd44-dddd-dddd-dddd-dddddddddddd';
    const screenOnTts = jest.fn();
    const staleScreenOnTts = jest.fn();
    const session = new RealtimeSession({
      conversationId: 'conv-x',
      queryClient: qc,
    });
    const activeOwner = Symbol('active-screen-owner');
    const staleOwner = Symbol('stale-screen-owner');

    session.attachUiCallbacks({ onTtsSegment: screenOnTts }, activeOwner);
    session.attachUiCallbacks({ onTtsSegment: staleScreenOnTts }, staleOwner);

    await session.connect();
    await new Promise((r) => setImmediate(r));

    const ws = MockWebSocket.instances[0]!;
    ws.simulateMessage({
      type: 'tts_audio',
      conversation_id: 'conv-x',
      data: {
        audio_base64: 'ZmFrZS1tcDM=',
        audio_url: 'https://example.com/tts-d.mp3',
        index: 0,
        total: 1,
        assistant_message_id: aid,
        manual: true,
      },
      timestamp: new Date().toISOString(),
    });

    expect(screenOnTts).toHaveBeenCalledTimes(1);
    expect(staleScreenOnTts).not.toHaveBeenCalled();
    session.dispose();
  });

  it('replays topic_suggestions after attachUiCallbacks with owner (warmup path)', async () => {
    const onTopic = jest.fn();
    const owner = Symbol('ui');
    const session = new RealtimeSession({
      conversationId: 'conv-x',
      queryClient: qc,
    });

    await session.connect();
    await new Promise((r) => setImmediate(r));

    const ws = MockWebSocket.instances[0]!;
    ws.simulateMessage({
      type: 'topic_suggestions',
      conversation_id: 'conv-x',
      data: {
        reason: 'opening',
        stage: 'childhood',
        suggestions: [
          {
            id: 'place',
            label: 'where you grew up',
            text: "I'd like to talk about where you grew up.",
          },
        ],
      },
      timestamp: new Date().toISOString(),
    });

    expect(onTopic).not.toHaveBeenCalled();

    session.attachUiCallbacks(
      {
        onTopicSuggestions: onTopic,
        onStreamingText: () => {},
        onTtsSegment: () => {},
        onError: () => {},
        onStateChange: () => {},
      },
      owner,
    );

    expect(onTopic).toHaveBeenCalledTimes(1);
    expect(onTopic.mock.calls[0]![0].suggestions[0]!.id).toBe('place');

    session.dispose();
  });
});