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

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

export interface TaskRunRuntime {
  sessionId: string;
  runId: string | null;
  thinking: boolean;
  controller: AbortController | null;
}

type Listener = () => void;

/** 按 sessionId 隔离长任务流状态，对齐 useTurnRuntime 模式 */
export function useTaskRunStream(activeSessionId: string) {
  const storeRef = useRef({
    sessions: new Map<string, TaskRunRuntime>(),
    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 getRuntime = useCallback(
    (sessionId: string): TaskRunRuntime => {
      const store = storeRef.current;
      if (!store.sessions.has(sessionId)) {
        store.sessions.set(sessionId, {
          sessionId,
          runId: null,
          thinking: false,
          controller: null,
        });
      }
      return store.sessions.get(sessionId)!;
    },
    [],
  );

  const snapshot = useSyncExternalStore(
    subscribe,
    () => getRuntime(activeSessionId),
    () => getRuntime(activeSessionId),
  );

  const setThinking = useCallback(
    (sessionId: string, thinking: boolean) => {
      const rt = getRuntime(sessionId);
      rt.thinking = thinking;
      notify();
    },
    [getRuntime, notify],
  );

  const setRunId = useCallback(
    (sessionId: string, runId: string | null) => {
      const rt = getRuntime(sessionId);
      rt.runId = runId;
      notify();
    },
    [getRuntime, notify],
  );

  const setController = useCallback(
    (sessionId: string, controller: AbortController | null) => {
      const rt = getRuntime(sessionId);
      rt.controller = controller;
      notify();
    },
    [getRuntime, notify],
  );

  const abortSession = useCallback(
    (sessionId: string) => {
      const rt = getRuntime(sessionId);
      rt.controller?.abort();
      rt.controller = null;
      rt.thinking = false;
      rt.runId = null;
      notify();
    },
    [getRuntime, notify],
  );

  const abortAllSessions = useCallback(() => {
    for (const rt of storeRef.current.sessions.values()) {
      rt.controller?.abort();
      rt.controller = null;
      rt.thinking = false;
    }
    notify();
  }, [notify]);

  return {
    runtime: snapshot,
    isThinking: snapshot.thinking,
    activeTaskRunId: snapshot.runId,
    taskStreamAbortRef: { current: snapshot.controller },
    setThinking,
    setRunId,
    setController,
    abortSession,
    abortAllSessions,
  };
}