ReactでリアルタイムチャットUIを作る — WebSocketクライアントとOptimistic UI

❌ コンポーネントに直接書く:
  - ページ遷移のたびにWebSocket接続が切れる
  - 再接続ロジックがUIロジックと混ざって複雑になる
  - テストが困難

✅ サービスクラスに分離:
  - コンポーネントのライフサイクルから独立
  - 接続管理の責務が明確に分離
  - シングルトンで接続を使い回せる

type MessageHandler = (data: Record<string, unknown>) => void;

export class WebSocketService {
  private ws: WebSocket | null;
  private messageHandlers: Set<MessageHandler>;
  private reconnectAttempts: number;
  private maxReconnectAttempts: number;
  private reconnectDelay: number;
  private heartbeatInterval: ReturnType<typeof setInterval> | null;
  private reconnectTimer: ReturnType<typeof setTimeout> | null;
  private intentionalClose: boolean;
  private url: string;
  private getAccessToken: () => string | null;

  constructor(url: string, getAccessToken: () => string | null) {
    this.url = url;
    this.getAccessToken = getAccessToken;
    this.ws = null;
    this.messageHandlers = new Set();
    this.reconnectAttempts = 0;
    this.maxReconnectAttempts = 5;
    this.reconnectDelay = 1000;
    this.heartbeatInterval = null;
    this.reconnectTimer = null;
    this.intentionalClose = false;
  }

connect(): Promise<void> {
    // 既存のゾンビ接続をクリーンアップ
    if (this.ws) {
      this.ws.onclose = null; // 旧接続のcloseで再接続が走らないように
      this.ws.close();
      this.ws = null;
      this.stopHeartbeat();
    }

    this.intentionalClose = false;

    return new Promise((resolve, reject) => {
      const token = this.getAccessToken();
      if (!token) {
        reject(new Error('No access token available'));
        return;
      }

      // クエリパラメータでトークンを渡す（第4回のconnect.tsで受け取る）
      const wsUrl = `${this.url}?token=${token}`;
      this.ws = new WebSocket(wsUrl);

      this.ws.onopen = () => {
        this.reconnectAttempts = 0;  // 成功したらリセット
        this.startHeartbeat();
        resolve();
      };

      this.ws.onmessage = (event) => {
        const message = JSON.parse(event.data) as Record<string, unknown>;
        this.messageHandlers.forEach((handler) => handler(message));
      };

      this.ws.onclose = () => {
        this.stopHeartbeat();
        if (!this.intentionalClose) {
          this.attemptReconnect();  // 意図しない切断のみ再接続
        }
      };

      this.ws.onerror = (error) => reject(error);
    });
  }

private attemptReconnect(): void {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnect attempts reached');
      return;
    }

    this.reconnectAttempts++;
    const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);
    // 1回目: 1秒、2回目: 2秒、3回目: 4秒、4回目: 8秒、5回目: 16秒

    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect().catch(() => {});
    }, delay);
  }

再接続タイムライン:
──×─[1s]─再接続──×─[2s]─再接続──×─[4s]─再接続──✓ 成功
  切断     待機              待機              待機

private startHeartbeat(): void {
    this.stopHeartbeat();
    this.heartbeatInterval = setInterval(() => {
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ action: 'ping' }));
      }
    }, 30000); // 30秒ごと
  }

let wsService: WebSocketService | null = null;

export function getWebSocketService(
  url?: string,
  getAccessToken?: () => string | null
): WebSocketService {
  if (!wsService && url && getAccessToken) {
    wsService = new WebSocketService(url, getAccessToken);
  }
  if (!wsService) {
    throw new Error('WebSocketService not initialized');
  }
  return wsService;
}

import { create } from 'zustand';

export interface Message {
  messageId: string;
  chatRoomId: string;
  senderId: string;
  content: string;
  messageType: 'user' | 'system';
  readBy: string[];
  createdAt: string;
  timestamp: number;
}

export interface ChatRoom {
  chatRoomId: string;
  name?: string;
  participantIds: string[];
  type: 'direct' | 'group';
  activityId?: string;
  lastMessageAt?: string;
  lastMessage?: string;
  createdAt: string;
  unreadCount?: number;
}

interface ChatState {
  rooms: ChatRoom[];
  currentRoom: ChatRoom | null;
  messages: Message[];
  isConnected: boolean;
  setRooms: (rooms: ChatRoom[]) => void;
  setCurrentRoom: (room: ChatRoom | null) => void;
  setMessages: (messages: Message[]) => void;
  addMessage: (message: Message) => void;
  removeMessage: (messageId: string) => void;
  setConnected: (connected: boolean) => void;
}

