feat: move useAgentTimeoutIdStore away from zustand and put it at the session level

Author: 1egoman

packages/react/src/hooks/index.ts

@@ -58,6 +58,12 @@ export * from './useTextStream';
 export * from './useTranscriptions';
 export * from './useSequentialRoomConnectDisconnect';
 export * from './useSession';
-export * from './useAgent';
+export {
+  type AgentState,
+  AgentEvent,
+  type AgentCallbacks,
+  type UseAgentReturn,
+  useAgent,
+} from './useAgent';
 export * from './useEvents';
 export * from './useSessionMessages';

packages/react/src/hooks/useAgent.ts

@@ -10,7 +10,6 @@ import {
 import type TypedEventEmitter from 'typed-emitter';
 import { EventEmitter } from 'events';
 import * as React from 'react';
-import { create } from 'zustand';
 import { ParticipantAgentAttributes, TrackReference } from '@livekit/components-core';
 
 import { useParticipantTracks } from './useParticipantTracks';
@@ -172,91 +171,73 @@ const generateDerivedStateValues = <State extends AgentState>(state: State) =>
     isAvailable: State extends 'listening' | 'thinking' | 'speaking' ? true : false;
   };
 
-const useAgentTimeoutIdStore = create<{
+/** Internal hook used by useSession to store global agent state */
+export const useAgentTimeoutIdStore = (): {
   agentTimeoutFailureReason: string | null;
   startAgentTimeout: (agentConnectTimeoutMilliseconds?: number) => void;
   clearAgentTimeout: () => void;
   updateAgentTimeoutState: (agentState: AgentState) => void;
   updateAgentTimeoutParticipantExists: (agentParticipantExists: boolean) => void;
-  internal: {
-    agentTimeoutId: ReturnType<typeof setTimeout> | null;
-    agentState: AgentState;
-    agentParticipantExists: boolean;
-  };
-}>((set, get) => {
+} => {
+  const [agentTimeoutFailureReason, setAgentTimeoutFailureReason] = React.useState<string | null>(
+    null,
+  );
+  const [agentTimeoutId, setAgentTimeoutId] = React.useState<ReturnType<typeof setTimeout> | null>(
+    null,
+  );
+
+  const agentStateRef = React.useRef<AgentState>('connecting');
+  const agentParticipantExistsRef = React.useRef(false);
+
   const startAgentConnectedTimeout = (agentConnectTimeoutMilliseconds?: number) => {
     return setTimeout(() => {
-      const {
-        internal: { agentState, agentParticipantExists },
-      } = get();
-      if (!agentParticipantExists) {
-        set((old) => ({ ...old, agentTimeoutFailureReason: 'Agent did not join the room.' }));
+      if (!agentParticipantExistsRef.current) {
+        setAgentTimeoutFailureReason('Agent did not join the room.');
         return;
       }
 
-      const { isAvailable } = generateDerivedStateValues(agentState);
+      const { isAvailable } = generateDerivedStateValues(agentStateRef.current);
       if (!isAvailable) {
-        set((old) => ({
-          ...old,
-          agentTimeoutFailureReason: 'Agent connected but did not complete initializing.',
-        }));
+        setAgentTimeoutFailureReason('Agent connected but did not complete initializing.');
         return;
       }
     }, agentConnectTimeoutMilliseconds ?? DEFAULT_AGENT_CONNECT_TIMEOUT_MILLISECONDS);
   };
 
   return {
-    agentTimeoutFailureReason: null,
-    startAgentTimeout: (agentConnectTimeoutMilliseconds?: number) => {
-      set((old) => {
-        if (old.internal.agentTimeoutId) {
-          clearTimeout(old.internal.agentTimeoutId);
-        }
-
-        return {
-          ...old,
-          agentTimeoutFailureReason: null,
-          internal: {
-            ...old.internal,
-            agentTimeoutId: startAgentConnectedTimeout(agentConnectTimeoutMilliseconds),
-            agentState: 'connecting',
-            agentParticipantExists: false,
-          },
-        };
-      });
-    },
-    clearAgentTimeout: () => {
-      set((old) => {
-        if (old.internal.agentTimeoutId) {
-          clearTimeout(old.internal.agentTimeoutId);
+    agentTimeoutFailureReason,
+    startAgentTimeout: React.useCallback(
+      (agentConnectTimeoutMilliseconds?: number) => {
+        if (agentTimeoutId) {
+          clearTimeout(agentTimeoutId);
         }
-        return {
-          ...old,
-          agentTimeoutFailureReason: null,
-          internal: {
-            ...old.internal,
-            agentTimeoutId: null,
-            agentState: 'connecting',
-            agentParticipantExists: false,
-          },
-        };
-      });
-    },
 
-    updateAgentTimeoutState: (agentState: AgentState) => {
-      set((old) => ({ ...old, internal: { ...old.internal, agentState: agentState } }));
-    },
-    updateAgentTimeoutParticipantExists: (agentParticipantExists: boolean) => {
-      set((old) => ({ ...old, internal: { ...old.internal, agentParticipantExists } }));
-    },
+        setAgentTimeoutFailureReason(null);
+        setAgentTimeoutId(startAgentConnectedTimeout(agentConnectTimeoutMilliseconds));
+        agentStateRef.current = 'connecting';
+        agentParticipantExistsRef.current = false;
+      },
+      [agentTimeoutId],
+    ),
+    clearAgentTimeout: React.useCallback(() => {
+      if (agentTimeoutId) {
+        clearTimeout(agentTimeoutId);
+      }
 
-    internal: {
-      agentTimeoutId: null,
-      agentState: 'connecting',
-      agentParticipantExists: false,
-    },
+      setAgentTimeoutFailureReason(null);
+      setAgentTimeoutId(null);
+      agentStateRef.current = 'connecting';
+      agentParticipantExistsRef.current = false;
+    }, [agentTimeoutId]),
+
+    updateAgentTimeoutState: React.useCallback((agentState: AgentState) => {
+      agentStateRef.current = agentState;
+    }, []),
+    updateAgentTimeoutParticipantExists: React.useCallback((agentParticipantExists: boolean) => {
+      agentParticipantExistsRef.current = agentParticipantExists;
+    }, []),
   };
-});
+};
 
 type SessionStub = Pick<UseSessionReturn, 'connectionState' | 'room' | 'internal'>;
 
@@ -275,7 +256,15 @@ export function useAgent(session?: SessionStub): UseAgentReturn {
 
   const {
     room,
-    internal: { agentConnectTimeoutMilliseconds },
+    internal: {
+      agentConnectTimeoutMilliseconds,
+
+      agentTimeoutFailureReason,
+      startAgentTimeout,
+      clearAgentTimeout,
+      updateAgentTimeoutState,
+      updateAgentTimeoutParticipantExists,
+    },
   } = session;
 
   const emitter = React.useMemo(() => new EventEmitter() as TypedEventEmitter<AgentCallbacks>, []);
@@ -398,14 +387,6 @@ export function useAgent(session?: SessionStub): UseAgentReturn {
     };
   }, [room.localParticipant]);
 
-  const {
-    agentTimeoutFailureReason,
-    startAgentTimeout,
-    clearAgentTimeout,
-    updateAgentTimeoutState,
-    updateAgentTimeoutParticipantExists,
-  } = useAgentTimeoutIdStore();
-
   const failureReasons = React.useMemo(() => {
     return agentTimeoutFailureReason ? [agentTimeoutFailureReason] : [];
   }, [agentTimeoutFailureReason]);

packages/react/src/hooks/useSession.ts

@@ -14,7 +14,7 @@ import {
 import { EventEmitter } from 'events';
 
 import { useMaybeRoomContext } from '../context';
-import { useAgent } from './useAgent';
+import { AgentState, useAgent, useAgentTimeoutIdStore } from './useAgent';
 import { TrackReference } from '@livekit/components-core';
 import { useLocalParticipant } from './useLocalParticipant';
 
@@ -72,6 +72,12 @@ type SessionStateCommon = {
     emitter: TypedEventEmitter<SessionCallbacks>;
     tokenSource: TokenSourceConfigurable | TokenSourceFixed;
     agentConnectTimeoutMilliseconds?: number;
+
+    agentTimeoutFailureReason: string | null;
+    startAgentTimeout: (agentConnectTimeoutMilliseconds?: number) => void;
+    clearAgentTimeout: () => void;
+    updateAgentTimeoutState: (agentState: AgentState) => void;
+    updateAgentTimeoutParticipantExists: (agentParticipantExists: boolean) => void;
   };
 };
 
@@ -264,17 +270,45 @@ export function useSession(
     };
   }, [localParticipant, microphonePublication, microphonePublication?.isMuted]);
 
