Agent-System/web/src/hooks/useTurnRuntime.ts

import { useCallback, useRef, useSyncExternalStore } from 'react';

export type TurnRuntimeStatus =
  | 'pending'
  | 'streaming'
  | 'completed'
  | 'failed'
  | 'cancelled';

export interface TurnRuntime {
  clientTurnId: string;
  sessionId: string;
  userMsgId: string;
  asstId: string;
  status: TurnRuntimeStatus;
  controller: AbortController;
  startedAt: number;
}

type Listener = () => void;

/** 按 sessionId → clientTurnId 隔离流运行态 */
export function useTurnRuntime(activeSessionId: string) {
  const storeRef = useRef({
    turns: new Map<string, Map<string, TurnRuntime>>(),
    listeners: new Set<Listener>(),
  });

  const subscribe = useCallback((listener: Listener) => {
    storeRef.current.listeners.add(listener);
    return () => {
      storeRef.current.listeners.delete(listener);
    };
  }, []);

  const notify = useCallback(() => {
    for (const l of storeRef.current.listeners) {
      l();
    }
  }, []);

  const getSessionMap = useCallback((sessionId: string) => {
    const store = storeRef.current;
    if (!store.turns.has(sessionId)) {
      store.turns.set(sessionId, new Map());
    }
    return store.turns.get(sessionId)!;
  }, []);

  const registerTurn = useCallback(
    (runtime: TurnRuntime) => {
      const sessionMap = getSessionMap(runtime.sessionId);
      // 同 session 新 turn 开始前，仅中止该 session 内仍在 pending/streaming 的旧 turn
      for (const [turnId, existing] of sessionMap) {
        if (
          turnId !== runtime.clientTurnId &&
          (existing.status === 'pending' || existing.status === 'streaming')
        ) {
          existing.controller.abort();
          existing.status = 'cancelled';
        }
      }
      sessionMap.set(runtime.clientTurnId, runtime);
      notify();
    },
    [getSessionMap, notify],
  );

  const updateTurnStatus = useCallback(
    (sessionId: string, clientTurnId: string, status: TurnRuntimeStatus) => {
      const rt = getSessionMap(sessionId).get(clientTurnId);
      if (rt) {
        rt.status = status;
        notify();
      }
    },
    [getSessionMap, notify],
  );

  const removeTurn = useCallback(
    (sessionId: string, clientTurnId: string) => {
      getSessionMap(sessionId).delete(clientTurnId);
      notify();
    },
    [getSessionMap, notify],
  );

  const abortTurn = useCallback(
    (sessionId: string, clientTurnId: string) => {
      const rt = getSessionMap(sessionId).get(clientTurnId);
      if (rt) {
        rt.controller.abort();
        rt.status = 'cancelled';
        notify();
      }
    },
    [getSessionMap, notify],
  );

  const abortActiveSessionTurns = useCallback(() => {
    const sessionMap = getSessionMap(activeSessionId);
    for (const rt of sessionMap.values()) {
      if (rt.status === 'pending' || rt.status === 'streaming') {
        rt.controller.abort();
        rt.status = 'cancelled';
      }
    }
    notify();
  }, [activeSessionId, getSessionMap, notify]);

  const isSessionThinking = useCallback(
    (sessionId: string) => {
      const sessionMap = storeRef.current.turns.get(sessionId);
      if (!sessionMap) return false;
      for (const rt of sessionMap.values()) {
        if (rt.status === 'pending' || rt.status === 'streaming') return true;
      }
      return false;
    },
    [],
  );

  const isThinking = useSyncExternalStore(
    subscribe,
    () => isSessionThinking(activeSessionId),
    () => false,
  );

  const getTurn = useCallback(
    (sessionId: string, clientTurnId: string) =>
      storeRef.current.turns.get(sessionId)?.get(clientTurnId),
    [],
  );

  return {
    registerTurn,
    updateTurnStatus,
    removeTurn,
    abortTurn,
    abortActiveSessionTurns,
    isThinking,
    isSessionThinking,
    getTurn,
  };
}