⚛️ refactor: Move the Steer Atoms to Jotai in the Chat Package (#17135)

Author: berry-13

client/src/Providers/MessagePartsHostContext.tsx

@@ -1,7 +1,12 @@
 import { useCallback, useMemo } from 'react';
 import { useAtomValue, useSetAtom } from 'jotai';
 import { useToastContext } from '@librechat/client';
-import { setDefaultMessagePartsHost } from '@librechat/chat/react';
+import { setDefaultMessagePartsHost, pendingSteersByConvoId } from '@librechat/chat/react';
+import {
+  liveAppliedSteerFamily,
+  liveAppliedSteerIdsAtom,
+  escalatingSteerFamily,
+} from '@librechat/chat/react';
 import {
   useRecoilCallback,
   useRecoilState,
@@ -17,11 +22,6 @@ import type {
   MessagePartsHost,
 } from '~/hooks/Chat/contract';
 import type { Artifact } from '~/common';
-import {
-  liveAppliedSteerFamily,
-  liveAppliedSteerIdsAtom,
-  escalatingSteerFamily,
-} from '~/store/steer';
 import store, { ptcTraceByToolCallId, ptcTraceKey, sandboxStartingByToolCallId } from '~/store';
 import { useMessageContext, useFileMapContext, useAgentsMapContext } from '~/Providers';
 import { getAppPartRenderers } from '~/components/Chat/Messages/Content/Parts/slots';
@@ -124,8 +124,7 @@ export const appMessagePartsHost: MessagePartsHost = {
     ),
   useToolArtifactClaim: (artifactId) => useRecoilState(store.toolArtifactClaim(artifactId)),
   useArtifactPanel: useAppArtifactPanel,
-  usePendingSteers: (conversationId) =>
-    useRecoilValue(store.pendingSteersByConvoId(conversationId)),
+  usePendingSteers: (conversationId) => useAtomValue(pendingSteersByConvoId(conversationId)),
   useSteerEscalating: (conversationId) => useAtomValue(escalatingSteerFamily(conversationId)),
   usePaneConversationId: (index) => useRecoilValue(store.conversationIdByIndex(index)),
   useLiveAppliedSteer: useAppLiveAppliedSteer,

client/src/components/Chat/Input/Composer/Queue.tsx

@@ -1,11 +1,15 @@
 import { memo, useId, useRef, useMemo, useState, useCallback } from 'react';
 import { useAtomValue } from 'jotai';
-import { useRecoilValue } from 'recoil';
 import { createPortal } from 'react-dom';
 import * as Ariakit from '@ariakit/react';
 import { useDrag, useDrop } from 'react-dnd';
 import { AnimatePresence, motion, useReducedMotion } from 'framer-motion';
 import { claimQueuedIntent, releaseQueuedIntent, hasQueuedIntent } from '@librechat/chat';
+import {
+  escalatingSteerFamily,
+  revealedQueuedTurnFamily,
+  pendingSteersByConvoId,
+} from '@librechat/chat/react';
 import {
   Button,
   IconButton,
@@ -31,12 +35,10 @@ import type { RestoreToComposer } from '~/Providers/ComposerRestoreContext';
 import type { SteeringControls } from '~/hooks/Chat/useSteering';
 import type { QueuedMessage } from '~/hooks/Chat/queue';
 import { useQueuedTurnPortal } from '~/components/Chat/Steering/QueuedTurnPortal';
-import { escalatingSteerFamily, revealedQueuedTurnFamily } from '~/store/steer';
 import EscalateNowButton from '~/components/Chat/Input/EscalateNowButton';
 import { queuedMessagesByConvoId } from '~/hooks/Chat/queue';
 import { useLocalize } from '~/hooks';
 import { cn } from '~/utils';
-import store from '~/store';
 
 const DRAG_TYPE = 'queued-message';
 
@@ -733,7 +735,7 @@ function Queue({
      `aria-describedby` at whichever copy the document happened to keep. */
   const reorderHintId = useId();
   const queued = useAtomValue(queuedMessagesByConvoId(steering.queueKey));
-  const pendingSteers = useRecoilValue(store.pendingSteersByConvoId(conversationId));
+  const pendingSteers = useAtomValue(pendingSteersByConvoId(conversationId));
   const escalating = useAtomValue(escalatingSteerFamily(conversationId));
   /* Only one interrupt can be unresolved at a time: a second arm would seal the
      same run twice. The escalating flag covers an arm's round trip, before its

client/src/components/Chat/Input/Composer/__tests__/Queue.spec.tsx

@@ -5,6 +5,7 @@ import { getDefaultStore, useSetAtom } from 'jotai';
 import { HTML5Backend } from 'react-dnd-html5-backend';
 import { ReasoningEffort } from 'librechat-data-provider';
 import { hasQueuedIntent, releaseQueuedIntent } from '@librechat/chat';
+import { revealedQueuedTurnFamily, resetSteerFamilies } from '@librechat/chat/react';
 import { act, render, screen, within, waitFor, fireEvent } from '@testing-library/react';
 import type { SteeringControls } from '~/hooks/Chat/useSteering';
 import type { QueuedMessage } from '~/hooks/Chat/queue';
@@ -13,7 +14,6 @@ import {
   useQueuedTurnPortal,
 } from '~/components/Chat/Steering/QueuedTurnPortal';
 import { queuedMessagesByConvoId, resetQueueFamilies } from '~/hooks/Chat/queue';
-import { revealedQueuedTurnFamily } from '~/store/steer';
 import Queue from '../Queue';
 jest.mock('~/hooks', () => ({
   useLocalize: () => (key: string, options?: Record<string, string | number>) => {
@@ -201,7 +201,10 @@ function renderQueue(
   );
 }
 
-beforeEach(() => resetQueueFamilies());
+beforeEach(() => {
+  resetQueueFamilies();
+  resetSteerFamilies();
+});
 
 describe('Queue', () => {
   beforeEach(() => jest.clearAllMocks());

client/src/components/Chat/Input/__tests__/ChatForm.queryParams.spec.tsx

@@ -8,6 +8,10 @@ import { QueryKeys, EModelEndpoint } from 'librechat-data-provider';
 import { act, render, screen, fireEvent } from '@testing-library/react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { MemoryRouter, useNavigate, useLocation } from 'react-router-dom';
+import {
+  activeGenerationCreatedAtByConvoId,
+  activeGenerationProtocolVersionByConvoId,
+} from '@librechat/chat/react';
 import type { TConversation } from 'librechat-data-provider';
 import type { Transport } from '~/hooks/Chat/contract';
 import type { ChatFormValues } from '~/common';
@@ -167,6 +171,8 @@ function mountComposer(
   const jotaiStore = createStore();
   if (liveRun) {
     jotaiStore.set(showStopButtonByIndex(0), true);
+    jotaiStore.set(activeGenerationCreatedAtByConvoId(conversation.conversationId ?? ''), 1000);
+    jotaiStore.set(activeGenerationProtocolVersionByConvoId(conversation.conversationId ?? ''), 2);
   }
   const wrapper = ({ children }: { children: React.ReactNode }) => (
     <QueryClientProvider client={queryClient}>
@@ -175,14 +181,6 @@ function mountComposer(
           initializeState={({ set }) => {
             if (liveRun) {
               set(store.isSubmittingFamily(0), true);
-              set(
-                store.activeGenerationCreatedAtByConvoId(conversation.conversationId ?? ''),
-                1000,
-              );
-              set(
-                store.activeGenerationProtocolVersionByConvoId(conversation.conversationId ?? ''),
-                2,
-              );
             }
             if (speechSettingsInitialized) {
               set(store.engineSTT, 'external');

client/src/components/Chat/Input/__tests__/QueuedTurnPortal.spec.tsx

@@ -3,13 +3,13 @@ import { RecoilRoot } from 'recoil';
 import { EModelEndpoint } from 'librechat-data-provider';
 import { act, render, screen } from '@testing-library/react';
 import { Provider as JotaiProvider, createStore } from 'jotai';
+import { revealedQueuedTurnFamily } from '@librechat/chat/react';
 import type { TConversation } from 'librechat-data-provider';
 import {
   QueuedTurnPortalProvider,
   useQueuedTurnPortal,
 } from '~/components/Chat/Steering/QueuedTurnPortal';
 import PendingTurn from '~/components/Chat/Messages/PendingTurn';
-import { revealedQueuedTurnFamily } from '~/store/steer';
 import { ChatContext } from '~/Providers';
 
 jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key }));

client/src/components/Chat/Messages/Content/ApprovalContext.tsx

@@ -7,9 +7,9 @@ import {
   useRef,
   useState,
 } from 'react';
-import { useRecoilValue } from 'recoil';
-import { useAtom, useStore } from 'jotai';
 import { Constants } from 'librechat-data-provider';
+import { useAtom, useStore, useAtomValue } from 'jotai';
+import { activeGenerationCreatedAtByConvoId } from '@librechat/chat/react';
 import type { Agents } from 'librechat-data-provider';
 import type { AskAnswerStatus } from '~/components/Chat/ask/state';
 import {
@@ -22,7 +22,6 @@ import { askSubmitStatusAtom } from '~/components/Chat/ask/state';
 import { resolveAskUserQuestionPart } from '~/utils/approval';
 import { ChatContext } from '~/Providers/ChatContext';
 import { useGetEphemeralAgent } from '~/store/agents';
-import store from '~/store';
 
 /** Shared read/write for the ask-answer submit status. */
 export function useAskSubmitStatus(): {
@@ -438,8 +437,8 @@ export function useResumeSubmit() {
   /** Ask status lives in Jotai so it works from the composer (outside the
    *  provider); tool-approval status stays on the context. */
   const { setAskStatus } = useAskSubmitStatus();
-  const activeGenerationCreatedAt = useRecoilValue(
-    store.activeGenerationCreatedAtByConvoId(conversation?.conversationId ?? Constants.NEW_CONVO),
+  const activeGenerationCreatedAt = useAtomValue(
+    activeGenerationCreatedAtByConvoId(conversation?.conversationId ?? Constants.NEW_CONVO),
   );
 
   const buildResumeFields = useCallback((): ResumeAgentFields | null => {

client/src/components/Chat/Messages/Content/Parts/__tests__/PendingSteers.test.tsx

@@ -2,6 +2,7 @@ import React from 'react';
 import { RecoilRoot } from 'recoil';
 import { MemoryRouter } from 'react-router-dom';
 import { Provider as JotaiProvider, createStore } from 'jotai';
+import { pendingSteersByConvoId } from '@librechat/chat/react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { Constants, ContentTypes, QueryKeys } from 'librechat-data-provider';
 import { render, screen, fireEvent, waitFor, act } from '@testing-library/react';
@@ -86,6 +87,7 @@ function renderPending(
   });
   queryClient.setQueryData([QueryKeys.messages, CONVO_ID], messages ?? []);
   const jotaiStore = createStore();
+  jotaiStore.set(pendingSteersByConvoId(CONVO_ID), steers);
   if (activeSiblingIndex != null) {
     jotaiStore.set(
       siblingIdxFamily(
@@ -110,7 +112,6 @@ function renderPending(
                   conversationId: CONVO_ID,
                 } as TConversation);
               }
-              set(store.pendingSteersByConvoId(CONVO_ID), steers);
             }}
           >
             <PendingSteers conversationId={CONVO_ID} index={pane?.index} fullWidth={fullWidth} />

client/src/components/Chat/Messages/Content/Parts/__tests__/SteerPart.test.tsx

@@ -2,11 +2,11 @@ import React from 'react';
 import { RecoilRoot } from 'recoil';
 import { QueryKeys } from 'librechat-data-provider';
 import { Provider, createStore, useAtomValue } from 'jotai';
+import { liveAppliedSteerIdsAtom } from '@librechat/chat/react';
 import { render, screen, fireEvent } from '@testing-library/react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import type { TMessage } from 'librechat-data-provider';
 import { STEER_ICON } from '~/components/Chat/Steering/identity';
-import { liveAppliedSteerIdsAtom } from '~/store/steer';
 import SteerPart from '../SteerPart';
 import store from '~/store';
 

client/src/components/Chat/Messages/Content/__tests__/ApprovalContext.test.tsx

@@ -2,10 +2,9 @@ import React from 'react';
 import { RecoilRoot } from 'recoil';
 import { act, renderHook } from '@testing-library/react';
 import { Provider as JotaiProvider, createStore } from 'jotai';
-import type { MutableSnapshot } from 'recoil';
+import { activeGenerationCreatedAtByConvoId } from '@librechat/chat/react';
 import ApprovalProvider, { useApprovalContext, useResumeSubmit } from '../ApprovalContext';
 import { ChatContext } from '~/Providers/ChatContext';
-import store from '~/store';
 
 const mockApprovalMutate = jest.fn();
 const mockAskMutate = jest.fn();
@@ -29,17 +28,15 @@ const chatContextValue = {
 
 const createWrapper = (generationCreatedAt: number | null) =>
   function ResumeWrapper({ children }: { children: React.ReactNode }) {
-    const jotaiStore = React.useRef(createStore()).current;
-    const initializeState = (snapshot: MutableSnapshot) => {
+    const [jotaiStore] = React.useState(() => {
+      const created = createStore();
       if (generationCreatedAt != null) {
-        snapshot.set(
-          store.activeGenerationCreatedAtByConvoId('conversation-1'),
-          generationCreatedAt,
-        );
+        created.set(activeGenerationCreatedAtByConvoId('conversation-1'), generationCreatedAt);
       }
-    };
+      return created;
+    });
     return (
-      <RecoilRoot initializeState={initializeState}>
+      <RecoilRoot>
         <JotaiProvider store={jotaiStore}>
           <ChatContext.Provider value={chatContextValue}>
             <ApprovalProvider>{children}</ApprovalProvider>

client/src/components/Chat/Messages/HoverButtons.tsx

@@ -2,6 +2,7 @@ import React, { useState, useMemo, useCallback, memo } from 'react';
 import { Copy, Check } from 'lucide';
 import { useAtomValue } from 'jotai';
 import { useRecoilState } from 'recoil';
+import { revealedQueuedTurnFamily } from '@librechat/chat/react';
 import { getSpeechText, findMessageById, isUserInitiatedCompaction } from 'librechat-data-provider';
 import {
   Button,
@@ -15,7 +16,6 @@ import type { TConversation, TMessage, TFeedback } from 'librechat-data-provider
 import { useMessagesIsSubmitting, useOptionalMessagesOperations } from '~/Providers';
 import { useGenerationsByLatest, useLocalize } from '~/hooks';
 import { hasEditablePart } from './Content/editableParts';
-import { revealedQueuedTurnFamily } from '~/store/steer';
 import { Fork } from '~/components/Conversations';
 import { hoverButtonClasses } from './styles';
 import MessageAudio from './MessageAudio';