+  const {
+    agentTimeoutFailureReason,
+    startAgentTimeout,
+    clearAgentTimeout,
+    updateAgentTimeoutState,
+    updateAgentTimeoutParticipantExists,
+  } = useAgentTimeoutIdStore();
+
+  const sessionInternal: UseSessionReturn['internal'] = React.useMemo(
+    () => ({
+      emitter,
+      tokenSource,
+      agentConnectTimeoutMilliseconds,
+
+      agentTimeoutFailureReason,
+      startAgentTimeout,
+      clearAgentTimeout,
+      updateAgentTimeoutState,
+      updateAgentTimeoutParticipantExists,
+    }),
+    [
+      emitter,
+      agentConnectTimeoutMilliseconds,
+      tokenSource,
+      agentTimeoutFailureReason,
+      startAgentTimeout,
+      clearAgentTimeout,
+      updateAgentTimeoutState,
+      updateAgentTimeoutParticipantExists,
+    ],
+  );
+
   const conversationState = React.useMemo(():
     | SessionStateConnecting
     | SessionStateConnected
     | SessionStateDisconnected => {
     const common: SessionStateCommon = {
       room,
-      internal: {
-        emitter,
-        tokenSource,
-        agentConnectTimeoutMilliseconds,
-      },
+      internal: sessionInternal,
     };
 
     switch (roomConnectionState) {
@@ -320,8 +354,7 @@ export function useSession(
         };
     }
   }, [
-    tokenSource,
-    agentConnectTimeoutMilliseconds,
+    sessionInternal,
     room,
     emitter,
     roomConnectionState,
@@ -387,12 +420,9 @@ export function useSession(
       () => ({
         connectionState: conversationState.connectionState,
         room,
-        internal: {
-          emitter,
-          tokenSource,
-        },
+        internal: sessionInternal,
       }),
-      [conversationState, emitter, room, tokenSource],
+      [conversationState, room, sessionInternal],
     ),
   );