export const useChatStore = create<ChatState>((set) => ({
  rooms: [],
  currentRoom: null,
  messages: [],
  isConnected: false,
  // ... アクション定義
}));

1. ユーザーAがメッセージを送信
2. Optimistic UIで即座にローカルstateに追加 (messageId: "local-123")
3. サーバーで処理され、WebSocketでブロードキャスト
4. ユーザーA自身もそのブロードキャストを受信 (messageId: "local-123")
5. → 同じメッセージが2つ表示される！

addMessage: (message) =>
  set((state) => {
    const isDuplicate = state.messages.some(
      (m) =>
        // messageIdが一致（Optimistic追加 + WebSocket受信の重複）
        m.messageId === message.messageId ||
        // 同じ送信者・同じ内容・5秒以内（IDが異なるケース対策）
        (m.senderId === message.senderId &&
          m.content === message.content &&
          Math.abs(m.timestamp - message.timestamp) < 5000)
    );
    if (isDuplicate) return state; // 何も変更しない
    return { messages: [...state.messages, message] };
  }),

通常のUI:                    Optimistic UI:
送信ボタン押す                送信ボタン押す
  ↓                           ↓
サーバーに送信 (200ms)        ① 即座にローカルstateに追加（見た目に表示）
  ↓                           ② 同時にサーバーに送信
サーバーから応答                 ↓ (200ms)
  ↓                          サーバーから応答
画面に表示                      ↓
                             ✓ 成功 → そのまま（何もしない）
合計: 200ms+                  ✗ 失敗 → ローカルから削除（ロールバック）

                             合計: 0ms（体感）

const handleSendMessage = useCallback(async (content: string) => {
  const wsService = getWebSocketService(WS_URL, () => idToken);

  // 接続チェック・再接続
  if (!wsService.isConnected()) {
    try {
      await wsService.connect();
    } catch {
      setSendError('接続に失敗しました。ページを更新してください。');
      return;
    }
  }

  const { userId } = useAuthStore.getState();
  const messageId = crypto.randomUUID();
  const now = new Date().toISOString();

  // ① Optimistic: 即座にローカルに追加
  addMessage({
    messageId,
    chatRoomId: chatRoomId!,
    senderId: userId || '',
    content,
    messageType: 'user',
    readBy: [],
    createdAt: now,
    timestamp: Date.now(),
  });

  try {
    // ② サーバーに送信
    wsService.sendMessage('sendMessage', {
      chatRoomId,
      content,
    });
  } catch (err) {
    // ③ 失敗時: ローカルから削除（ロールバック）
    removeMessage(messageId);
    setSendError('メッセージの送信に失敗しました。');
  }
}, [chatRoomId, idToken, addMessage, removeMessage]);

useEffect(() => {
  if (!chatRoomId) return;

  const handleVisibilityChange = async () => {
    if (document.hidden) {
      hiddenAtRef.current = Date.now();
      return;
    }

    // タブが再表示された
    const hiddenMs = hiddenAtRef.current
      ? Date.now() - hiddenAtRef.current
      : 0;
    const wsService = getWebSocketService(WS_URL, () => idToken);

    // 30秒以上非表示だった or 接続が切れている場合に再接続
    if (hiddenMs > 30000 || !wsService.isConnected()) {
      setIsReconnecting(true);
      try {
        await wsService.connect();
        setConnected(true);
      } catch (e) {
        console.error('Reconnect failed:', e);
      } finally {
        setIsReconnecting(false);
      }
    }
  };

  document.addEventListener('visibilitychange', handleVisibilityChange);
  return () => document.removeEventListener('visibilitychange', handleVisibilityChange);
}, [chatRoomId, idToken, setConnected]);

useEffect(() => {
  if (!chatRoomId) return;

  const wsService = getWebSocketService(WS_URL, () => idToken);
  wsService.connect().then(() => setConnected(true));

  const unsubscribe = wsService.onMessage((message) => {
    if (message.type === 'message' && message.data) {
      const data = message.data as Record<string, unknown>;

      // 現在開いているルームのメッセージのみ追加
      if (data.chatRoomId === chatRoomId) {
        addMessage({
          messageId: data.messageId as string,
          chatRoomId: data.chatRoomId as string,
          senderId: data.senderId as string,
          content: data.content as string,
          messageType: data.messageType === 'system' ? 'system' : 'user',
          createdAt: data.createdAt as string,
          readBy: (data.readBy as string[]) || [],
          timestamp: typeof data.timestamp === 'number'
            ? data.timestamp
            : Date.now(),
        });
      }
    }
  });

  return () => {
    unsubscribe();
    wsService.disconnect();
    setConnected(false);
  };
}, [chatRoomId, idToken, addMessage, setConnected]);